Skip to content

Latest commit

 

History

History
414 lines (328 loc) · 11.4 KB

File metadata and controls

414 lines (328 loc) · 11.4 KB

GraphRAG Architecture

Design Principles

  1. Infrastructure, not product - No business logic, just building blocks
  2. Interface-first - Define contracts, implementations are pluggable
  3. Context everywhere - Cancellation, timeouts, tracing via context.Context
  4. Fail loud - Explicit errors, no silent degradation
  5. Code-based config - Functional options, no external config files

System Layers

Product Layer (consumers)
  - task-keeper, future products
  - Memory tiers, Scope isolation, ADR storage
                    |
          uses as library
                    |
GraphRAG Library (pkg/graphrag)
  - Retrieval Pipeline (Stage interface, Pipeline builder, Scoring/combination)
  - Core Interfaces (GraphStore, VectorStore, Embedder, CacheStore, Reranker, Indexer)
  - Implementations (Neo4j, Qdrant, Redis, Voyage, OpenAI, Cohere, Mock, Sidecar)
                    |
       gRPC, Bolt, HTTP
                    |
External Services
  - Neo4j (7687), Qdrant (6334), Redis (6379), Voyage API, OpenAI API, Cohere API

Core Interfaces

GraphStore

Manages documents and their relationships in a graph database.

type GraphStore interface {
    // Document operations
    Store(ctx context.Context, doc *Document) error
    Get(ctx context.Context, id string) (*Document, error)
    Delete(ctx context.Context, id string) error

    // Relationship operations
    FindRelated(ctx context.Context, id string, depth int) ([]*Document, error)
    CreateRelation(ctx context.Context, from, to string, relType string) error
    DeleteRelation(ctx context.Context, from, to string, relType string) error

    // Graph algorithms
    ComputePageRank(ctx context.Context) (map[string]float64, error)

    // Lifecycle
    Close() error
}

Implementations: store/neo4j, store/mock

VectorStore

Manages vector embeddings for similarity search.

type VectorStore interface {
    // Vector operations
    Upsert(ctx context.Context, doc *Document, vector []float32) error
    Search(ctx context.Context, vector []float32, limit int, filter Filter) ([]SearchResult, error)
    Delete(ctx context.Context, id string) error

    // Lifecycle
    Close() error
}

Implementations: store/qdrant, store/mock

Embedder

Converts text to vector embeddings.

type Embedder interface {
    // Batch embedding (for documents)
    Embed(ctx context.Context, texts []string) ([][]float32, error)

    // Single embedding (for queries, may have different input_type)
    EmbedQuery(ctx context.Context, query string) ([]float32, error)

    // Metadata
    Dimension() int
    Provider() string
}

Implementations: embedding/voyage, embedding/openai, embedding/cache, embedding/sidecar, embedding/mock

Reranker

Reorders documents by semantic relevance to query.

type Reranker interface {
    Rerank(ctx context.Context, query string, docs []string, topK int) ([]Result, error)
    Provider() string
}

type Result struct {
    Index    int
    Score    float64
    Document string
}

Implementations: rerank/voyage, rerank/cohere, rerank/sidecar, rerank/mock


Document Model

type Document struct {
    ID        string            // Unique identifier
    Content   string            // Main text content
    Metadata  map[string]any    // Flexible metadata
    CreatedAt time.Time
    UpdatedAt time.Time

    // Set during retrieval, not stored
    Score     float64           // Combined relevance score
    Signals   map[string]float64 // Individual signal scores
}

Design decisions:

  • No scope/tenant fields - that's product-layer concern
  • No memory tier - that's product-layer concern
  • Metadata is flexible - consumers add their fields
  • Score/Signals are ephemeral - only exist during retrieval

Retrieval Pipeline

Pipeline Architecture

Query -> [Stage 1] -> [Stage 2] -> ... -> [Stage N] -> Scorer -> Results
            |            |                  |
        Candidates   Expanded         Reranked

Stage Interface

type Stage interface {
    Name() string
    Execute(ctx context.Context, input StageInput) (StageOutput, error)
}

type StageInput struct {
    Query      string
    Candidates []*ScoredDocument
    Context    map[string]any  // Stage-specific data
}

type StageOutput struct {
    Candidates []*ScoredDocument
    Context    map[string]any
}

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

Pre-built Stages

Stage Input Output Signal
Semantic Query Top-K similar docs (ID only) semantic: 0.0-1.0
Hydrate Candidates Full documents from GraphStore -
GraphExpansion Candidates + Related docs graph: 0.0-1.0
TemporalDecay Candidates Same (scored) temporal: 0.0-1.0
PageRank Candidates Same (scored) pagerank: 0.0-1.0
Access Candidates Same (scored) access: 0.0-1.0
Rerank Candidates Reordered top-K rerank: 0.0-1.0

Scoring

Final score = weighted sum of signals:

type WeightedScorer struct {
    Semantic  float64  // e.g., 0.30
    Graph     float64  // e.g., 0.25
    Temporal  float64  // e.g., 0.15
    Access    float64  // e.g., 0.10
    PageRank  float64  // e.g., 0.10
    Rerank    float64  // e.g., 0.10
}

func (s WeightedScorer) Score(signals map[string]float64) float64 {
    return s.Semantic*signals["semantic"] +
           s.Graph*signals["graph"] +
           s.Temporal*signals["temporal"] +
           s.Access*signals["access"] +
           s.PageRank*signals["pagerank"] +
           s.Rerank*signals["rerank"]
}

Configuration Pattern

All configuration via functional options:

// Backend configuration
graphStore, err := neo4j.New(
    neo4j.WithURI("bolt://localhost:7687"),
    neo4j.WithAuth("neo4j", "password"),
    neo4j.WithMaxConnectionPoolSize(50),
    neo4j.WithConnectionTimeout(30 * time.Second),
)

// Embedder with retry
embedder, err := voyage.New(
    voyage.WithAPIKey(apiKey),
    voyage.WithModel("voyage-3.5"),
    voyage.WithRetry(5, 100*time.Millisecond),  // maxAttempts, baseDelay
    voyage.WithMaxBackoff(10 * time.Second),
    voyage.WithTimeout(30 * time.Second),
)

Why code-based config?

  1. Type-safe - compiler catches errors
  2. Testable - can inject test values
  3. Explicit - no hidden defaults or env var magic
  4. IDE support - autocomplete, go to definition

Error Handling

Sentinel Errors

var (
    // Storage errors
    ErrNotFound       = errors.New("graphrag: document not found")
    ErrAlreadyExists  = errors.New("graphrag: document already exists")
    ErrConnectionFail = errors.New("graphrag: connection failed")

    // Provider errors
    ErrRateLimited    = errors.New("graphrag: rate limited")
    ErrAuthFailed     = errors.New("graphrag: authentication failed")
    ErrInvalidInput   = errors.New("graphrag: invalid input")

    // Pipeline errors
    ErrStageFailure   = errors.New("graphrag: stage execution failed")
    ErrNoCandidates   = errors.New("graphrag: no candidates found")
    ErrPipelineEmpty  = errors.New("graphrag: pipeline has no stages")
)

Error Wrapping

// Implementations wrap with context
return fmt.Errorf("neo4j: failed to store document %s: %w", doc.ID, ErrConnectionFail)

// Consumers check with errors.Is
if errors.Is(err, graphrag.ErrNotFound) {
    // Handle not found
}

Retry Strategy

All external providers implement retry with exponential backoff:

type RetryConfig struct {
    MaxAttempts       int
    BaseBackoff       time.Duration
    MaxBackoff        time.Duration
    BackoffMultiplier float64
    RetryableErrors   []error  // Which errors to retry
}

// Default: 3 attempts, 1s base, 30s max, 2x multiplier
func DefaultRetryConfig() RetryConfig {
    return RetryConfig{
        MaxAttempts:       3,
        BaseBackoff:       1 * time.Second,
        MaxBackoff:        30 * time.Second,
        BackoffMultiplier: 2.0,
        RetryableErrors:   []error{ErrRateLimited, ErrConnectionFail},
    }
}

Features:

  • Exponential backoff with jitter
  • Respects Retry-After header (429 responses)
  • Context-aware (respects deadlines)
  • Distinguishes retryable vs fatal errors

Context Usage

All operations accept context for:

  1. Cancellation - Stop long-running operations
  2. Timeouts - Prevent hanging requests
  3. Tracing - OpenTelemetry span propagation
  4. Values - Request-scoped data (use sparingly)
// With timeout
ctx, cancel := context.WithTimeout(ctx, 30*time.Second)
defer cancel()

results, err := pipeline.Execute(ctx, query, 10)

// With tracing
ctx, span := tracer.Start(ctx, "graphrag.retrieve")
defer span.End()

Concurrency Model

  • Stores are thread-safe - Share across goroutines
  • Embedders are thread-safe - Internal connection pooling
  • Pipelines are thread-safe - Stateless execution
  • Documents are value types - Copy, don't share
// Safe: shared store
store, _ := neo4j.New(...)
go func() { store.Get(ctx, "id1") }()
go func() { store.Get(ctx, "id2") }()

// Safe: shared pipeline
pipeline := retrieval.NewPipeline().Build()
go func() { pipeline.Execute(ctx, "query1", 10) }()
go func() { pipeline.Execute(ctx, "query2", 10) }()

Package Structure

pkg/graphrag/
├── doc.go              # Package documentation
├── document.go         # Document type
├── errors.go           # Sentinel errors
├── context.go          # Context helpers
├── options.go          # Common option types
├── metrics.go          # Observability interface
│
├── store/              # Storage interfaces
│   ├── graph.go        # GraphStore interface
│   ├── vector.go       # VectorStore interface
│   ├── cache.go        # CacheStore interface
│   ├── neo4j/          # Neo4j implementation
│   ├── qdrant/         # Qdrant implementation
│   ├── redis/          # Redis implementation
│   └── mock/           # Mock for testing
│
├── embedding/          # Embedding interfaces
│   ├── embedder.go     # Embedder interface
│   ├── types.go        # Request/response types
│   ├── voyage/         # Voyage AI implementation
│   ├── openai/         # OpenAI implementation
│   ├── cache/          # Caching wrapper
│   ├── sidecar/        # Local model sidecar
│   └── mock/           # Mock for testing
│
├── rerank/             # Reranking interfaces
│   ├── reranker.go     # Reranker interface
│   ├── voyage/         # Voyage implementation
│   ├── cohere/         # Cohere implementation
│   ├── sidecar/        # Local model sidecar
│   └── mock/           # Mock for testing
│
├── retrieval/          # Retrieval pipeline
│   ├── pipeline.go     # Pipeline builder
│   ├── stage.go        # Stage interface
│   ├── scorer.go       # Scoring strategies
│   ├── types.go        # StageInput/Output types
│   ├── stages/         # Pre-built stages
│   └── presets/        # Common configurations
│
├── indexing/           # Document ingestion
│   ├── indexer.go      # Indexer interface
│   ├── chunker/        # Text chunking
│   └── enricher/       # Context enrichment
│
└── server/             # HTTP API server