33 lines
1.0 KiB
Python
33 lines
1.0 KiB
Python
import httpx2
|
|
from openai import OpenAI
|
|
|
|
from utils import AgentMessage
|
|
|
|
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 |