示例

真实场景的样例代码集中在两个位置:

• docs/examples.md —— RAG 管道、智能体记忆、各种索引、事务、批量插入、嵌入缓存。
• examples/ —— 可直接 python examples/<name>.py 运行的 Python 脚本。

精选片段

最小 RAG 循环

import numpy as np
from pistadb import PistaDB, Metric, Index

with PistaDB("rag.pst", dim=1536, metric=Metric.COSINE, index=Index.HNSW) as db:
    # 1. 把文本分块灌入索引(先调用你的嵌入模型)
    for chunk_id, vec, text in your_chunks:
        db.insert(chunk_id, vec, label=text[:255])

    # 2. 查询时检索
    q_vec = embed(user_question)
    hits = db.search(q_vec, k=5)

    # 3. 拼进 LLM Prompt
    context = "\n---\n".join(h.label for h in hits)
    answer = llm(f"Context:\n{context}\n\nQuestion: {user_question}")

事务(ACID 风格)

with db.transaction() as tx:
    tx.insert(1, vec1)
    tx.delete(99)
    tx.update(2, vec2)
    # 块内抛出任何异常都会触发完整回滚。

批量插入

# 多线程 ring-buffer 入库 —— 跑满 SSD 吞吐
db.batch_insert(ids, vectors, labels=labels, num_threads=8)

嵌入缓存

from pistadb import EmbeddingCache

cache = EmbeddingCache("openai_embed.pcc", dim=1536, capacity=1_000_000)
vec = cache.get_or_compute(text, lambda t: openai_embed(t))

一个持久化的 LRU 缓存,避免跨运行重复调用嵌入模型。

更多

包含 HNSW 调参、IVF 训练、SQ 内存压缩、ScaNN 两阶段搜索、智能体记忆、Milvus 兼容 schema API 在内的完整代码,见仓库中的 docs/examples.md。