chatbot.py

import json
import os
from typing import List, Dict, Tuple
from transformers import pipeline
from knowledge_base import KnowledgeBase
from retriever import Retriever

class ITMOChatbot:
    def __init__(self):
        self.knowledge_base = KnowledgeBase()
        self.retriever = Retriever()
        self.generator = None
        self.max_history_turns = 3
        self.max_context_tokens = 1200
        self.relevance_threshold = 0.38
        
        # Ленивая загрузка генеративной модели
        self._load_generator()
    
    def _load_generator(self):
        try:
            print('Загрузка генеративной модели...')
            self.generator = pipeline(
                'text2text-generation',
                model='cointegrated/rut5-base-multitask',
                device=-1  # CPU
            )
            print('Генеративная модель загружена успешно')
        except Exception as e:
            print(f'Ошибка загрузки генеративной модели: {e}')
            self.generator = None
    
    def chat(self, message: str, history: List[List[str]]) -> Tuple[str, float]:
        if not message.strip():
            return '', 0.0
        
        # Проверка релевантности
        if not self.knowledge_base.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_answer(message, context, history)
        
        # Расчет релевантности
        relevance_score = self._calculate_relevance_score(message, context)
        
        return response, relevance_score
    
    def recommend_courses(self, profile: Dict) -> str:
        semester = profile.get('semester')
        if not semester:
            return 'Пожалуйста, укажите семестр для получения рекомендаций.'
        
        try:
            semester = int(semester)
        except ValueError:
            return 'Пожалуйста, выберите корректный семестр.'
        
        # Получение курсов для семестра
        courses = self.knowledge_base.get_courses_by_semester(semester)
        
        if not courses:
            return f'К сожалению, не найдено курсов для {semester} семестра.'
        
        # Генерация рекомендаций с помощью LLM
        recommendations = self._generate_recommendations(profile, courses)
        
        return recommendations
    
    def _get_context(self, message: str) -> List[Dict]:
        try:
            # Сначала пробуем RAG поиск
            if self.retriever.index:
                results = self.retriever.retrieve(message, k=6, threshold=0.35)
                if results:
                    return results
            
            # Fallback на простой поиск
            return self.knowledge_base.search_courses(message)
            
        except Exception as e:
            print(f'Ошибка получения контекста: {e}')
            return self.knowledge_base.search_courses(message)
    
    def _generate_answer(self, message: str, context: List[Dict], history: List[List[str]]) -> str:
        if not self.generator:
            return self._fallback_answer(context)
        
        try:
            # Подготовка промпта
            prompt = self._build_prompt(message, context, history)
            
            # Генерация ответа
            response = self.generator(
                prompt,
                max_new_tokens=200,
                temperature=0.4,
                do_sample=True,
                pad_token_id=self.generator.tokenizer.eos_token_id
            )
            
            answer = response[0]['generated_text'].strip()
            
            # Очистка ответа от лишних элементов
            if answer.startswith('Ответ:'):
                answer = answer[6:].strip()
            elif answer.startswith('Бот:'):
                answer = answer[4:].strip()
            
            # Убираем лишние скобки и форматирование
            if answer.startswith('[[') and answer.endswith(']]'):
                try:
                    import ast
                    parsed = ast.literal_eval(answer)
                    if isinstance(parsed, list) and len(parsed) > 0 and isinstance(parsed[0], list) and len(parsed[0]) > 1:
                        answer = parsed[0][1]
                except:
                    answer = self._fallback_answer(context)
            
            # Проверяем, что ответ не пустой и не содержит технических деталей
            if answer and len(answer) > 10 and not answer.startswith('['):
                return answer
            else:
                return self._fallback_answer(context)
            
        except Exception as e:
            print(f'Ошибка генерации ответа: {e}')
            return self._fallback_answer(context)
    
    def _generate_recommendations(self, profile: Dict, courses: List[Dict]) -> str:
        if not self.generator:
            return self._fallback_recommendations(profile, courses)
        
        try:
            # Подготовка промпта для рекомендаций
            prompt = self._build_recommendations_prompt(profile, courses)
            
            # Генерация рекомендаций
            response = self.generator(
                prompt,
                max_new_tokens=400,
                temperature=0.5,
                do_sample=True,
                pad_token_id=self.generator.tokenizer.eos_token_id
            )
            
            recommendations = response[0]['generated_text'].strip()
            
            # Очистка ответа
            if recommendations.startswith('Рекомендации:'):
                recommendations = recommendations[14:].strip()
            
            # Проверяем качество ответа
            if recommendations and len(recommendations) > 20 and not recommendations.startswith('['):
                return recommendations
            else:
                return self._fallback_recommendations(profile, courses)
            
        except Exception as e:
            print(f'Ошибка генерации рекомендаций: {e}')
            return self._fallback_recommendations(profile, courses)
    
    def _build_prompt(self, message: str, context: List[Dict], history: List[List[str]]) -> str:
        # Системные инструкции
        system_prompt = '''Ты - помощник для абитуриентов магистратур ITMO. Отвечай на вопросы о программах и курсах на основе предоставленного контекста. Отвечай кратко, дружелюбно и по делу. Если информации недостаточно, скажи об этом прямо. НЕ используй скобки или специальное форматирование в ответе.'''
        
        # История диалога (последние 3 хода)
        history_text = ''
        if history:
            recent_history = history[-self.max_history_turns:]
            for user_msg, bot_msg in recent_history:
                history_text += f'Пользователь: {user_msg}\nБот: {bot_msg}\n\n'
        
        # Контекст
        context_text = 'Информация о курсах:\n'
        for i, item in enumerate(context, 1):
            context_text += f'{i}. {item["name"]} ({item["semester"]} семестр, {item["credits"]} кредитов)\n'
            if item.get('short_desc'):
                context_text += f'   {item["short_desc"]}\n'
            context_text += '\n'
        
        # Полный промпт
        full_prompt = f'{system_prompt}\n\n{history_text}{context_text}Пользователь: {message}\nБот:'
        
        return full_prompt
    
    def _build_recommendations_prompt(self, profile: Dict, courses: List[Dict]) -> str:
        # Описание профиля студента
        profile_text = f'''Студент с профилем:
- Опыт программирования: {profile.get('programming_experience', 2)}/5
- Уровень математики: {profile.get('math_level', 2)}/4
- Интересы: {", ".join(profile.get('interests', []))}
- Целевой семестр: {profile.get('semester')}

Доступные курсы в {profile.get('semester')} семестре:'''

        # Список курсов
        courses_text = ''
        for i, course in enumerate(courses[:10], 1):  # Топ-10 курсов
            tags = ', '.join(course.get('tags', []))
            courses_text += f'{i}. {course["name"]} ({course["credits"]} кредитов)\n'
            if course.get('short_desc'):
                courses_text += f'   Описание: {course["short_desc"]}\n'
            courses_text += f'   Теги: {tags}\n\n'
        
        # Инструкции для рекомендаций
        instructions = '''Для такого студента с такими навыками какие из курсов подойдут? 
Выбери 3-5 наиболее подходящих курсов и объясни почему они подходят для этого профиля.
Учитывай уровень сложности, интересы и опыт студента. Отвечай на русском языке.'''
        
        full_prompt = f'{profile_text}\n\n{courses_text}\n{instructions}\n\nРекомендации:'
        
        return full_prompt
    
    def _fallback_answer(self, context: List[Dict]) -> str:
        if not context:
            return 'К сожалению, не нашел релевантной информации в учебных планах ITMO.'
        
        response = 'Найденная информация:\n\n'
        for i, item in enumerate(context, 1):
            response += f'{i}. {item["name"]} ({item["semester"]} семестр, {item["credits"]} кредитов)\n'
            if item.get('short_desc'):
                response += f'   {item["short_desc"]}\n'
            response += '\n'
        
        return response
    
    def _fallback_recommendations(self, profile: Dict, courses: List[Dict]) -> str:
        semester = profile.get('semester')
        interests = profile.get('interests', [])
        programming_exp = profile.get('programming_experience', 2)
        math_level = profile.get('math_level', 2)
        
        # Простая логика рекомендаций
        recommendations = []
        for course in courses[:5]:
            score = 0
            why_reasons = []
            
            # Оценка по интересам
            matching_tags = [tag for tag in interests if tag in course.get('tags', [])]
            if matching_tags:
                score += 2
                why_reasons.append(f'соответствует вашим интересам: {", ".join(matching_tags)}')
            
            # Оценка по опыту программирования
            if programming_exp >= 3 and any(tag in course.get('tags', []) for tag in ['ml', 'dl', 'systems']):
                score += 1
                why_reasons.append('подходит для вашего уровня программирования')
            
            # Оценка по математике
            if math_level >= 3 and any(tag in course.get('tags', []) for tag in ['math', 'stats', 'dl']):
                score += 1
                why_reasons.append('соответствует вашему уровню математики')
            
            if score > 0:
                recommendations.append({
                    'name': course['name'],
                    'credits': course['credits'],
                    'why': '; '.join(why_reasons) if why_reasons else 'курс из учебного плана программы'
                })
        
        if not recommendations:
            # Если нет подходящих, показываем все курсы
            for course in courses[:3]:
                recommendations.append({
                    'name': course['name'],
                    'credits': course['credits'],
                    'why': 'курс из учебного плана программы'
                })
        
        result = f'🎯 Рекомендуемые курсы для {semester} семестра:\n\n'
        for i, rec in enumerate(recommendations, 1):
            result += f'{i}. {rec["name"]} ({rec["credits"]} кредитов)\n'
            result += f'   {rec["why"]}\n\n'
        
        return result
    
    def _get_irrelevant_response(self) -> str:
        return '''Похоже, вопрос не относится к магистратурам ITMO и их учебным планам. 

Попробуйте спросить, например:
• "Какие дисциплины по NLP в 1 семестре программы ИИ?"
• "Расскажи о программе AI Product"
• "Какие курсы по машинному обучению есть в программе ИИ?"
• "Сколько кредитов за дисциплину 'Глубокое обучение'?"'''
    
    def _calculate_relevance_score(self, message: str, context: List[Dict]) -> float:
        if not context:
            return 0.0
        
        # Простой расчет на основе количества найденных результатов
        return min(len(context) / 6.0, 1.0)