diff --git a/examples/test_chromadb.py b/examples/test_chromadb.py index 5293f05d..99090a6d 100644 --- a/examples/test_chromadb.py +++ b/examples/test_chromadb.py @@ -17,7 +17,9 @@ if not os.path.exists(WORKING_DIR): # ChromaDB Configuration CHROMADB_USE_LOCAL_PERSISTENT = False # Local PersistentClient Configuration -CHROMADB_LOCAL_PATH = os.environ.get("CHROMADB_LOCAL_PATH", os.path.join(WORKING_DIR, "chroma_data")) +CHROMADB_LOCAL_PATH = os.environ.get( + "CHROMADB_LOCAL_PATH", os.path.join(WORKING_DIR, "chroma_data") +) # Remote HttpClient Configuration CHROMADB_HOST = os.environ.get("CHROMADB_HOST", "localhost") CHROMADB_PORT = int(os.environ.get("CHROMADB_PORT", 8000)) diff --git a/lightrag/kg/chroma_impl.py b/lightrag/kg/chroma_impl.py index 7b7642d6..cb3b59f1 100644 --- a/lightrag/kg/chroma_impl.py +++ b/lightrag/kg/chroma_impl.py @@ -67,7 +67,9 @@ class ChromaVectorDBStorage(BaseVectorStorage): if "token_authn" in auth_provider: headers = { - config.get("auth_header_name", "X-Chroma-Token"): auth_credentials + config.get( + "auth_header_name", "X-Chroma-Token" + ): auth_credentials } elif "basic_authn" in auth_provider: auth_credentials = config.get("auth_credentials", "admin:admin") @@ -154,7 +156,9 @@ class ChromaVectorDBStorage(BaseVectorStorage): embedding = await self.embedding_func([query]) results = self._collection.query( - query_embeddings=embedding.tolist() if not isinstance(embedding, list) else embedding, + query_embeddings=embedding.tolist() + if not isinstance(embedding, list) + else embedding, n_results=top_k * 2, # Request more results to allow for filtering include=["metadatas", "distances", "documents"], ) diff --git a/lightrag/types.py b/lightrag/types.py index 35036d2f..5e3d2948 100644 --- a/lightrag/types.py +++ b/lightrag/types.py @@ -1,7 +1,7 @@ from __future__ import annotations from pydantic import BaseModel -from typing import List, Dict, Any, Optional +from typing import Any, Optional class GPTKeywordExtractionFormat(BaseModel):