Merge pull request #26 from bucolucas/enhance-persona-tool

Enhance Persona Tool

path/to/PersonaTool.py

@@ -0,0 +1,4 @@
+class PersonaTool:
+    def get_functions(self):
+        # Implementation to return functions in JSON format
+        pass

tools/persona_tool.py

@@ -0,0 +1,58 @@
+import openai
+import json
+from tools.base_tool import BaseTool
+
+class PersonaTool(BaseTool):
+    def __init__(self, api_key: str):
+        super().__init__()
+        openai.api_key = api_key
+
+    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")