Files

65 lines
1.7 KiB
Python

"""Base classes for AI provider integration."""
from __future__ import annotations
from abc import ABC, abstractmethod
from dataclasses import dataclass, field
from datetime import datetime, timezone
from typing import Any
@dataclass
class AIResult:
"""Result from an AI provider call."""
success: bool
content: str = ""
model: str = ""
provider: str = ""
tokens_used: int = 0
duration_ms: int = 0
error: str = ""
metadata: dict[str, Any] = field(default_factory=dict)
class AIProvider(ABC):
"""Abstract base class for AI providers.
All AI providers (Yandex GPT, GigaChat, etc.) must implement this.
"""
name: str = ""
model: str = ""
timeout: int = 10
max_retries: int = 3
def __init__(self, api_key: str = "", **kwargs: Any):
self.api_key = api_key
for key, value in kwargs.items():
setattr(self, key, value)
self._initialize()
def _initialize(self) -> None:
"""Post-initialization hook for provider-specific setup."""
@abstractmethod
async def analyze(self, prompt: str, **kwargs: Any) -> AIResult:
"""Send a prompt to the AI provider and return the result.
Args:
prompt: The formatted prompt to send
**kwargs: Additional provider-specific parameters
Returns:
AIResult with the response content
"""
def format_prompt(
self, system_prompt: str, user_prompt: str, **kwargs: Any
) -> str:
"""Format a prompt with system instructions and user message.
Override in provider if the API uses separate system/user messages.
"""
return f"{system_prompt}\n\n{user_prompt}"