zleap-sag

Local-first memory & knowledge engine for AI agents. Ingest documents, extract an event/entity graph with an LLM, and retrieve over it — fully local by default, progressive to production databases.

Distribution zleap-sag · import zleap.sag · Python ≥ 3.11 · MIT

Highlights

• Zero-infra by default — embedded SQLite + LanceDB (built-in BM25), no services to run. Existing default paths remain ./.zleap/sag.db and ./.zleap/lancedb/.
• Progressive to production — swap in MySQL / PostgreSQL / OceanBase (relational) and Elasticsearch / pgvector (vector) by changing config only; pipeline code is unchanged.
• Freely composable storage — combine any relational backend with LanceDB, Elasticsearch, pgvector, or OceanBase Vector. A physical single-database deployment shares only the target; relation and vector still use separate tables, runtimes, pools, and transactions.
• Explicit five-stage API — Parse → Chunking → Index → Extract → Search, with immutable, JSON-serializable hand-off contracts and no "last run" engine state.
• Profile-driven retrieval — five explicit profiles: vector, atomic, full_expand, pruned_expand_llm, and pruned_expand_rff.

Quick Start

1. Install

pip install zleap-sag        # runs as-is: embedded SQLite + LanceDB, no services needed

Add an extra only when you switch that backend on (quote the brackets in shells):

pip install "zleap-sag[es]"        # Elasticsearch vector store
pip install "zleap-sag[mysql]"     # MySQL or OceanBase       (aiomysql driver)
pip install "zleap-sag[postgres]"  # PostgreSQL, incl. pgvector single-DB
pip install "zleap-sag[summary]"   # optional Sumy article summaries
pip install "zleap-sag[tables]"    # CSV/XLSX via MarkItDown, including multi-sheet workbooks
pip install "zleap-sag[all]"       # all backends + litellm

2. Run

Point it at any OpenAI-compatible LLM + embedding endpoint. Storage defaults to ./.zleap/ (SQLite + LanceDB); the schema is created automatically on first run.

import asyncio
from zleap.sag import DataEngine, EngineConfig
from zleap.sag.config import LLMConfig, EmbeddingConfig
from zleap.sag.pipeline import SearchOptions, SearchRequest, SearchScope


async def main():
    config = EngineConfig(
        storage_mode="normal",  # required: normal | lite
        llm=LLMConfig(api_key="sk-...", base_url="https://your-gateway/v1", model="qwen3.6-flash"),
        embedding=EmbeddingConfig(model="bge-large-en-v1.5"),
    )
    async with DataEngine(config) as engine:  # start(): create local tables on first run
        chunk_ref = await engine.ingest("your_document.md")
        event_ref = await engine.extract(chunk_ref)
        result = await engine.search(
            SearchRequest(
                query="Who founded Acme?",
                scope=SearchScope(
                    data_source_ids=(event_ref.data_source_id,),
                    source_ids=(event_ref.source_id,),
                ),
                options=SearchOptions(strategy="full_expand", top_k=5, return_type="event"),
            )
        )
        for event in result.events:
            print(event.content[:200])


asyncio.run(main())

Storage lands in ./.zleap/ — add it to your .gitignore. storage_mode selects the schema contract but does not rewrite data_dir; use an explicit different data_dir if you intentionally maintain a second embedded store. Runnable scripts: examples/.

The two-command example stores the extracted events and reusable search scope in one JSON file:

python examples/05_extract_and_search.py extract article.md --env-file .env
python examples/05_extract_and_search.py search "Who led the round?" \
  --scope extracted_events.json --env-file .env

Both modes store EventEntity embeddings in the formal event_entity_vectors collection—never as relation-database bytes. One vector record is keyed by the event_entity.id and carries event_id, entity_id, source metadata, description, status, and timestamps. Normal uses the configured embedding dimension; Lite stores and queries the first 128 float values. The physical column is provider-native (dense_vector, vector(N), OceanBase VECTOR(N), or a fixed-size Arrow float32 list in LanceDB).

Configuration

Two ways to configure — pick one

zleap-sag always builds a single EngineConfig. Provide it by parameter injection or environment variables — they are alternatives, not layered.

Does setting env vars remove the need to pass keys? Only if you call EngineConfig.from_env(). Plain EngineConfig(...) never reads the environment — you pass keys there explicitly. Don't mix the two: use from_env() or inject params.

Environment-variable path, minimal set (zero-infra):

export OPENAI_API_KEY=sk-...
export SAG_STORAGE_MODE=normal                     # required: normal | lite
export OPENAI_BASE_URL=https://your-gateway/v1     # optional; defaults to OpenAI
export LLM_MODEL=qwen3.6-flash
export EMBEDDING_MODEL=bge-large-en-v1.5

config = EngineConfig.from_env()  # or EngineConfig.from_env(env_file=".env")

Required vs optional

• Required: storage_mode (normal or lite), llm, and embedding.
• Optional (all default to the local stack): storage backend, data_dir (./.zleap), rerank, language, log_level.

Two gotchas

• Separate embedding endpoint? Set EmbeddingConfig(base_url=..., api_key=...) (or EMBEDDING_BASE_URL / EMBEDDING_API_KEY). If omitted, embedding reuses the LLM's key and URL.
• Fixed-dimension embedding model? Use EmbeddingConfig(model="BAAI/bge-m3", schema_dimensions=1024). Storage uses 1024 dimensions and validates returned vectors; the API request omits dimensions. Set request_dimensions only if the API supports and requires that parameter. The two settings are independent; omitting either leaves it unset. Environment equivalents are EMBEDDING_SCHEMA_DIMENSIONS and EMBEDDING_REQUEST_DIMENSIONS (empty, none, or null means unset). Migration: dimensions / EMBEDDING_DIMENSIONS have been removed and are rejected. To retain the old combined behavior, explicitly set both new values to the old dimension.

Full variable reference (storage, backends, rerank): .env.example.

Custom entity types (optional)

Extraction keeps only entities whose type is defined. A generic set (person, organization, location, product, event, time, …) is seeded automatically. To add domain types, declare them in config — they are seeded on schema init, idempotently: existing types are skipped, only new ones are added.

config = EngineConfig(
    storage_mode="normal",
    entity_types=[
        "contract",
        "invoice",
        "party",
    ],  # str, or EntityTypeConfig(type=..., description=...)
    llm=LLMConfig(...),
    embedding=EmbeddingConfig(...),
)

Storage backends

Switch backends by changing EngineConfig only — the ingest/extract/search code is identical. Spin up local backends for testing with make up (docker compose).

Initialization

• Local SQLite — start() creates the schema and seeds default entity types automatically. Nothing to call.
• Production backends — call await engine.init_schema() before start() once. It only creates completely missing tables in the selected mode and never alters or migrates an existing table.
• Production mutations — multi-worker deployments must use process_source, delete_source, and delete_data_source with caller-owned OperationContext values. Direct ingest → extract calls are stage-level APIs for local workflows and custom orchestration; they do not provide the durable fence, lease, checkpoint, generation switch, or exact replay contract.

Example — MySQL + Elasticsearch storage wiring

from zleap.sag import DataEngine
from zleap.sag.pipeline import SearchOptions, SearchRequest, SearchScope
from zleap.sag.config import (
    EmbeddingConfig,
    ElasticsearchVectorConfig,
    EngineConfig,
    LLMConfig,
    RelationalConfig,
)

config = EngineConfig(
    storage_mode="normal",
    relational=RelationalConfig(
        provider="mysql", host="localhost", user="sag2", password="sag2", database="sag2"
    ),
    vector=ElasticsearchVectorConfig(hosts=["http://localhost:9200"]),
    llm=LLMConfig(api_key="sk-...", base_url="https://your-gateway/v1", model="qwen3.6-flash"),
    embedding=EmbeddingConfig(model="bge-large-en-v1.5"),
)

engine = DataEngine(config)
await engine.init_schema()  # once, before start(), for production backends
async with engine:
    # Stage-level example. A production worker uses process_source(), shown below.
    chunk_ref = await engine.ingest("your_document.md")
    event_ref = await engine.extract(chunk_ref)
    result = await engine.search(
        SearchRequest(
            query="Who founded Acme?",
            scope=SearchScope(data_source_ids=(event_ref.data_source_id,)),
            options=SearchOptions(strategy="full_expand", top_k=5, return_type="event"),
        )
    )

Cancelling stage-level extraction

Pass a request-local CancellationToken to engine.extract(..., cancellation=token) and call token.cancel() from the pause handler. The same token reaches the ExtractAdapter and built-in extractor. Cancellation stops and drains active and queued generation tasks, skips the batch commit, and raises PipelineCancelledError with an extract/cancelled observer event. A commit that has already started is allowed to finish so its atomic write/rollback boundary is preserved. Cancelling a client request does not guarantee that a remote LLM provider stops server-side work.

Custom ExtractAdapter.extract implementations must accept the keyword-only cancellation: CancellationToken | None = None argument and honor it during generation, for example with await cancellation.run(generate_batch). Raise asyncio.CancelledError when token-driven cancellation interrupts generation.

To verify with real LLM and embedding requests, configure the root .env as for integration tests and run from the repository root:

.venv/bin/python -m pytest packages/sag/tests/integration/test_extraction_cancellation_live.py -q -s

This test uses synthetic text and temporary SQLite/LanceDB storage. It verifies successful extraction, cancellation after two HTTP request bodies have been sent, preservation of the previous events, and successful extraction with a fresh token. It writes a result.json report under the pytest temporary directory printed at completion. Missing model configuration skips the test; running it consumes API usage.

Production mutation lifecycle

from hashlib import sha256

from zleap.sag.operations import OperationContext, ProcessSourceRequest
from zleap.sag.pipeline import SourceDescriptor, TextSource

data_source_id = "11111111-1111-1111-1111-111111111111"
source_id = "22222222-2222-2222-2222-222222222222"
markdown = "# Acme\nAcme was founded by Jane."

request = ProcessSourceRequest(
    context=OperationContext(
        operation_id="publish-job-42-attempt-1",
        idempotency_key="publish-job-42",
        request_digest=sha256(markdown.encode()).hexdigest(),
        fence_scope=data_source_id,
        fence_token=7,
        owner_id="knowledge-worker-3",
    ),
    source=TextSource(
        text=markdown,
        descriptor=SourceDescriptor(
            data_source_id=data_source_id,
            source_id=source_id,
            source_type="article",
        ),
    ),
)

async with DataEngine(config, data_source_id=data_source_id) as engine:
    result = await engine.process_source(request)
    if result.status == "failed":
        raise RuntimeError(f"{result.failure_code}: retryable={result.retryable}")
    # An uncertain caller response is recovered with get_operation_status(operation_id)
    # or by replaying the exact same request.

Resume extraction within the same durable task

process_source() persists each completed Chunk's extraction outcome in the existing SAGOperation.work_payload JSON. No database table or column migration is required. Exact replay of the same request, including after an Engine/process restart, reuses successful results (including valid empty results) and retries unfinished or retryable failed Chunks. Input content, model settings, prompts and entity contracts are fingerprinted; changes reject reuse. Chunk records are unchanged. Results stay invisible until the complete generation is committed and published.

from zleap.sag.pipeline import CancellationToken

cancel = CancellationToken()
# The application's pause action calls cancel.cancel(). Cancellation propagates
# as asyncio.CancelledError, preserves durable progress and releases the lease.
result = await engine.process_source(request, cancellation=cancel)

state = await engine.get_operation_status(request.context.operation_id)
progress = state.extraction_progress
# progress contains resume_token (= operation_id), succeeded_chunk_ids,
# failed_chunk_ids, pending_chunk_ids, retryable_chunk_ids and token_usage.

# After a pause/retryable failure, replay the same request with a fresh token.
result = await engine.process_source(request, cancellation=CancellationToken())

extraction_progress exposes counts/IDs, not document or event payloads. Its generation_attempts counts completed Chunk attempts across deliveries; token_usage accumulates provider-reported prompt/completion/total tokens from structured extraction calls, including repair retries and parent summaries. Cache reuse adds no tokens. Missing usage, failed/cancelled deliveries, or crash recovery set complete=False: totals are known usage, not a complete billing claim. Custom LLM adapters bypassing the built-in structured client cannot supply complete usage. Transport retries without returned usage cannot be measured exactly.

Checkpoint updates are serialized and merge the latest JSON under the task lease. An in-flight checkpoint transaction finishes before cancellation returns. A model response lost before its result checkpoint may be generated again; model calls are not exactly-once. Checkpoint write failures abort preparation and remain retryable. Per-Chunk progress remains at stage indexed; only the completed batch advances to extract_prepared. Parent summaries may be regenerated until that whole-batch checkpoint exists; commit retries after it reuse the prepared batch.

This is same-task recovery via process_source, not incremental publication or a resume feature for plain extract(). Existing terminal error policies and namespace fences still apply: non-retryable failures follow cleanup, and a superseding operation can invalidate the old task. Do not create a new operation or raise its fence to resume. Large tasks rewrite growing JSON values; this version does not add batching or an external result store.

Run the live model/embedding resume checks with the repository's .env credentials:

.venv/bin/python -m pytest packages/sag/tests/integration/test_extraction_resume_live.py -q -s

These use isolated SQLite/LanceDB storage, checkpoint the first Chunk, then cancel the next request after its HTTP body has been sent. Both a cancellation token and direct task cancellation are covered. A new Engine reopens the same storage and must reuse that first Chunk, finish the remaining Chunks and parent summary, and replay the completed operation without new model calls. Each case writes a result.json report. These checks do not simulate SIGKILL or frontend button clicks.

Example — physical single database

# One PostgreSQL for relational + vector (pip install "zleap-sag[postgres]")
from zleap.sag.config import PgVectorConfig, PostgresConnectionConfig

config = EngineConfig(
    storage_mode="normal",
    relational=RelationalConfig(
        provider="postgres",
        host="localhost",
        port=5432,
        user="sag2",
        password="sag2",
        database="sag2",
    ),
    vector=PgVectorConfig(
        connection=PostgresConnectionConfig(
            host="localhost", port=5432, user="sag2", password="sag2", database="sag2"
        )
    ),
    llm=LLMConfig(...),
    embedding=EmbeddingConfig(...),
)

# One OceanBase for SQL + vector (pip install "zleap-sag[mysql]", OceanBase ≥ 4.3.3)
from zleap.sag.config import OceanBaseConnectionConfig, OceanBaseVectorConfig

config = EngineConfig(
    storage_mode="lite",
    relational=RelationalConfig(
        provider="oceanbase", host="localhost", port=2881, user="root", password="", database="sag2"
    ),
    vector=OceanBaseVectorConfig(
        connection=OceanBaseConnectionConfig(
            host="localhost", port=2881, user="root", password="", database="sag2"
        )
    ),
    llm=LLMConfig(...),
    embedding=EmbeddingConfig(...),
)

OceanBase ANN indexes need the tenant setting ob_vector_memory_limit_percentage > 0; if unset, the engine falls back to exact vector search automatically.

Matching connection values above mean “same physical database”; they do not trigger runtime reuse. To deploy PostgreSQL + Elasticsearch or OceanBase + Elasticsearch, keep the relational config and replace only vector with ElasticsearchVectorConfig(...).

How it works

The engine exposes five independent stages. Each output is passed explicitly to the next stage:

• parse(SourceInput, ParseOptions) → ParsedSource; built-in Markdown, HTML, and plain text parsing, plus automatic CSV/XLSX → Markdown conversion with the [tables] extra.
• chunk(ParsedSource, ChunkOptions) → deterministic in-memory ChunkSet.
• index(ChunkSet, SourceDescriptor, IndexOptions) → ChunkSetRef.
• extract(ChunkSetRef | PersistedChunkSelector, ExtractionOptions) → EventSetRef.
• load_events(EventSetRef) → portable EventDetail records scoped to that exact source.
• search(SearchRequest) → SearchResult; scope always contains non-empty data_source_ids, optional source_ids/source_types/creator_ids, and optional timezone-aware time bounds. creator_ids requires an injected SearchScopeResolver.

Time bounds filter stored event start_time / end_time; chunk searches use their associated events to apply this restriction. Built-in extraction currently leaves event times unset, including parent summaries, rather than deriving them from document creation times, message timestamps, or extraction time. Events with both times unset do not match a scope with either time bound. Omit both bounds to search these events; the bounds do not select documents or messages by their source timestamps.

For Markdown documents, ChunkOptions(strategy="heading_strict") creates one chunk for every non-empty heading block. Consecutive headings with the same text remain separate, and a heading block is never split again by sentences or max_tokens; chunk text preserves the source block after trimming its outer whitespace. This matches SAG-Benchmark's heading_strict corpus semantics.

In standard and overlap modes, valid Markdown tables are recognized as TABLE blocks. Rows stay intact, every split table chunk repeats the complete header, and table chunks never mix with neighboring prose or a different table. ArticleSection keeps one header evidence record that all chunks from the same table reference. CSV and XLSX files enter this exact path after MarkItDown normalization; parse() remains storage-free and works before engine.start().

ingest() remains a Parse → Chunking → Index convenience method, but does not save its result on the engine. The no-argument extract() and implicit-scope search() forms are removed. SearchOptions.strategy is required. Old names multi/multi1/hopllm/multi_es are no longer accepted by the public dispatcher.

Search profiles and typed overrides

Multi-hop expansion is tuned through graph.max_hops on the same strategy (e.g. full_expand + graph=GraphSearchOptions(max_hops=2)); there is no separate multi-hop strategy name.

The built-in search runtime has dedicated Vector, Atomic, and Production Executors. Production performs direct Event-vector plus Entity lexical/vector recall; both precise and fast run bounded one-hop graph expansion. The precise variants (pruned_expand_llm / pruned_expand_rerank) rank by LLM or an external rerank model respectively; the fast variant (pruned_expand_rff) uses deterministic RRF, and full_expand also runs on Production. Atomic invokes its concrete searcher directly; there is no second SAGSearcher dispatcher. Hosts may replace the engine-local Production Executor without importing host code into this package. The Atomic executor honors score thresholds, candidate limits, graph hops, and selection.enabled; unsupported query rewrite, custom Selection prompts, rationale output, or non-vector ranking fail explicitly instead of being silently ignored.

Profiles provide stable defaults. Requests may override supported behavior through typed sub-options; there is no strategy_options dictionary:

from zleap.sag.pipeline import (
    GraphSearchOptions,
    RankingOptions,
    SearchOptions,
    SearchOutputOptions,
)

options = SearchOptions(
    strategy="pruned_expand_llm",
    top_k=10,
    graph=GraphSearchOptions(enabled=True, max_hops=2),
    ranking=RankingOptions(rerank_threshold=0.6),
    output=SearchOutputOptions(return_graph=True),
)

graph.enabled controls retrieval expansion; output.return_graph only controls whether the computed graph is returned. Search does not filter content by default; an optional EngineConfig.guard provider can filter queries and returned search results.

Extract options and atomic batches

Extract has one public strategy, sag_extract, and one response schema. Customize behavior with typed options instead of an unvalidated strategy_options dictionary:

from zleap.sag.pipeline import ExtractionLimits, ExtractionOptions

event_ref = await engine.extract(
    chunk_ref,
    ExtractionOptions(
        background="Preserve monetary values in their original currency.",
        guidance_rules=("Ignore headers and footers.",),
        limits=ExtractionLimits(
            max_events_per_chunk=20,
            min_entities_per_event=1,
            max_entities_per_event=20,
        ),
        max_retries=5,
        enable_parent_summary=True,
    ),
)

Each Chunk gets one initial generation plus up to five validation-repair generations. Transport retries remain controlled by LLMConfig.max_retries. The built-in Extract adapter commits atomically per source: if any Chunk exhausts its retries, or an enabled parent summary fails, no new events or vectors are written and the previous source snapshot remains active. Failure details are available on ExtractionBatchFailure.failures; EventSetRef represents only a fully committed success.

For ordinary rich extraction, out-of-range integer references are removed before local validation (for 128 inputs, [1, 129] becomes [1]). Valid IDs keep their original order; duplicates and invalid types are not repaired. An all-invalid list becomes empty and fails the required-reference check, following the configured retry/failure or drop_event policy. All other event constraints still apply. Parent summaries still require empty references and do not use this repair.

To retain valid events from a mixed response, set ExtractionOptions(on_contract_violation="drop_event") (default: "raise"). Event schema/semantic violations are dropped without regenerating the Chunk; entity types are filtered before truncation and the minimum entity count check. Valid children of a dropped event are promoted in source order. Depth limits use the original tree depth, and event-count limits still apply to all survivors. Invalid envelopes/JSON and parent-summary requests remain strict and repairable. If all events are dropped, extraction succeeds with no events; like an ordinary empty result, this can replace the source's previous snapshot with an empty one. In drop mode, events_seen, events_kept, events_dropped, and events_promoted in extraction stats describe successful validation passes (before later parsing, grounding checks, or persistence); validation_issue_counts includes drop reasons. Bounded warning samples contain original event paths without event text. A zero events_seen distinguishes an originally empty response from an all-dropped one. The validator also accepts an ExtractionValidationReport via report= for auditing. Built-in LLM clients still send the full generation schema, while deferring local schema validation to the event validator. Custom LLM adapters used with ordinary rich extraction or drop_event must accept validate_response_schema=False, consume it locally (not forward it to the provider), and return the parsed response without whole-response rejection. Source commits remain atomic over the retained results; this option does not change storage transaction or parent-summary failure semantics.

enable_article_summary=True is valid only for text/article sources and requires pip install "zleap-sag[summary]". Sumy is imported lazily, never downloads NLTK data at runtime, and the extractive summary is cached by source_version.

DataEngine is an async context manager: async with DataEngine(config) as engine runs start() on enter and aclose() on exit. Stage contracts are imported from zleap.sag.pipeline; durable production contracts come from zleap.sag.operations, zleap.sag.queries, zleap.sag.records, and zleap.sag.maintenance. Parse and Chunking need no started storage; use parse(TextSource(...)) followed by chunk(parsed), or chunk_text(...). See MIGRATING_PIPELINE.md for breaking API changes.

Notes

• Multiple DataEngine instances may coexist in one process. Each owns its relational and vector resources; closing one does not close another instance's pools or client.
• pruned_expand_llm / pruned_expand_rerank / pruned_expand_rff resolve to the registered production Executor. The LLM and rerank variants need real lexical retrieval (LanceDB or Elasticsearch). Algorithms are registered as SearchExecutorSpec values before start() and the engine-local catalog is frozen after startup. Use register_search_algorithm(...) for a complete versioned contract; Missing capabilities or executors fail explicitly unless the request declares a valid fallback_strategy.
• All engine errors derive from SagError — catch it at the boundary.

Links

Examples · API Reference · Changelog · Contributing · Config reference