rag/rag_integration.py

"""
RAG Integration - Unified interface for ChromaDB queries across all agents
Provides context-aware retrieval for nutrition, exercise, health tips, etc.
"""

from typing import Dict, Any, List, Optional
from rag.query_engine import query_with_context
import json


class RAGIntegration:
    """
    Unified RAG interface for all agents
    Retrieves relevant health knowledge from ChromaDB
    """
    
    def __init__(self):
        """Initialize RAG integration"""
        self.nutrition_prompt = """You are a nutrition expert. Use the provided documents to answer nutrition questions.
Focus on: dietary recommendations, macro/micronutrients, meal planning, food groups."""
        
        self.exercise_prompt = """You are a fitness expert. Use the provided documents to answer exercise questions.
Focus on: workout routines, exercise techniques, fitness progression, safety guidelines."""
        
        self.health_prompt = """You are a health consultant. Use the provided documents to answer health questions.
Focus on: health tips, disease prevention, wellness practices, lifestyle recommendations."""
    
    def query_nutrition(self, query: str) -> Dict[str, Any]:
        """
        Query nutrition knowledge from ChromaDB
        
        Args:
            query: Nutrition question
        
        Returns:
            Dict with answer, sources, and metadata
        """
        return query_with_context(query, self.nutrition_prompt)
    
    def query_exercise(self, query: str) -> Dict[str, Any]:
        """
        Query exercise knowledge from ChromaDB
        
        Args:
            query: Exercise question
        
        Returns:
            Dict with answer, sources, and metadata
        """
        return query_with_context(query, self.exercise_prompt)
    
    def query_health(self, query: str) -> Dict[str, Any]:
        """
        Query general health knowledge from ChromaDB
        
        Args:
            query: Health question
        
        Returns:
            Dict with answer, sources, and metadata
        """
        return query_with_context(query, self.health_prompt)
    
    def query_generic(self, query: str, context: str = "") -> Dict[str, Any]:
        """
        Generic query with custom context
        
        Args:
            query: Question
            context: Custom system prompt context
        
        Returns:
            Dict with answer, sources, and metadata
        """
        prompt = f"""You are a helpful health assistant. Use the provided documents to answer questions.
{context}"""
        return query_with_context(query, prompt)
    
    def extract_answer(self, result: Dict[str, Any]) -> str:
        """Extract just the answer text from RAG result"""
        return result.get('answer', '')
    
    def extract_sources(self, result: Dict[str, Any]) -> List[Dict[str, Any]]:
        """Extract source documents from RAG result"""
        return result.get('source_docs', [])
    
    def extract_metadata(self, result: Dict[str, Any]) -> Dict[str, Any]:
        """Extract metadata (time, tokens, status) from RAG result"""
        return result.get('metadata', {})
    
    def format_response_with_sources(self, result: Dict[str, Any]) -> str:
        """
        Format RAG response (answer only, no sources)
        
        Args:
            result: RAG query result
        
        Returns:
            Answer text only
        """
        answer = result.get('answer', '')
        return answer
    
    def is_success(self, result: Dict[str, Any]) -> bool:
        """Check if RAG query was successful"""
        metadata = result.get('metadata', {})
        return metadata.get('status') == 'SUCCESS'
    
    def get_query_time(self, result: Dict[str, Any]) -> float:
        """Get query execution time in seconds"""
        metadata = result.get('metadata', {})
        return metadata.get('time_s', 0.0)
    
    def get_token_count(self, result: Dict[str, Any]) -> int:
        """Get token count used in query"""
        metadata = result.get('metadata', {})
        return metadata.get('tokens', 0)


# Global RAG instance
_rag_instance = None


def get_rag_integration() -> RAGIntegration:
    """Get or create global RAG integration instance"""
    global _rag_instance
    if _rag_instance is None:
        _rag_instance = RAGIntegration()
    return _rag_instance


# Convenience functions for direct access
def query_nutrition(query: str) -> Dict[str, Any]:
    """Query nutrition knowledge"""
    rag = get_rag_integration()
    return rag.query_nutrition(query)


def query_exercise(query: str) -> Dict[str, Any]:
    """Query exercise knowledge"""
    rag = get_rag_integration()
    return rag.query_exercise(query)


def query_health(query: str) -> Dict[str, Any]:
    """Query health knowledge"""
    rag = get_rag_integration()
    return rag.query_health(query)


def query_generic(query: str, context: str = "") -> Dict[str, Any]:
    """Generic query with custom context"""
    rag = get_rag_integration()
    return rag.query_generic(query, context)