main.py

"""
Multi-Agent RAG System with Agent2Agent Communication
"""

import os
from agent import ResearcherAgent, SummarizerAgent, AnswerAgent
from a2a_protocol import AgentMessage, create_message
from state import MultiAgentState
from langgraph.graph import StateGraph, END
from dotenv import load_dotenv
from langchain_ollama import ChatOllama, OllamaEmbeddings
from langchain_community.vectorstores import Chroma
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_community.document_loaders import TextLoader, DirectoryLoader

load_dotenv()

# Initialize LLM
llm = ChatOllama(model="llama3", temperature=0)
embeddings = OllamaEmbeddings(model="nomic-embed-text")

# Load and process documents
def load_and_process_documents(documents_dir: str = "./multi_agent_rag/documents"):
    """
    Load documents, split into chunks, and create vector store.
    This is the RAG setup phase.
    """
    print("Loading documents...")
    
    # Load documents
    if os.path.exists(documents_dir):
        loader = DirectoryLoader(
            documents_dir,
            glob="**/*.txt",
            loader_cls=TextLoader
        )
        documents = loader.load()
    else:
        print(f"Documents directory '{documents_dir}' not found. Creating sample...")
        # Create a sample document
        os.makedirs(documents_dir, exist_ok=True)
        with open(f"{documents_dir}/sample.txt", "w") as f:
            f.write("""
            Artificial Intelligence and Machine Learning
            
            Artificial Intelligence (AI) is a branch of computer science that aims to create 
            intelligent machines capable of performing tasks that typically require human intelligence.
            
            Machine Learning (ML) is a subset of AI that enables systems to learn and improve 
            from experience without being explicitly programmed. It uses algorithms to analyze 
            data, identify patterns, and make predictions.
            
            Deep Learning is a subset of machine learning that uses neural networks with multiple 
            layers to learn complex patterns in data. It has been particularly successful in 
            image recognition, natural language processing, and speech recognition.
            
            Key applications of AI and ML include:
            - Natural Language Processing (NLP)
            - Computer Vision
            - Recommendation Systems
            - Autonomous Vehicles
            - Healthcare Diagnostics
            """)
        loader = DirectoryLoader(
            documents_dir,
            glob="**/*.txt",
            loader_cls=TextLoader
        )
        documents = loader.load()
    
    print(f"Loaded {len(documents)} documents")
    
    # Split documents into chunks
    text_splitter = RecursiveCharacterTextSplitter(
        chunk_size=500,
        chunk_overlap=50
    )
    chunks = text_splitter.split_documents(documents)
    print(f"Split into {len(chunks)} chunks")
    
    # Create vector store (this is where embeddings are stored)
    vector_store = Chroma.from_documents(
        documents=chunks,
        embedding=embeddings,
        persist_directory="./multi_agent_rag/chroma_db"
    )
    print("Created vector store")
    
    return vector_store

vector_store = load_and_process_documents()

# Create agents
researcher = ResearcherAgent("researcher", vector_store, llm)
summarizer = SummarizerAgent("summarizer", llm)
answerer = AnswerAgent("answerer", llm)

# Define workflow nodes with Agent2Agent Protocol
def retrieve_node(state: MultiAgentState) -> MultiAgentState:
    """Node: Researcher agent retrieves documents using A2A protocol"""
    print(f"\n[Researcher Agent] Retrieving documents...")

    # Create request message to researcher agent
    request = create_message(
        from_agent="workflow",
        to_agent="researcher",
        message_type="request",
        task="retrieve_documents",
        data={"query": state['user_query']}
    )

    # Researcher processes the message
    response = researcher.receive_message(request)

    # Update state from response
    # documents = researcher.retrieve_documents(state['user_query'])
    documents = response.get('data', {}).get('documents', [])

    state['retrieved_documents'] = documents
    state['retrieved_context'] = "\n\n".join(documents)
    state['agent_messages'].append(request)
    state['agent_messages'].append(response)
    state['current_agent'] = "researcher"
    state['workflow_status'] = "retrieving"

    print(f"[A2A Protocol] Retrieved {response['data'].get('count', 0)} documents")
    return state

def summarize_node(state: MultiAgentState) -> MultiAgentState:
    """Node: Summarizer agent summarizes documents using A2A protocol"""
    print(f"\n[Summarizer Agent] Summarizing documents...")

    # Create request message to summarizer agent
    request = create_message(
        from_agent="workflow",
        to_agent="summarizer",
        message_type="request",
        task="summarize_documents",
        data={"documents": state['retrieved_documents']}
    )

    # Summarizer processes the message
    # summary = summarizer.summarize(state['retrieved_documents'])
    response = summarizer.receive_message(request)
    summary = response.get('data', {}).get('summary', '')
    
    # Update state from response
    state['summary'] = summary
    state['agent_messages'].append(request)
    state['agent_messages'].append(response)
    state['current_agent'] = "summarizer"
    state['workflow_status'] = "summarizing"

    print(f"[A2A Protocol] Summarized {response['data'].get('source_count', 0)} documents")
    return state

def answer_node(state: MultiAgentState) -> MultiAgentState:
    """Node: Answer agent generates answer using A2A protocol"""
    print(f"\n[Answer Agent] Generating answer...")

    context = state.get('summary', state.get('retrieved_context', ''))

    # Create request message to answer agent
    request = create_message(
        from_agent="workflow",
        to_agent="answerer",
        message_type="request",
        task="generate_answer",
        data={"query": state['user_query'], "context": context}
    )

    # Answer agent processes the message
    response = answerer.receive_message(request)
    # answer = answerer.generate_answer(state['user_query'], context)
    answer = response.get('data', {}).get('answer', '')

    # Update state from response
    state['final_answer'] = answer
    state['agent_messages'].append(request)
    state['agent_messages'].append(response)
    state['current_agent'] = "answerer"
    state['workflow_status'] = "complete"

    print(f"[A2A Protocol] Answer generated")
    return state

# Create workflow
def create_multi_agent_rag_workflow():
    """Create the multi-agent RAG workflow"""
    workflow = StateGraph(MultiAgentState)
    # Add nodes
    workflow.add_node("retrieve", retrieve_node)
    workflow.add_node("summarize", summarize_node)
    workflow.add_node("answer", answer_node)

    # Define the flow
    workflow.set_entry_point("retrieve")
    workflow.add_edge("retrieve", "summarize")
    workflow.add_edge("summarize", "answer")
    workflow.add_edge("answer", END)

    return workflow.compile()

def main():
    """Main function to run the multi-agent RAG system"""
    print("=" * 80)
    print("Multi-Agent RAG System with Agent2Agent Communication")
    print("=" * 80)

    agent = create_multi_agent_rag_workflow()

    # Interactive loop
    print("\nAsk questions (type 'quit' to exit):")
    print("-" * 80)

    while True:
        query = input("\nYour question: ").strip()
        if query.lower() in ['quit', 'exit', 'q']:
            print("Goodbye!")
            break
        if not query:
            continue
        initial_state = {
            "user_query": query,
            "agent_messages": [],
            "retrieved_context": "",
            "retrieved_documents": [],
            "summary": "",
            "final_answer": "",
            "current_agent": "initial",
            "workflow_status": "initializing",
        }
        result = agent.invoke(initial_state)
        print("\n" + "-" * 80)
        print("Answer:")
        print(result["final_answer"])
        print("-" * 80)
        print("Workflow Status:", result["workflow_status"])
        print("Documents Retrieved:", len(result["retrieved_documents"]))
        print("Agent Messages Exchanged:", len(result["agent_messages"]))
        print("-" * 80)
        if result["agent_messages"]:
            print("\nAgent Communication Log:")
            for msg in result["agent_messages"]:
                print(f"  {msg['from_agent']} → {msg['to_agent']}: {msg['task']} ({msg['message_type']})")
        print("-" * 80)

if __name__ == "__main__":
    main()