File size: 2,250 Bytes
ad469c9 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 | import argparse
from pathlib import Path
from typing import List
import chromadb
def _parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(description="Local ChromaDB persistence demo")
parser.add_argument(
"--path",
default="chroma_data",
help="Local persistence directory",
)
parser.add_argument(
"--collection",
default="knowledge_base",
help="Collection name",
)
parser.add_argument(
"--query",
default="Tell me about vector stores",
help="Query text",
)
return parser.parse_args()
def _seed_collection(collection: chromadb.Collection) -> None:
documents = [
"Chroma is a lightweight, open-source vector database built for AI.",
"Python is a high-level programming language used extensively in data science.",
"The celestial body closest to Earth is the Moon.",
]
metadatas = [
{"category": "tech", "source": "docs"},
{"category": "tech", "source": "wiki"},
{"category": "science", "source": "space-facts"},
]
ids = ["doc1", "doc2", "doc3"]
collection.add(documents=documents, metadatas=metadatas, ids=ids)
def main() -> None:
args = _parse_args()
persist_path = Path(args.path).resolve()
persist_path.mkdir(parents=True, exist_ok=True)
print(f"Using local Chroma persistence at: {persist_path}")
client = chromadb.PersistentClient(path=str(persist_path))
collection = client.get_or_create_collection(name=args.collection)
if collection.count() == 0:
print("Seeding collection with sample documents...")
_seed_collection(collection)
print(f"Collection '{args.collection}' has {collection.count()} documents.")
results = collection.query(query_texts=[args.query], n_results=2)
print("\n--- Search Results ---")
for doc, meta, distance in zip(
results["documents"][0],
results["metadatas"][0],
results["distances"][0],
):
print(f"Matched Document: {doc}")
print(f"Metadata: {meta}")
print(f"Distance Score (Lower is better): {distance:.4f}")
print()
print("----------------------")
if __name__ == "__main__":
main()
|