145 lines
5.5 KiB
Python
145 lines
5.5 KiB
Python
# rag/agents/roles.py
|
||
"""
|
||
Определение базового класса агента и конкретных ролей.
|
||
"""
|
||
|
||
import logging
|
||
from abc import ABC, abstractmethod
|
||
from typing import Optional, List, Dict, Any
|
||
|
||
from rag.services.giga_client import GigaClient
|
||
from rag.config_models import AppConfig
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
|
||
class BaseAgent(ABC):
|
||
"""
|
||
Базовый абстрактный класс для агента.
|
||
Каждый агент имеет системный промпт, может использовать GigaChat.
|
||
"""
|
||
|
||
def __init__(
|
||
self,
|
||
name: str,
|
||
system_prompt: str,
|
||
giga: GigaClient,
|
||
config: AppConfig,
|
||
temperature: Optional[float] = None,
|
||
):
|
||
self.name = name
|
||
self.system_prompt = system_prompt
|
||
self.giga = giga
|
||
self.config = config
|
||
self.temperature = temperature or getattr(config, 'agent_temperature', 0.3)
|
||
|
||
@abstractmethod
|
||
async def process(
|
||
self,
|
||
query: str,
|
||
context: str,
|
||
history: List[Dict[str, str]],
|
||
tools: Optional[List[Dict[str, Any]]] = None,
|
||
) -> Dict[str, Any]:
|
||
"""
|
||
Обрабатывает запрос и возвращает ответ или запрос на действие.
|
||
|
||
Возвращает словарь с ключами:
|
||
'type': 'answer' или 'action'
|
||
'content': если 'answer' -> текст ответа; если 'action' -> {'tool': ..., 'params': ...}
|
||
"""
|
||
pass
|
||
|
||
|
||
class MethodologistAgent(BaseAgent):
|
||
"""Агент-методолог, специализируется на стандартах, ГОСТах, методиках."""
|
||
|
||
async def process(
|
||
self,
|
||
query: str,
|
||
context: str,
|
||
history: List[Dict[str, str]],
|
||
tools: Optional[List[Dict[str, Any]]] = None,
|
||
) -> Dict[str, Any]:
|
||
# Формируем промпт с учётом системного промпта и контекста
|
||
# Здесь можно использовать prompt_builder, но для простоты составим вручную
|
||
messages = []
|
||
if self.system_prompt:
|
||
messages.append({"role": "system", "content": self.system_prompt})
|
||
messages.extend(history)
|
||
messages.append({"role": "user", "content": f"Контекст:\n{context}\n\nВопрос: {query}"})
|
||
|
||
try:
|
||
response = await self.giga.chat(
|
||
history=messages,
|
||
query="", # уже включено в history
|
||
system_prompt=None,
|
||
file_id=None,
|
||
temperature=self.temperature,
|
||
)
|
||
# В простейшем случае возвращаем ответ как финальный
|
||
return {"type": "answer", "content": response}
|
||
except Exception as e:
|
||
logger.error(f"Ошибка в MethodologistAgent: {e}")
|
||
return {"type": "answer", "content": f"Извините, произошла ошибка: {str(e)}"}
|
||
|
||
|
||
class BoxSolutionAgent(BaseAgent):
|
||
"""Агент для генерации коробочных решений."""
|
||
|
||
async def process(
|
||
self,
|
||
query: str,
|
||
context: str,
|
||
history: List[Dict[str, str]],
|
||
tools: Optional[List[Dict[str, Any]]] = None,
|
||
) -> Dict[str, Any]:
|
||
messages = []
|
||
if self.system_prompt:
|
||
messages.append({"role": "system", "content": self.system_prompt})
|
||
messages.extend(history)
|
||
messages.append({"role": "user", "content": f"Контекст:\n{context}\n\nЗапрос на коробочное решение: {query}"})
|
||
|
||
try:
|
||
response = await self.giga.chat(
|
||
history=messages,
|
||
query="",
|
||
system_prompt=None,
|
||
file_id=None,
|
||
temperature=self.temperature,
|
||
)
|
||
return {"type": "answer", "content": response}
|
||
except Exception as e:
|
||
logger.error(f"Ошибка в BoxSolutionAgent: {e}")
|
||
return {"type": "answer", "content": f"Извините, произошла ошибка: {str(e)}"}
|
||
|
||
|
||
class PersonalAssistantAgent(BaseAgent):
|
||
"""Персональный ассистент (бронирование, календарь, командировки)."""
|
||
|
||
async def process(
|
||
self,
|
||
query: str,
|
||
context: str,
|
||
history: List[Dict[str, str]],
|
||
tools: Optional[List[Dict[str, Any]]] = None,
|
||
) -> Dict[str, Any]:
|
||
# Для персонального ассистента можно использовать ReAct или просто генерацию
|
||
messages = []
|
||
if self.system_prompt:
|
||
messages.append({"role": "system", "content": self.system_prompt})
|
||
messages.extend(history)
|
||
messages.append({"role": "user", "content": f"Контекст:\n{context}\n\nЗапрос: {query}"})
|
||
|
||
try:
|
||
response = await self.giga.chat(
|
||
history=messages,
|
||
query="",
|
||
system_prompt=None,
|
||
file_id=None,
|
||
temperature=self.temperature,
|
||
)
|
||
return {"type": "answer", "content": response}
|
||
except Exception as e:
|
||
logger.error(f"Ошибка в PersonalAssistantAgent: {e}")
|
||
return {"type": "answer", "content": f"Извините, произошла ошибка: {str(e)}"} |