79 lines
2.5 KiB
Plaintext
79 lines
2.5 KiB
Plaintext
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 |