metadata.py

from pydantic import BaseModel
from typing import Literal, Union, List, Dict


AllowedOps = Literal["$in", "$eq", "$gt", "$gte", "$lt", "$lte", "$ne", "$nin"]


class MetadataFilter(BaseModel):
    metadata_field: str
    metadata_search_operator: AllowedOps
    metadata_value: Union[str, int, float, List[Union[str, int, float]]]


class MetadataWhereClause(BaseModel):
    filters: List[MetadataFilter]

    def to_chroma_where(self) -> Dict:
        """Convert list of MetadataFilter into Chroma-compatible where clause with AND logic."""
        if not self.filters:
            return {}
        if len(self.filters) == 1:
            f = self.filters[0]
            return {f.metadata_field: {f.metadata_search_operator: f.metadata_value}}
        return {
            "$and": [
                {f.metadata_field: {f.metadata_search_operator: f.metadata_value}}
                for f in self.filters
            ]
        }