agent.py

from abc import ABC, abstractmethod
from typing import Any, Dict
from a2a_protocol import AgentMessage, create_message

class BaseAgent(ABC):
    """Base class for all agents"""
    def __init__(self, name: str, llm=None):
        self.name = name
        self.llm = llm
        self.message_history = []

    def receive_message(self, message: AgentMessage) -> AgentMessage:
        """Receive a message from another agent"""
        self.message_history.append(message)
        return self.process_message(message)

    def send_message(self,
        to_agent: str,
        task: str,
        data: Dict[str, Any],
        message_type: str = "request"
        ) -> AgentMessage:
        """Create and Send a message to another agent"""
        return create_message(
            from_agent=self.name,
            to_agent=to_agent,
            message_type=message_type,
            task=task,
            data=data,
        )

    @abstractmethod
    def process_message(self, message: AgentMessage) -> AgentMessage:
        """Process a message and return a response"""
        pass






class ResearcherAgent(BaseAgent):
    """Agent responsible for retrieving relevant documents"""
    def __init__(self, name: str, vector_store, llm=None):
        super().__init__(name, llm)
        self.vector_store = vector_store

    def process_message(self, message: AgentMessage) -> AgentMessage:
        """Process a message and return a response"""
        query = message["data"].get("query","")
        # retrieve relevant documents using RAG
        doc_texts = self.retrieve(query)
        response_data = {
            "documents": doc_texts,
            "count": len(doc_texts),
            "query": query,
        }
        return self.send_message(
            to_agent=message["from_agent"], # send back to the sender
            task="retrieval_complete",
            data=response_data,
            message_type="response",
        )

    def retrieve(self, query):
        """Retrieve relevant documents from the vector store"""
        relevant_docs = self.vector_store.similarity_search(query, k=5)
        return [doc.page_content for doc in relevant_docs]






class SummarizerAgent(BaseAgent):
    """Agent responsible for summarizing documents"""
    def __init__(self, name: str, llm=None):
        super().__init__(name, llm)

    def process_message(self, message: AgentMessage) -> AgentMessage:
        """Process a message and return a summarized response"""
        documents = message["data"].get("documents", [])
        summary = self.summarize(documents)
        response_data = {
            "summary": summary,
            "source_count": len(documents),
        }
        return self.send_message(
            to_agent=message["from_agent"],
            task="summarization_complete",
            data=response_data,
            message_type="response",
        )

    def summarize(self, documents: list[str]) -> str:
        """Summarize a list of documents"""
        combined_text = "\n\n".join(documents)
        prompt = f"""Summarize the following documents, focusing on the key points and main ideas.
        {combined_text}
        Summary:
        """
        if self.llm:
            response = self.llm.invoke(prompt)
            return response.content
        else:
            return f"No LLM provided, summary of {len(documents)} documents"






class AnswerAgent(BaseAgent):
    """Agent responsible for generating final answers"""
    def __init__(self, name: str, llm=None):
        super().__init__(name, llm)

    def process_message(self, message: AgentMessage) -> AgentMessage:
        """Process a message and return an answer"""
        query = message["data"].get("query", "")
        context = message["data"].get("context", "")
        answer = self.generate_answer(query, context)
        response_data = {
            "answer": answer,
            "query": query
        }
        return self.send_message(
            to_agent=message["from_agent"],
            task="answer_generation_complete",
            data=response_data,
            message_type="response",
        )

    def generate_answer(self, query: str, context: str) -> str:
        """Generating an answer"""
        prompt = f"""Answer the question based on the context provided.
        Context:
        {context}
        Question:
        {query}
        Answer based on the context. If the context doesn't have enough information, say so."""
        
        if self.llm:
            response = self.llm.invoke(prompt)
            return response.content
        else:
            return f"No LLM provided, answer generation requires LLM."