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"])