Update AI_Agent/coordinator.py

AI_Agent/coordinator.py

@@ -3,27 +3,31 @@ 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
-from AI_Agent.chains.task_assigner_chain import TaskAssignerChain
 
 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."
         ])
-        self.llm = HuggingFaceAdapter(model_name="tiiuae/falcon-7b-instruct")
+        
+        # Use the CPU-friendly Google Gemma 3n model
+        self.llm = HuggingFaceAdapter(model_name="google/gemma-3n-E2B-it")
+        
+        # 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)
-        self.task_assigner_chain = TaskAssignerChain(self.llm)
-
 
     async def run_task(self, brief: str):
         # Step 1: Task decomposition
@@ -55,7 +59,6 @@ class Coordinator:
 
 
 if __name__ == "__main__":
-    import asyncio
     coord = Coordinator()
     brief = "Build a task management app with user authentication and task sharing"
     result = asyncio.run(coord.run_task(brief))
@@ -63,4 +66,3 @@ if __name__ == "__main__":
     print("Decomposed tasks:\n", result["decomposition"]["tasks_text"])
     print("Assigned tasks:\n", result["assignment"]["assigned_tasks_text"])
     print("Final synthesis:\n", result["synthesis"]["answer"])
-