From 53a6f826889661faae5f80f4bfd264faaa1e663d Mon Sep 17 00:00:00 2001 From: itsmvd Date: Mon, 23 Dec 2024 10:46:56 +0000 Subject: [PATCH] linter stuff --- timesketch/lib/llms/aistudio.py | 23 ++++++++++++----------- 1 file changed, 12 insertions(+), 11 deletions(-) diff --git a/timesketch/lib/llms/aistudio.py b/timesketch/lib/llms/aistudio.py index b4a1097bd1..13665e631a 100644 --- a/timesketch/lib/llms/aistudio.py +++ b/timesketch/lib/llms/aistudio.py @@ -26,6 +26,7 @@ except ImportError: has_required_deps = False + class AIStudio(interface.LLMProvider): """AI Studio LLM provider.""" @@ -49,7 +50,6 @@ def __init__(self, config: dict): genai.configure(api_key=self._api_key) self.model = genai.GenerativeModel(model_name=self._model_name) - def generate(self, prompt: str, response_schema: Optional[dict] = None) -> str: """ Generate text using the Google AI Studio service. @@ -61,18 +61,18 @@ def generate(self, prompt: str, response_schema: Optional[dict] = None) -> str: Returns: The generated text as a string (or parsed data if response_schema is provided). """ - + generation_config = genai.GenerationConfig( temperature=self._temperature, top_p=self._top_p, - top_k = self._top_k, - max_output_tokens = self._max_output_tokens, + top_k=self._top_k, + max_output_tokens=self._max_output_tokens, ) - + if response_schema: - generation_config.response_mime_type = "application/json" - generation_config.response_schema = response_schema - + generation_config.response_mime_type = "application/json" + generation_config.response_schema = response_schema + response = self.model.generate_content( contents=prompt, generation_config=generation_config, @@ -81,9 +81,10 @@ def generate(self, prompt: str, response_schema: Optional[dict] = None) -> str: if response_schema: try: return json.loads(response.text) - except Exception as e: - print(f"Error processing JSON response: {e}") - return f"Error processing JSON: {e}. Raw response: {response.text}" + except Exception as error: + raise ValueError( + f"Error JSON parsing text: {response.text}: {e}" + ) from error return response.text