feat: add ChromaClient for managing embeddings with ChromaDB

- Updated pyproject.toml to include chromadb as a dependency.
- Implemented ChromaClient class for handling text embeddings and interactions with ChromaDB.
- Added methods for upserting texts and querying embeddings.
- Integrated OpenAI's embedding model for generating text embeddings.

pyproject.toml

@@ -16,7 +16,8 @@ dependencies = [
     "alembic (>=1.18.4,<2.0.0)",
     "aiosqlite (>=0.22.1,<0.23.0)",
     "pydantic-ai-slim (>=1.67.0,<2.0.0)",
-    "openai (>=2.26.0,<3.0.0)"
+    "openai (>=2.26.0,<3.0.0)",
+    "chromadb (>=1.5.5,<2.0.0)"
 ]
 
 

src/utils/_chroma_client.py

@@ -0,0 +1,102 @@
+import asyncio
+import uuid
+import os
+import chromadb
+from chromadb.config import Settings
+from openai import AsyncOpenAI
+
+
+class ChromaClient:
+    def __init__(self):
+        self.collection_name = os.getenv("CHROMA_COLLECTION_NAME", "default_collection")
+        self.persist_directory = os.getenv("CHROMA_PERSIST_DIRECTORY", "./chroma_db")
+        self.use_persistent = (
+            os.getenv("CHROMA_USE_PERSISTENT", "true").lower() == "true"
+        )
+        self.chroma_host = os.getenv("CHROMA_HOST", "localhost")
+        self.chroma_port = int(os.getenv("CHROMA_PORT", "8000"))
+        self.embedding_model = os.getenv(
+            "OPENAI_EMBEDDING_MODEL", "text-embedding-3-small"
+        )
+
+        self.openai = AsyncOpenAI()
+        self.client = None
+        self.collection = None
+
+    async def __aenter__(self):
+        if self.use_persistent:
+            self.client = chromadb.PersistentClient(path=self.persist_directory)
+        else:
+            self.client = chromadb.HttpClient(
+                host=self.chroma_host,
+                port=self.chroma_port,
+                settings=Settings(anonymized_telemetry=False),
+            )
+
+        self.collection = self.client.get_or_create_collection(
+            name=self.collection_name,
+            metadata={"hnsw:space": "cosine"},
+        )
+        return self
+
+    async def __aexit__(self, exc_type, exc_value, traceback):
+        self.client = None
+        self.collection = None
+
+    async def _get_text_embedding(self, text: str) -> list[float]:
+        response = await self.openai.embeddings.create(
+            input=text,
+            model=self.embedding_model,
+        )
+        return response.data[0].embedding
+
+    async def _get_batch_embeddings(self, texts: list[str]) -> list[list[float]]:
+        response = await self.openai.embeddings.create(
+            input=texts,
+            model=self.embedding_model,
+        )
+        return [item.embedding for item in response.data]
+
+    async def upsert(self, texts: list[str], metadatas: list[dict] = None):
+        if not texts:
+            return
+
+        if metadatas is None:
+            metadatas = [{} for _ in texts]
+
+        if len(texts) != len(metadatas):
+            raise ValueError("texts and metadatas must have the same length")
+
+        ids = [meta.pop("id", str(uuid.uuid4())) for meta in metadatas]
+        embeddings = await self._get_batch_embeddings(texts)
+
+        loop = asyncio.get_event_loop()
+        await loop.run_in_executor(
+            None,
+            lambda: self.collection.upsert(
+                ids=ids,
+                embeddings=embeddings,
+                documents=texts,
+                metadatas=metadatas,
+            ),
+        )
+
+    async def query(self, query: str, n_results: int = 5) -> dict:
+        query_embedding = await self._get_text_embedding(query)
+
+        loop = asyncio.get_event_loop()
+        results = await loop.run_in_executor(
+            None,
+            lambda: self.collection.query(
+                query_embeddings=[query_embedding],
+                n_results=n_results,
+                include=["documents", "metadatas", "distances"],
+            ),
+        )
+
+        return {
+            "ids": results["ids"][0],
+            "documents": results["documents"][0],
+            "metadatas": results["metadatas"][0],
+            "distances": results["distances"][0],
+        }