import os
from openai import OpenAI
from selfmemory.llms.base import LLMBase
from selfmemory.llms.configs import BaseLlmConfig
class OpenAIStructuredLLM(LLMBase):
def __init__(self, config: BaseLlmConfig | None = None):
super().__init__(config)
if not self.config.model:
self.config.model = "gpt-4o-2024-08-06"
api_key = self.config.api_key or os.getenv("OPENAI_API_KEY")
base_url = (
self.config.openai_base_url
or os.getenv("OPENAI_API_BASE")
or "https://api.openai.com/v1"
)
self.client = OpenAI(api_key=api_key, base_url=base_url)
def generate_response(
self,
messages: list[dict[str, str]],
response_format: str | None = None,
tools: list[dict] | None = None,
tool_choice: str = "auto",
) -> str:
"""
Generate a response based on the given messages using OpenAI.
Args:
messages (List[Dict[str, str]]): A list of dictionaries, each containing a 'role' and 'content' key.
response_format (Optional[str]): The desired format of the response. Defaults to None.
Returns:
str: The generated response.
"""
params = {
"model": self.config.model,
"messages": messages,
"temperature": self.config.temperature,
}
if response_format:
params["response_format"] = response_format
if tools:
params["tools"] = tools
params["tool_choice"] = tool_choice
response = self.client.beta.chat.completions.parse(**params)
return response.choices[0].message.content
