weave_logger.py

#!/usr/bin/env python3

import weave
import os
from typing import Dict, List, Any, Optional
from datetime import datetime

class BhagavadGitaWeaveLogger:
    def __init__(self, project_name: str = "bhagavad-gita-qa"):
        """Initialize Weave logging for the Bhagavad Gita bot"""
        self.project_name = project_name
        self.init_weave()
    
    def init_weave(self):
        """Initialize Weave with WandB project"""
        try:
            # Initialize Weave normally
            weave.init(self.project_name)
            
            # Disable automatic DSPy tracing to prevent multiple traces per session
            import dspy
            if hasattr(dspy.settings, 'trace') and dspy.settings.trace:
                dspy.settings.trace = None
                print("✅ Disabled DSPy automatic tracing")
            
            print("✅ Weave initialized (minimal mode)")
        except Exception as e:
            print(f"⚠️ Weave initialization warning: {e}")
    
    @weave.op(name="Bhagavad_Gita_Session")
    def log_complete_session(
        self, 
        user_problem: str,
        session_id: str,
        problem_analysis: Dict[str, str],
        matched_contexts: List[Dict[str, Any]],
        contextualized_teachings: List[Dict[str, str]],
        final_synthesis: Dict[str, str],
        status: str = "success"
    ) -> Dict[str, Any]:
        """Single comprehensive operation that logs entire user session"""
        
        # Extract similarity scores and chapters
        similarity_scores = [ctx.get("similarity_score", 0.0) for ctx in matched_contexts if ctx.get("similarity_score")]
        chapters_used = list(set([ctx.get("chapter_num", 0) for ctx in matched_contexts if ctx.get("chapter_num")]))
        
        # Create comprehensive session data
        session_data = {
            "session_id": session_id,
            "timestamp": datetime.now().isoformat(),
            "status": status,
            
            # Input
            "user_input": {
                "original_problem": user_problem,
                "problem_length": len(user_problem),
                "word_count": len(user_problem.split())
            },
            
            # Analysis
            "problem_analysis": {
                "key_themes": problem_analysis.get("key_themes", ""),
                "emotional_state": problem_analysis.get("emotional_state", ""),  
                "core_question": problem_analysis.get("core_question", "")
            },
            
            # Context Matching
            "context_matching": {
                "num_verses_found": len(matched_contexts),
                "avg_similarity": sum(similarity_scores) / len(similarity_scores) if similarity_scores else 0.0,
                "max_similarity": max(similarity_scores) if similarity_scores else 0.0,
                "min_similarity": min(similarity_scores) if similarity_scores else 0.0,
                "chapters_covered": chapters_used,
                "num_chapters": len(chapters_used),
                "verses_details": [
                    {
                        "verse_id": f"{ctx.get('chapter_num', 0)}.{ctx.get('verse_num', 0)}",
                        "similarity_score": ctx.get("similarity_score", 0.0),
                        "sanskrit_original": ctx.get("sanskrit_text", ""),
                        "english_translation": ctx.get("english_text", ""),
                        "chapter": ctx.get("chapter_num", 0),
                        "verse": ctx.get("verse_num", 0)
                    }
                    for ctx in matched_contexts
                ]
            },
            
            # Generated Guidance
            "generated_guidance": {
                "contextualized_teachings": [
                    {
                        "verse_reference": teaching.get("verse_reference", ""),
                        "contextual_guidance": teaching.get("contextual_guidance", ""),
                        "practical_application": teaching.get("practical_application", ""),
                        "reflection_question": teaching.get("reflection_question", ""),
                        "guidance_length": len(teaching.get("contextual_guidance", ""))
                    }
                    for teaching in contextualized_teachings
                ],
                "final_response": final_synthesis.get("final_response", ""),
                "closing_blessing": final_synthesis.get("closing_blessing", ""),
                "total_response_length": len(final_synthesis.get("final_response", "")) + len(final_synthesis.get("closing_blessing", "")),
                "avg_guidance_length": sum([len(t.get("contextual_guidance", "")) for t in contextualized_teachings]) / len(contextualized_teachings) if contextualized_teachings else 0
            }
        }
        
        # Set comprehensive attributes for easy filtering and analysis
        weave.attributes({
            # Core identifiers
            "session_id": session_id,
            "status": status,
            
            # Problem categorization
            "primary_theme": problem_analysis.get("key_themes", "").split(",")[0].strip() if problem_analysis.get("key_themes") else "unknown",
            "primary_emotion": problem_analysis.get("emotional_state", "").split(",")[0].strip() if problem_analysis.get("emotional_state") else "unknown",
            
            # Context quality metrics
            "verses_found": len(matched_contexts),
            "avg_similarity": session_data["context_matching"]["avg_similarity"],
            "max_similarity": session_data["context_matching"]["max_similarity"],
            "chapters_used": len(chapters_used),
            "multi_chapter": len(chapters_used) > 1,
            
            # Response quality metrics
            "response_length": session_data["generated_guidance"]["total_response_length"],
            "has_blessing": bool(final_synthesis.get("closing_blessing", "")),
            "avg_teaching_length": session_data["generated_guidance"]["avg_guidance_length"],
            
            # Processing success indicators
            "high_relevance": session_data["context_matching"]["max_similarity"] > 0.7 if similarity_scores else False,
            "good_coverage": len(matched_contexts) >= 3,
            "comprehensive_response": session_data["generated_guidance"]["total_response_length"] > 500
        })
        
        return session_data
    
    @weave.op(name="User_Feedback")
    def log_user_feedback(
        self, 
        session_id: str, 
        feedback_type: str, 
        comments: str = ""
    ) -> Dict[str, Any]:
        """Log user feedback as separate minimal operation"""
        feedback_data = {
            "session_id": session_id,
            "timestamp": datetime.now().isoformat(),
            "feedback_type": feedback_type,  # "positive" or "negative"
            "comments": comments,
            "has_comments": bool(comments.strip()),
            "comments_length": len(comments.strip()) if comments else 0,
            "sentiment_score": 1 if feedback_type == "positive" else -1
        }
        
        weave.attributes({
            "session_id": session_id,
            "feedback_type": feedback_type,
            "has_comments": bool(comments.strip()),
            "sentiment": feedback_type
        })
        
        return feedback_data
    
    @weave.op(name="Session_Error")
    def log_error(
        self, 
        session_id: str, 
        error_step: str, 
        error_message: str, 
        user_input: str = ""
    ) -> Dict[str, Any]:
        """Log errors as minimal separate operations"""
        error_data = {
            "session_id": session_id,
            "timestamp": datetime.now().isoformat(),
            "error_step": error_step,
            "error_message": error_message,
            "user_input": user_input[:200] if user_input else "",  # Truncate for logging
            "status": "error"
        }
        
        weave.attributes({
            "session_id": session_id,
            "error_step": error_step,
            "status": "error",
            "has_user_input": bool(user_input)
        })
        
        return error_data