from dataclasses import dataclass
from typing import Literal

import httpx2
from openai import OpenAI


@dataclass
class AgentMessage:
    content: str
    """Content of the message"""

    role: Literal["system", "assistant", "user"]
    """Who sent the message"""

class Agent:
    """Perform operations with timeline events using OpenAI-compatible API"""
    def __init__(self, *, model: str, base_url: str | None, api_key: str, **kwargs) -> None:
        self._client = OpenAI(
            base_url=base_url,
            api_key=api_key,
            http_client=httpx2.Client(verify=False),
            **kwargs
        )
        self._model = model
        

    def completion(self, messages: list[AgentMessage], **kwargs) -> str:
        """Generate a completion for specified messages."""
        messages_raw = []
        for m in messages:
            messages_raw.append(
                {
                    "role": m.role,
                    "content": m.content
                }
            )
        response = self._client.chat.completions.create(
            model=self._model,
            messages=messages_raw,
            **kwargs
        )
        return response.choices[0].message.content # type: ignore