tessl/pypi-dspy

A declarative framework for programming foundation models through compositional Python code, enabling modular AI systems with automated optimization algorithms that synthesize examples, generate instructions, and finetune weights based on user-defined metrics.

Overview

Eval results

Files

Retrieval

Name: tessl/pypi-dspy
Author: tessl

Retrieval components for fetching relevant passages from document stores and computing text embeddings for similarity search. Essential for building RAG (Retrieval-Augmented Generation) systems.

Capabilities

Retrieve Module

Retrieval module that queries a retrieval model to fetch relevant passages.

class Retrieve:
    """
    Retrieval module for fetching passages.

    Queries the configured retrieval model (rm) to fetch relevant
    passages based on input queries. Core component for RAG systems.
    """

    def __init__(self, k: int = 3, callbacks: list = None):
        """
        Initialize retrieval module.

        Args:
            k (int): Number of passages to retrieve (default: 3)
            callbacks (list | None): Optional callbacks for monitoring
        """
        pass

    def forward(self, query: str, k: int = None, **kwargs):
        """
        Retrieve passages for query.

        Args:
            query (str): Search query
            k (int | None): Override number of passages to retrieve
            **kwargs: Additional retrieval parameters

        Returns:
            Prediction with 'passages' field containing list of strings
        """
        pass

    def __call__(self, query: str, k: int = None, **kwargs):
        """
        Execute retrieval (same as forward).

        Args:
            query (str): Search query
            k (int | None): Number of passages
            **kwargs: Additional parameters

        Returns:
            Prediction with passages
        """
        pass

    def dump_state(self) -> dict:
        """
        Serialize state.

        Returns:
            Dictionary representation of state
        """
        pass

    def load_state(self, state: dict):
        """
        Load serialized state.

        Args:
            state (dict): State dictionary
        """
        pass

Usage:

import dspy

# Configure with retrieval model
colbert = dspy.ColBERTv2(url="http://localhost", port=8893)
dspy.configure(
    lm=dspy.LM('openai/gpt-4o-mini'),
    rm=colbert
)

# Basic retrieval
retrieve = dspy.Retrieve(k=5)
result = retrieve(query="What is machine learning?")
print(len(result.passages))  # 5
for passage in result.passages:
    print(passage)

# Override k per call
result = retrieve(query="What is deep learning?", k=10)
print(len(result.passages))  # 10

# Use in RAG system
class RAG(dspy.Module):
    def __init__(self, k=3):
        super().__init__()
        self.retrieve = dspy.Retrieve(k=k)
        self.generate = dspy.ChainOfThought("context, question -> answer")

    def forward(self, question):
        context = self.retrieve(query=question).passages
        return self.generate(context=context, question=question)

rag = RAG(k=5)
result = rag(question="What is artificial intelligence?")
print(result.answer)

Embeddings

Embedding-based retrieval using vector similarity.

class Embeddings:
    """
    Embedding-based retrieval.

    Uses text embeddings and vector similarity for document retrieval.
    Alternative to ColBERT for simpler embedding-based search.
    """

    def __init__(
        self,
        embedder,
        documents: list = None,
        k: int = 3
    ):
        """
        Initialize embeddings retrieval.

        Args:
            embedder: Embedder instance for computing embeddings
            documents (list | None): List of documents to search
            k (int): Default number of results (default: 3)
        """
        pass

    def __call__(self, query: str, k: int = None):
        """
        Retrieve documents by embedding similarity.

        Args:
            query (str): Search query
            k (int | None): Number of documents to return

        Returns:
            List of top-k most similar documents
        """
        pass

Usage:

import dspy
from sentence_transformers import SentenceTransformer

# Set up embedder
model = SentenceTransformer("all-MiniLM-L6-v2")
embedder = dspy.Embedder(model.encode)

# Prepare documents
documents = [
    "Machine learning is a subset of AI.",
    "Deep learning uses neural networks.",
    "Natural language processing analyzes text.",
    # ... more documents
]

# Create embeddings retriever
embeddings = dspy.Embeddings(
    embedder=embedder,
    documents=documents,
    k=3
)

# Retrieve similar documents
results = embeddings(query="What is ML?", k=5)
for doc in results:
    print(doc)

# Use as retrieval model
dspy.configure(
    lm=dspy.LM('openai/gpt-4o-mini'),
    rm=embeddings
)

# Now Retrieve module uses embeddings
retrieve = dspy.Retrieve(k=3)
result = retrieve(query="neural networks")

ColBERTv2 Retriever

ColBERTv2 dense retrieval model for high-quality passage retrieval.

class ColBERTv2:
    """
    ColBERTv2 retrieval model for dense passage retrieval.

    Connects to a ColBERT server for retrieving passages using
    dense token-level representations.
    """

    def __init__(self, url: str = "http://0.0.0.0", port: int = None, **kwargs):
        """
        Initialize ColBERT retriever.

        Args:
            url (str): ColBERT server URL (default: "http://0.0.0.0")
            port (int | str | None): Server port
            **kwargs: Additional ColBERT-specific parameters
        """
        pass

    def __call__(self, query: str, k: int = 10, **kwargs):
        """
        Retrieve documents for query.

        Args:
            query (str): Search query
            k (int): Number of documents to retrieve
            **kwargs: Additional parameters

        Returns:
            List of retrieved passages
        """
        pass

Usage:

import dspy

# Initialize ColBERT retriever
colbert = dspy.ColBERTv2(url="http://localhost", port=8893)

# Configure as retrieval model
dspy.configure(
    lm=dspy.LM('openai/gpt-4o-mini'),
    rm=colbert
)

# Use with Retrieve module
retrieve = dspy.Retrieve(k=5)
result = retrieve(query="What is machine learning?")
print(result.passages)

WeaviateRM

Weaviate vector database retrieval model.

class WeaviateRM:
    """
    Weaviate vector database retrieval model.

    Connects to a Weaviate instance for semantic search
    using vector embeddings.
    """

    def __init__(
        self,
        weaviate_collection_name: str,
        weaviate_client=None,
        weaviate_collection_text_key: str = "content",
        k: int = 3
    ):
        """
        Initialize Weaviate retrieval model.

        Args:
            weaviate_collection_name (str): Name of Weaviate collection
            weaviate_client: Weaviate client instance
            weaviate_collection_text_key (str): Field name for text content
            k (int): Default number of results
        """
        pass

    def __call__(self, query: str, k: int = None, **kwargs):
        """
        Retrieve documents from Weaviate.

        Args:
            query (str): Search query
            k (int | None): Number of documents
            **kwargs: Additional parameters

        Returns:
            List of retrieved passages
        """
        pass

Usage:

import dspy
import weaviate

# Set up Weaviate client
client = weaviate.Client("http://localhost:8080")

# Create Weaviate retriever
weaviate_rm = dspy.WeaviateRM(
    weaviate_collection_name="Documents",
    weaviate_client=client,
    weaviate_collection_text_key="content",
    k=5
)

# Configure
dspy.configure(
    lm=dspy.LM('openai/gpt-4o-mini'),
    rm=weaviate_rm
)

# Use for retrieval
retrieve = dspy.Retrieve(k=5)
result = retrieve(query="machine learning concepts")

DatabricksRM

Databricks Vector Search retrieval model.

class DatabricksRM:
    """
    Databricks Vector Search retrieval model.

    Connects to Databricks Vector Search for enterprise-scale
    semantic search capabilities.
    """

    def __init__(
        self,
        databricks_index_name: str,
        databricks_endpoint: str = None,
        databricks_token: str = None,
        columns: list = None,
        k: int = 3
    ):
        """
        Initialize Databricks Vector Search retriever.

        Args:
            databricks_index_name (str): Vector search index name
            databricks_endpoint (str | None): Databricks endpoint URL
            databricks_token (str | None): Authentication token
            columns (list | None): Columns to return from index
            k (int): Default number of results
        """
        pass

    def __call__(self, query: str, k: int = None, **kwargs):
        """
        Retrieve from Databricks Vector Search.

        Args:
            query (str): Search query
            k (int | None): Number of results
            **kwargs: Additional parameters

        Returns:
            List of retrieved passages
        """
        pass

Usage:

import dspy

# Create Databricks retriever
databricks_rm = dspy.DatabricksRM(
    databricks_index_name="my_index",
    databricks_endpoint=os.getenv("DATABRICKS_ENDPOINT"),
    databricks_token=os.getenv("DATABRICKS_TOKEN"),
    columns=["content", "source"],
    k=10
)

# Configure
dspy.configure(
    lm=dspy.LM('openai/gpt-4o-mini'),
    rm=databricks_rm
)

# Use for retrieval
retrieve = dspy.Retrieve(k=5)
result = retrieve(query="data engineering best practices")

Retrieval Patterns

Basic RAG

Simple retrieval-augmented generation:

import dspy

# Configure
dspy.configure(
    lm=dspy.LM('openai/gpt-4o-mini'),
    rm=my_retrieval_model
)

# RAG module
class SimpleRAG(dspy.Module):
    def __init__(self, k=3):
        super().__init__()
        self.retrieve = dspy.Retrieve(k=k)
        self.generate = dspy.ChainOfThought("context, question -> answer")

    def forward(self, question):
        context = self.retrieve(query=question).passages
        return self.generate(context=context, question=question)

# Use
rag = SimpleRAG(k=5)
result = rag(question="What is quantum computing?")

Multi-Hop RAG

Retrieve multiple times for complex questions:

import dspy

class MultiHopRAG(dspy.Module):
    """Multi-hop retrieval for complex questions."""

    def __init__(self, k=3, hops=2):
        super().__init__()
        self.k = k
        self.hops = hops
        self.retrieve = dspy.Retrieve(k=k)
        self.generate_query = dspy.Predict("context, question -> search_query")
        self.generate_answer = dspy.ChainOfThought("context, question -> answer")

    def forward(self, question):
        all_passages = []

        # First retrieval
        context = self.retrieve(query=question).passages
        all_passages.extend(context)

        # Additional hops
        for hop in range(self.hops - 1):
            # Generate follow-up query based on current context
            query_result = self.generate_query(
                context="\n".join(all_passages),
                question=question
            )

            # Retrieve with new query
            new_passages = self.retrieve(query=query_result.search_query).passages
            all_passages.extend(new_passages)

        # Generate final answer with all context
        return self.generate_answer(
            context="\n".join(all_passages),
            question=question
        )

multi_rag = MultiHopRAG(k=3, hops=3)
result = multi_rag(question="Who was the president when the iPhone was released?")

Hybrid Retrieval

Combine multiple retrieval methods:

import dspy

class HybridRetrieval(dspy.Module):
    """Combine keyword and semantic retrieval."""

    def __init__(self, k=5):
        super().__init__()
        self.keyword_retrieve = dspy.Retrieve(k=k)
        self.semantic_retrieve = dspy.Retrieve(k=k)
        self.rerank = dspy.Predict("passages, query -> top_passages: list")
        self.generate = dspy.ChainOfThought("context, question -> answer")

    def forward(self, question):
        # Keyword retrieval (e.g., ColBERT)
        keyword_passages = self.keyword_retrieve(query=question).passages

        # Semantic retrieval (e.g., embeddings)
        with dspy.context(rm=semantic_rm):
            semantic_passages = self.semantic_retrieve(query=question).passages

        # Combine and deduplicate
        all_passages = list(set(keyword_passages + semantic_passages))

        # Rerank passages
        reranked = self.rerank(
            passages=all_passages,
            query=question
        )

        # Generate answer
        return self.generate(
            context="\n".join(reranked.top_passages[:5]),
            question=question
        )

Self-Retrieval

Generate queries dynamically:

import dspy

class SelfRetrievingRAG(dspy.Module):
    """Generate retrieval queries dynamically."""

    def __init__(self, k=3):
        super().__init__()
        self.generate_queries = dspy.Predict(
            "question -> search_queries: list"
        )
        self.retrieve = dspy.Retrieve(k=k)
        self.synthesize = dspy.ChainOfThought(
            "question, contexts -> answer"
        )

    def forward(self, question):
        # Generate multiple search queries
        queries_result = self.generate_queries(question=question)
        queries = queries_result.search_queries

        # Retrieve for each query
        all_passages = []
        for query in queries:
            passages = self.retrieve(query=query).passages
            all_passages.extend(passages)

        # Deduplicate
        unique_passages = list(set(all_passages))

        # Synthesize answer from all contexts
        return self.synthesize(
            question=question,
            contexts="\n\n".join(unique_passages)
        )

Retrieval with Filtering

Filter retrieved passages before generation:

import dspy

class FilteredRAG(dspy.Module):
    """Filter retrieved passages by relevance."""

    def __init__(self, k=5, top_k=3):
        super().__init__()
        self.retrieve = dspy.Retrieve(k=k)
        self.judge = dspy.Predict(
            "passage, question -> relevant: bool, score: float"
        )
        self.generate = dspy.ChainOfThought("context, question -> answer")
        self.top_k = top_k

    def forward(self, question):
        # Retrieve candidates
        passages = self.retrieve(query=question).passages

        # Judge each passage
        scored_passages = []
        for passage in passages:
            judgment = self.judge(passage=passage, question=question)
            if judgment.relevant:
                scored_passages.append((passage, judgment.score))

        # Sort by score and take top_k
        scored_passages.sort(key=lambda x: -x[1])
        top_passages = [p for p, s in scored_passages[:self.top_k]]

        # Generate answer
        return self.generate(
            context="\n".join(top_passages),
            question=question
        )

Custom Retrieval Model

Implement custom retrieval:

import dspy

class CustomRetriever:
    """Custom retrieval implementation."""

    def __init__(self, index_path: str):
        self.index_path = index_path
        # Load your index/database
        self.index = self._load_index()

    def _load_index(self):
        # Load your search index
        pass

    def __call__(self, query: str, k: int = 3):
        """
        Retrieve passages for query.

        Args:
            query: Search query
            k: Number of results

        Returns:
            List of passage strings
        """
        # Your retrieval logic
        results = self._search(query, k)
        return results

    def _search(self, query, k):
        # Implement search logic
        pass

# Use custom retriever
retriever = CustomRetriever("/path/to/index")
dspy.configure(
    lm=dspy.LM('openai/gpt-4o-mini'),
    rm=retriever
)

# Works with Retrieve module
retrieve = dspy.Retrieve(k=5)
result = retrieve(query="custom search")

ColBERT Setup

Set up ColBERTv2 retrieval server:

import dspy

# Initialize ColBERT retriever
colbert = dspy.ColBERTv2(
    url="http://localhost",
    port=8893
)

# Configure as retrieval model
dspy.configure(
    lm=dspy.LM('openai/gpt-4o-mini'),
    rm=colbert
)

# Use in modules
retrieve = dspy.Retrieve(k=5)
result = retrieve(query="machine learning papers")

# Direct usage
passages = colbert(query="deep learning", k=10)

Embedding-Based Retrieval

Build retrieval with embeddings:

import dspy
import numpy as np
from sentence_transformers import SentenceTransformer

# Load embedding model
model = SentenceTransformer("all-MiniLM-L6-v2")
embedder = dspy.Embedder(model.encode)

# Prepare document collection
documents = [
    "Document 1 text...",
    "Document 2 text...",
    # ... more documents
]

# Compute document embeddings
doc_embeddings = embedder(documents)

# Custom retriever with embeddings
class EmbeddingRetriever:
    def __init__(self, documents, embeddings, embedder):
        self.documents = documents
        self.embeddings = embeddings
        self.embedder = embedder

    def __call__(self, query, k=3):
        # Embed query
        query_embedding = self.embedder([query])[0]

        # Compute similarities
        similarities = np.dot(self.embeddings, query_embedding)

        # Get top-k
        top_indices = np.argsort(similarities)[-k:][::-1]
        return [self.documents[i] for i in top_indices]

# Create and configure
retriever = EmbeddingRetriever(documents, doc_embeddings, embedder)
dspy.configure(
    lm=dspy.LM('openai/gpt-4o-mini'),
    rm=retriever
)

# Use with RAG
rag = SimpleRAG(k=5)
result = rag(question="What is AI?")

Install with Tessl CLI