用 ChromaDB 存向量、喂上下文：给 LLM 加一层可检索的记忆

pip install chromadb

import chromadb

# 内存模式——适合快速实验，进程结束数据消失
client = chromadb.Client()

# 持久化模式——数据写入指定目录，下次启动还在
# client = chromadb.PersistentClient(path="./chroma_data")

collection = client.get_or_create_collection(
    name="product_docs",
    metadata={"hnsw:space": "cosine"}  # 可选：指定距离度量，默认 L2
)

# 添加文档——ChromaDB 自动用内置模型生成嵌入
docs = [
    "AcmeChat 支持最多 500 人同时在线的视频会议，延迟低于 200ms。",
    "AcmeChat 的端到端加密基于 AES-256，密钥由会议发起者生成。",
    "AcmeDocs 提供实时协同编辑，支持 Markdown 和富文本，历史版本保留 90 天。",
    "AcmeDocs 的权限体系分三级：只读、评论、编辑，可按文件夹粒度配置。",
]

collection.add(
    documents=docs,
    ids=["doc1", "doc2", "doc3", "doc4"],       # 每条文档必须唯一
    metadatas=[                                  # 可选：附加元数据，用于过滤
        {"product": "AcmeChat", "feature": "video"},
        {"product": "AcmeChat", "feature": "security"},
        {"product": "AcmeDocs", "feature": "collab"},
        {"product": "AcmeDocs", "feature": "permission"},
    ]
)

print(f"已存入 {collection.count()} 条文档")

results = collection.query(
    query_texts=["AcmeChat 的安全性怎么样？"],
    n_results=2,                                # 返回最相关的 2 条
    where={"product": "AcmeChat"},              # 只在 AcmeChat 的文档里找
)

for doc, dist in zip(results["documents"][0], results["distances"][0]):
    print(f"[距离: {dist:.4f}] {doc}")

[距离: 0.32] AcmeChat 的端到端加密基于 AES-256，密钥由会议发起者生成。
[距离: 0.58] AcmeChat 支持最多 500 人同时在线的视频会议，延迟低于 200ms。

import chromadb
import openai

openai_client = openai.OpenAI()  # 自动读取 OPENAI_API_KEY 环境变量

def openai_embed(texts: list[str]) -> list[list[float]]:
    resp = openai_client.embeddings.create(
        model="text-embedding-3-small",
        input=texts,
    )
    return [item.embedding for item in resp.data]

chroma_client = chromadb.PersistentClient(path="./chroma_openai")

# 自定义嵌入函数——告诉 ChromaDB 用我们的函数而不是内置模型
from chromadb.utils.embedding_functions import EmbeddingFunction

class OpenAIEmbedding(EmbeddingFunction):
    def __call__(self, input: list[str]) -> list[list[float]]:
        return openai_embed(input)

ef = OpenAIEmbedding()

collection = chroma_client.get_or_create_collection(
    name="product_docs_openai",
    embedding_function=ef,
)

# 后续 add / query 和之前完全一样，ChromaDB 会自动调用 OpenAIEmbedding
collection.add(
    documents=docs,
    ids=["doc1", "doc2", "doc3", "doc4"],
    metadatas=[
        {"product": "AcmeChat", "feature": "video"},
        {"product": "AcmeChat", "feature": "security"},
        {"product": "AcmeDocs", "feature": "collab"},
        {"product": "AcmeDocs", "feature": "permission"},
    ]
)

results = collection.query(
    query_texts=["AcmeChat 的安全性怎么样？"],
    n_results=2,
    where={"product": "AcmeChat"},
)

for doc, dist in zip(results["documents"][0], results["distances"][0]):
    print(f"[距离: {dist:.4f}] {doc}")

import chromadb
import openai

# ---- 1. 初始化 ChromaDB ----
chroma_client = chromadb.PersistentClient(path="./chroma_data")
collection = chroma_client.get_collection(name="product_docs")

# ---- 2. 检索相关文档 ----
def retrieve_context(question: str, n: int = 3) -> str:
    results = collection.query(
        query_texts=[question],
        n_results=n,
    )
    # 把多条文档拼成一段上下文
    return "\n\n".join(results["documents"][0])

# ---- 3. 构造 Prompt 并调用 LLM ----
openai_client = openai.OpenAI()

def rag_answer(question: str) -> str:
    context = retrieve_context(question)
    prompt = f"""你是一个产品助手。请只根据以下上下文回答问题，如果上下文没有相关信息，就说"我不确定"。

上下文：
{context}

问题：{question}
回答："""

    resp = openai_client.chat.completions.create(
        model="gpt-4o-mini",
        messages=[{"role": "user", "content": prompt}],
        temperature=0.2,  # 低温度，减少幻觉
    )
    return resp.choices[0].message.content

# ---- 4. 测试 ----
print(rag_answer("AcmeChat 最多支持多少人同时在线？"))
print(rag_answer("AcmeDocs 能编辑 Excel 吗？"))

AcmeChat 支持最多 500 人同时在线的视频会议。
我不确定。

import chromadb

def chunk_text(text: str, chunk_size: int = 400, overlap: int = 50) -> list[str]:
    chunks = []
    start = 0
    while start < len(text):
        end = start + chunk_size
        chunks.append(text[start:end])
        start += chunk_size - overlap
    return chunks

# 读取长文档
with open("product_handbook.md", "r", encoding="utf-8") as f:
    full_text = f.read()

chunks = chunk_text(full_text)

client = chromadb.PersistentClient(path="./chroma_data")
collection = client.get_or_create_collection(name="handbook")

collection.add(
    documents=chunks,
    ids=[f"chunk_{i}" for i in range(len(chunks))],
    metadatas=[{"source": "product_handbook.md", "chunk_index": i} for i in range(len(chunks))],
)

print(f"已分段存入 {len(chunks)} 个片段")