codekingpro/portable-devtools

1"""Utility functions for managing collection statistics.
2
3This module provides standalone functions for enabling, disabling, and retrieving
4statistics for ChromaDB collections. These functions work with the attached function
5system to automatically compute metadata value frequencies.
6
7Example:
8    >>> from chromadb.utils.statistics import attach_statistics_function, get_statistics
9    >>> import chromadb
10    >>>
11    >>> client = chromadb.Client()
12    >>> collection = client.get_or_create_collection("my_collection")
13    >>>
14    >>> # Attach statistics function with output collection name
15    >>> attach_statistics_function(collection, "my_collection_statistics")
16    >>>
17    >>> # Add some data
18    >>> collection.add(
19    ...     ids=["id1", "id2"],
20    ...     documents=["doc1", "doc2"],
21    ...     metadatas=[{"category": "A"}, {"category": "B"}]
22    ... )
23    >>>
24    >>> # Get statistics from the named output collection
25    >>> stats = get_statistics(collection, "my_collection_statistics")
26    >>> print(stats)
27"""
28
29from typing import TYPE_CHECKING, Optional, Dict, Any, cast, Tuple
30from collections import defaultdict
31from chromadb.api.types import OneOrMany, Where, maybe_cast_one_to_many
32from chromadb.api.functions import STATISTICS_FUNCTION
33
34if TYPE_CHECKING:
35    from chromadb.api.models.Collection import Collection
36    from chromadb.api.models.AttachedFunction import AttachedFunction
37
38
39def get_statistics_fn_name(collection: "Collection") -> str:
40    """Generate the default name for the statistics attached function.
41
42    Args:
43        collection: The collection to generate the name for
44
45    Returns:
46        str: The statistics function name
47    """
48    return f"{collection.name}_stats"
49
50
51def attach_statistics_function(
52    collection: "Collection", stats_collection_name: str
53) -> Tuple["AttachedFunction", bool]:
54    """Attach statistics collection function to a collection.
55
56    This attaches the statistics function which will automatically compute
57    and update metadata value frequencies whenever records are added, updated,
58    or deleted.
59
60    Args:
61        collection: The collection to enable statistics for
62        stats_collection_name: Name of the collection where statistics will be stored.
63
64    Returns:
65        Tuple of (AttachedFunction, created) where created is True if newly created,
66        False if already existed (idempotent request)
67
68    Example:
69        >>> attached_fn, created = attach_statistics_function(collection, "my_collection_statistics")
70        >>> if created:
71        ...     print("Statistics function newly attached")
72        >>> collection.add(ids=["id1"], documents=["doc1"], metadatas=[{"key": "value"}])
73        >>> # Statistics are automatically computed
74        >>> stats = get_statistics(collection, "my_collection_statistics")
75    """
76    return collection.attach_function(
77        function=STATISTICS_FUNCTION,
78        name=get_statistics_fn_name(collection),
79        output_collection=stats_collection_name,
80        params=None,
81    )
82
83
84def get_statistics_fn(collection: "Collection") -> "AttachedFunction":
85    """Get the statistics attached function for a collection.
86
87    Args:
88        collection: The collection to get the statistics function for
89
90    Returns:
91        AttachedFunction: The statistics function
92
93    Raises:
94        NotFoundError: If statistics are not enabled
95        AssertionError: If the attached function is not a statistics function
96    """
97    af = collection.get_attached_function(get_statistics_fn_name(collection))
98    assert (
99        af.function_name == "statistics"
100    ), "Attached function is not a statistics function"
101    return af
102
103
104def detach_statistics_function(
105    collection: "Collection", delete_stats_collection: bool = False
106) -> bool:
107    """Detach statistics collection function from a collection.
108
109    Args:
110        collection: The collection to disable statistics for
111        delete_stats_collection: If True, also delete the statistics output collection.
112                                  Defaults to False.
113
114    Returns:
115        bool: True if successful
116
117    Example:
118        >>> detach_statistics_function(collection, delete_stats_collection=True)
119    """
120    attached_fn = get_statistics_fn(collection)
121    return collection.detach_function(
122        attached_fn.name, delete_output_collection=delete_stats_collection
123    )
124
125
126def get_statistics(
127    collection: "Collection",
128    stats_collection_name: str,
129    keys: Optional[OneOrMany[str]] = None,
130) -> Dict[str, Any]:
131    """Get the current statistics for a collection.
132
133    Statistics include frequency counts for all metadata key-value pairs,
134    as well as a summary with the total record count.
135
136    Args:
137        collection: The collection to get statistics for
138        stats_collection_name: Name of the statistics collection to read from.
139        keys: Optional metadata key(s) to filter statistics for. Can be a single key
140              string or a list of keys. If provided, only returns statistics for
141              those specific keys.
142
143    Returns:
144        Dict[str, Any]: A dictionary with the structure:
145            {
146                "statistics": {
147                    "key1": {
148                        "value1": {"count": count, ...},
149                        "value2": {"count": count, ...}
150                    },
151                    "key2": {...},
152                    ...
153                },
154                "summary": {
155                    "total_count": count
156                }
157            }
158
159    Example:
160        >>> attach_statistics_function(collection, "my_collection_statistics")
161        >>> collection.add(
162        ...     ids=["id1", "id2"],
163        ...     documents=["doc1", "doc2"],
164        ...     metadatas=[{"category": "A", "score": 10}, {"category": "B", "score": 10}]
165        ... )
166        >>> # Wait for statistics to be computed
167        >>> stats = get_statistics(collection, "my_collection_statistics")
168        >>> print(stats)
169        {
170            "statistics": {
171                "category": {
172                    "A": {"count": 1},
173                    "B": {"count": 1}
174                },
175                "score": {
176                    "10": {"count": 2}
177                }
178            },
179            "summary": {
180                "total_count": 2
181            }
182        }
183
184    Raises:
185        ValueError: If more than 30 keys are provided in the keys filter.
186    """
187    # Normalize keys to list
188    keys_list = maybe_cast_one_to_many(keys)
189
190    # Validate keys count to avoid issues with large $in queries
191    MAX_KEYS = 30
192    if keys_list is not None and len(keys_list) > MAX_KEYS:
193        raise ValueError(
194            f"Too many keys provided: {len(keys_list)}. "
195            f"Maximum allowed is {MAX_KEYS} keys per request. "
196            "Consider calling get_statistics multiple times with smaller key batches."
197        )
198
199    # Import here to avoid circular dependency
200    from chromadb.api.models.Collection import Collection
201
202    # Get the statistics output collection model from the server
203    stats_collection_model = collection._client.get_collection(
204        name=stats_collection_name,
205        tenant=collection.tenant,
206        database=collection.database,
207    )
208
209    # Wrap it in a Collection object to access get/query methods
210    stats_collection = Collection(
211        client=collection._client,
212        model=stats_collection_model,
213        embedding_function=None,  # Statistics collections don't need embedding functions
214        data_loader=None,
215    )
216
217    # Get all statistics records by paginating through the stats collection
218    stats: Dict[str, Dict[str, Dict[str, int]]] = defaultdict(lambda: defaultdict(dict))
219    summary: Dict[str, Any] = {}
220
221    offset = 0
222    # When filtering by keys, also include "summary" entries to get total_count
223    where_filter: Optional[Where] = (
224        cast(Where, {"key": {"$in": keys_list + ["summary"]}})
225        if keys_list is not None
226        else None
227    )
228
229    while True:
230        page = stats_collection.get(
231            include=["metadatas"], offset=offset, where=where_filter
232        )
233
234        metadatas = page.get("metadatas") or []
235        if not metadatas:
236            break
237
238        for metadata in metadatas:
239            if metadata is None:
240                continue
241
242            meta_key = metadata.get("key")
243            value = metadata.get("value")
244            value_label = metadata.get("value_label")
245            value_type = metadata.get("type")
246            count = metadata.get("count")
247
248            if (
249                meta_key is not None
250                and value is not None
251                and value_type is not None
252                and count is not None
253            ):
254                if meta_key == "summary":
255                    if value == "total_count":
256                        summary["total_count"] = count
257                else:
258                    # Prioritize value_label if present, otherwise use value
259                    stats_key = value_label if value_label is not None else value
260                    assert isinstance(meta_key, str)
261                    assert isinstance(stats_key, str)
262                    assert isinstance(count, int)
263                    stats[meta_key][stats_key]["count"] = count
264
265        # Advance to next page using the actual number of items returned
266        offset += len(metadatas)
267
268    result = {"statistics": dict(stats)}
269    if summary:
270        result["summary"] = summary
271
272    return result
273