Fixes for anthropic

2024-08-18 08:55:22 -05:00
parent 5134312cab
commit 3d521f4790
3 changed files with 17 additions and 11 deletions
@@ -0,0 +1,79 @@
+import os
+import openai
+import json
+from tools.base_tool import BaseTool
+
+class PersonaTool(BaseTool):
+
+
+    def __init__(self):
+        super().__init__()
+        self.api_key = os.environ.get("OPENAI_API_KEY")
+
+    GPT_4O = "gpt-4o"
+    GPT_4O_MINI = "gpt-4o-mini"
+
+    def generate_response(self, persona_description: str, query: str) -> str:
+        """
+        Makes a call to the OpenAI API using the persona as a system prompt.
+        
+        Parameters:
+            persona_description (str): Description of the persona.
+            query (str): Query to be processed.
+        
+        Returns:
+            str: The response generated by the OpenAI API.
+        """
+        response = openai.ChatCompletion.create(
+            model="gpt-3.5-turbo",
+            messages=[
+                {"role": "system", "content": persona_description},
+                {"role": "user", "content": query}
+            ]
+        )
+        return response.choices[0].message['content']
+
+    def get_functions(self):
+        return json.dumps({
+            "functions": [
+                {
+                    "name": "generate_response",
+                    "description": "Generates a response based on a persona description and a user query.",
+                    "parameters": {
+                        "type": "object",
+                        "properties": {
+                            "persona_description": {
+                                "type": "string",
+                                "description": "Description of the persona."
+                            },
+                            "query": {
+                                "type": "string",
+                                "description": "User's query to be processed."
+                            }
+                        },
+                        "required": ["persona_description", "query"]
+                    }
+                }
+            ]
+        })
+
+    def execute(self, function_name, **kwargs):
+        if function_name == "generate_response":
+            return self.generate_response(kwargs.get("persona_description"), kwargs.get("query"))
+        else:
+            raise ValueError(f"Function {function_name} not found")
+        
+    def get_chat_response(client, messages, model):
+        
+        model_max_tokens = {
+            GPT_4O: 4096,
+            GPT_4O_MINI: 16384
+        }
+
+        response = client.chat.completions.create(
+            model=model,
+            messages=messages,
+            function_call="none",
+            max_tokens=model_max_tokens[model]
+        )
+        return response