Skip to content

Commit

Permalink
fix: chroma examples error due to None parsing in flattens (#215)
Browse files Browse the repository at this point in the history
  • Loading branch information
mhordynski authored Nov 29, 2024
1 parent ffbcc68 commit 7f82f95
Show file tree
Hide file tree
Showing 3 changed files with 14 additions and 9 deletions.
4 changes: 2 additions & 2 deletions examples/document-search/chroma_otel.py
Original file line number Diff line number Diff line change
Expand Up @@ -51,7 +51,7 @@

import asyncio

from chromadb import PersistentClient
from chromadb import EphemeralClient
from opentelemetry import trace
from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import OTLPSpanExporter
from opentelemetry.sdk.resources import SERVICE_NAME, Resource
Expand Down Expand Up @@ -102,7 +102,7 @@ async def main() -> None:
model="text-embedding-3-small",
)
vector_store = ChromaVectorStore(
client=PersistentClient("./chroma"),
client=EphemeralClient(),
index_name="jokes",
)
document_search = DocumentSearch(
Expand Down
Original file line number Diff line number Diff line change
@@ -1,9 +1,9 @@
from typing import Any

SimpleTypes = str | int | float | bool | None

def flatten_dict(
input_dict: dict[str, Any], parent_key: str = "", sep: str = "."
) -> dict[str, str | int | float | bool]:

def flatten_dict(input_dict: dict[str, Any], parent_key: str = "", sep: str = ".") -> dict[str, SimpleTypes]:
"""
Recursively flatten a nested dictionary and lists, converting non-primitive types to strings.
Expand All @@ -18,7 +18,7 @@ def flatten_dict(
Raises:
ValueError: If the dictionary cannot be safely flattened due to the presence of the separator in the dict key.
"""
items: dict[str, str | int | float | bool] = {}
items: dict[str, SimpleTypes] = {}
for k, v in input_dict.items():
if sep in k:
raise ValueError(f"Separator '{sep}' found in key '{parent_key}' Cannot flatten dictionary safely.")
Expand All @@ -36,9 +36,9 @@ def flatten_dict(
if isinstance(item, dict):
items = {**items, **flatten_dict(item, list_key, sep=sep)}
else:
items[list_key] = item if isinstance(item, str | int | float | bool) else str(item)
items[list_key] = item if isinstance(item, SimpleTypes) else str(item)
else:
items[new_key] = v if isinstance(v, str | int | float | bool) else str(v)
items[new_key] = v if isinstance(v, SimpleTypes) else str(v)

return items

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -80,7 +80,7 @@ async def store(self, entries: list[VectorStoreEntry]) -> None:
metadatas = [entry.metadata for entry in entries]

# Flatten metadata
flattened_metadatas = [flatten_dict(metadata) for metadata in metadatas]
flattened_metadatas = [self._flatten_metadata(metadata) for metadata in metadatas]

metadatas = (
flattened_metadatas
Expand Down Expand Up @@ -180,3 +180,8 @@ async def list(
)
for id, metadata, embedding, document in zip(ids, metadatas, embeddings, documents, strict=True)
]

@staticmethod
def _flatten_metadata(metadata: dict) -> dict:
"""Flattens the metadata dictionary. Removes any None values as they are not supported by ChromaDB."""
return {k: v for k, v in flatten_dict(metadata).items() if v is not None}

0 comments on commit 7f82f95

Please sign in to comment.