Add new file
This commit is contained in:
145
rag/agents/roles.py
Normal file
145
rag/agents/roles.py
Normal file
@@ -0,0 +1,145 @@
|
||||
# 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)}"}
|
||||
Reference in New Issue
Block a user