metadata.py

from pydantic import BaseModel
from typing import Literal, Optional, 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]
#     conditional_operator: Literal["$and", "$or"] = "$and"

#     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 {
#             self.conditional_operator: [
#                 {f.metadata_field: {f.metadata_search_operator: f.metadata_value}}
#                 for f in self.filters
#             ]
#         }

class MetadataWhereClause(BaseModel):
    filters: Optional[List["MetadataFilter"]] = None
    groups: Optional[List["MetadataWhereClause"]] = None
    conditional_operator: Literal["$and", "$or"] = "$and"

    def to_chroma_where(self) -> Dict:
        parts = []

        # Handle direct filters
        if self.filters:
            for f in self.filters:
                parts.append({f.metadata_field: {f.metadata_search_operator: f.metadata_value}})

        # Handle nested groups
        if self.groups:
            for g in self.groups:
                parts.append(g.to_chroma_where())

        if not parts:
            return {}

        if len(parts) == 1:
            return parts[0]

        # More than one part → wrap with conditional operator
        return {self.conditional_operator: parts}