chatbot.py

import re
import json
from typing import List, Dict, Tuple
from knowledge_base import KnowledgeBase
from retriever import Retriever

class ITMOChatbot:
    def __init__(self):
        self.kb = KnowledgeBase()
        self.retriever = Retriever()
        self.max_history_turns = 3
        self.max_context_tokens = 1200
        self.relevance_threshold = 0.38
        
        try:
            from transformers import pipeline
            self.generator = pipeline('text2text-generation', model='cointegrated/rut5-base-multitask')
        except Exception as e:
            print(f'Генеративная модель не загружена: {e}')
            self.generator = None
    
    def chat(self, message: str, history: list) -> Tuple[str, float]:
        if not message.strip():
            return 'Пожалуйста, задайте вопрос.', 0.0
        
        if not self.kb.is_itmo_query(message):
            return self._get_irrelevant_response(), 0.0
        
        context = self._get_context(message)
        if not context:
            return 'К сожалению, не нашел релевантной информации в учебных планах ITMO.', 0.0
        
        response = self._generate_response(message, history, context)
        relevance_score = self._calculate_relevance_score(message, context)
        
        return response, relevance_score
    
    def recommend_courses(self, profile: dict) -> str:
        if not profile.get('semester'):
            return 'Пожалуйста, укажите целевой семестр для получения рекомендаций.'
        
        recommendations = self.kb.recommend(profile)
        if not recommendations:
            return 'К сожалению, не удалось найти подходящие курсы для вашего профиля.'
        
        result = '🎯 Рекомендуемые курсы (из официальных учебных планов ITMO):\n\n'
        for i, rec in enumerate(recommendations[:7], 1):
            result += f'{i}. {rec["name"]} ({rec["semester"]} семестр, {rec["credits"]} кредитов)\n'
            result += f'   {rec["why"]}\n\n'
        
        return result
    
    def _get_context(self, message: str) -> List[Dict]:
        try:
            results = self.retriever.retrieve(message, k=6, threshold=0.35)
            # Преобразуем результаты в нужный формат
            formatted_results = []
            for result in results:
                course_id = result.get('course_id')
                if course_id:
                    course = self.kb.get_course_by_id(course_id)
                    if course:
                        course['score'] = result.get('score', 0.0)
                        formatted_results.append(course)
            return formatted_results
        except Exception as e:
            print(f'Ошибка при получении контекста: {e}')
            return []
    
    def _generate_response(self, message: str, history: list, context: List[Dict]) -> str:
        if not context:
            return 'В предоставленных данных об этом не сказано.'
        
        prompt = self._build_prompt(message, history, context)
        
        if self.generator:
            try:
                response = self.generator(
                    prompt,
                    max_new_tokens=180,
                    temperature=0.4,
                    do_sample=True
                )[0]['generated_text']
                return response.strip()
            except Exception as e:
                print(f'Ошибка генерации: {e}')
        
        return self._fallback_response(context)
    
    def _build_prompt(self, message: str, history: list, context: List[Dict]) -> str:
        system_prompt = 'Отвечай только по контексту (ниже). Если недостаточно данных — прямо скажи: "в предоставленных данных об этом не сказано". Отвечай кратко и по делу.'
        
        history_text = ''
        if history:
            recent_history = history[-self.max_history_turns:]
            for turn in recent_history:
                history_text += f'Пользователь: {turn[0]}\nБот: {turn[1]}\n'
        
        context_text = 'Контекст:\n'
        for item in context:
            context_text += f'- {item["name"]} ({item["semester"]} семестр, {item["credits"]} кредитов): {item["short_desc"]}\n'
        
        prompt = f'{system_prompt}\n\n{history_text}Контекст:\n{context_text}\nВопрос: {message}'
        
        if len(prompt) > self.max_context_tokens * 4:
            prompt = prompt[:self.max_context_tokens * 4]
        
        return prompt
    
    def _fallback_response(self, context: List[Dict]) -> str:
        if not context:
            return 'В предоставленных данных об этом не сказано.'
        
        courses = []
        for item in context[:3]:
            courses.append(f'{item["name"]} ({item["semester"]} семестр, {item["credits"]} кредитов)')
        
        return f'Найденные курсы: {", ".join(courses)}. Для более подробной информации обратитесь к официальным учебным планам ITMO.'
    
    def _calculate_relevance_score(self, message: str, context: List[Dict]) -> float:
        if not context:
            return 0.0
        
        scores = [item.get('score', 0.0) for item in context]
        return sum(scores) / len(scores) if scores else 0.0
    
    def _get_irrelevant_response(self) -> str:
        return '''Похоже, вопрос не относится к магистратурам ITMO и их учебным планам. 



Попробуйте спросить, например:

• "Какие дисциплины по NLP в 1 семестре программы ИИ?"

• "Расскажи о программе AI Product"

• "Какие курсы по машинному обучению есть в программе ИИ?"

• "Сколько кредитов за дисциплину 'Глубокое обучение'?"'''