Spaces:
Sleeping
Sleeping
File size: 6,740 Bytes
068aa4e |
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 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 |
"""
Agent Orchestrator
Connects all agents using LangGraph workflow
"""
import os
from dotenv import load_dotenv
from langgraph.graph import StateGraph, START, END
from typing import TypedDict, List, Dict
load_dotenv()
class AgentState(TypedDict):
"""State passed between agents"""
original_query: str
reformulated_query: str
retrieved_documents: List[Dict]
synthesized_answer: str
validation_result: Dict
final_answer: str
metadata: Dict
class AgentOrchestrator:
def __init__(self, rag_system):
"""Initialize orchestrator with RAG system"""
print("🔗 Initializing Agent Orchestrator...\n")
self.rag = rag_system
self.workflow = self._build_workflow()
print("✅ Agent Orchestrator ready!\n")
def _build_workflow(self):
"""Build LangGraph workflow"""
workflow = StateGraph(AgentState)
# Define nodes
workflow.add_node("query_understanding", self._query_understanding_node)
workflow.add_node("retrieval", self._retrieval_node)
workflow.add_node("synthesis", self._synthesis_node)
workflow.add_node("validation", self._validation_node)
workflow.add_node("finalize", self._finalize_node)
# Define edges
workflow.add_edge(START, "query_understanding")
workflow.add_edge("query_understanding", "retrieval")
workflow.add_edge("retrieval", "synthesis")
workflow.add_edge("synthesis", "validation")
workflow.add_edge("validation", "finalize")
workflow.add_edge("finalize", END)
return workflow.compile()
def _query_understanding_node(self, state: AgentState) -> AgentState:
"""Query Understanding Agent Node"""
print("\n" + "=" * 70)
print("🧠 AGENT 1: QUERY UNDERSTANDING")
print("=" * 70)
original_query = state["original_query"]
reformulated_query = self.rag.query_agent.reformulate_query(original_query)
state["reformulated_query"] = reformulated_query
state["metadata"]["query_understanding_time"] = 0
return state
def _retrieval_node(self, state: AgentState) -> AgentState:
"""Multi-Source Retrieval Agent Node"""
print("\n" + "=" * 70)
print("🔍 AGENT 2: MULTI-SOURCE RETRIEVAL")
print("=" * 70)
reformulated_query = state["reformulated_query"]
retrieved_results = self.rag.retrieval_agent.retrieve(reformulated_query, top_k=5)
# Convert to document format
documents = []
for result in retrieved_results:
documents.append({
'content': result['content'],
'source': result.get('source', 'unknown'),
'score': result['score']
})
state["retrieved_documents"] = documents
state["metadata"]["num_documents_retrieved"] = len(documents)
return state
def _synthesis_node(self, state: AgentState) -> AgentState:
"""Synthesis Agent Node"""
print("\n" + "=" * 70)
print("🧬 AGENT 3: SYNTHESIS")
print("=" * 70)
original_query = state["original_query"]
documents = state["retrieved_documents"]
synthesized_answer = self.rag.synthesis_agent.synthesize(
original_query,
documents
)
state["synthesized_answer"] = synthesized_answer
return state
def _validation_node(self, state: AgentState) -> AgentState:
"""Validation Agent Node"""
print("\n" + "=" * 70)
print("✅ AGENT 4: VALIDATION")
print("=" * 70)
answer = state["synthesized_answer"]
documents = state["retrieved_documents"]
validation_result = self.rag.validation_agent.validate(answer, documents)
state["validation_result"] = validation_result
return state
def _finalize_node(self, state: AgentState) -> AgentState:
"""Finalize and format response"""
print("\n" + "=" * 70)
print("📋 FINALIZATION")
print("=" * 70 + "\n")
state["final_answer"] = state["synthesized_answer"]
return state
def run(self, query: str) -> Dict:
"""Run complete agent orchestration workflow"""
print("\n" + "=" * 80)
print("🚀 MULTI-AGENT ORCHESTRATION WORKFLOW")
print("=" * 80)
print(f"\nINPUT QUERY: {query}\n")
# Initialize state
initial_state = {
"original_query": query,
"reformulated_query": "",
"retrieved_documents": [],
"synthesized_answer": "",
"validation_result": {},
"final_answer": "",
"metadata": {}
}
# Run workflow
final_state = self.workflow.invoke(initial_state)
# Format and display results
self._display_results(final_state)
return final_state
def _display_results(self, state: AgentState):
"""Display final results"""
print("\n" + "=" * 80)
print("🎯 FINAL RESULTS")
print("=" * 80 + "\n")
print("ORIGINAL QUERY:")
print(f" {state['original_query']}\n")
print("REFORMULATED QUERY:")
print(f" {state['reformulated_query']}\n")
print("ANSWER:")
print("-" * 80)
print(state['final_answer'])
print("-" * 80 + "\n")
validation = state['validation_result']
print("VALIDATION:")
print(f" Status: {'✅ VALID' if validation['is_valid'] else '⚠️ NEEDS REVIEW'}")
print(f" Confidence: {validation['confidence']}%\n")
print("SOURCES:")
for i, doc in enumerate(state['retrieved_documents'], 1):
print(f" {i}. {doc['source']} (relevance: {doc['score']:.2%})")
print("\n" + "=" * 80 + "\n")
# Test the orchestrator
if __name__ == "__main__":
from rag_system import RAGSystem
api_key = os.getenv("GROQ_API_KEY")
# Initialize RAG system
print("Initializing RAG System...")
rag = RAGSystem(groq_api_key=api_key)
# Initialize orchestrator
orchestrator = AgentOrchestrator(rag)
# Test queries
test_queries = [
"How do I create a FastAPI endpoint?",
"What is the leave policy?",
"Tell me about remote work"
]
for query in test_queries:
result = orchestrator.run(query)
print("\n" + "=" * 80 + "\n")
|