bhagavad_gita_bot.py

import os
import uuid
from typing import Dict, List, Any, Optional
import json
from pathlib import Path
from data_loader import BhagavadGitaDataLoader
from semantic_search import SemanticSearch
from dspy_signatures import ProblemAnalyzer, TeachingContextualizer, WisdomSynthesizer
from weave_logger import BhagavadGitaWeaveLogger
import dspy

class BhagavadGitaBot:
    def __init__(self, cache_dir: str = "cache"):
        self.cache_dir = Path(cache_dir)
        self.cache_dir.mkdir(exist_ok=True)
        
        self.data_loader = BhagavadGitaDataLoader(cache_dir)
        self.semantic_search = SemanticSearch(cache_dir=cache_dir)
        
        self.problem_analyzer = ProblemAnalyzer()
        self.teaching_contextualizer = TeachingContextualizer()
        self.wisdom_synthesizer = WisdomSynthesizer()
        
        # Initialize Weave logger AFTER other components to avoid tracking their initialization
        self.logger = BhagavadGitaWeaveLogger()
        
        self._ensure_data_loaded()
    
    def _ensure_data_loaded(self):
        try:
            self.semantic_search._compute_embeddings()
        except Exception as e:
            print(f"Warning: Could not precompute embeddings: {e}")
            print("Embeddings will be computed on first search.")
    
    def analyze_problem(self, user_problem: str) -> Dict[str, str]:
        try:
            analysis = self.problem_analyzer(user_problem)
            return {
                'key_themes': analysis.key_themes,
                'emotional_state': analysis.emotional_state,
                'core_question': analysis.core_question
            }
        except Exception as e:
            print(f"Error in problem analysis: {e}")
            return {
                'key_themes': "guidance, wisdom",
                'emotional_state': "seeking help",
                'core_question': user_problem
            }
    
    def find_relevant_verses(self, problem_analysis: Dict[str, str], top_k: int = 3) -> List[Dict[str, Any]]:
        themes = [theme.strip() for theme in problem_analysis['key_themes'].split(',')]
        emotional_keywords = [emotion.strip() for emotion in problem_analysis['emotional_state'].split(',')]
        
        query_text = f"{problem_analysis['core_question']} {' '.join(themes)} {' '.join(emotional_keywords)}"
        
        try:
            verses = self.semantic_search.get_contextual_verses(
                problem_description=query_text,
                emotion_keywords=emotional_keywords,
                top_k=top_k
            )
            return verses
        except Exception as e:
            print(f"Error in finding verses: {e}")
            return []
    
    def contextualize_verse(self, user_problem: str, verse_data: Dict[str, Any], emotional_state: str) -> Dict[str, str]:
        try:
            contextualized = self.teaching_contextualizer(
                user_problem=user_problem,
                verse_data=verse_data,
                user_emotional_state=emotional_state
            )
            return {
                'contextual_guidance': contextualized.contextual_guidance,
                'practical_application': contextualized.practical_application,
                'reflection_question': contextualized.reflection_question,
                'verse_reference': f"Bhagavad Gita {verse_data['chapter_num']}.{verse_data['verse_num']}",
                'sanskrit_text': verse_data['sanskrit_text'],
                'english_translation': verse_data['english_text']
            }
        except Exception as e:
            print(f"Error in contextualizing verse: {e}")
            return {
                'contextual_guidance': f"The wisdom from Chapter {verse_data['chapter_num']}, Verse {verse_data['verse_num']} offers guidance for your situation.",
                'practical_application': "Reflect on this teaching and how it applies to your current circumstances.",
                'reflection_question': "How might this ancient wisdom guide your next steps?",
                'verse_reference': f"Bhagavad Gita {verse_data['chapter_num']}.{verse_data['verse_num']}",
                'sanskrit_text': verse_data['sanskrit_text'],
                'english_translation': verse_data['english_text']
            }
    
    def synthesize_response(self, user_problem: str, contextualized_teachings: List[Dict[str, str]]) -> Dict[str, str]:
        teachings_text = "\n\n".join([
            f"**{teaching['verse_reference']}**\n{teaching['contextual_guidance']}\n*Application: {teaching['practical_application']}*"
            for teaching in contextualized_teachings
        ])
        
        core_message = f"Based on {len(contextualized_teachings)} relevant verses from the Bhagavad Gita"
        
        try:
            synthesis = self.wisdom_synthesizer(
                user_problem=user_problem,
                contextual_teachings=teachings_text,
                core_message=core_message
            )
            return {
                'final_response': synthesis.final_response,
                'closing_blessing': synthesis.closing_blessing
            }
        except Exception as e:
            print(f"Error in synthesis: {e}")
            return {
                'final_response': "The Bhagavad Gita offers timeless wisdom for your situation. Consider these teachings carefully and apply them with patience and understanding.",
                'closing_blessing': "May you find peace and clarity on your spiritual journey."
            }
    
    def get_guidance(self, user_problem: str, include_verses: bool = True, session_id: Optional[str] = None) -> Dict[str, Any]:
        # Generate session ID if not provided
        if not session_id:
            session_id = str(uuid.uuid4())[:8]
        
        try:
            print(f"Analyzing problem: {user_problem[:100]}...")
            problem_analysis = self.analyze_problem(user_problem)
            
            print("Finding relevant verses...")
            relevant_verses = self.find_relevant_verses(problem_analysis, top_k=3)
            
            if not relevant_verses:
                # Log error
                self.logger.log_error(session_id, "semantic_search", "No relevant verses found", user_problem)
                return {
                    'error': 'Could not find relevant verses. Please try rephrasing your problem.',
                    'problem_analysis': problem_analysis,
                    'session_id': session_id
                }
            
            print(f"Contextualizing {len(relevant_verses)} verses...")
            contextualized_teachings = []
            for verse in relevant_verses:
                contextual_teaching = self.contextualize_verse(
                    user_problem=user_problem,
                    verse_data=verse,
                    emotional_state=problem_analysis['emotional_state']
                )
                contextualized_teachings.append(contextual_teaching)
            
            print("Synthesizing final response...")
            synthesis = self.synthesize_response(user_problem, contextualized_teachings)
            
            # Log everything in ONE comprehensive trace
            self.logger.log_complete_session(
                user_problem=user_problem,
                session_id=session_id,
                problem_analysis=problem_analysis,
                matched_contexts=relevant_verses,
                contextualized_teachings=contextualized_teachings,
                final_synthesis=synthesis,
                status="success"
            )
            
            response = {
                'problem_analysis': problem_analysis,
                'main_response': synthesis['final_response'],
                'closing_blessing': synthesis['closing_blessing'],
                'relevant_verses': contextualized_teachings if include_verses else len(contextualized_teachings),
                'session_id': session_id
            }
            
            return response
            
        except Exception as e:
            # Log error
            self.logger.log_error(session_id, "get_guidance", str(e), user_problem)
            return {
                'error': f'An error occurred: {str(e)}',
                'session_id': session_id
            }
    
    def save_session(self, session_data: Dict[str, Any], filename: str = None):
        if filename is None:
            from datetime import datetime
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
            filename = f"session_{timestamp}.json"
        
        session_file = self.cache_dir / filename
        with open(session_file, 'w', encoding='utf-8') as f:
            json.dump(session_data, f, indent=2, ensure_ascii=False)
        
        print(f"Session saved to {session_file}")
        return str(session_file)