mantisdk 0.1.0__py3-none-any.whl

mantisdk/algorithm/gepa/lib/adapters/generic_rag_adapter/generic_rag_adapter.py

@@ -0,0 +1,496 @@
+# Copyright (c) 2025 Lakshya A Agrawal and the GEPA contributors
+# https://github.com/gepa-ai/gepa
+
+from typing import Any, TypedDict
+
+from mantisdk.algorithm.gepa.lib.adapters.generic_rag_adapter.evaluation_metrics import RAGEvaluationMetrics
+from mantisdk.algorithm.gepa.lib.adapters.generic_rag_adapter.rag_pipeline import RAGPipeline
+from mantisdk.algorithm.gepa.lib.adapters.generic_rag_adapter.vector_store_interface import VectorStoreInterface
+from mantisdk.algorithm.gepa.lib.core.adapter import EvaluationBatch, GEPAAdapter
+
+
+class RAGDataInst(TypedDict):
+    """
+    Data instance for RAG evaluation and optimization.
+
+    This TypedDict defines the structure for training and validation examples
+    used in RAG system optimization with GEPA.
+
+    Attributes:
+            query (str): User query or question to be answered
+        ground_truth_answer (str): Expected/correct answer for evaluation
+            relevant_doc_ids (List[str]): List of document IDs that should ideally
+            be retrieved for this query (used for retrieval evaluation)
+            metadata (Dict[str, Any]): Additional context, tags, or configuration
+            specific to this example (e.g., difficulty level, category)
+
+    Example:
+            .. code-block:: python
+
+            data_inst = RAGDataInst(
+                query="What is machine learning?",
+                ground_truth_answer="Machine learning is a subset of AI...",
+                relevant_doc_ids=["doc_001", "doc_042"],
+                metadata={"category": "AI", "difficulty": "beginner"}
+            )
+    """
+
+    query: str
+    ground_truth_answer: str
+    relevant_doc_ids: list[str]
+    metadata: dict[str, Any]
+
+
+class RAGTrajectory(TypedDict):
+    """
+    Detailed trajectory capturing all RAG pipeline execution steps.
+
+    This TypedDict captures the complete execution trace of the RAG pipeline,
+    providing visibility into each step for analysis and optimization.
+
+    Attributes:
+            original_query (str): Original user query as provided
+        reformulated_query (str): Query after reformulation step (if enabled)
+            retrieved_docs (List[Dict[str, Any]]): Documents retrieved from vector store
+            with their content, metadata, and similarity scores
+            synthesized_context (str): Context after document synthesis step
+            generated_answer (str): Final answer generated by the LLM
+            execution_metadata (Dict[str, Any]): Pipeline execution metadata including
+            retrieval metrics, generation metrics, token counts, and performance data
+
+    Note:
+            Trajectories are only captured when capture_traces=True is passed to
+        the evaluate() method, as they can be memory-intensive for large batches.
+    """
+
+    original_query: str
+    reformulated_query: str
+    retrieved_docs: list[dict[str, Any]]
+    synthesized_context: str
+    generated_answer: str
+    execution_metadata: dict[str, Any]
+
+
+class RAGOutput(TypedDict):
+    """
+    Final output from RAG system execution.
+
+    This TypedDict represents the final result of RAG pipeline execution,
+    containing both the generated answer and associated metadata.
+
+    Attributes:
+            final_answer (str): The generated answer from the RAG system
+            confidence_score (float): Estimated confidence in the answer (0.0 to 1.0)
+            based on retrieval quality and generation metrics
+            retrieved_docs (List[Dict[str, Any]]): Documents that were retrieved
+            and used for answer generation
+        total_tokens (int): Estimated total token usage for the pipeline execution
+
+    Example:
+            .. code-block:: python
+
+            output = RAGOutput(
+                final_answer="Machine learning is a method of data analysis...",
+                confidence_score=0.87,
+                retrieved_docs=[{"content": "...", "score": 0.9}],
+                total_tokens=450
+            )
+    """
+
+    final_answer: str
+    confidence_score: float
+    retrieved_docs: list[dict[str, Any]]
+    total_tokens: int
+
+
+class GenericRAGAdapter(GEPAAdapter[RAGDataInst, RAGTrajectory, RAGOutput]):
+    """
+    Generic GEPA adapter for RAG system optimization with pluggable vector stores.
+
+    This adapter enables GEPA's evolutionary prompt optimization to work with any
+    vector store implementation through the VectorStoreInterface. It provides
+    comprehensive evaluation of both retrieval and generation quality.
+
+    Optimizable Components:
+            - Query reformulation prompts: Improve query understanding and reformulation
+            - Context synthesis prompts: Optimize document combination and summarization
+            - Answer generation prompts: Enhance final answer quality and formatting
+            - Reranking criteria: Improve document relevance ordering
+
+    Evaluation Metrics:
+            - Retrieval Quality: Precision, recall, F1, mean reciprocal rank (MRR)
+            - Generation Quality: Token F1, BLEU score, faithfulness, answer relevance
+        - Combined Score: Weighted combination for overall system performance
+
+    Vector Store Support:
+            Works with any vector store implementing VectorStoreInterface, including:
+            ChromaDB, Weaviate, Qdrant, Pinecone, Milvus, and custom implementations.
+
+    Example:
+            .. code-block:: python
+
+            from mantisdk.algorithm.gepa.lib.adapters.generic_rag_adapter import GenericRAGAdapter, ChromaVectorStore
+            import gepa
+
+            vector_store = ChromaVectorStore.create_local("./kb", "docs")
+            adapter = GenericRAGAdapter(vector_store=vector_store, llm_model="gpt-4")
+
+            result = gepa.optimize(
+                seed_candidate={"answer_generation": "Answer based on context:"},
+                trainset=train_data,
+                valset=val_data,
+                adapter=adapter,
+                max_metric_calls=50
+            )
+            print(result.best_candidate)  # Optimized prompts
+    """
+
+    def __init__(
+        self,
+        vector_store: VectorStoreInterface,
+        llm_model,
+        embedding_model: str = "text-embedding-3-small",
+        embedding_function=None,
+        rag_config: dict[str, Any] | None = None,
+        failure_score: float = 0.0,
+    ):
+        """
+        Initialize the GenericRAGAdapter for RAG system optimization.
+
+        Args:
+            vector_store: Vector store implementation (ChromaDB, Weaviate, etc.)
+                Must implement VectorStoreInterface for similarity search operations.
+            llm_model: LLM client for text generation. Can be:
+                - String model name (uses litellm for inference)
+                - Callable that takes messages and returns response text
+                - Any object with a callable interface for LLM inference
+            embedding_model: Model name for text embeddings (default: "text-embedding-3-small").
+                Used when embedding_function is not provided.
+            embedding_function: Optional custom embedding function that takes text
+                and returns List[float]. If None, uses default litellm embeddings.
+            rag_config: RAG pipeline configuration dictionary. Keys include:
+                - "retrieval_strategy": "similarity", "hybrid", or "vector"
+                - "top_k": Number of documents to retrieve (default: 5)
+                - "retrieval_weight": Weight for retrieval in combined score (default: 0.3)
+                - "generation_weight": Weight for generation in combined score (default: 0.7)
+                - "hybrid_alpha": Semantic vs keyword balance for hybrid search (default: 0.5)
+                - "filters": Default metadata filters for retrieval
+            failure_score: Score assigned when evaluation fails (default: 0.0)
+
+        Example:
+            .. code-block:: python
+
+                vector_store = WeaviateVectorStore.create_local(collection_name="docs")
+                adapter = GenericRAGAdapter(
+                    vector_store=vector_store,
+                    llm_model="gpt-4",
+                    rag_config={
+                        "retrieval_strategy": "hybrid",
+                        "top_k": 5,
+                        "hybrid_alpha": 0.7
+                    }
+                )
+        """
+        self.vector_store = vector_store
+        self.rag_pipeline = RAGPipeline(
+            vector_store=vector_store,
+            llm_client=llm_model,
+            embedding_model=embedding_model,
+            embedding_function=embedding_function,
+        )
+        self.evaluator = RAGEvaluationMetrics()
+        self.config = rag_config or self._default_config()
+        self.failure_score = failure_score
+
+    def evaluate(
+        self,
+        batch: list[RAGDataInst],
+        candidate: dict[str, str],
+        capture_traces: bool = False,
+    ) -> EvaluationBatch[RAGTrajectory, RAGOutput]:
+        """
+        Evaluate RAG system performance on a batch of query-answer examples.
+
+        This method runs the complete RAG pipeline on each example in the batch,
+        evaluating both retrieval and generation quality using the provided
+        prompt components.
+
+        Args:
+            batch: List of RAG evaluation examples, each containing:
+                - query: Question to answer
+                - ground_truth_answer: Expected correct answer
+                - relevant_doc_ids: Documents that should be retrieved
+                - metadata: Additional context for evaluation
+            candidate: Dictionary mapping prompt component names to their text.
+                Supported components:
+                - "query_reformulation": Prompt for improving user queries
+                - "context_synthesis": Prompt for combining retrieved documents
+                - "answer_generation": Prompt for generating final answers
+                - "reranking_criteria": Criteria for reordering retrieved documents
+            capture_traces: If True, capture detailed execution trajectories
+                for each example. Required for reflective dataset generation but
+                increases memory usage.
+
+        Returns:
+            EvaluationBatch containing:
+            - outputs: List of RAGOutput for each example
+            - scores: List of combined quality scores (higher = better)
+            - trajectories: List of detailed execution traces (if capture_traces=True)
+
+        Raises:
+            Exception: Individual example failures are caught and assigned failure_score.
+                Only systemic failures (e.g., vector store unavailable) raise exceptions.
+
+        Example:
+            .. code-block:: python
+
+                prompts = {
+                    "answer_generation": "Answer the question based on this context:"
+                }
+                result = adapter.evaluate(
+                    batch=validation_data,
+                    candidate=prompts,
+                    capture_traces=True
+                )
+                avg_score = sum(result.scores) / len(result.scores)
+                print(f"Average RAG performance: {avg_score:.3f}")
+        """
+        outputs: list[RAGOutput] = []
+        scores: list[float] = []
+        trajectories: list[RAGTrajectory] | None = [] if capture_traces else None
+
+        for data_inst in batch:
+            try:
+                # Execute RAG pipeline with candidate prompts
+                rag_result = self.rag_pipeline.execute_rag(
+                    query=data_inst["query"], prompts=candidate, config=self.config
+                )
+
+                # Evaluate retrieval quality
+                retrieval_metrics = self.evaluator.evaluate_retrieval(
+                    rag_result["retrieved_docs"], data_inst["relevant_doc_ids"]
+                )
+
+                # Evaluate generation quality
+                generation_metrics = self.evaluator.evaluate_generation(
+                    rag_result["generated_answer"], data_inst["ground_truth_answer"], rag_result["synthesized_context"]
+                )
+
+                # Calculate combined score
+                overall_score = self.evaluator.combined_rag_score(
+                    retrieval_metrics,
+                    generation_metrics,
+                    retrieval_weight=self.config.get("retrieval_weight", 0.3),
+                    generation_weight=self.config.get("generation_weight", 0.7),
+                )
+
+                # Prepare output
+                output = RAGOutput(
+                    final_answer=rag_result["generated_answer"],
+                    confidence_score=generation_metrics.get("answer_confidence", 0.5),
+                    retrieved_docs=rag_result["retrieved_docs"],
+                    total_tokens=rag_result["metadata"]["total_tokens"],
+                )
+
+                outputs.append(output)
+                scores.append(overall_score)
+
+                # Capture trajectory if requested
+                if capture_traces:
+                    trajectory = RAGTrajectory(
+                        original_query=rag_result["original_query"],
+                        reformulated_query=rag_result["reformulated_query"],
+                        retrieved_docs=rag_result["retrieved_docs"],
+                        synthesized_context=rag_result["synthesized_context"],
+                        generated_answer=rag_result["generated_answer"],
+                        execution_metadata={
+                            **rag_result["metadata"],
+                            "retrieval_metrics": retrieval_metrics,
+                            "generation_metrics": generation_metrics,
+                            "overall_score": overall_score,
+                        },
+                    )
+                    trajectories.append(trajectory)
+
+            except Exception as e:
+                # Handle individual example failure
+                error_output = RAGOutput(
+                    final_answer=f"Error: {e!s}", confidence_score=0.0, retrieved_docs=[], total_tokens=0
+                )
+
+                outputs.append(error_output)
+                scores.append(self.failure_score)
+
+                if capture_traces:
+                    error_trajectory = RAGTrajectory(
+                        original_query=data_inst["query"],
+                        reformulated_query=data_inst["query"],
+                        retrieved_docs=[],
+                        synthesized_context="",
+                        generated_answer=f"Error: {e!s}",
+                        execution_metadata={"error": str(e)},
+                    )
+                    trajectories.append(error_trajectory)
+
+        return EvaluationBatch(outputs=outputs, scores=scores, trajectories=trajectories)
+
+    def make_reflective_dataset(
+        self,
+        candidate: dict[str, str],
+        eval_batch: EvaluationBatch[RAGTrajectory, RAGOutput],
+        components_to_update: list[str],
+    ) -> dict[str, list[dict[str, Any]]]:
+        """
+        Generate reflective dataset for evolutionary prompt optimization.
+
+        This method analyzes the evaluation results and creates training examples
+        that GEPA's proposer can use to improve the specified prompt components.
+        Each component gets a tailored dataset with input-output pairs and feedback.
+
+        Args:
+            candidate: Current prompt components that were evaluated
+            eval_batch: Evaluation results from evaluate() with capture_traces=True.
+                Must contain trajectories for analysis.
+            components_to_update: List of component names to generate improvement
+                suggestions for. Must be subset of candidate.keys().
+
+        Returns:
+            Dictionary mapping component names to their reflective datasets.
+            Each dataset is a list of examples with structure:
+            - "Inputs": Input data for the component (query, docs, etc.)
+            - "Generated Outputs": What the component currently produces
+            - "Feedback": Analysis of performance and suggestions for improvement
+
+        Example:
+            .. code-block:: python
+
+                reflective_data = adapter.make_reflective_dataset(
+                    candidate=current_prompts,
+                    eval_batch=evaluation_results,  # with trajectories
+                    components_to_update=["answer_generation", "context_synthesis"]
+                )
+                print(reflective_data["answer_generation"][0]["Feedback"])
+                # Output: "The generated answer lacks specific details from the context..."
+
+        Note:
+            This method requires eval_batch to have been created with
+            capture_traces=True, otherwise trajectories will be None.
+        """
+        reflective_data: dict[str, list[dict[str, Any]]] = {}
+
+        for component in components_to_update:
+            component_examples = []
+
+            # Process each trajectory to create examples for this component
+            for traj, output, score in zip(
+                eval_batch.trajectories or [], eval_batch.outputs, eval_batch.scores, strict=False
+            ):
+                example = self._create_component_example(component, traj, output, score, candidate)
+                if example:
+                    component_examples.append(example)
+
+            # Only include components that have examples
+            if component_examples:
+                reflective_data[component] = component_examples
+
+        return reflective_data
+
+    def _create_component_example(
+        self, component_name: str, trajectory: RAGTrajectory, output: RAGOutput, score: float, candidate: dict[str, str]
+    ) -> dict[str, Any] | None:
+        """Create a reflective example for a specific component."""
+
+        if component_name == "query_reformulation":
+            return {
+                "Inputs": {
+                    "original_query": trajectory["original_query"],
+                    "current_prompt": candidate.get(component_name, ""),
+                },
+                "Generated Outputs": trajectory["reformulated_query"],
+                "Feedback": self._generate_query_reformulation_feedback(trajectory, score),
+            }
+
+        elif component_name == "context_synthesis":
+            return {
+                "Inputs": {
+                    "query": trajectory["original_query"],
+                    "retrieved_docs": [doc["content"] for doc in trajectory["retrieved_docs"]],
+                    "current_prompt": candidate.get(component_name, ""),
+                },
+                "Generated Outputs": trajectory["synthesized_context"],
+                "Feedback": self._generate_context_synthesis_feedback(trajectory, score),
+            }
+
+        elif component_name == "answer_generation":
+            return {
+                "Inputs": {
+                    "query": trajectory["original_query"],
+                    "context": trajectory["synthesized_context"],
+                    "current_prompt": candidate.get(component_name, ""),
+                },
+                "Generated Outputs": trajectory["generated_answer"],
+                "Feedback": self._generate_answer_generation_feedback(trajectory, output, score),
+            }
+
+        elif component_name == "reranking_criteria":
+            return {
+                "Inputs": {
+                    "query": trajectory["original_query"],
+                    "documents": [doc["content"] for doc in trajectory["retrieved_docs"]],
+                    "current_criteria": candidate.get(component_name, ""),
+                },
+                "Generated Outputs": "Document ranking applied",
+                "Feedback": self._generate_reranking_feedback(trajectory, score),
+            }
+
+            return None
+
+    def _generate_query_reformulation_feedback(self, trajectory: RAGTrajectory, score: float) -> str:
+        """Generate feedback for query reformulation component."""
+        if score > 0.7:
+            return f"Good query reformulation. The reformulated query '{trajectory['reformulated_query']}' helped retrieve relevant documents and generated a good answer."
+        else:
+            return f"The query reformulation from '{trajectory['original_query']}' to '{trajectory['reformulated_query']}' may not have improved retrieval. Consider making the reformulated query more specific or preserving key terms."
+
+    def _generate_context_synthesis_feedback(self, trajectory: RAGTrajectory, score: float) -> str:
+        """Generate feedback for context synthesis component."""
+        if score > 0.7:
+            return "Context synthesis worked well - the synthesized context effectively supported answer generation."
+        else:
+            return "Context synthesis could be improved. The synthesized context may not have highlighted the most relevant information or may have been too verbose/concise."
+
+    def _generate_answer_generation_feedback(self, trajectory: RAGTrajectory, output: RAGOutput, score: float) -> str:
+        """Generate feedback for answer generation component."""
+        if score > 0.7:
+            return f"Good answer generation. The generated answer '{trajectory['generated_answer']}' was accurate and well-supported by the context."
+        else:
+            return f"Answer generation needs improvement. The generated answer '{trajectory['generated_answer']}' may not be fully accurate or well-supported by the provided context."
+
+    def _generate_reranking_feedback(self, trajectory: RAGTrajectory, score: float) -> str:
+        """Generate feedback for reranking criteria component."""
+        if score > 0.7:
+            return "Document reranking appears to have helped surface more relevant documents for answer generation."
+        else:
+            return "Document reranking may not have improved relevance. Consider adjusting the criteria to better prioritize documents that contain the answer."
+
+    def _default_config(self) -> dict[str, Any]:
+        """
+        Get default configuration for RAG pipeline.
+
+        Returns:
+            Dictionary with default RAG configuration parameters:
+            - retrieval_strategy: "similarity" (semantic search)
+            - top_k: 5 (number of documents to retrieve)
+            - retrieval_weight: 0.3 (30% weight for retrieval metrics)
+            - generation_weight: 0.7 (70% weight for generation metrics)
+            - hybrid_alpha: 0.5 (balanced semantic/keyword for hybrid search)
+            - filters: None (no metadata filtering by default)
+        """
+        return {
+            "retrieval_strategy": "similarity",
+            "top_k": 5,
+            "retrieval_weight": 0.3,
+            "generation_weight": 0.7,
+            "hybrid_alpha": 0.5,
+            "filters": None,
+        }