| import argparse |
| import os |
| from typing import List, Tuple |
|
|
| import chromadb |
| from chromadb.config import Settings |
|
|
|
|
| def _parse_args() -> argparse.Namespace: |
| parser = argparse.ArgumentParser(description="ChromaDB client demo") |
| parser.add_argument( |
| "--collection", |
| default=os.getenv("CHROMA_COLLECTION", "knowledge_base"), |
| help="Collection name", |
| ) |
| parser.add_argument( |
| "--query", |
| default=os.getenv("CHROMA_QUERY", "Tell me about vector stores"), |
| help="Query text", |
| ) |
| parser.add_argument( |
| "--n-results", |
| type=int, |
| default=int(os.getenv("CHROMA_N_RESULTS", "2")), |
| help="Number of results to return", |
| ) |
| return parser.parse_args() |
|
|
|
|
| def _build_client() -> chromadb.Client: |
| host = os.getenv("CHROMA_HOST", "maxxcarl-emb1024.hf.space") |
| port = int(os.getenv("CHROMA_PORT", "443")) |
| ssl = os.getenv("CHROMA_SSL", "true").lower() in {"1", "true", "yes", "y"} |
| http_path = os.getenv("CHROMA_HTTP_PATH", "/chroma") |
|
|
| settings = Settings( |
| chroma_server_host=host, |
| chroma_server_http_port=port, |
| chroma_server_ssl=ssl, |
| chroma_server_http_path=http_path, |
| ) |
| return chromadb.Client(settings) |
|
|
|
|
| def _seed_collection(collection: chromadb.Collection) -> Tuple[List[str], List[str]]: |
| 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) |
| return ids, documents |
|
|
|
|
| def main() -> None: |
| args = _parse_args() |
| client = _build_client() |
|
|
| print("Connecting to Chroma...") |
| collection = client.get_or_create_collection(name=args.collection) |
| print(f"Using collection: {args.collection}") |
|
|
| if collection.count() == 0: |
| print("Seeding collection with sample documents...") |
| ids, _ = _seed_collection(collection) |
| print(f"Added {len(ids)} documents.") |
| else: |
| print(f"Collection already has {collection.count()} documents.") |
|
|
| print(f"Query: {args.query}") |
| results = collection.query(query_texts=[args.query], n_results=args.n_results) |
|
|
| 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() |
|
|
