src/openai_agent.py

import os
import json
from openai import OpenAI


class AgentResponse:
    def __init__(self, response):
        self.content = response

    def __str__(self):
        return self.content

    def __repr__(self):
        return self.content

    def lower(self):
        self.content = self.content.lower()
        return self

    def is_equal(self, condition):
        return self.content == condition

    def as_boolean(self):
        try:
            return bool(self.content)
        except ValueError:
            raise ValueError(f"Cannot convert {self.content} to boolean")

    def as_json(self):
        try:
            return json.loads(self.content)
        except json.JSONDecodeError:
            raise ValueError(f"Cannot convert {self.content} to JSON")


class Agent:
    def __init__(self, prompt_folder: str = "prompts"):
        organization = os.getenv("OPENAI_ORGANIZATION")
        project = os.getenv("OPENAI_PROJECT")
        api_key = os.getenv("OPENAI_API_KEY")

        if not organization or not project or not api_key:
            raise ValueError(
                "Please set the OPENAI_ORGANIZATION, OPENAI_PROJECT, and OPENAI_API_KEY environment variables"
            )
        self.client = OpenAI(
            organization=organization,
            project=project,
            api_key=api_key,
        )
        self.prompt_folder = prompt_folder

    def request(self, query, model="gpt-4o-mini", max_tokens=300) -> AgentResponse:
        response = self.client.chat.completions.create(
            model=model,
            messages=[
                {
                    "role": "user",
                    "content": [
                        {
                            "type": "text",
                            "text": query,
                        }
                    ],
                }
            ],
            max_tokens=max_tokens,
        )
        return AgentResponse(response.choices[0].message.content)

    def read_prompt(self, prompt_name: str):
        prompt_file = f"{self.prompt_folder}/{prompt_name}.txt"

        if not os.path.exists(prompt_file):
            raise FileNotFoundError(f"Prompt file not found at {prompt_file}")

        with open(prompt_file, "r") as file:
            return file.read()

    def image_request(
        self, image: str, query: str, max_tokens: int = 300, force_json: bool = False
    ) -> AgentResponse:
        try:
            response = self.client.chat.completions.create(
                model="gpt-4o-mini",
                messages=[
                    {
                        "role": "user",
                        "content": [
                            {
                                "type": "text",
                                "text": query,
                            },
                            {
                                "type": "image_url",
                                "image_url": {"url": f"data:image/jpg;base64,{image}"},
                            },
                        ],
                    }
                ],
                max_tokens=max_tokens,
                response_format={"type": "json_object"} if force_json else None,
            )
            result = AgentResponse(response.choices[0].message.content)
            return result
        except Exception as e:
            raise