Skip to content

ARCHITECTURE.md

Latest commit

 

History

History
582 lines (470 loc) · 21.7 KB

ARCHITECTURE.md

File metadata and controls

582 lines (470 loc) · 21.7 KB
Raw

HaystackBench — Architecture Document

Version: 1.0 Last Updated: 2025-01-01

1. System Overview

HaystackBench is a Python package with three user-facing interfaces:

  1. CLI (haystackbench command) — for power users and CI/CD
  2. Web UI (FastAPI + SvelteKit) — for interactive use and visualization
  3. Python API — for programmatic integration into research pipelines

These interfaces all consume the same core library (haystackbench/), which is designed to be fully usable as a standalone Python package independent of the UI.

┌─────────────────────────────────────────────────────────────────┐
│                        User Interfaces                           │
│   CLI (typer)          Web UI (SvelteKit)    Python API          │
└──────────────┬─────────────────┬──────────────────┬────────────┘
               │                 │                  │
               └─────────────────▼──────────────────┘
                          FastAPI Backend
                    (ui/backend/main.py)
                               │
               ┌───────────────▼───────────────────┐
               │        Core Library                │
               │    haystackbench/                  │
               │                                    │
               │  ┌─────────┐   ┌───────────────┐  │
               │  │  Tasks  │   │    Config     │  │
               │  └────┬────┘   └───────┬───────┘  │
               │       │                │           │
               │  ┌────▼────────────────▼────────┐  │
               │  │        Core Engine            │  │
               │  │  needle ─ haystack ─ injector │  │
               │  │  evaluator ─ runner           │  │
               │  └────┬──────────────┬───────────┘  │
               │       │              │              │
               │  ┌────▼────┐  ┌──────▼──────┐      │
               │  │Provider │  │  Storage    │      │
               │  │Registry │  │  (SQLite)   │      │
               │  └────┬────┘  └─────────────┘      │
               └───────┼────────────────────────────┘
                       │
          ┌────────────┼────────────────┐
          ▼            ▼                ▼
     Anthropic       OpenAI          Ollama ...
     API             API             (local)

2. Module Structure & Responsibilities

2.1 haystackbench/core/

The heart of the library. All modules here are pure Python with no web dependencies.

needle.py

Responsibility: Generate all types of needles (synthetic facts to be inserted).

Key classes:

  • Needle — dataclass holding template, filled text, value, entity name, attribute
  • NeedleConfig — Pydantic model for needle generation parameters
  • NeedleFactory — factory class with methods:
    • generate_single() → Needle
    • generate_batch(count) → List[Needle]
    • generate_relational_chain(depth, structure) → RelationalNeedleSet
    • generate_kv_pairs(count) → List[KVNeedle]

haystack.py

Responsibility: Source, construct, and manage background text (haystacks).

Key classes:

  • Haystack — dataclass holding raw text, token count, source metadata
  • HaystackConfig — Pydantic model for haystack parameters
  • HaystackBuilder — builds haystacks from various sources
  • HaystackSource — enum of supported sources

Haystack sources are implemented as plugins in haystackbench/haystacks/:

  • standard/paul_graham.py — downloads and caches PG essays
  • standard/wikipedia.py — fetches Wikipedia articles via API
  • standard/arxiv.py — fetches arXiv abstracts
  • generators/synthetic.py — generates synthetic paragraphs

injector.py

Responsibility: Insert needles at specific depth positions within a haystack.

Key classes:

  • InjectedDocument — dataclass holding the final text, metadata about injection points
  • Injector — performs depth-based injection with paragraph boundary enforcement

Key logic:

  • Converts depth % → token position → finds nearest paragraph boundary
  • Handles multi-needle injection with configurable distribution strategies
  • Validates that injected document length matches target token count

evaluator.py

Responsibility: Score model responses against expected answers.

Key classes:

  • EvaluationResult — score (0.0–1.0), is_correct bool, method used, explanation
  • Evaluator — routes scoring to the appropriate method:
    • score_exact() — string normalization + substring check
    • score_semantic() — sentence-transformer cosine similarity
    • score_llm_judge() — async call to judge model
    • score_numeric() — parse and compare numbers

runner.py

Responsibility: Orchestrate full experiment execution (async, concurrent, with checkpointing).

Key classes:

  • TestRunner — main orchestrator
    • run_experiment(config, provider) → ExperimentResult
    • estimate_cost(config) → CostEstimate
    • _run_cell(config, cell) → ResultRecord

Key behaviors:

  • Builds the full grid of (context_length, depth, num_needles, trial, model) tuples
  • Depth-sweep deduplication: for uniform and bookends distributions with num_needles > 1, needle placement is identical across all depths — _build_cells() collapses the depth axis to a single value (0.5), eliminating redundant cells
  • Context window pre-flight: before each API call, checks the model's max_context_tokens; cells that exceed it raise a non-retriable error and are recorded as skipped
  • Respects max_concurrency via asyncio.Semaphore
  • Saves each ResultRecord to SQLite immediately after completion
  • On resume, queries SQLite for existing results and skips those cells
  • Emits progress events via callback for WebSocket streaming
  • Needle generation is seeded from (experiment_id, model, trial) so the same needle is used across all (context_length × depth) cells for a given trial — holding the needle constant while varying the haystack

2.2 haystackbench/tasks/

Each task type encapsulates the logic for generating prompts and interpreting results for that task.

base.py

Abstract base class BaseTask:

class BaseTask(ABC):
    @abstractmethod
    def build_prompt(self, document: InjectedDocument, needle_set: ...) -> str

    @abstractmethod
    def build_system_prompt(self) -> str

    @abstractmethod
    def score_response(self, response: str, evaluator: Evaluator) -> EvaluationResult

    @property
    @abstractmethod
    def task_type(self) -> TaskType

Task Implementations

Each file implements one BaseTask subclass:

File Class Task Type
single_retrieval.py SingleRetrievalTask S-RT
multi_retrieval.py MultiRetrievalTask M-RT
multi_reasoning.py MultiReasoningTask M-RS
ancestral_trace.py AncestralTraceTask ATC
counting.py CountingTask COUNTING
key_value.py KeyValueTask KEY_VALUE
conflicting.py ConflictingNeedlesTask CONFLICTING

Task Registry (__init__.py)

TASK_REGISTRY: Dict[TaskType, Type[BaseTask]] = {
    TaskType.S_RT: SingleRetrievalTask,
    ...
}

def get_task(task_type: TaskType, config: TaskConfig) -> BaseTask:
    return TASK_REGISTRY[task_type](config)

2.3 haystackbench/providers/

Each provider wraps one LLM API vendor.

base.py

class BaseProvider(ABC):
    @abstractmethod
    async def complete(self, prompt, model, max_tokens, temperature, system_prompt) -> ProviderResponse

    @abstractmethod
    def list_models(self) -> List[ModelInfo]

    @abstractmethod
    def validate_credentials(self) -> bool

ProviderResponse dataclass:

  • content: str — model response text
  • prompt_tokens: int
  • completion_tokens: int
  • latency_ms: int
  • model: str — exact model version returned by API

Provider Registry

PROVIDER_REGISTRY = {
    "anthropic": AnthropicProvider,
    "openai": OpenAIProvider,
    "google": GoogleProvider,
    ...
}

Rate Limiting

Each provider has a configurable requests_per_minute limit. The runner enforces this via token bucket within the provider wrapper.

2.4 haystackbench/config/

schema.py

All Pydantic v2 models:

  • ExperimentConfig — top-level experiment configuration, including depth_distribution and custom_needles
  • NeedleConfig — needle generation parameters
  • HaystackConfig — haystack source and construction parameters
  • ModelConfig — provider + model ID + per-model overrides
  • ScoringConfig — scoring method and judge model
  • CustomNeedleSpec — needle text, expected answer, optional label and custom question; embedded in ExperimentConfig.custom_needles so experiments are self-contained

defaults.py

Default values for all config fields. Centralized so they're easy to find and modify.

presets.py

Named preset configurations returned as ExperimentConfig instances:

  • get_preset(name: str) -> ExperimentConfig

Presets: quick, standard, full, large_context, needlebench_32k, needlebench_128k, needlebench_1m

Preset Context lengths Depths Cells/model
quick 4K, 16K, 65K 5 15
standard 4K–131K (6) 11 66
full 4K–1M (9) 21 567
large_context 128K–1M (4) 11 44
needlebench_1m 4K–1M (7) 11 77

2.5 haystackbench/storage/

credentials.py

class CredentialStore:
    def save(self, provider: str, api_key: str) -> None
    def load(self, provider: str) -> Optional[str]
    def delete(self, provider: str) -> None
    def list_providers(self) -> List[str]

Uses keyring library. Falls back to os.environ[f"{PROVIDER.upper()}_API_KEY"] if keyring is unavailable.

results.py

class ResultStore:
    async def init_db(self) -> None          # Create tables if not exist
    async def save_experiment(self, config) -> str  # Returns experiment ID
    async def save_result(self, result: ResultRecord) -> None
    async def load_results(self, experiment_id: str) -> List[ResultRecord]
    async def load_experiment(self, experiment_id: str) -> ExperimentRecord
    async def list_experiments(self) -> List[ExperimentSummary]
    async def update_experiment_status(self, id, status, completed_cells) -> None
    async def update_experiment_summary(self, id, accuracy, cells_correct,
                                        elapsed_seconds, cost_usd) -> None

The experiments table carries summary columns (accuracy, cells_correct, elapsed_seconds, cost_usd) that are populated by update_experiment_summary() after a run completes. The config_json column is denormalised in GET /api/experiments/{id} to expose task_type, haystack_source, and models directly.

Cell results carry a metadata_json blob with: needle_texts, expected_answer, actual_depths, question, and a truncated haystack_text — enabling the cell detail view without re-running.

custom_models.py

class CustomModelStore:
    async def init_db(self) -> None
    async def list_models(self) -> List[dict]
    async def add_model(self, provider, model_id, display_name,
                        max_context_tokens, ...) -> dict
    async def delete_model(self, provider, model_id) -> bool
    async def context_window_map(self) -> dict[str, dict[str, int]]

Stores user-defined models in the custom_models table (shared DB). context_window_map() is used by cost estimation and context-guardrail logic to extend the static provider model lists.

custom_needles.py

class CustomNeedleStore:
    async def init_db(self) -> None
    async def list_needles(self) -> List[dict]
    async def add_needle(self, needle_text, expected_answer,
                         name, question) -> dict
    async def update_needle(self, needle_id, ...) -> bool
    async def delete_needle(self, needle_id) -> bool
    async def get_needles_by_ids(self, ids: List[int]) -> List[dict]

Stores reusable needle/answer pairs in the custom_needles table. When an experiment is created with custom needle IDs, the runner embeds the full CustomNeedleSpec objects into ExperimentConfig.custom_needles so the experiment remains self-contained.

cache.py

class ResponseCache:
    async def get(self, cache_key: str) -> Optional[ProviderResponse]
    async def set(self, cache_key: str, response: ProviderResponse) -> None
    def make_key(self, model: str, prompt: str) -> str  # SHA256 hash

Cache key: SHA256(model + prompt). Stored in SQLite as a separate table.

2.6 haystackbench/viz/

All visualization modules produce both:

  1. Static images (PNG/SVG) via matplotlib
  2. JSON-serializable data for web UI (ECharts format)

heatmap.py

def generate_heatmap(
    results: List[ResultRecord],
    title: str,
    colormap: str = 'RdYlGn',
    annotate: bool = True,
    output_path: Optional[Path] = None
) -> Figure

report.py

def generate_html_report(
    experiment: ExperimentRecord,
    results: List[ResultRecord],
    output_path: Path
) -> None

Generates a self-contained HTML file with all charts as base64-embedded SVGs and a sortable data table.

2.7 haystackbench/cli.py

Built with typer. Command groups:

  • haystackbench setup — credential wizard
  • haystackbench models — model management
  • haystackbench run — experiment execution
  • haystackbench results — results management
  • haystackbench haystack — haystack utilities
  • haystackbench needle — needle utilities
  • haystackbench serve — start web UI server

Uses rich for tables, progress bars, and colored output.

2.8 ui/backend/

FastAPI application with:

  • Lifespan events: database initialization, credential store setup
  • CORS: configured for local SvelteKit dev server
  • WebSocket manager: tracks active experiment runs, broadcasts progress events
  • Background tasks: experiment runs execute as asyncio tasks, not blocking the event loop

Route organization:

  • routes/models.py — provider and credential management
  • routes/tests.py — experiment CRUD and execution control
  • routes/results.py — results retrieval and export; includes GET /{id}/cells for per-cell detail
  • routes/haystacks.py — haystack management
  • routes/settings.py — custom model and custom needle CRUD (/api/settings/models, /api/settings/needles)
  • websocket.py — WebSocket connection manager

2.9 ui/frontend/

SvelteKit application:

Store structure (src/stores/):

  • experiments.ts — experiment list and current experiment state
  • results.ts — loaded results and active filters
  • models.ts — provider connection status and model lists
  • ws.ts — WebSocket connection and event handlers

Component hierarchy:

App.svelte
├── Layout.svelte (nav, sidebar)
├── pages/
│   ├── Dashboard.svelte
│   ├── Models.svelte
│   ├── NewExperiment.svelte (wizard)
│   │   ├── steps/Step1TaskType.svelte
│   │   ├── steps/Step2Needles.svelte
│   │   ├── steps/Step3Haystack.svelte
│   │   ├── steps/Step4Grid.svelte
│   │   ├── steps/Step5Models.svelte
│   │   └── steps/Step6Review.svelte
│   ├── ExperimentProgress.svelte
│   ├── Results.svelte
│   │   ├── HeatmapTab.svelte
│   │   ├── ChartsTab.svelte
│   │   └── DrilldownModal.svelte
│   ├── ResultsLibrary.svelte
│   └── Settings.svelte
└── components/ (shared)
    ├── HeatmapChart.svelte (ECharts wrapper)
    ├── RadarChart.svelte
    ├── LineChart.svelte
    ├── ProgressGrid.svelte (live updating)
    └── Toast.svelte

3. Data Flow

3.1 Experiment Execution Flow

User configures experiment (UI/CLI/YAML)
        │
        ▼
ExperimentConfig (Pydantic validation)
        │
        ▼
TestRunner.run_experiment()
        │
        ├── Build experiment grid (all cells)
        ├── Load existing results from SQLite (for resume)
        ├── Filter out completed cells
        │
        ▼ (for each remaining cell, concurrently)
TestRunner._run_cell(cell)
        │
        ├── Pre-flight: check cell.context_length ≤ model.max_context_tokens
        ├── NeedleFactory.generate_*() — or load from CustomNeedleSpec pool
        ├── HaystackBuilder.build() → Haystack (cached)
        ├── Task.compute_depths() → depth list (respects depth_distribution)
        ├── Injector.inject_multiple() → InjectedDocument
        ├── Task.build_prompt() → str
        ├── ResponseCache.get() → hit? return cached
        ├── Provider.complete() → ProviderResponse
        ├── ResponseCache.set() → cache response
        ├── Evaluator.score_response() → EvaluationResult (auto/numeric/substring)
        ├── ResultStore.save_result() → persist (with metadata: needle, expected, depths)
        └── progress_callback(progress_event) → WebSocket broadcast

3.2 Results Query Flow

User requests results (UI/CLI)
        │
        ▼
ResultStore.load_results(experiment_id)
        │
        ▼
List[ResultRecord]
        │
        ├── viz/heatmap.py → matplotlib Figure → PNG/SVG
        ├── viz/radar.py → spider chart
        ├── viz/report.py → HTML report
        └── Direct JSON/CSV export

4. Threading & Concurrency Model

  • Main event loop: asyncio (single-threaded)
  • Concurrent API calls: asyncio.gather() with semaphore (max_concurrency)
  • Blocking operations: tokenization, file I/O run in asyncio.get_event_loop().run_in_executor(None, ...)
  • matplotlib rendering: synchronous, called from sync context after async data collection
  • FastAPI: fully async, no sync route handlers
  • WebSocket broadcasting: fan-out to all connected clients via asyncio.gather()

5. Configuration Loading Priority

YAML config file
    ↓ (overrides)
CLI flags
    ↓ (overrides)
Environment variables
    ↓ (overrides)
Code defaults (config/defaults.py)

6. Error Handling Strategy

Error Category Behavior
Provider API error (5xx) Exponential backoff, retry up to N times, then mark cell as failed
Rate limit (429) Respect Retry-After header, then continue
Credential missing Raise CredentialError with actionable message
Model not available Raise ModelNotFoundError with list of available models
Context length exceeded Pre-flight check skips cell before API call; non-retriable; recorded in skipped_by_model in cost estimate
Network timeout Retry with backoff
Invalid config Pydantic ValidationError with field-level error messages
SQLite error Log error, mark experiment as failed, preserve existing data

7. Extension Points

Adding a New Provider

  1. Create haystackbench/providers/myprovider.py
  2. Subclass BaseProvider, implement 3 abstract methods
  3. Register in haystackbench/providers/__init__.py
  4. Add to ModelInfo list with context window sizes
  5. Write a test in tests/unit/test_providers.py

Adding a New Task Type

  1. Create haystackbench/tasks/mytask.py
  2. Subclass BaseTask, implement abstract methods
  3. Register in haystackbench/tasks/__init__.py
  4. Add to TaskType enum in haystackbench/config/schema.py
  5. Add preset config if applicable
  6. Write a test in tests/unit/test_tasks.py

Adding a New Haystack Source

  1. Create haystackbench/haystacks/standard/mysource.py
  2. Implement fetch(target_tokens, **kwargs) -> str
  3. Register in HaystackSource enum
  4. Handle caching via HaystackCache

8. Dependencies

Core Dependencies

Package Version Purpose
pydantic >=2.0 Config validation
aiosqlite >=0.19 Async SQLite
tiktoken >=0.5 Token counting
keyring >=24.0 Credential storage
httpx >=0.25 Async HTTP client
rich >=13.0 CLI formatting
typer >=0.9 CLI framework
pyyaml >=6.0 YAML config
matplotlib >=3.7 Chart generation
seaborn >=0.12 Statistical visualization

Provider SDKs (optional, installed with extras)

Package Provider
anthropic Anthropic
openai OpenAI
google-generativeai Google
mistralai Mistral
cohere Cohere
groq Groq

UI Dependencies (optional)

Package Purpose
fastapi Backend framework
uvicorn ASGI server
websockets WebSocket support

Optional Enhancement Dependencies

Package Purpose
sentence-transformers Semantic scoring
numpy Numeric operations