Spaces:
No application file
No application file
File size: 3,353 Bytes
b325aad |
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 |
import time
from typing import List
from src.config.settings import settings
from src.agenticRAG.models.state import AgentState
from src.agenticRAG.models.schemas import QueryRequest, QueryResponse
from src.agenticRAG.graph.builder import GraphBuilder
from loguru import logger
class AgenticRAGSystem:
"""Main AgenticRAG system"""
def __init__(self):
# Validate settings
settings.validate()
# Create graph
self.app = GraphBuilder.create_graph()
logger.info("AgenticRAG system initialized successfully")
def process_query(self, query: str) -> QueryResponse:
"""Process a single query"""
start_time = time.time()
try:
# Initialize state
initial_state = AgentState(user_query=query)
# Run the graph
final_state = self.app.invoke(initial_state)
# Calculate processing time
processing_time = time.time() - start_time
# Create response
response = QueryResponse(
query=final_state.user_query,
upgraded_query=final_state.upgraded_query,
route_taken=final_state.route_decision,
response=final_state.final_response,
metadata=final_state.metadata,
processing_time=processing_time
)
logger.info(f"Query processed successfully in {processing_time:.2f}s")
return response
except Exception as e:
logger.error(f"Error processing query: {e}")
raise
def process_batch(self, queries: List[str]) -> List[QueryResponse]:
"""Process multiple queries"""
responses = []
for query in queries:
try:
response = self.process_query(query)
responses.append(response)
except Exception as e:
logger.error(f"Error processing query '{query}': {e}")
return responses
def agenticRAGResponse(query: str) -> QueryResponse:
"""Function to get response for a single query"""
system = AgenticRAGSystem()
return system.process_query(query)
def main():
"""Main function"""
# Initialize system
system = AgenticRAGSystem()
# Test queries
test_queries = [
"What is machine learning?",
"Latest news about AI",
"Write a poem about spring"
]
# Process queries
for query in test_queries:
print(f"\n{'='*50}")
print(f"Query: {query}")
print(f"{'='*50}")
try:
response = system.process_query(query)
print(f"Original Query: {response.query}")
print(f"Upgraded Query: {response.upgraded_query}")
print(f"Route Taken: {response.route_taken}")
print(f"Response: {response.response}")
print(f"Processing Time: {response.processing_time:.2f}s")
print(f"Metadata: {response.metadata}")
except Exception as e:
print(f"Error: {e}")
if __name__ == "__main__":
main() |