core/agent.py - master - angrybricks/dementor

core/agent.py
132 строки5 KB
MaksPV
Убрал опасные данные и сделал скрипт для развёртывания БД
30 сен 2025, 20:52
30 сен 2025, 20:521adc81f
"""
Упрощенный агент для помощи в учёбе
"""

import os
from typing import Optional

from dotenv import load_dotenv
from langchain_gigachat.chat_models import GigaChat
from langchain_core.messages import HumanMessage
from langgraph.prebuilt import create_react_agent
from langgraph.checkpoint.memory import MemorySaver

from modules.database.tools import DatabaseTools, DatabaseConfig
from modules.learning_materials.tools import MaterialsTools
from modules.quiz_generation.tools import QuizTools


class StudyAssistantAgent:
    def __init__(self, credentials: Optional[str] = None, student_id: int = 1):
        load_dotenv()

        # Инициализация GigaChat
        credentials = credentials or os.getenv("GIGACHAT_CREDENTIALS")
        if not credentials:
            raise ValueError("GIGACHAT_CREDENTIALS не задан")

        # Параметры модели и SSL из переменных окружения
        model_name = os.getenv("GIGACHAT_MODEL")
        verify_ssl = os.getenv("GIGACHAT_VERIFY_SSL", "false").lower() in {"1", "true", "yes"}
        base_url = os.getenv("GIGACHAT_BASE_URL")
        if model_name and base_url:
            self.giga = GigaChat(
                credentials=credentials,
                model=model_name,
                base_url=base_url,
                verify_ssl_certs=verify_ssl
            )
        elif model_name:
            self.giga = GigaChat(
                credentials=credentials,
                model=model_name,
                verify_ssl_certs=verify_ssl
            )
        elif base_url:
            self.giga = GigaChat(
                credentials=credentials,
                base_url=base_url,
                verify_ssl_certs=verify_ssl
            )
        else:
            self.giga = GigaChat(
                credentials=credentials,
                verify_ssl_certs=verify_ssl
            )

        # Конфигурация БД из переменных окружения
        db_config = DatabaseConfig(
            host=os.getenv("DB_HOST", "localhost"),
            port=int(os.getenv("DB_PORT", "3306")),
            database=os.getenv("DB_NAME", ""),
            user=os.getenv("DB_USER", ""),
            password=os.getenv("DB_PASSWORD", "")
        )

        # Инициализация модулей
        self.database = DatabaseTools(db_config)
        self.quiz = QuizTools(db_config)
        self.materials = MaterialsTools(db_config)

        # Сбор инструментов
        self.tools = []
        self.tools.extend(self.database.get_tools())
        self.tools.extend(self.quiz.get_tools())
        self.tools.extend(self.materials.get_tools())

        # Создание агента
        self.giga_with_functions = self.giga.bind_functions(self.tools)

        self.agent_executor = create_react_agent(
            self.giga_with_functions,
            self.tools,
            checkpointer=MemorySaver(),
            prompt=self._get_simple_prompt(student_id),
        )

    def _get_simple_prompt(self, student_id: int) -> str:
        return f"""
        Ты - ИИ помощник для учёбы. Твои основные задачи:
        
        1. ПОЛУЧЕНИЕ МАТЕРИАЛОВ:
           - Используй get_educational_materials(student_id={student_id}) для получения учебных материалов
           - Показывай список доступных материалов с ID и названиями
        
        2. ОБЪЯСНЕНИЕ МАТЕРИАЛОВ:
           - Используй explain_material для получения объяснений и ссылок на PDF
           - Всегда показывай прямую ссылку на PDF файл
        
        3. ВИКТОРИНЫ:
           - generate_quiz - создание теста по материалу
           - answer_question - проверка ответов
           - Задавай по одному вопросу за раз
        
        4. FAQ:
           - Используй get_faq для поиска ответов на вопросы
        
        Будь полезным и конкретным! Всегда показывай ссылки на PDF когда это возможно.
        """

    def chat(self, message: str, thread_id: str = "default") -> str:
        config = {"configurable": {"thread_id": thread_id}, "recursion_limit": 10}

        try:
            response = self.agent_executor.invoke(
                {"messages": [HumanMessage(content=message)]},
                config=config
            )
            return response['messages'][-1].content
        except Exception as e:
            return f"Ошибка: {str(e)}"


# Пример использования
if __name__ == "__main__":
    agent = StudyAssistantAgent()

    print("Упрощенный агент готов!")
    print("Примеры запросов:")
    print("- 'Покажи мои учебные материалы'")
    print("- 'Объясни материал про математику'")
    print("- 'Создай викторину по последней теме'")
    print("- 'Найди ответы на вопросы по физике'")