Implement get_functions to return OpenAI function schema as a JSON string

tools/persona_tool.py

@@ -1,11 +1,11 @@
 import openai
+import json
 from tools.base_tool import BaseTool
 
 class PersonaTool(BaseTool):
-    def __init__(self):
+    def __init__(self, api_key: str):
         super().__init__()
-        # Initialize OpenAI API key
+        openai.api_key = api_key
-        openai.api_key = "YOUR_OPENAI_API_KEY"
 
     def generate_response(self, persona_description: str, query: str) -> str:
         """
@@ -18,16 +18,41 @@ class PersonaTool(BaseTool):
         Returns:
             str: The response generated by the OpenAI API.
         """
-        try:
+        response = openai.ChatCompletion.create(
-            response = openai.ChatCompletion.create(
+            model="gpt-3.5-turbo",
-                model="gpt-3.5-turbo",  # Specify the model
+            messages=[
-                messages=[
+                {"role": "system", "content": persona_description},
-                    {"role": "system", "content": persona_description},
+                {"role": "user", "content": query}
-                    {"role": "user", "content": query},
+            ]
-                ],
+        )
-                max_tokens=150  # Adjust token limit as needed
+        return response.choices[0].message['content']
-            )
-            return response['choices'][0]['message']['content']
 
-        except Exception as e:
+    def get_functions(self):
-            return f"An error occurred: {str(e)}"
+        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")