Spaces:
Runtime error
Runtime error
File size: 5,083 Bytes
eeb0f9c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 |
"""
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)
|