Spaces:
Sleeping
Sleeping
File size: 2,818 Bytes
1e00e3b 9f08e65 1e00e3b 1c6c663 1e00e3b 9f08e65 1e00e3b 9f08e65 33f0ebb 9f08e65 1e00e3b 9f08e65 1e00e3b 1c6c663 1e00e3b 1c6c663 1e00e3b 1c6c663 1e00e3b 1c6c663 1e00e3b 1c6c663 1e00e3b 1c6c663 1e00e3b 1c6c663 1e00e3b 1c6c663 1e00e3b 1c6c663 1e00e3b |
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 |
# AI_Agent/coordinator.py
from AI_Agent.chains.retrieval_chain import RetrievalChain
from AI_Agent.chains.reasoning_chain import ReasoningChain
from AI_Agent.chains.synthesis_chain import SynthesisChain
from AI_Agent.chains.task_decomposer_chain import TaskDecomposerChain
from AI_Agent.chains.task_assigner_chain import TaskAssignerChain
from AI_Agent.llm_adapters.hf_adapter import HuggingFaceAdapter
from AI_Agent.vector_store.vector_store import SimpleVectorStore
import asyncio
class Coordinator:
def __init__(self):
# Initialize vector store with sample documents
self.vector_store = SimpleVectorStore()
self.vector_store.add_documents([
"Doc 1: Notes on building retrieval chains.",
"Doc 2: Sample LLM reasoning methods.",
"Doc 3: Multi-chain orchestration tips."
])
# Use the CPU-friendly Google Gemma 3n model
self.llm = HuggingFaceAdapter(model_name="EleutherAI/gpt-neo-125M")
# Initialize all chains
self.task_decomposer_chain = TaskDecomposerChain(self.llm)
self.task_assigner_chain = TaskAssignerChain(self.llm)
self.retrieval_chain = RetrievalChain(self.vector_store)
self.reasoning_chain = ReasoningChain(self.llm)
self.synthesis_chain = SynthesisChain(self.llm)
async def run_task(self, brief: str):
# Step 1: Task decomposition
decomposition_result = await self.task_decomposer_chain.run(brief)
tasks_text = decomposition_result["tasks_text"]
# Step 2: Task assignment
assignment_result = await self.task_assigner_chain.run(tasks_text)
assigned_tasks_text = assignment_result["assigned_tasks_text"]
# Step 3: Retrieval
retrieval_result = await self.retrieval_chain.run(brief)
contexts = retrieval_result["contexts"]
# Step 4: Reasoning
reasoning_result = await self.reasoning_chain.run(brief, contexts)
reasoning_text = reasoning_result["reasoning"]
# Step 5: Synthesis
synthesis_result = await self.synthesis_chain.run(brief, reasoning_text, contexts)
return {
"decomposition": decomposition_result,
"assignment": assignment_result,
"retrieval": retrieval_result,
"reasoning": reasoning_result,
"synthesis": synthesis_result
}
if __name__ == "__main__":
coord = Coordinator()
brief = "Build a task management app with user authentication and task sharing"
result = asyncio.run(coord.run_task(brief))
print("Decomposed tasks:\n", result["decomposition"]["tasks_text"])
print("Assigned tasks:\n", result["assignment"]["assigned_tasks_text"])
print("Final synthesis:\n", result["synthesis"]["answer"])
|