modules/database/tools.py - master - angrybricks/dementor

modules/database/tools.py
161 строка5 KB
MaksPV
УПРОСТИЛ ВСЁ
30 сен 2025, 19:46
30 сен 2025, 19:4601b3746
"""
Упрощенные инструменты для работы с базой данных
"""

from typing import List, Dict, Optional
from datetime import date
import mysql.connector
from pydantic import BaseModel, Field, PrivateAttr
from langchain_core.tools import BaseTool


class DatabaseConfig(BaseModel):
    host: str
    user: str
    password: str
    database: str
    port: int = 3306


class GetMaterialsInput(BaseModel):
    student_id: int = Field(description="ID студента")


class GetMaterialsOutput(BaseModel):
    materials: List[Dict] = Field(description="Учебные материалы")
    status: str = Field(description="Статус")
    message: str = Field(description="Сообщение")


class GetFAQInput(BaseModel):
    question: Optional[str] = Field(default=None, description="Текст вопроса")
    class_name: Optional[str] = Field(default=None, description="Предмет")


class GetFAQOutput(BaseModel):
    qa_pairs: List[Dict] = Field(description="Вопросы и ответы")
    status: str = Field(description="Статус")
    message: str = Field(description="Сообщение")


class BaseDBTool(BaseTool):
    _db_config: DatabaseConfig = PrivateAttr()

    def __init__(self, db_config: DatabaseConfig, **kwargs):
        super().__init__(**kwargs)
        self._db_config = db_config

    def get_connection(self):
        try:
            return mysql.connector.connect(
                host=self._db_config.host,
                user=self._db_config.user,
                password=self._db_config.password,
                database=self._db_config.database,
                port=self._db_config.port
            )
        except Exception as e:
            print(f"Ошибка подключения: {e}")
            return None


class GetMaterialsTool(BaseDBTool):
    name: str = "get_educational_materials"
    description: str = "Получает учебные материалы студента из базы данных"
    args_schema: type[BaseModel] = GetMaterialsInput

    def _run(self, student_id: int) -> GetMaterialsOutput:
        try:
            connection = self.get_connection()
            if not connection:
                return GetMaterialsOutput(
                    materials=[], status="error", message="Ошибка подключения"
                )

            cursor = connection.cursor(dictionary=True)

            query = """
            SELECT em.* 
            FROM educational_material em
            JOIN edu_material_for_group eg ON em.id = eg.edmat_id
            WHERE eg.group_id = (SELECT group_id FROM student WHERE id = %s)
            """

            cursor.execute(query, (student_id,))
            results = cursor.fetchall()

            # Конвертируем даты в строки
            for result in results:
                for key, value in result.items():
                    if isinstance(value, (date)):
                        result[key] = value.isoformat()

            cursor.close()
            connection.close()

            return GetMaterialsOutput(
                materials=results,
                status="success",
                message=f"Найдено материалов: {len(results)}"
            )

        except Exception as e:
            return GetMaterialsOutput(
                materials=[], status="error", message=f"Ошибка: {str(e)}"
            )


class GetFAQTool(BaseDBTool):
    name: str = "get_faq"
    description: str = "Ищет вопросы и ответы в базе знаний"
    args_schema: type[BaseModel] = GetFAQInput

    def _run(self, question: Optional[str] = None, class_name: Optional[str] = None) -> GetFAQOutput:
        try:
            connection = self.get_connection()
            if not connection:
                return GetFAQOutput(
                    qa_pairs=[], status="error", message="Ошибка подключения"
                )

            cursor = connection.cursor(dictionary=True)

            query = "SELECT * FROM faq_and_answers WHERE 1=1"
            params = []

            if question:
                query += " AND question LIKE %s"
                params.append(f"%{question}%")

            if class_name:
                query += " AND class = %s"
                params.append(class_name)

            cursor.execute(query, params)
            results = cursor.fetchall()

            cursor.close()
            connection.close()

            return GetFAQOutput(
                qa_pairs=results,
                status="success",
                message=f"Найдено вопросов: {len(results)}"
            )

        except Exception as e:
            return GetFAQOutput(
                qa_pairs=[], status="error", message=f"Ошибка: {str(e)}"
            )


class DatabaseTools:
    def __init__(self, db_config: DatabaseConfig):
        self.db_config = db_config
        self.tools = [
            GetMaterialsTool(db_config=db_config),
            GetFAQTool(db_config=db_config),
        ]

    def get_tools(self):
        return self.tools