This commit is contained in:
Timothy J. Baek
2024-10-03 20:51:21 -07:00
parent 2fc07fd6a2
commit 57360b7a61
2 changed files with 6 additions and 7 deletions

View File

@@ -68,19 +68,18 @@ class ChromaClient:
def query(
self, collection_name: str, filter: dict, limit: int = 1
) -> Optional[SearchResult]:
) -> Optional[GetResult]:
# Query the items from the collection based on the filter.
collection = self.client.get_collection(name=collection_name)
if collection:
result = collection.query(
result = collection.get(
where=filter,
n_results=limit,
limit=limit,
)
return SearchResult(
return GetResult(
**{
"ids": result["ids"],
"distances": result["distances"],
"documents": result["documents"],
"metadatas": result["metadatas"],
}