nikita
642542ad10
- Added different AIs for ASR and for structure builder - Forcing models to use JSON from now on
44 lines
1.2 KiB
Python
44 lines
1.2 KiB
Python
from dataclasses import dataclass
|
|
from typing import Literal
|
|
from pydantic import BaseModel
|
|
|
|
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.parse(
|
|
model=self._model,
|
|
messages=messages_raw,
|
|
**kwargs
|
|
)
|
|
return response.choices[0].message.content # type: ignore |