chromadb/telemetry/product/__init__.py

from abc import abstractmethod
import os
from typing import ClassVar, Dict, Any
import uuid
import chromadb
from chromadb.config import Component
from pathlib import Path
from enum import Enum

TELEMETRY_WHITELISTED_SETTINGS = [
    "chroma_api_impl",
    "is_persistent",
    "chroma_server_ssl_enabled",
]


class ServerContext(Enum):
    NONE = "None"
    FASTAPI = "FastAPI"


class ProductTelemetryEvent:
    max_batch_size: ClassVar[int] = 1
    batch_size: int

    def __init__(self, batch_size: int = 1):
        self.batch_size = batch_size

    @property
    def properties(self) -> Dict[str, Any]:
        return self.__dict__

    @property
    def name(self) -> str:
        return self.__class__.__name__

    # A batch key is used to determine whether two events can be batched together.
    # If a TelemetryEvent's max_batch_size > 1, batch_key() and batch() MUST be
    # implemented.
    # Otherwise they are ignored.
    @property
    def batch_key(self) -> str:
        return self.name

    def batch(self, other: "ProductTelemetryEvent") -> "ProductTelemetryEvent":
        raise NotImplementedError


class ProductTelemetryClient(Component):
    USER_ID_PATH = str(Path.home() / ".cache" / "chroma" / "telemetry_user_id")
    UNKNOWN_USER_ID = "UNKNOWN"
    SERVER_CONTEXT: ServerContext = ServerContext.NONE
    _curr_user_id = None

    @abstractmethod
    def capture(self, event: ProductTelemetryEvent) -> None:
        pass

    @property
    def context(self) -> Dict[str, Any]:
        chroma_version = chromadb.__version__
        settings = chromadb.get_settings()
        telemetry_settings = {}
        for whitelisted in TELEMETRY_WHITELISTED_SETTINGS:
            telemetry_settings[whitelisted] = settings[whitelisted]

        self._context = {
            "chroma_version": chroma_version,
            "server_context": self.SERVER_CONTEXT.value,
            **telemetry_settings,
        }
        return self._context

    @property
    def user_id(self) -> str:
        if self._curr_user_id:
            return self._curr_user_id

        # File access may fail due to permissions or other reasons. We don't want to
        # crash so we catch all exceptions.
        try:
            if not os.path.exists(self.USER_ID_PATH):
                os.makedirs(os.path.dirname(self.USER_ID_PATH), exist_ok=True)
                with open(self.USER_ID_PATH, "w") as f:
                    new_user_id = str(uuid.uuid4())
                    f.write(new_user_id)
                self._curr_user_id = new_user_id
            else:
                with open(self.USER_ID_PATH, "r") as f:
                    self._curr_user_id = f.read()
        except Exception:
            self._curr_user_id = self.UNKNOWN_USER_ID
        return self._curr_user_id