Added code samples for AI-Agents
This commit is contained in:
13
agents/__init__.py
Normal file
13
agents/__init__.py
Normal file
@@ -0,0 +1,13 @@
|
||||
"""ИИ-агенты для проекта Новая Планета."""
|
||||
from agents.chat_agent import ChatAgent
|
||||
from agents.gigachat_client import GigaChatClient
|
||||
from agents.recommendation_engine import RecommendationEngine
|
||||
from agents.schedule_generator import ScheduleGenerator
|
||||
|
||||
__all__ = [
|
||||
"GigaChatClient",
|
||||
"ScheduleGenerator",
|
||||
"ChatAgent",
|
||||
"RecommendationEngine",
|
||||
]
|
||||
|
||||
115
agents/chat_agent.py
Normal file
115
agents/chat_agent.py
Normal file
@@ -0,0 +1,115 @@
|
||||
"""ИИ-агент для чата 'Планета Земля'."""
|
||||
from typing import List, Optional
|
||||
from uuid import UUID
|
||||
|
||||
from models.gigachat_types import GigaChatMessage
|
||||
from prompts.persona import EARTH_PERSONA
|
||||
|
||||
from agents.gigachat_client import GigaChatClient
|
||||
from services.cache_service import CacheService
|
||||
|
||||
|
||||
class ChatAgent:
|
||||
"""ИИ-агент для общения с детьми и родителями."""
|
||||
|
||||
def __init__(self, gigachat: GigaChatClient, cache: CacheService):
|
||||
self.gigachat = gigachat
|
||||
self.cache = cache
|
||||
|
||||
async def chat(
|
||||
self,
|
||||
user_id: UUID,
|
||||
message: str,
|
||||
conversation_id: Optional[str] = None,
|
||||
model: str = "GigaChat-2-Lite",
|
||||
) -> tuple[str, int]:
|
||||
"""
|
||||
Отправить сообщение и получить ответ.
|
||||
|
||||
Args:
|
||||
user_id: ID пользователя
|
||||
message: Текст сообщения
|
||||
conversation_id: ID разговора (для контекста)
|
||||
model: Модель GigaChat
|
||||
|
||||
Returns:
|
||||
(ответ, количество использованных токенов)
|
||||
"""
|
||||
# Загружаем контекст из кэша
|
||||
context_messages = []
|
||||
if conversation_id:
|
||||
cached_context = await self.cache.get_context(str(conversation_id))
|
||||
context_messages = [
|
||||
GigaChatMessage(role=msg["role"], content=msg["content"])
|
||||
for msg in cached_context
|
||||
]
|
||||
|
||||
# Добавляем системный промпт в начало
|
||||
system_message = GigaChatMessage(role="system", content=EARTH_PERSONA)
|
||||
if not context_messages or context_messages[0].role != "system":
|
||||
context_messages.insert(0, system_message)
|
||||
|
||||
# Добавляем текущее сообщение пользователя
|
||||
context_messages.append(GigaChatMessage(role="user", content=message))
|
||||
|
||||
# Отправляем запрос
|
||||
response = await self.gigachat.chat_with_response(
|
||||
message=message,
|
||||
context=context_messages,
|
||||
model=model,
|
||||
temperature=0.7,
|
||||
max_tokens=1500,
|
||||
)
|
||||
|
||||
assistant_message = response.choices[0].message.content
|
||||
tokens_used = response.usage.total_tokens
|
||||
|
||||
# Сохраняем в контекст
|
||||
if conversation_id:
|
||||
await self.cache.add_message(str(conversation_id), "user", message)
|
||||
await self.cache.add_message(str(conversation_id), "assistant", assistant_message)
|
||||
|
||||
return assistant_message, tokens_used
|
||||
|
||||
async def chat_with_context(
|
||||
self,
|
||||
user_id: UUID,
|
||||
message: str,
|
||||
context: Optional[List[dict]] = None,
|
||||
model: str = "GigaChat-2-Lite",
|
||||
) -> tuple[str, int]:
|
||||
"""
|
||||
Отправить сообщение с явным контекстом.
|
||||
|
||||
Args:
|
||||
user_id: ID пользователя
|
||||
message: Текст сообщения
|
||||
context: Явный контекст разговора
|
||||
model: Модель GigaChat
|
||||
|
||||
Returns:
|
||||
(ответ, количество использованных токенов)
|
||||
"""
|
||||
context_messages = [GigaChatMessage(role="system", content=EARTH_PERSONA)]
|
||||
|
||||
if context:
|
||||
for msg in context:
|
||||
context_messages.append(
|
||||
GigaChatMessage(role=msg["role"], content=msg["content"])
|
||||
)
|
||||
|
||||
context_messages.append(GigaChatMessage(role="user", content=message))
|
||||
|
||||
response = await self.gigachat.chat_with_response(
|
||||
message=message,
|
||||
context=context_messages,
|
||||
model=model,
|
||||
temperature=0.7,
|
||||
max_tokens=1500,
|
||||
)
|
||||
|
||||
assistant_message = response.choices[0].message.content
|
||||
tokens_used = response.usage.total_tokens
|
||||
|
||||
return assistant_message, tokens_used
|
||||
|
||||
124
agents/gigachat_client.py
Normal file
124
agents/gigachat_client.py
Normal file
@@ -0,0 +1,124 @@
|
||||
"""Клиент для работы с GigaChat API."""
|
||||
import json
|
||||
from typing import List, Optional
|
||||
|
||||
import aiohttp
|
||||
|
||||
from models.gigachat_types import GigaChatMessage, GigaChatRequest, GigaChatResponse
|
||||
from services.token_manager import TokenManager
|
||||
|
||||
|
||||
class GigaChatClient:
|
||||
"""Клиент для взаимодействия с GigaChat API."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
token_manager: TokenManager,
|
||||
base_url: Optional[str] = None,
|
||||
):
|
||||
self.token_manager = token_manager
|
||||
self.base_url = base_url or "https://gigachat.devices.sberbank.ru/api/v1"
|
||||
self._session: Optional[aiohttp.ClientSession] = None
|
||||
|
||||
async def _get_session(self) -> aiohttp.ClientSession:
|
||||
"""Получить HTTP сессию (lazy initialization)."""
|
||||
if self._session is None or self._session.closed:
|
||||
self._session = aiohttp.ClientSession()
|
||||
return self._session
|
||||
|
||||
async def chat(
|
||||
self,
|
||||
message: str,
|
||||
context: Optional[List[GigaChatMessage]] = None,
|
||||
model: str = "GigaChat-2",
|
||||
temperature: float = 0.7,
|
||||
max_tokens: int = 2000,
|
||||
) -> str:
|
||||
"""
|
||||
Отправить сообщение в GigaChat.
|
||||
|
||||
Args:
|
||||
message: Текст сообщения
|
||||
context: История сообщений
|
||||
model: Модель GigaChat (GigaChat-2, GigaChat-2-Lite, GigaChat-2-Pro, GigaChat-2-Max)
|
||||
temperature: Температура генерации
|
||||
max_tokens: Максимальное количество токенов
|
||||
|
||||
Returns:
|
||||
Ответ от модели
|
||||
"""
|
||||
messages = context or []
|
||||
messages.append(GigaChatMessage(role="user", content=message))
|
||||
|
||||
request = GigaChatRequest(
|
||||
model=model,
|
||||
messages=messages,
|
||||
temperature=temperature,
|
||||
max_tokens=max_tokens,
|
||||
)
|
||||
|
||||
response = await self._make_request(request)
|
||||
return response.choices[0].message.content
|
||||
|
||||
async def chat_with_response(
|
||||
self,
|
||||
message: str,
|
||||
context: Optional[List[GigaChatMessage]] = None,
|
||||
model: str = "GigaChat-2",
|
||||
temperature: float = 0.7,
|
||||
max_tokens: int = 2000,
|
||||
) -> GigaChatResponse:
|
||||
"""
|
||||
Отправить сообщение и получить полный ответ.
|
||||
|
||||
Args:
|
||||
message: Текст сообщения
|
||||
context: История сообщений
|
||||
model: Модель GigaChat
|
||||
temperature: Температура генерации
|
||||
max_tokens: Максимальное количество токенов
|
||||
|
||||
Returns:
|
||||
Полный ответ от API
|
||||
"""
|
||||
messages = context or []
|
||||
messages.append(GigaChatMessage(role="user", content=message))
|
||||
|
||||
request = GigaChatRequest(
|
||||
model=model,
|
||||
messages=messages,
|
||||
temperature=temperature,
|
||||
max_tokens=max_tokens,
|
||||
)
|
||||
|
||||
return await self._make_request(request)
|
||||
|
||||
async def _make_request(self, request: GigaChatRequest) -> GigaChatResponse:
|
||||
"""Выполнить запрос к API."""
|
||||
token = await self.token_manager.get_token()
|
||||
session = await self._get_session()
|
||||
|
||||
headers = {
|
||||
"Authorization": f"Bearer {token}",
|
||||
"Content-Type": "application/json",
|
||||
}
|
||||
|
||||
url = f"{self.base_url}/chat/completions"
|
||||
|
||||
async with session.post(
|
||||
url,
|
||||
headers=headers,
|
||||
json=request.model_dump(exclude_none=True),
|
||||
) as response:
|
||||
if response.status != 200:
|
||||
error_text = await response.text()
|
||||
raise Exception(f"GigaChat API error: {response.status} - {error_text}")
|
||||
|
||||
data = await response.json()
|
||||
return GigaChatResponse(**data)
|
||||
|
||||
async def close(self):
|
||||
"""Закрыть HTTP сессию."""
|
||||
if self._session and not self._session.closed:
|
||||
await self._session.close()
|
||||
|
||||
130
agents/recommendation_engine.py
Normal file
130
agents/recommendation_engine.py
Normal file
@@ -0,0 +1,130 @@
|
||||
"""Рекомендательная система для заданий (MVP-1)."""
|
||||
from typing import Dict, List, Optional
|
||||
|
||||
import numpy as np
|
||||
from sklearn.feature_extraction.text import TfidfVectorizer
|
||||
from sklearn.metrics.pairwise import cosine_similarity
|
||||
|
||||
|
||||
class RecommendationEngine:
|
||||
"""Простая рекомендательная система на основе TF-IDF."""
|
||||
|
||||
def __init__(self):
|
||||
self.vectorizer = TfidfVectorizer(max_features=100, stop_words="english")
|
||||
self.task_vectors = None
|
||||
self.tasks = []
|
||||
|
||||
def fit(self, tasks: List[Dict]):
|
||||
"""
|
||||
Обучить модель на исторических данных.
|
||||
|
||||
Args:
|
||||
tasks: Список заданий с полями: title, description, category, completed
|
||||
"""
|
||||
self.tasks = tasks
|
||||
|
||||
# Создаем текстовые описания для векторизации
|
||||
texts = []
|
||||
for task in tasks:
|
||||
text = f"{task.get('title', '')} {task.get('description', '')} {task.get('category', '')}"
|
||||
texts.append(text)
|
||||
|
||||
if texts:
|
||||
self.task_vectors = self.vectorizer.fit_transform(texts)
|
||||
|
||||
def recommend(
|
||||
self,
|
||||
preferences: List[str],
|
||||
completed_tasks: Optional[List[str]] = None,
|
||||
top_k: int = 5,
|
||||
) -> List[Dict]:
|
||||
"""
|
||||
Рекомендовать задания на основе предпочтений.
|
||||
|
||||
Args:
|
||||
preferences: Предпочтения пользователя
|
||||
completed_tasks: Список уже выполненных заданий (для исключения)
|
||||
top_k: Количество рекомендаций
|
||||
|
||||
Returns:
|
||||
Список рекомендованных заданий
|
||||
"""
|
||||
if not self.tasks or self.task_vectors is None:
|
||||
return []
|
||||
|
||||
# Векторизуем предпочтения
|
||||
preferences_text = " ".join(preferences)
|
||||
preference_vector = self.vectorizer.transform([preferences_text])
|
||||
|
||||
# Вычисляем схожесть
|
||||
similarities = cosine_similarity(preference_vector, self.task_vectors)[0]
|
||||
|
||||
# Исключаем уже выполненные задания
|
||||
if completed_tasks:
|
||||
for i, task in enumerate(self.tasks):
|
||||
if task.get("title") in completed_tasks or task.get("id") in completed_tasks:
|
||||
similarities[i] = -1
|
||||
|
||||
# Получаем топ-K индексов
|
||||
top_indices = np.argsort(similarities)[::-1][:top_k]
|
||||
top_indices = [idx for idx in top_indices if similarities[idx] > 0]
|
||||
|
||||
return [self.tasks[idx] for idx in top_indices]
|
||||
|
||||
def recommend_by_category(
|
||||
self,
|
||||
category: str,
|
||||
completed_tasks: Optional[List[str]] = None,
|
||||
top_k: int = 3,
|
||||
) -> List[Dict]:
|
||||
"""
|
||||
Рекомендовать задания по категории.
|
||||
|
||||
Args:
|
||||
category: Категория заданий
|
||||
completed_tasks: Выполненные задания
|
||||
top_k: Количество рекомендаций
|
||||
|
||||
Returns:
|
||||
Список рекомендованных заданий
|
||||
"""
|
||||
category_tasks = [task for task in self.tasks if task.get("category") == category]
|
||||
|
||||
if completed_tasks:
|
||||
category_tasks = [
|
||||
task
|
||||
for task in category_tasks
|
||||
if task.get("title") not in completed_tasks
|
||||
and task.get("id") not in completed_tasks
|
||||
]
|
||||
|
||||
# Сортируем по популярности (можно добавить поле rating)
|
||||
return category_tasks[:top_k]
|
||||
|
||||
def get_popular_tasks(self, top_k: int = 10) -> List[Dict]:
|
||||
"""
|
||||
Получить популярные задания.
|
||||
|
||||
Args:
|
||||
top_k: Количество заданий
|
||||
|
||||
Returns:
|
||||
Список популярных заданий
|
||||
"""
|
||||
# Простая эвристика: задания, которые чаще выполняются
|
||||
task_scores: Dict[str, float] = {}
|
||||
|
||||
for task in self.tasks:
|
||||
task_id = task.get("id") or task.get("title")
|
||||
if task.get("completed", False):
|
||||
task_scores[task_id] = task_scores.get(task_id, 0) + 1
|
||||
|
||||
# Сортируем по популярности
|
||||
sorted_tasks = sorted(
|
||||
self.tasks,
|
||||
key=lambda t: task_scores.get(t.get("id") or t.get("title"), 0),
|
||||
reverse=True,
|
||||
)
|
||||
|
||||
return sorted_tasks[:top_k]
|
||||
|
||||
168
agents/schedule_generator.py
Normal file
168
agents/schedule_generator.py
Normal file
@@ -0,0 +1,168 @@
|
||||
"""Генератор расписаний с использованием GigaChat."""
|
||||
import json
|
||||
from typing import List, Optional
|
||||
|
||||
from models.gigachat_types import GigaChatMessage
|
||||
from models.schedule import Schedule, Task
|
||||
from prompts.schedule_prompts import SCHEDULE_GENERATION_PROMPT
|
||||
|
||||
from agents.gigachat_client import GigaChatClient
|
||||
|
||||
|
||||
class ScheduleGenerator:
|
||||
"""Генератор расписаний для детей с РАС."""
|
||||
|
||||
def __init__(self, gigachat: GigaChatClient):
|
||||
self.gigachat = gigachat
|
||||
|
||||
async def generate(
|
||||
self,
|
||||
child_age: int,
|
||||
preferences: List[str],
|
||||
date: str,
|
||||
existing_tasks: Optional[List[str]] = None,
|
||||
model: str = "GigaChat-2-Pro",
|
||||
) -> Schedule:
|
||||
"""
|
||||
Сгенерировать расписание.
|
||||
|
||||
Args:
|
||||
child_age: Возраст ребенка
|
||||
preferences: Предпочтения ребенка
|
||||
date: Дата расписания
|
||||
existing_tasks: Существующие задания для учета
|
||||
model: Модель GigaChat
|
||||
|
||||
Returns:
|
||||
Объект расписания
|
||||
"""
|
||||
preferences_str = ", ".join(preferences) if preferences else "не указаны"
|
||||
|
||||
prompt = SCHEDULE_GENERATION_PROMPT.format(
|
||||
age=child_age,
|
||||
preferences=preferences_str,
|
||||
date=date,
|
||||
)
|
||||
|
||||
if existing_tasks:
|
||||
prompt += f"\n\nУчти существующие задания: {', '.join(existing_tasks)}"
|
||||
|
||||
# Используем более высокую температуру для разнообразия
|
||||
response_text = await self.gigachat.chat(
|
||||
message=prompt,
|
||||
model=model,
|
||||
temperature=0.8,
|
||||
max_tokens=3000,
|
||||
)
|
||||
|
||||
# Парсим JSON из ответа
|
||||
schedule_data = self._parse_json_response(response_text)
|
||||
|
||||
# Создаем объект Schedule
|
||||
tasks = [
|
||||
Task(
|
||||
title=task_data["title"],
|
||||
description=task_data.get("description"),
|
||||
duration_minutes=task_data["duration_minutes"],
|
||||
category=task_data.get("category", "обучение"),
|
||||
)
|
||||
for task_data in schedule_data.get("tasks", [])
|
||||
]
|
||||
|
||||
return Schedule(
|
||||
title=schedule_data.get("title", f"Расписание на {date}"),
|
||||
date=date,
|
||||
tasks=tasks,
|
||||
)
|
||||
|
||||
async def update(
|
||||
self,
|
||||
existing_schedule: Schedule,
|
||||
user_request: str,
|
||||
model: str = "GigaChat-2-Pro",
|
||||
) -> Schedule:
|
||||
"""
|
||||
Обновить существующее расписание.
|
||||
|
||||
Args:
|
||||
existing_schedule: Текущее расписание
|
||||
user_request: Запрос на изменение
|
||||
model: Модель GigaChat
|
||||
|
||||
Returns:
|
||||
Обновленное расписание
|
||||
"""
|
||||
from prompts.schedule_prompts import SCHEDULE_UPDATE_PROMPT
|
||||
|
||||
schedule_json = existing_schedule.model_dump_json()
|
||||
|
||||
prompt = SCHEDULE_UPDATE_PROMPT.format(
|
||||
existing_schedule=schedule_json,
|
||||
user_request=user_request,
|
||||
)
|
||||
|
||||
response_text = await self.gigachat.chat(
|
||||
message=prompt,
|
||||
model=model,
|
||||
temperature=0.7,
|
||||
max_tokens=3000,
|
||||
)
|
||||
|
||||
schedule_data = self._parse_json_response(response_text)
|
||||
|
||||
tasks = [
|
||||
Task(
|
||||
title=task_data["title"],
|
||||
description=task_data.get("description"),
|
||||
duration_minutes=task_data["duration_minutes"],
|
||||
category=task_data.get("category", "обучение"),
|
||||
)
|
||||
for task_data in schedule_data.get("tasks", [])
|
||||
]
|
||||
|
||||
return Schedule(
|
||||
id=existing_schedule.id,
|
||||
title=schedule_data.get("title", existing_schedule.title),
|
||||
date=existing_schedule.date,
|
||||
tasks=tasks,
|
||||
user_id=existing_schedule.user_id,
|
||||
)
|
||||
|
||||
def _parse_json_response(self, response_text: str) -> dict:
|
||||
"""
|
||||
Извлечь JSON из ответа модели.
|
||||
|
||||
Args:
|
||||
response_text: Текст ответа
|
||||
|
||||
Returns:
|
||||
Распарсенный JSON
|
||||
"""
|
||||
# Пытаемся найти JSON в ответе
|
||||
response_text = response_text.strip()
|
||||
|
||||
# Удаляем markdown код блоки если есть
|
||||
if response_text.startswith("```json"):
|
||||
response_text = response_text[7:]
|
||||
if response_text.startswith("```"):
|
||||
response_text = response_text[3:]
|
||||
if response_text.endswith("```"):
|
||||
response_text = response_text[:-3]
|
||||
|
||||
response_text = response_text.strip()
|
||||
|
||||
try:
|
||||
return json.loads(response_text)
|
||||
except json.JSONDecodeError:
|
||||
# Если не удалось распарсить, пытаемся найти JSON объект в тексте
|
||||
start_idx = response_text.find("{")
|
||||
end_idx = response_text.rfind("}") + 1
|
||||
|
||||
if start_idx >= 0 and end_idx > start_idx:
|
||||
try:
|
||||
return json.loads(response_text[start_idx:end_idx])
|
||||
except json.JSONDecodeError:
|
||||
pass
|
||||
|
||||
raise ValueError(f"Не удалось распарсить JSON из ответа: {response_text[:200]}")
|
||||
|
||||
Reference in New Issue
Block a user