import openai
from tools.base_tool import BaseTool

class PersonaTool(BaseTool):
    def __init__(self):
        super().__init__()
        # Initialize OpenAI API key
        openai.api_key = "YOUR_OPENAI_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.
        """
        try:
            response = openai.ChatCompletion.create(
                model="gpt-3.5-turbo",  # Specify the model
                messages=[
                    {"role": "system", "content": persona_description},
                    {"role": "user", "content": query},
                ],
                max_tokens=150  # Adjust token limit as needed
            )
            return response['choices'][0]['message']['content']
        
        except Exception as e:
            return f"An error occurred: {str(e)}"