AI_Agent/coordinator.py

# 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.llm_adapters.hf_adapter import HuggingFaceAdapter
from AI_Agent.vector_store.vector_store import SimpleVectorStore
import asyncio

class Coordinator:
    def __init__(self):
        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")
        self.task_decomposer_chain = TaskDecomposerChain(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):
        decomposition_result = await self.task_decomposer_chain.run(brief)
        tasks_text = decomposition_result["tasks_text"]

        retrieval_result = await self.retrieval_chain.run(brief)
        contexts = retrieval_result["contexts"]

        reasoning_result = await self.reasoning_chain.run(brief, contexts)
        reasoning_text = reasoning_result["reasoning"]

        synthesis_result = await self.synthesis_chain.run(brief, reasoning_text, contexts)

        return {
            "decomposition": decomposition_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("Final synthesis:\n", result["synthesis"]["answer"])