🐛 fix(AnthropicLLMEngine): Use fix_busted_json to fix badly formatted json by claude

src/engines/LLMEngine/AnthropicLLMEngine.py

@@ -1,5 +1,6 @@
 import anthropic
 import gradio as gr
+import fix_busted_json
 import orjson
 
 from .BaseLLMEngine import BaseLLMEngine
@@ -18,9 +19,8 @@ class AnthropicLLMEngine(BaseLLMEngine):
 
     def __init__(self, options: list) -> None:
         self.model = options[0]
-        self.client = anthropic.Anthropic(
+        api_key = self.retrieve_setting(identifier="anthropic_api_key")["api_key"]
-            api_key="YourAnthropicAPIKeyHere"
+        self.client = anthropic.Anthropic(api_key=api_key)
-        )  # Ensure API key is securely managed
         super().__init__()
 
     def generate(
@@ -34,27 +34,35 @@ class AnthropicLLMEngine(BaseLLMEngine):
         frequency_penalty: float = 0,
         presence_penalty: float = 0,
     ) -> str | dict:
-        prompt = f"""{anthropic.HUMAN_PROMPT} {system_prompt} {anthropic.HUMAN_PROMPT} {chat_prompt} {anthropic.AI_PROMPT}"""
+        tries = 0
+        while tries < 2:
+            messages = [
+                {"role": "user", "content": chat_prompt},
+            ]
             if json_mode:
                 # anthropic does not officially support JSON mode, but we can bias the output towards a JSON-like format
-            prompt += " {"
+                messages.append({"role": "assistant", "content": "{"})
-        # noinspection PyArgumentList
+            response: anthropic.types.Message = self.client.messages.create(
-        response: anthropic.types.Completion = self.client.completions.create(
+                max_tokens=max_tokens,
-            max_tokens_to_sample=max_tokens,
+                messages=messages,
-            prompt=prompt,
                 model=self.model,
                 top_p=top_p,
                 temperature=temperature,
-            frequency_penalty=frequency_penalty,
+                system=system_prompt,
             )
 
-        content = response.completion
+            content = response.content[0].text
             if json_mode:
-            # we add back the opening curly brace wich is not included in the response since it is in the prompt
                 content = "{" + content
                 # we remove everything after the last closing curly brace
                 content = content[: content.rfind("}") + 1]
-            return orjson.loads(content)
+                content = content.replace("\n", "")
+                try:
+                    returnable = fix_busted_json.repair_json(content)
+                    returnable = orjson.loads(returnable)
+                    return returnable
+                except Exception as e: # noqa wait for library to imlement pep https://peps.python.org/pep-0352/ (Required Superclass for Exceptions
+                    tries += 1
             else:
                 return content
 
@@ -68,3 +76,21 @@ class AnthropicLLMEngine(BaseLLMEngine):
                 value=ANTHROPIC_POSSIBLE_MODELS[0],
             )
         ]
+
+    @classmethod
+    def get_settings(cls):
+        current_api_key = cls.retrieve_setting(identifier="anthropic_api_key")
+        current_api_key = current_api_key["api_key"] if current_api_key else ""
+        api_key_input = gr.Textbox(
+            label="Anthropic API Key",
+            type="password",
+            value=current_api_key,
+        )
+        save = gr.Button("Save")
+
+        def save_api_key(api_key: str):
+            cls.store_setting(identifier="anthropic_api_key", data={"api_key": api_key})
+            gr.Info("API key saved successfully.")
+            return gr.update(value=api_key)
+
+        save.click(save_api_key, inputs=[api_key_input])