dialogue/agents.py

from dataclasses import dataclass, field
from typing import List, Dict
import anthropic
import aiohttp
import json
from datetime import datetime
import os

@dataclass
class Message:
    content: str
    timestamp: datetime = field(default_factory=datetime.now)
    metadata: Dict = field(default_factory=dict)

    def to_dict(self):
        return {
            'content': self.content,
            'timestamp': self.timestamp.isoformat(),
            'metadata': self.metadata
        }

class BaseAgent:
    def __init__(self, name: str, system_prompt: str):
        self.name = name
        self.system_prompt = system_prompt
        self.messages: List[Message] = []
    
    def add_message(self, message: Message):
        self.messages.append(message)
    
    def get_context(self, context_depth: int = 3) -> List[Dict]:
        recent_messages = self.messages[-context_depth:] if len(self.messages) > context_depth else self.messages
        return [{"role": "assistant" if i % 2 == 0 else "user", 
                "content": msg.content} 
                for i, msg in enumerate(recent_messages)]

class ClaudeAgent(BaseAgent):
    def __init__(self, name: str, system_prompt: str):
        super().__init__(name, system_prompt)
        self.client = anthropic.Client(api_key=os.getenv("ANTHROPIC_API_KEY"))

    async def generate_response(self, prompt: str) -> str:
        messages = self.get_context(context_depth=3)
        messages.append({"role": "user", "content": prompt})
        
        try:
            response = self.client.messages.create(
                model="claude-3-5-sonnet-20241022",
                max_tokens=1000,
                system=self.system_prompt,
                messages=messages
            )
            return response.content[0].text
        except Exception as e:
            return f"Ошибка API Claude: {str(e)}"

class GrokAgent(BaseAgent):
    def __init__(self, name: str, system_prompt: str):
        super().__init__(name, system_prompt)
        self.api_key = os.getenv("XAI_API_KEY")
        self.api_url = "https://api.x.ai/v1/chat/completions"

    async def generate_response(self, prompt: str) -> str:
        messages = self.get_context(context_depth=3)
        messages.append({"role": "user", "content": prompt})
        
        headers = {
            "Authorization": f"Bearer {self.api_key}",
            "Content-Type": "application/json"
        }
        
        payload = {
            "messages": messages,
            "model": "grok-beta",
            "stream": False,
            "temperature": 1.0
        }
        
        try:
            async with aiohttp.ClientSession() as session:
                async with session.post(
                    self.api_url, 
                    headers=headers,
                    json=payload
                ) as response:
                    if response.status == 200:
                        data = await response.json()
                        return data['choices'][0]['message']['content']
                    else:
                        return f"Ошибка API Grok: {response.status}"
        except Exception as e:
            return f"Ошибка при обращении к API Grok: {str(e)}"