From 8eaad30ce0a4fe326dff7855e723d52639be4a34 Mon Sep 17 00:00:00 2001 From: itsmvd Date: Mon, 23 Dec 2024 11:04:53 +0000 Subject: [PATCH] update interface.py --- timesketch/lib/llms/interface.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/timesketch/lib/llms/interface.py b/timesketch/lib/llms/interface.py index 743b0253c3..014e0887b4 100644 --- a/timesketch/lib/llms/interface.py +++ b/timesketch/lib/llms/interface.py @@ -14,6 +14,7 @@ """Interface for LLM providers.""" import string +from typing import Optional from flask import current_app @@ -82,13 +83,13 @@ def prompt_from_template(self, template: str, kwargs: dict) -> str: formatter = string.Formatter() return formatter.format(template, **kwargs) - def generate(self, prompt: str) -> str: + def generate(self, prompt: str, response_schema: Optional[dict] = None) -> str: """Generate a response from the LLM provider. Args: prompt: The prompt to generate a response for. - temperature: The temperature to use for the response. - stream: Whether to stream the response. + response_schema: An optional JSON schema to define the expected + response format. Returns: The generated response.