114 lines
3.9 KiB
Python
114 lines
3.9 KiB
Python
import base64
|
|
import mimetypes
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
from typing import Any, Literal, cast
|
|
|
|
import httpx2
|
|
from openai import OpenAI
|
|
from openai.types.chat import (
|
|
ChatCompletionContentPartParam,
|
|
ChatCompletionMessageParam,
|
|
)
|
|
|
|
|
|
@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
|
|
|
|
@staticmethod
|
|
def _raw_content(response: Any) -> str:
|
|
"""Read message content without triggering SDK-side schema validation."""
|
|
try:
|
|
content = response.http_response.json()["choices"][0]["message"]["content"]
|
|
except (AttributeError, KeyError, IndexError, TypeError) as error:
|
|
raise RuntimeError("Agent returned an invalid response") from error
|
|
if not isinstance(content, str):
|
|
raise RuntimeError("Agent returned an empty response")
|
|
return content
|
|
|
|
def completion(self, messages: list[AgentMessage], **kwargs) -> str:
|
|
"""Generate a completion for specified messages."""
|
|
messages_raw: list[ChatCompletionMessageParam] = []
|
|
for m in messages:
|
|
messages_raw.append(
|
|
cast(ChatCompletionMessageParam, {
|
|
"role": m.role,
|
|
"content": m.content
|
|
})
|
|
)
|
|
response = self._client.chat.completions.with_raw_response.parse(
|
|
model=self._model,
|
|
messages=messages_raw,
|
|
**kwargs
|
|
)
|
|
return self._raw_content(response)
|
|
|
|
def completion_with_images(
|
|
self,
|
|
messages: list[AgentMessage],
|
|
image_paths: list[str],
|
|
*,
|
|
detail: Literal["low", "high", "auto"] = "auto",
|
|
**kwargs,
|
|
) -> str:
|
|
"""Generate a completion with local images attached to the last message."""
|
|
if not messages or messages[-1].role != "user":
|
|
raise ValueError("The last message must be a user message")
|
|
if not image_paths:
|
|
raise ValueError("At least one image is required")
|
|
|
|
messages_raw: list[ChatCompletionMessageParam] = [
|
|
cast(ChatCompletionMessageParam, {
|
|
"role": message.role,
|
|
"content": message.content,
|
|
})
|
|
for message in messages[:-1]
|
|
]
|
|
content: list[ChatCompletionContentPartParam] = [
|
|
{"type": "text", "text": messages[-1].content}
|
|
]
|
|
for index, image_path in enumerate(image_paths, start=1):
|
|
path = Path(image_path)
|
|
mime_type = mimetypes.guess_type(path.name)[0] or "image/jpeg"
|
|
if not mime_type.startswith("image/"):
|
|
raise ValueError(f"Unsupported image type: {image_path}")
|
|
encoded = base64.b64encode(path.read_bytes()).decode("ascii")
|
|
content.append({
|
|
"type": "text",
|
|
"text": f"Изображение {index}",
|
|
})
|
|
content.append({
|
|
"type": "image_url",
|
|
"image_url": {
|
|
"url": f"data:{mime_type};base64,{encoded}",
|
|
"detail": detail,
|
|
},
|
|
})
|
|
messages_raw.append({
|
|
"role": "user",
|
|
"content": content,
|
|
})
|
|
|
|
response = self._client.chat.completions.with_raw_response.parse(
|
|
model=self._model,
|
|
messages=messages_raw,
|
|
**kwargs,
|
|
)
|
|
return self._raw_content(response)
|