mantisdk 0.1.0__py3-none-any.whl

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

@@ -0,0 +1,238 @@
+# Copyright (c) 2025 Lakshya A Agrawal and the GEPA contributors
+# https://github.com/gepa-ai/gepa
+
+from typing import Any, Callable
+
+from mantisdk.algorithm.gepa.lib.adapters.generic_rag_adapter.vector_store_interface import VectorStoreInterface
+
+
+class RAGPipeline:
+    """
+    Generic RAG pipeline that works with any vector store.
+
+    This pipeline orchestrates the full RAG process:
+    1. Query reformulation (optional)
+    2. Document retrieval
+    3. Document reranking (optional)
+    4. Context synthesis
+    5. Answer generation
+    """
+
+    def __init__(
+        self,
+        vector_store: VectorStoreInterface,
+        llm_client,
+        embedding_model: str = "text-embedding-3-small",
+        embedding_function: Callable[[str], list[float]] | None = None,
+    ):
+        """
+        Initialize the RAG pipeline.
+
+        Args:
+            vector_store: Vector store interface implementation
+            llm_client: LLM client for generation (should have a callable interface)
+            embedding_model: Model name for embeddings (if using default embedding function)
+            embedding_function: Optional custom embedding function
+        """
+        self.vector_store = vector_store
+        self.llm_client = llm_client
+        self.embedding_model = embedding_model
+        self.embedding_function = embedding_function
+
+        # Initialize default embedding function if none provided
+        if self.embedding_function is None:
+            self.embedding_function = self._default_embedding_function
+
+    def execute_rag(
+        self,
+        query: str,
+        prompts: dict[str, str],
+        config: dict[str, Any],
+    ) -> dict[str, Any]:
+        """
+        Execute the full RAG pipeline with given prompts and configuration.
+
+        Args:
+            query: User query
+            prompts: Dictionary of prompt templates for different stages
+            config: Configuration parameters for retrieval and generation
+
+        Returns:
+            Dictionary containing all pipeline outputs and metadata
+        """
+        # Step 1: Query reformulation (if enabled)
+        reformulated_query = query
+        if "query_reformulation" in prompts and prompts["query_reformulation"].strip():
+            reformulated_query = self._reformulate_query(query, prompts["query_reformulation"])
+
+        # Step 2: Retrieval
+        retrieved_docs = self._retrieve_documents(reformulated_query, config)
+
+        # Step 3: Reranking (if enabled)
+        if "reranking_criteria" in prompts and prompts["reranking_criteria"].strip():
+            retrieved_docs = self._rerank_documents(retrieved_docs, query, prompts["reranking_criteria"], config)
+
+        # Step 4: Context synthesis
+        context = self._synthesize_context(retrieved_docs, query, prompts.get("context_synthesis", ""))
+
+        # Step 5: Answer generation
+        answer = self._generate_answer(query, context, prompts.get("answer_generation", ""))
+
+        return {
+            "original_query": query,
+            "reformulated_query": reformulated_query,
+            "retrieved_docs": retrieved_docs,
+            "synthesized_context": context,
+            "generated_answer": answer,
+            "metadata": {
+                "retrieval_count": len(retrieved_docs),
+                "total_tokens": self._estimate_token_count(context + answer),
+                "vector_store_type": self.vector_store.get_collection_info().get("vector_store_type", "unknown"),
+            },
+        }
+
+    def _reformulate_query(self, query: str, reformulation_prompt: str) -> str:
+        """Reformulate the user query using the provided prompt."""
+        messages = [
+            {"role": "system", "content": reformulation_prompt},
+            {"role": "user", "content": f"Original query: {query}"},
+        ]
+
+        try:
+            if callable(self.llm_client):
+                response = self.llm_client(messages)
+            else:
+                # Assume it's a litellm-style client
+                response = self.llm_client.completion(messages=messages).choices[0].message.content
+
+            return response.strip() if response else query
+        except Exception:
+            # Fallback to original query if reformulation fails
+            return query
+
+    def _retrieve_documents(self, query: str, config: dict[str, Any]) -> list[dict[str, Any]]:
+        """Retrieve documents using the configured strategy."""
+        retrieval_strategy = config.get("retrieval_strategy", "similarity")
+        k = config.get("top_k", 5)
+        filters = config.get("filters", None)
+
+        if retrieval_strategy == "similarity":
+            return self.vector_store.similarity_search(query, k=k, filters=filters)
+        elif retrieval_strategy == "hybrid":
+            if self.vector_store.supports_hybrid_search():
+                alpha = config.get("hybrid_alpha", 0.5)
+                return self.vector_store.hybrid_search(query, k=k, alpha=alpha, filters=filters)
+            else:
+                # Fallback to similarity search
+                return self.vector_store.similarity_search(query, k=k, filters=filters)
+        elif retrieval_strategy == "vector":
+            # Use pre-computed embedding
+            query_vector = self.embedding_function(query)
+            return self.vector_store.vector_search(query_vector, k=k, filters=filters)
+        else:
+            raise ValueError(f"Unknown retrieval strategy: {retrieval_strategy}")
+
+    def _rerank_documents(
+        self, documents: list[dict[str, Any]], query: str, reranking_prompt: str, config: dict[str, Any]
+    ) -> list[dict[str, Any]]:
+        """Rerank documents based on relevance criteria."""
+        if not documents:
+            return documents
+
+        # For simplicity, we'll use a prompt-based reranking approach
+        # In production, you might use dedicated reranking models
+        try:
+            doc_texts = [f"Document {i + 1}: {doc['content']}" for i, doc in enumerate(documents)]
+            doc_context = "\n\n".join(doc_texts)
+
+            messages = [
+                {"role": "system", "content": reranking_prompt},
+                {
+                    "role": "user",
+                    "content": f"Query: {query}\n\nDocuments:\n{doc_context}\n\nPlease rank these documents by relevance (return document numbers in order, e.g., '3,1,4,2,5'):",
+                },
+            ]
+
+            if callable(self.llm_client):
+                response = self.llm_client(messages)
+            else:
+                response = self.llm_client.completion(messages=messages).choices[0].message.content
+
+            # Parse the ranking response
+            ranking_str = response.strip()
+            rankings = [int(x.strip()) - 1 for x in ranking_str.split(",") if x.strip().isdigit()]
+
+            # Reorder documents based on ranking
+            if len(rankings) == len(documents):
+                return [documents[i] for i in rankings if 0 <= i < len(documents)]
+        except Exception:
+            pass
+
+        # Return original order if reranking fails
+        return documents
+
+    def _synthesize_context(self, documents: list[dict[str, Any]], query: str, synthesis_prompt: str) -> str:
+        """Synthesize retrieved documents into coherent context."""
+        if not documents:
+            return ""
+
+        if not synthesis_prompt.strip():
+            # Default: simple concatenation
+            contexts = []
+            for i, doc in enumerate(documents):
+                contexts.append(f"[Document {i + 1}] {doc['content']}")
+            return "\n\n".join(contexts)
+
+        # Use LLM for context synthesis
+        doc_texts = [doc["content"] for doc in documents]
+        doc_context = "\n\n".join(f"Document {i + 1}: {text}" for i, text in enumerate(doc_texts))
+
+        messages = [
+            {"role": "system", "content": synthesis_prompt},
+            {"role": "user", "content": f"Query: {query}\n\nRetrieved Documents:\n{doc_context}"},
+        ]
+
+        try:
+            if callable(self.llm_client):
+                response = self.llm_client(messages)
+            else:
+                response = self.llm_client.completion(messages=messages).choices[0].message.content
+
+            return response.strip() if response else doc_context
+        except Exception:
+            # Fallback to simple concatenation
+            return doc_context
+
+    def _generate_answer(self, query: str, context: str, generation_prompt: str) -> str:
+        """Generate the final answer using the query and synthesized context."""
+        if not generation_prompt.strip():
+            generation_prompt = "You are a helpful assistant. Answer the user's question based on the provided context."
+
+        messages = [
+            {"role": "system", "content": generation_prompt},
+            {"role": "user", "content": f"Context:\n{context}\n\nQuestion: {query}"},
+        ]
+
+        try:
+            if callable(self.llm_client):
+                response = self.llm_client(messages)
+            else:
+                response = self.llm_client.completion(messages=messages).choices[0].message.content
+
+            return response.strip() if response else "I couldn't generate an answer based on the provided context."
+        except Exception as e:
+            return f"Error generating answer: {e!s}"
+
+    def _default_embedding_function(self, text: str) -> list[float]:
+        """Default embedding function using litellm."""
+        try:
+            import litellm
+
+            response = litellm.embedding(model=self.embedding_model, input=text)
+            return response.data[0].embedding
+        except Exception as e:
+            raise RuntimeError(f"Failed to generate embeddings: {e!s}") from e
+
+    def _estimate_token_count(self, text: str) -> int:
+        """Rough estimate of token count (4 chars per token approximation)."""
+        return len(text) // 4

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

@@ -0,0 +1,212 @@
+# Copyright (c) 2025 Lakshya A Agrawal and the GEPA contributors
+# https://github.com/gepa-ai/gepa
+
+from abc import ABC, abstractmethod
+from typing import Any
+
+
+class VectorStoreInterface(ABC):
+    """
+    Abstract interface for vector store operations in RAG systems.
+
+    This interface defines the core operations needed for retrieval-augmented generation,
+    enabling GEPA to work with any vector store implementation (ChromaDB, Weaviate, Qdrant,
+    Pinecone, Milvus, etc.) through a unified API.
+
+    The interface supports:
+    - Semantic similarity search
+    - Vector-based search with pre-computed embeddings
+    - Hybrid search (semantic + keyword, where supported)
+    - Metadata filtering and collection introspection
+
+    Implementing this interface allows your vector store to be used with GEPA's
+    evolutionary prompt optimization for RAG systems.
+
+    Example:
+        .. code-block:: python
+
+            class MyVectorStore(VectorStoreInterface):
+                def similarity_search(self, query, k=5, filters=None):
+                    # Your implementation
+                    return documents
+
+            vector_store = MyVectorStore()
+            adapter = GenericRAGAdapter(vector_store=vector_store)
+    """
+
+    @abstractmethod
+    def similarity_search(
+        self,
+        query: str,
+        k: int = 5,
+        filters: dict[str, Any] | None = None,
+    ) -> list[dict[str, Any]]:
+        """
+        Search for documents semantically similar to the query text.
+
+        This method performs semantic similarity search using the vector store's
+        default embedding model or configured vectorizer.
+
+        Args:
+            query: Text query to search for similar documents
+            k: Maximum number of documents to return (default: 5)
+            filters: Optional metadata filters to constrain search.
+                Format: {"key": "value"} or {"key": {"$op": value}}
+
+        Returns:
+            List of documents ordered by similarity score (highest first).
+            Each document is a dictionary with:
+            - "content" (str): The document text content
+            - "metadata" (dict): Document metadata including any doc_id
+            - "score" (float): Similarity score between 0.0 and 1.0 (higher = more similar)
+
+        Raises:
+            NotImplementedError: Must be implemented by concrete vector store classes
+
+        Example:
+            .. code-block:: python
+
+                results = vector_store.similarity_search(
+                    query="machine learning algorithms",
+                    k=3,
+                    filters={"category": "AI"}
+                )
+                print(results[0]["content"])  # Most similar document
+        """
+        pass
+
+    @abstractmethod
+    def vector_search(
+        self,
+        query_vector: list[float],
+        k: int = 5,
+        filters: dict[str, Any] | None = None,
+    ) -> list[dict[str, Any]]:
+        """
+        Search using a pre-computed query embedding vector.
+
+        This method allows direct vector similarity search when you already have
+        the query embedding, avoiding the need for additional embedding computation.
+
+        Args:
+            query_vector: Pre-computed embedding vector for the query.
+                Must match the dimensionality of vectors in the collection.
+            k: Maximum number of documents to return (default: 5)
+            filters: Optional metadata filters to constrain search
+
+        Returns:
+            List of documents ordered by vector similarity (highest first).
+            Same format as similarity_search().
+
+        Raises:
+            NotImplementedError: Must be implemented by concrete vector store classes
+            ValueError: If query_vector dimensions don't match collection
+
+        Example:
+            .. code-block:: python
+
+                import numpy as np
+                query_vector = embedding_model.encode("machine learning")
+                results = vector_store.vector_search(query_vector.tolist(), k=5)
+        """
+        pass
+
+    def hybrid_search(
+        self,
+        query: str,
+        k: int = 5,
+        alpha: float = 0.5,
+        filters: dict[str, Any] | None = None,
+    ) -> list[dict[str, Any]]:
+        """
+        Hybrid semantic + keyword search combining vector and text-based matching.
+
+        This method combines semantic similarity (vector search) with keyword-based
+        search (like BM25) to leverage both approaches. The alpha parameter controls
+        the balance between the two search methods.
+
+        Args:
+            query: Text query to search for
+            k: Maximum number of documents to return (default: 5)
+            alpha: Weight for semantic vs keyword search (default: 0.5)
+                - 0.0 = pure keyword/BM25 search
+                - 1.0 = pure semantic/vector search
+                - 0.5 = balanced hybrid search
+            filters: Optional metadata filters to constrain search
+
+        Returns:
+            List of documents ordered by hybrid similarity score (highest first).
+            Same format as similarity_search().
+
+        Note:
+            If hybrid search is not supported by the vector store implementation,
+            this method falls back to similarity_search() with a warning.
+            Use supports_hybrid_search() to check availability.
+
+        Example:
+            .. code-block:: python
+
+                # Balanced hybrid search
+                results = vector_store.hybrid_search("AI algorithms", alpha=0.5)
+
+                # More semantic-focused
+                results = vector_store.hybrid_search("AI algorithms", alpha=0.8)
+        """
+        # Default fallback implementation
+        return self.similarity_search(query, k, filters)
+
+    @abstractmethod
+    def get_collection_info(self) -> dict[str, Any]:
+        """
+        Get metadata and statistics about the vector store collection.
+
+        This method provides introspection capabilities for the collection,
+        returning key information about its configuration and contents.
+
+        Returns:
+            Dictionary containing collection metadata with keys:
+            - "name" (str): Collection/index name
+            - "document_count" (int): Total number of documents
+            - "dimension" (int): Vector embedding dimension (0 if unknown)
+            - "vector_store_type" (str): Type of vector store (e.g., "chromadb", "weaviate")
+            - Additional store-specific metadata as available
+
+        Raises:
+            NotImplementedError: Must be implemented by concrete vector store classes
+
+        Example:
+            .. code-block:: python
+
+                info = vector_store.get_collection_info()
+                print(f"Collection {info['name']} has {info['document_count']} documents")
+                print(f"Vector dimension: {info['dimension']}")
+        """
+        pass
+
+    def get_embedding_dimension(self) -> int:
+        """
+        Get the embedding dimension of the vector store.
+
+        Returns:
+            Dimension of the vectors in the collection
+        """
+        info = self.get_collection_info()
+        return info.get("dimension", 0)
+
+    def supports_hybrid_search(self) -> bool:
+        """
+        Check if the vector store supports hybrid search.
+
+        Returns:
+            True if hybrid search is supported, False otherwise
+        """
+        return False
+
+    def supports_metadata_filtering(self) -> bool:
+        """
+        Check if the vector store supports metadata filtering.
+
+        Returns:
+            True if metadata filtering is supported, False otherwise
+        """
+        return True

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

@@ -0,0 +1,2 @@
+# Copyright (c) 2025 Lakshya A Agrawal and the GEPA contributors
+# https://github.com/gepa-ai/gepa

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

@@ -0,0 +1,196 @@
+# Copyright (c) 2025 Lakshya A Agrawal and the GEPA contributors
+# https://github.com/gepa-ai/gepa
+
+from typing import Any
+
+from mantisdk.algorithm.gepa.lib.adapters.generic_rag_adapter.vector_store_interface import VectorStoreInterface
+
+
+class ChromaVectorStore(VectorStoreInterface):
+    """
+    ChromaDB implementation of the VectorStoreInterface.
+
+    ChromaDB is an open-source embedding database that's easy to use
+    and perfect for local development and prototyping.
+    """
+
+    def __init__(self, client, collection_name: str, embedding_function=None):
+        """
+        Initialize ChromaVectorStore.
+
+        Args:
+            client: ChromaDB client instance
+            collection_name: Name of the collection to use
+            embedding_function: Optional embedding function for text queries
+        """
+        import importlib.util
+
+        if importlib.util.find_spec("chromadb") is None:
+            raise ImportError("ChromaDB is required for ChromaVectorStore. Install with: pip install litellm chromadb")
+
+        self.client = client
+        self.collection_name = collection_name
+        self.embedding_function = embedding_function
+
+        # Get or create the collection
+        try:
+            self.collection = self.client.get_collection(name=collection_name, embedding_function=embedding_function)
+        except Exception:
+            # Collection doesn't exist, create it
+            self.collection = self.client.create_collection(name=collection_name, embedding_function=embedding_function)
+
+    def similarity_search(
+        self,
+        query: str,
+        k: int = 5,
+        filters: dict[str, Any] | None = None,
+    ) -> list[dict[str, Any]]:
+        """Search for documents similar to the query text."""
+        # Convert filters to ChromaDB format if provided
+        where_clause = self._convert_filters(filters) if filters else None
+
+        results = self.collection.query(
+            query_texts=[query], n_results=k, where=where_clause, include=["documents", "metadatas", "distances"]
+        )
+
+        return self._format_results(results)
+
+    def vector_search(
+        self,
+        query_vector: list[float],
+        k: int = 5,
+        filters: dict[str, Any] | None = None,
+    ) -> list[dict[str, Any]]:
+        """Search using a pre-computed query vector."""
+        where_clause = self._convert_filters(filters) if filters else None
+
+        results = self.collection.query(
+            query_embeddings=[query_vector],
+            n_results=k,
+            where=where_clause,
+            include=["documents", "metadatas", "distances"],
+        )
+
+        return self._format_results(results)
+
+    def get_collection_info(self) -> dict[str, Any]:
+        """Get metadata about the ChromaDB collection."""
+        count = self.collection.count()
+
+        # Try to get a sample document to determine embedding dimension
+        dimension = 0
+        if count > 0:
+            sample = self.collection.peek(limit=1)
+            embeddings = sample.get("embeddings")
+            if embeddings is not None and len(embeddings) > 0:
+                dimension = len(embeddings[0])
+
+        return {
+            "name": self.collection_name,
+            "document_count": count,
+            "dimension": dimension,
+            "vector_store_type": "chromadb",
+        }
+
+    def supports_metadata_filtering(self) -> bool:
+        """ChromaDB supports metadata filtering."""
+        return True
+
+    def _convert_filters(self, filters: dict[str, Any]) -> dict[str, Any]:
+        """
+        Convert generic filters to ChromaDB where clause format.
+
+        Generic format: {"key": "value", "key2": {"$gt": 5}}
+        ChromaDB format: {"key": {"$eq": "value"}, "key2": {"$gt": 5}}
+        """
+        chroma_filters = {}
+
+        for key, value in filters.items():
+            if isinstance(value, dict):
+                # Already in operator format
+                chroma_filters[key] = value
+            else:
+                # Convert to equality operator
+                chroma_filters[key] = {"$eq": value}
+
+        return chroma_filters
+
+    def _format_results(self, results) -> list[dict[str, Any]]:
+        """Convert ChromaDB results to standardized format."""
+        documents = []
+
+        if not results["documents"] or not results["documents"][0]:
+            return documents
+
+        docs = results["documents"][0]
+        metadatas = results.get("metadatas", [None] * len(docs))[0] or [{}] * len(docs)
+        distances = results.get("distances", [0.0] * len(docs))[0]
+
+        for doc, metadata, distance in zip(docs, metadatas, distances, strict=False):
+            # Convert distance to similarity score (higher is better)
+            # ChromaDB uses cosine distance, so similarity = 1 - distance
+            distance_val = self._extract_distance_value(distance)
+            similarity_score = max(0.0, 1.0 - distance_val)
+
+            documents.append({"content": doc, "metadata": metadata or {}, "score": similarity_score})
+
+        return documents
+
+    def _extract_distance_value(self, distance) -> float:
+        """
+        Helper to extract a scalar float from a distance value returned by ChromaDB.
+        Handles numpy scalars, single-element lists/arrays, and plain floats.
+        """
+        try:
+            if hasattr(distance, "item"):  # numpy scalar
+                return distance.item()
+            elif hasattr(distance, "__len__") and not isinstance(distance, str | bytes) and len(distance) == 1:
+                return float(distance[0])
+            else:
+                return float(distance)
+        except (TypeError, ValueError, IndexError) as e:
+            import logging
+
+            logging.warning(f"Unexpected distance format: {type(distance)}, value: {distance}, error: {e}")
+            return 0.0  # Default fallback
+
+    @classmethod
+    def create_local(cls, persist_directory: str, collection_name: str, embedding_function=None) -> "ChromaVectorStore":
+        """
+        Create a ChromaVectorStore with local persistence.
+
+        Args:
+            persist_directory: Directory to persist the database
+            collection_name: Name of the collection
+            embedding_function: Optional embedding function
+
+        Returns:
+            ChromaVectorStore instance
+        """
+        try:
+            import chromadb
+        except ImportError as e:
+            raise ImportError("ChromaDB is required. Install with: pip install litellm chromadb") from e
+
+        client = chromadb.PersistentClient(path=persist_directory)
+        return cls(client, collection_name, embedding_function)
+
+    @classmethod
+    def create_memory(cls, collection_name: str, embedding_function=None) -> "ChromaVectorStore":
+        """
+        Create a ChromaVectorStore in memory (for testing).
+
+        Args:
+            collection_name: Name of the collection
+            embedding_function: Optional embedding function
+
+        Returns:
+            ChromaVectorStore instance
+        """
+        try:
+            import chromadb
+        except ImportError as e:
+            raise ImportError("ChromaDB is required. Install with: pip install litellm chromadb") from e
+
+        client = chromadb.Client()
+        return cls(client, collection_name, embedding_function)