agents.py

"""
Multi-Agent Document Q&A Workflow
=================================
Orchestrates three specialized agents (Retrieval, Summarization, Response Generation)
through LangChain, with vector embeddings and local LLM inference via Ollama.

Each agent has a defined role and tool access. The coordinator routes queries
through the pipeline: Retrieval → Summarization → Response Generation.
"""

from __future__ import annotations

import os
import logging
from dataclasses import dataclass, field
from typing import List, Optional

from langchain_community.llms import Ollama
from langchain_community.embeddings import OllamaEmbeddings
from langchain_community.vectorstores import FAISS
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_core.prompts import PromptTemplate
from langchain_core.documents import Document

logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(name)s] %(message)s")
logger = logging.getLogger("multi-agent-qa")


# ---------------------------------------------------------------------------
# Configuration
# ---------------------------------------------------------------------------

@dataclass
class AgentConfig:
    """Central configuration for the multi-agent pipeline."""
    ollama_base_url: str = os.getenv("OLLAMA_BASE_URL", "http://localhost:11434")
    model_name: str = os.getenv("OLLAMA_MODEL", "llama3.2")
    embedding_model: str = os.getenv("EMBEDDING_MODEL", "llama3.2")
    chunk_size: int = 500
    chunk_overlap: int = 50
    top_k: int = 5
    temperature: float = 0.1
    vectorstore_path: str = "vectorstore"


# ---------------------------------------------------------------------------
# Document Ingestion
# ---------------------------------------------------------------------------

class DocumentIngester:
    """Loads and chunks documents, then builds a FAISS vector store."""

    def __init__(self, config: AgentConfig):
        self.config = config
        self.splitter = RecursiveCharacterTextSplitter(
            chunk_size=config.chunk_size,
            chunk_overlap=config.chunk_overlap,
            separators=["\n\n", "\n", ". ", " ", ""],
        )
        self.embeddings = OllamaEmbeddings(
            base_url=config.ollama_base_url,
            model=config.embedding_model,
        )
        self.vectorstore: Optional[FAISS] = None

    def ingest_texts(self, texts: List[str], metadatas: Optional[List[dict]] = None) -> FAISS:
        """Split raw texts into chunks and build the vector store."""
        docs: List[Document] = []
        for i, text in enumerate(texts):
            meta = metadatas[i] if metadatas else {"source": f"doc_{i}"}
            chunks = self.splitter.create_documents([text], metadatas=[meta])
            docs.extend(chunks)
        logger.info("Created %d chunks from %d documents", len(docs), len(texts))

        self.vectorstore = FAISS.from_documents(docs, self.embeddings)
        logger.info("FAISS vector store built (%d vectors)", len(docs))
        return self.vectorstore

    def ingest_files(self, file_paths: List[str]) -> FAISS:
        """Read files from disk and ingest them."""
        texts, metas = [], []
        for path in file_paths:
            with open(path, "r", encoding="utf-8") as f:
                texts.append(f.read())
            metas.append({"source": os.path.basename(path)})
        return self.ingest_texts(texts, metas)

    def save(self, path: Optional[str] = None):
        """Persist the vector store to disk."""
        save_path = path or self.config.vectorstore_path
        if self.vectorstore:
            self.vectorstore.save_local(save_path)
            logger.info("Vector store saved to %s", save_path)

    def load(self, path: Optional[str] = None):
        """Load a persisted vector store."""
        load_path = path or self.config.vectorstore_path
        self.vectorstore = FAISS.load_local(
            load_path, self.embeddings, allow_dangerous_deserialization=True
        )
        logger.info("Vector store loaded from %s", load_path)


# ---------------------------------------------------------------------------
# Agent Definitions
# ---------------------------------------------------------------------------

class RetrievalAgent:
    """
    Agent 1: Retrieval
    -----------------
    Searches the vector store for chunks most relevant to the query.
    Returns ranked document chunks with similarity scores.
    """

    def __init__(self, vectorstore: FAISS, config: AgentConfig):
        self.vectorstore = vectorstore
        self.config = config
        logger.info("RetrievalAgent initialized (top_k=%d)", config.top_k)

    def retrieve(self, query: str) -> List[dict]:
        """Return top-k chunks with scores."""
        results = self.vectorstore.similarity_search_with_score(query, k=self.config.top_k)
        retrieved = []
        for doc, score in results:
            retrieved.append({
                "content": doc.page_content,
                "source": doc.metadata.get("source", "unknown"),
                "score": round(float(score), 4),
            })
        logger.info(
            "RetrievalAgent found %d chunks (best score: %.4f)",
            len(retrieved),
            retrieved[0]["score"] if retrieved else 0.0,
        )
        return retrieved


class SummarizationAgent:
    """
    Agent 2: Summarization
    ----------------------
    Takes retrieved chunks and produces a focused, condensed context
    that distills the key information relevant to the query.
    """

    PROMPT = PromptTemplate(
        input_variables=["query", "chunks"],
        template=(
            "You are a summarization agent. Your job is to distill retrieved document "
            "chunks into a focused summary that captures the key information relevant "
            "to the user's query.\n\n"
            "USER QUERY: {query}\n\n"
            "RETRIEVED CHUNKS:\n{chunks}\n\n"
            "Write a concise summary (3-5 sentences) that captures ONLY the information "
            "from the chunks that is relevant to answering the query. "
            "If the chunks do not contain relevant information, say so explicitly.\n\n"
            "SUMMARY:"
        ),
    )

    def __init__(self, config: AgentConfig):
        self.llm = Ollama(
            base_url=config.ollama_base_url,
            model=config.model_name,
            temperature=config.temperature,
        )
        self.chain = self.PROMPT | self.llm
        logger.info("SummarizationAgent initialized (model=%s)", config.model_name)

    def summarize(self, query: str, chunks: List[dict]) -> str:
        """Produce a query-focused summary from retrieved chunks."""
        chunks_text = "\n---\n".join(
            f"[Source: {c['source']} | Score: {c['score']}]\n{c['content']}"
            for c in chunks
        )
        summary = self.chain.invoke({"query": query, "chunks": chunks_text})
        logger.info("SummarizationAgent produced summary (%d chars)", len(summary))
        return summary.strip()


class ResponseAgent:
    """
    Agent 3: Response Generation
    ----------------------------
    Takes the summarized context and generates a final, well-structured
    answer to the user's question, citing sources where possible.
    """

    PROMPT = PromptTemplate(
        input_variables=["query", "summary", "sources"],
        template=(
            "You are a response-generation agent. Using the provided summary of "
            "relevant document content, generate a clear and accurate answer to "
            "the user's question.\n\n"
            "USER QUESTION: {query}\n\n"
            "DOCUMENT SUMMARY:\n{summary}\n\n"
            "SOURCES CONSULTED: {sources}\n\n"
            "INSTRUCTIONS:\n"
            "- Answer the question directly based on the summary.\n"
            "- If the summary indicates insufficient information, say so honestly.\n"
            "- Reference source documents when making specific claims.\n"
            "- Keep the answer focused and concise.\n\n"
            "ANSWER:"
        ),
    )

    def __init__(self, config: AgentConfig):
        self.llm = Ollama(
            base_url=config.ollama_base_url,
            model=config.model_name,
            temperature=config.temperature,
        )
        self.chain = self.PROMPT | self.llm
        logger.info("ResponseAgent initialized (model=%s)", config.model_name)

    def generate(self, query: str, summary: str, sources: List[str]) -> str:
        """Generate the final answer."""
        sources_str = ", ".join(sorted(set(sources))) if sources else "none"
        answer = self.chain.invoke({"query": query, "summary": summary, "sources": sources_str})
        logger.info("ResponseAgent generated answer (%d chars)", len(answer))
        return answer.strip()


# ---------------------------------------------------------------------------
# Coordinator — orchestrates the agent pipeline
# ---------------------------------------------------------------------------

@dataclass
class PipelineResult:
    """Captures the full trace of a query through the pipeline."""
    query: str
    retrieved_chunks: List[dict]
    summary: str
    answer: str
    sources: List[str]

    def __str__(self):
        header = f"Query: {self.query}\n{'=' * 60}"
        chunks_info = f"Retrieved {len(self.retrieved_chunks)} chunks from: {', '.join(self.sources)}"
        return f"{header}\n{chunks_info}\n\nSummary:\n{self.summary}\n\nAnswer:\n{self.answer}"


class MultiAgentCoordinator:
    """
    Orchestrates the three-agent pipeline:
      1. RetrievalAgent  → fetches relevant chunks from the vector store
      2. SummarizationAgent → distills chunks into focused context
      3. ResponseAgent → generates the final answer

    Maintains shared context between agents and logs the full trace.
    """

    def __init__(self, config: Optional[AgentConfig] = None):
        self.config = config or AgentConfig()
        self.ingester = DocumentIngester(self.config)
        self._retrieval_agent: Optional[RetrievalAgent] = None
        self._summarization_agent: Optional[SummarizationAgent] = None
        self._response_agent: Optional[ResponseAgent] = None
        logger.info("MultiAgentCoordinator created")

    def ingest(self, file_paths: Optional[List[str]] = None,
               texts: Optional[List[str]] = None,
               metadatas: Optional[List[dict]] = None):
        """Ingest documents — either from file paths or raw text strings."""
        if file_paths:
            self.ingester.ingest_files(file_paths)
        elif texts:
            self.ingester.ingest_texts(texts, metadatas)
        else:
            raise ValueError("Provide either file_paths or texts")

        # Initialize agents once the vector store is ready
        self._retrieval_agent = RetrievalAgent(self.ingester.vectorstore, self.config)
        self._summarization_agent = SummarizationAgent(self.config)
        self._response_agent = ResponseAgent(self.config)
        logger.info("All agents initialized and ready")

    def load_vectorstore(self, path: Optional[str] = None):
        """Load a previously saved vector store and initialize agents."""
        self.ingester.load(path)
        self._retrieval_agent = RetrievalAgent(self.ingester.vectorstore, self.config)
        self._summarization_agent = SummarizationAgent(self.config)
        self._response_agent = ResponseAgent(self.config)
        logger.info("Agents initialized from saved vector store")

    def query(self, question: str) -> PipelineResult:
        """
        Run a query through the full agent pipeline.

        Flow: question → RetrievalAgent → SummarizationAgent → ResponseAgent → answer
        """
        if not all([self._retrieval_agent, self._summarization_agent, self._response_agent]):
            raise RuntimeError("Call ingest() or load_vectorstore() before querying")

        logger.info("=" * 60)
        logger.info("QUERY: %s", question)
        logger.info("=" * 60)

        # Step 1: Retrieval Agent
        logger.info("Step 1/3: RetrievalAgent searching...")
        chunks = self._retrieval_agent.retrieve(question)

        # Step 2: Summarization Agent (with shared context from Step 1)
        logger.info("Step 2/3: SummarizationAgent distilling...")
        summary = self._summarization_agent.summarize(question, chunks)

        # Step 3: Response Agent (with shared context from Steps 1 & 2)
        sources = list({c["source"] for c in chunks})
        logger.info("Step 3/3: ResponseAgent generating answer...")
        answer = self._response_agent.generate(question, summary, sources)

        result = PipelineResult(
            query=question,
            retrieved_chunks=chunks,
            summary=summary,
            answer=answer,
            sources=sources,
        )
        logger.info("Pipeline complete")
        return result