cyclop/tools/persona_tool.prepy

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