Skip to content

PIPELINE.md

Latest commit

 

History

History
462 lines (349 loc) · 10.7 KB

PIPELINE.md

File metadata and controls

462 lines (349 loc) · 10.7 KB
Raw

Retrieval Pipeline

The retrieval pipeline is the core of GraphRAG. It orchestrates multiple stages to find the most relevant documents for a query.

Pipeline Concept

Query -> [Semantic] -> [Hydrate] -> [Graph] -> [Temporal] -> [Rerank] -> Score -> Results
             |            |            |            |           |
         50 docs      hydrated     +20 related    scored    reordered   top 10

Each stage:

  1. Receives candidates from previous stage
  2. May add, remove, or score documents
  3. Passes results to next stage
  4. Final scorer combines all signals

Building Pipelines

Fluent Builder API

pipeline := retrieval.NewPipeline().
    AddStage(stages.Semantic(vectorStore, embedder, 50)).
    AddStage(stages.Hydrate(graphStore)).
    AddStage(stages.GraphExpansion(graphStore, 2)).
    AddStage(stages.TemporalDecay(90 * 24 * time.Hour)).
    AddStage(stages.PageRank(graphStore)).
    AddStage(stages.Rerank(reranker, 20)).
    WithScoring(retrieval.WeightedScorer{
        Semantic: 0.30,
        Graph:    0.25,
        Temporal: 0.15,
        PageRank: 0.10,
        Rerank:   0.20,
    }).
    Build()

Using Presets

// Standard 7-stage pipeline
pipeline := presets.Standard(graphStore, vectorStore, embedder, reranker)

// Minimal vector-only
pipeline := presets.Minimal(vectorStore, embedder)

// No reranking (faster, lower quality)
pipeline := presets.NoRerank(graphStore, vectorStore, embedder)

Stage Interface

type Stage interface {
    // Name returns the stage identifier for logging/metrics
    Name() string

    // Execute processes candidates and returns updated set
    Execute(ctx context.Context, input StageInput) (StageOutput, error)
}

type StageInput struct {
    Query      string                 // Original query text
    Candidates []*ScoredDocument      // Documents from previous stage
    Context    map[string]any         // Stage-specific context
}

type StageOutput struct {
    Candidates []*ScoredDocument      // Updated document set
    Context    map[string]any         // Context for next stage
}

type ScoredDocument struct {
    Document *Document
    Signals  map[string]float64       // Named scores: semantic, graph, etc.
}

Pre-built Stages

1. Semantic Search

Vector similarity search using embeddings.

stage := stages.Semantic(vectorStore, embedder, 50)
Parameter Type Description
vectorStore VectorStore Vector database
embedder Embedder Embedding provider
limit int Max candidates to retrieve

Behavior:

  • Embeds query using embedder.EmbedQuery()
  • Searches vectorStore with cosine similarity
  • Returns top limit documents (ID only, not full content)
  • Sets semantic signal (0.0-1.0)

Options:

stages.Semantic(store, embedder, 50,
    stages.WithSemanticMinScore(0.6),  // Minimum similarity threshold
    stages.WithSemanticFilter(filter),  // Metadata filtering
)

Important: Semantic stage only populates document IDs for performance. Use Hydrate stage to load full content for downstream stages.

2. Hydrate

Loads full document content from GraphStore.

stage := stages.Hydrate(graphStore)
Parameter Type Description
graphStore GraphStore Graph database to fetch from

Behavior:

  • For each candidate, fetches full document from GraphStore
  • Preserves existing signals from previous stages
  • Skips documents that are already hydrated (have Content)
  • Removes documents that no longer exist

When to use: Add after Semantic stage if downstream stages need document content (e.g., Rerank).

3. Graph Expansion

Expands candidate set by traversing relationships.

stage := stages.GraphExpansion(graphStore, 2)
Parameter Type Description
graphStore GraphStore Graph database
depth int Relationship hops (1-3 recommended)

Behavior:

  • For each candidate, finds related documents up to depth hops
  • Adds discovered documents to candidate set
  • Sets graph signal based on relationship distance

Options:

stages.GraphExpansion(store, 2,
    stages.WithMaxExpansion(100),  // Cap total expansion
)

4. Temporal Decay

Scores documents by recency.

stage := stages.TemporalDecay(90 * 24 * time.Hour)
Parameter Type Description
halfLife time.Duration Time for score to decay 50%

Behavior:

  • Calculates age of each document
  • Applies exponential decay: score = 0.5^(age/halfLife)
  • Sets temporal signal (0.0-1.0)

Options:

stages.TemporalDecay(90 * 24 * time.Hour,
    stages.WithMinTemporalScore(0.1),  // Floor value
)

5. PageRank

Scores by graph centrality (structural importance).

stage := stages.PageRank(graphStore)
Parameter Type Description
graphStore GraphStore Graph database with PageRank support

Behavior:

  • Retrieves pre-computed PageRank scores
  • Normalizes to 0.0-1.0 range
  • Sets pagerank signal
  • Falls back to default if PageRank unavailable

Options:

stages.PageRank(store,
    stages.WithDefaultPageRank(0.5),  // Fallback value
)

6. Access Count

Scores by historical retrieval frequency.

stage := stages.Access()

Behavior:

  • Reads access_count from document metadata
  • Normalizes across candidate set
  • Sets access signal (0.0-1.0)
  • Higher access = higher score (frequently retrieved = trustworthy)

Options:

stages.Access(
    stages.WithAccessMetadataKey("access_count"),  // Metadata key
    stages.WithMaxAccessCount(1000),               // Normalization cap
)

7. Rerank

Cross-encoder reranking for precision.

stage := stages.Rerank(reranker, 20)
Parameter Type Description
reranker Reranker Reranking provider
topK int How many to rerank

Behavior:

  • Takes top topK candidates by current score
  • Sends query + documents to reranker
  • Reorders based on cross-encoder scores
  • Sets rerank signal (0.0-1.0)

Note: Requires documents to have Content - use Hydrate stage before Rerank.

Scoring Strategies

Weighted Scorer (Default)

Linear combination of signals:

scorer := retrieval.WeightedScorer{
    Semantic: 0.30,
    Graph:    0.25,
    Temporal: 0.15,
    PageRank: 0.10,
    Access:   0.10,
    Rerank:   0.10,
}

// Final score = 0.30*semantic + 0.25*graph + 0.15*temporal + 0.10*pagerank + 0.10*access + 0.10*rerank

Custom Scorer

Implement the Scorer interface:

type Scorer interface {
    Score(signals map[string]float64) float64
}

// Example: Prioritize recency
type RecencyBiasedScorer struct{}

func (s RecencyBiasedScorer) Score(signals map[string]float64) float64 {
    if signals["temporal"] < 0.3 {
        return signals["semantic"] * 0.5  // Penalize old docs
    }
    return signals["semantic"]*0.4 + signals["temporal"]*0.4 + signals["graph"]*0.2
}

Creating Custom Stages

type MyCustomStage struct {
    // dependencies
}

func (s *MyCustomStage) Name() string {
    return "my_custom"
}

func (s *MyCustomStage) Execute(ctx context.Context, input StageInput) (StageOutput, error) {
    var updated []*ScoredDocument

    for _, doc := range input.Candidates {
        // Your logic here
        score := calculateMyScore(doc)

        // Add signal
        doc.Signals["my_custom"] = score
        updated = append(updated, doc)
    }

    return StageOutput{
        Candidates: updated,
        Context:    input.Context,
    }, nil
}

// Use it
pipeline := retrieval.NewPipeline().
    AddStage(stages.Semantic(store, embedder, 50)).
    AddStage(&MyCustomStage{}).
    Build()

Pipeline Execution

Basic Execution

results, err := pipeline.Execute(ctx, "How does caching work?", 10)
if err != nil {
    return err
}

for _, r := range results {
    fmt.Printf("%.3f: %s\n", r.Score, r.Document.ID)
}

With Options

results, err := pipeline.Execute(ctx, query, 10,
    retrieval.WithMinScore(0.5),  // Minimum final score
)

Result Structure

type Result struct {
    Document *Document
    Score    float64              // Final combined score
    Signals  map[string]float64   // Individual signal scores
}

Observability

Logging

Each stage logs entry/exit with document counts:

level=DEBUG msg="stage started" stage=semantic candidates=0
level=DEBUG msg="stage completed" stage=semantic candidates=50 duration=123ms
level=DEBUG msg="stage started" stage=hydrate candidates=50
level=DEBUG msg="stage completed" stage=hydrate candidates=50 duration=45ms
level=DEBUG msg="stage started" stage=graph candidates=50
level=DEBUG msg="stage completed" stage=graph candidates=67 duration=45ms

Metrics

When enabled, records:

  • graphrag_stage_duration_seconds - Per-stage latency
  • graphrag_stage_candidates - Candidate count after each stage
  • graphrag_pipeline_duration_seconds - Total pipeline latency
  • graphrag_pipeline_results - Final result count

Tracing

Creates spans for pipeline execution:

[Pipeline: retrieve]
  +-- [Stage: semantic]
  +-- [Stage: hydrate]
  +-- [Stage: graph]
  +-- [Stage: temporal]
  +-- [Stage: pagerank]
  +-- [Stage: rerank]
  +-- [Scoring]

Performance Considerations

Stage Typical Latency Cost Driver
Semantic 50-200ms Embedding API call + vector search
Hydrate 10-50ms GraphStore lookup (batch)
GraphExpansion 10-50ms Neo4j traversal depth
TemporalDecay <1ms In-memory calculation
PageRank 10-30ms Neo4j lookup (pre-computed)
Access <1ms In-memory calculation
Rerank 100-500ms Reranker API call

Optimization tips:

  1. Limit semantic candidates - 50 is usually enough
  2. Use Hydrate strategically - Only if downstream stages need content
  3. Cap graph expansion - Depth 2 with max 100 expansion
  4. Skip PageRank if not needed - Use preset without it
  5. Batch reranking - Rerank top 20, not all 70+
  6. Enable caching - CacheStore for frequent queries

Error Handling

Pipeline fails fast on stage errors:

results, err := pipeline.Execute(ctx, query, 10)
if err != nil {
    switch {
    case errors.Is(err, graphrag.ErrConnectionFail):
        // Backend unavailable
    case errors.Is(err, graphrag.ErrRateLimited):
        // Provider rate limit
    case errors.Is(err, graphrag.ErrNoCandidates):
        // No results found
    default:
        // Unexpected error
    }
}

Individual stage failures are wrapped:

// Error includes stage context
"stage 'rerank' failed: graphrag: rate limited"