Added code samples for AI-Agents

This commit is contained in:
2025-12-17 20:22:46 +03:00
parent d66aed35d6
commit 0885618b25
29 changed files with 2007 additions and 0 deletions

13
agents/__init__.py Normal file
View 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
View 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
View 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()

View 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]

View 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]}")