app/services/chroma.py

import os
import re
import httpx
from .http_client import http_client

SLUG_STRIP_PATTERN = re.compile(r'-\d+$')

class ChromaSearcher:
    def __init__(self, token: str = None):
        self.api_key = token or os.environ.get("CHROMA_API_KEY", "")
        self.tenant = "99526d4b-48cf-4b20-896b-0947aa36d4ab"
        self.database = "QLESS2"
        self.collection_id = "8e924609-da78-4c85-87ed-95c05439f85e"
        self.url = f"https://api.trychroma.com/api/v2/tenants/{self.tenant}/databases/{self.database}/collections/{self.collection_id}/query"

    async def search(self, embedding: list[float]) -> tuple[str, float] | None:
        if not self.api_key:
            print("[ChromaSearcher] Warning: No Chroma API key found.")
            return None

        headers = {
            "x-chroma-token": self.api_key,
            "Content-Type": "application/json"
        }
        payload = {
            "query_embeddings": [embedding],
            "n_results": 1
        }

        client = http_client if http_client is not None else httpx.AsyncClient()
        close_client = http_client is None
        try:
            response = await client.post(self.url, headers=headers, json=payload, timeout=5.0)
            if response.status_code == 200:
                data = response.json()
                metadatas = data.get("metadatas", [])
                distances = data.get("distances", [])
 
                if metadatas and metadatas[0] and metadatas[0][0]:
                    item_meta = metadatas[0][0]
                    raw_name = item_meta.get("product_name", "Unknown Item")
                    slug = self._normalize_slug(str(raw_name))
                    distance = distances[0][0] if (distances and distances[0]) else 2.0
                    return slug, float(distance)
            else:
                status = response.status_code
                print(f"[ChromaSearcher] Query failed: {status} - {response.text}")
        except Exception as e:
            print(f"[ChromaSearcher] Error querying ChromaDB: {e}")
        finally:
            if close_client:
                await client.aclose()
        return None

    def _normalize_slug(self, slug: str) -> str:
        clean_slug = SLUG_STRIP_PATTERN.sub('', slug)
        mapping = {
            'roasted-almond-chocolate-bar-cadbury': 'dairy-milk-roast-almond-cadbury',
            'cadbury-dairy-milk-crispello': 'dairy-milk-crispello-cadbury',
            'fruit-and-nut-milk-chocolate-bar-cadbury': 'dairy-milk-chocolate-cadbury',
        }
        return mapping.get(clean_slug, clean_slug)