ultimos ajustes

src/memory/embeddings.py

@@ -0,0 +1,55 @@
+"""Embedding service — generates vectors via OpenAI API.
+
+Used for semantic search over the knowledge base.
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import Any
+
+from openai import AsyncOpenAI
+
+from ..config import settings
+
+logger = logging.getLogger(__name__)
+
+# Default model: cheap, fast, good enough for doc retrieval
+DEFAULT_MODEL = "text-embedding-3-small"
+DEFAULT_DIMENSIONS = 1536
+
+
+class EmbeddingService:
+    """Generates embeddings via OpenAI's embedding API."""
+
+    def __init__(
+        self,
+        api_key: str | None = None,
+        model: str = DEFAULT_MODEL,
+    ) -> None:
+        self._client = AsyncOpenAI(
+            api_key=api_key or settings.openai_api_key,
+        )
+        self._model = model
+
+    async def embed(self, text: str) -> list[float]:
+        """Generate embedding for a single text."""
+        # Truncate to ~8k tokens worth of text to stay within limits
+        text = text[:32_000]
+        response = await self._client.embeddings.create(
+            model=self._model,
+            input=text,
+        )
+        return response.data[0].embedding
+
+    async def embed_batch(self, texts: list[str]) -> list[list[float]]:
+        """Generate embeddings for multiple texts in one call."""
+        # Truncate each
+        texts = [t[:32_000] for t in texts]
+        response = await self._client.embeddings.create(
+            model=self._model,
+            input=texts,
+        )
+        # Sort by index to maintain order
+        sorted_data = sorted(response.data, key=lambda d: d.index)
+        return [d.embedding for d in sorted_data]