cocoindex-code 0.2.10__tar.gz → 0.2.11__tar.gz

{cocoindex_code-0.2.10 → cocoindex_code-0.2.11}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: cocoindex-code
-Version: 0.2.10
+Version: 0.2.11
 Summary: MCP server for indexing and querying codebases using CocoIndex
 Project-URL: Homepage, https://github.com/cocoindex-io/cocoindex-code
 Project-URL: Repository, https://github.com/cocoindex-io/cocoindex-code
@@ -345,6 +345,7 @@ embedding:
   provider: sentence-transformers                    # or "litellm"
   model: sentence-transformers/all-MiniLM-L6-v2
   device: mps                                        # optional: cpu, cuda, mps (auto-detected if omitted)
+  min_interval_ms: 300                               # optional: pace LiteLLM embedding requests to reduce 429s; defaults to 5 for LiteLLM
 
 envs:                                                # extra environment variables for the daemon
   OPENAI_API_KEY: your-key                           # only needed if not already in your shell environment
@@ -436,6 +437,7 @@ Set `OLLAMA_API_BASE` in `envs:` if your Ollama server is not at `http://localho
 ```yaml
 embedding:
   model: text-embedding-3-small
+  min_interval_ms: 300                               # optional: override the 5ms LiteLLM default
 envs:
   OPENAI_API_KEY: your-api-key
 ```

{cocoindex_code-0.2.10 → cocoindex_code-0.2.11}/README.md

@@ -306,6 +306,7 @@ embedding:
   provider: sentence-transformers                    # or "litellm"
   model: sentence-transformers/all-MiniLM-L6-v2
   device: mps                                        # optional: cpu, cuda, mps (auto-detected if omitted)
+  min_interval_ms: 300                               # optional: pace LiteLLM embedding requests to reduce 429s; defaults to 5 for LiteLLM
 
 envs:                                                # extra environment variables for the daemon
   OPENAI_API_KEY: your-key                           # only needed if not already in your shell environment
@@ -397,6 +398,7 @@ Set `OLLAMA_API_BASE` in `envs:` if your Ollama server is not at `http://localho
 ```yaml
 embedding:
   model: text-embedding-3-small
+  min_interval_ms: 300                               # optional: override the 5ms LiteLLM default
 envs:
   OPENAI_API_KEY: your-api-key
 ```

cocoindex_code-0.2.11/src/cocoindex_code/_version.py

@@ -0,0 +1,24 @@
+# file generated by vcs-versioning
+# don't change, don't track in version control
+from __future__ import annotations
+
+__all__ = [
+    "__version__",
+    "__version_tuple__",
+    "version",
+    "version_tuple",
+    "__commit_id__",
+    "commit_id",
+]
+
+version: str
+__version__: str
+__version_tuple__: tuple[int | str, ...]
+version_tuple: tuple[int | str, ...]
+commit_id: str | None
+__commit_id__: str | None
+
+__version__ = version = '0.2.11'
+__version_tuple__ = version_tuple = (0, 2, 11)
+
+__commit_id__ = commit_id = None

cocoindex_code-0.2.11/src/cocoindex_code/litellm_embedder.py

@@ -0,0 +1,126 @@
+"""LiteLLM embedder with optional request pacing and rate-limit retries."""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+import re
+import time
+from typing import Any
+
+import cocoindex as coco
+import numpy as np
+from cocoindex.ops.litellm import LiteLLMEmbedder, litellm
+from numpy.typing import NDArray
+
+logger = logging.getLogger(__name__)
+
+_RATE_LIMIT_DELAY_RE = re.compile(r"Please try again in ([0-9.]+)(ms|s)", re.IGNORECASE)
+_MAX_RATE_LIMIT_RETRIES = 6
+
+
+def _get_rate_limit_delay(exc: Exception, attempt: int) -> float | None:
+    message = str(exc)
+    if "rate limit" not in message.lower():
+        return None
+
+    match = _RATE_LIMIT_DELAY_RE.search(message)
+    if match is not None:
+        value = float(match.group(1))
+        unit = match.group(2).lower()
+        delay = value / 1000.0 if unit == "ms" else value
+    else:
+        delay = min(0.5 * (2**attempt), 10.0)
+
+    return min(delay + 0.1, 10.0)
+
+
+class PacedLiteLLMEmbedder(LiteLLMEmbedder):
+    """LiteLLM embedder that serializes requests and paces them when configured."""
+
+    def __init__(self, model: str, *, min_interval_ms: int | None = None, **kwargs: Any) -> None:
+        super().__init__(model, **kwargs)
+        self._min_request_interval_seconds = max(0.0, float(min_interval_ms or 0) / 1000.0)
+        self._request_lock: asyncio.Lock | None = None
+        self._next_request_at: float = 0.0
+
+    def _get_request_lock(self) -> asyncio.Lock:
+        if self._request_lock is None:
+            self._request_lock = asyncio.Lock()
+        return self._request_lock
+
+    async def _aembedding_with_rate_limit_retries(
+        self, *, model: str, input: list[str], **kwargs: Any
+    ) -> Any:
+        last_exc: Exception | None = None
+
+        for attempt in range(_MAX_RATE_LIMIT_RETRIES):
+            try:
+                return await litellm.aembedding(model=model, input=input, **kwargs)
+            except Exception as exc:  # noqa: BLE001
+                delay = _get_rate_limit_delay(exc, attempt)
+                last_exc = exc
+                if delay is None or attempt == _MAX_RATE_LIMIT_RETRIES - 1:
+                    raise
+
+                logger.warning(
+                    "Embedding rate limited for model %s, retrying in %.3fs (attempt %d/%d)",
+                    model,
+                    delay,
+                    attempt + 1,
+                    _MAX_RATE_LIMIT_RETRIES,
+                )
+                await asyncio.sleep(delay)
+
+        assert last_exc is not None
+        raise last_exc
+
+    async def run_embedding_request(self, *, input: list[str], **kwargs: Any) -> Any:
+        lock = self._get_request_lock()
+        async with lock:
+            now = time.monotonic()
+            if self._next_request_at > now:
+                await asyncio.sleep(self._next_request_at - now)
+
+            response = await self._aembedding_with_rate_limit_retries(
+                model=self._model,
+                input=input,
+                **kwargs,
+            )
+
+            now = time.monotonic()
+            if self._min_request_interval_seconds > 0:
+                self._next_request_at = now + self._min_request_interval_seconds
+            else:
+                self._next_request_at = now
+
+            return response
+
+    async def _get_dim(self) -> int:
+        if self._dim is not None:
+            return self._dim
+        async with self._get_lock():
+            if self._dim is not None:
+                return self._dim
+            response = await self.run_embedding_request(input=["hello"], **self._kwargs)
+            embedding = response.data[0]["embedding"]
+            self._dim = len(embedding)
+            return self._dim
+
+    @coco.fn.as_async(
+        batching=True,
+        max_batch_size=64,
+        memo=True,
+        version=1,
+        logic_tracking="self",
+    )
+    async def embed(
+        self,
+        texts: list[str],
+        input_type: str | None = None,
+    ) -> list[NDArray[np.float32]]:
+        kwargs = dict(self._kwargs)
+        if input_type is not None:
+            kwargs["input_type"] = input_type
+        response = await self.run_embedding_request(input=texts, **kwargs)
+        return [np.array(item["embedding"], dtype=np.float32) for item in response.data]

{cocoindex_code-0.2.10 → cocoindex_code-0.2.11}/src/cocoindex_code/settings.py

@@ -92,6 +92,7 @@ class EmbeddingSettings:
     model: str
     provider: str = "litellm"
     device: str | None = None
+    min_interval_ms: int | None = None
 
 
 @dataclass
@@ -351,6 +352,8 @@ def _user_settings_to_dict(settings: UserSettings) -> dict[str, Any]:
     }
     if settings.embedding.device is not None:
         emb["device"] = settings.embedding.device
+    if settings.embedding.min_interval_ms is not None:
+        emb["min_interval_ms"] = settings.embedding.min_interval_ms
     d["embedding"] = emb
     if settings.envs:
         d["envs"] = dict(settings.envs)
@@ -367,6 +370,8 @@ def _user_settings_from_dict(d: dict[str, Any]) -> UserSettings:
         emb_kwargs["provider"] = emb_dict["provider"]
     if "device" in emb_dict:
         emb_kwargs["device"] = emb_dict["device"]
+    if "min_interval_ms" in emb_dict:
+        emb_kwargs["min_interval_ms"] = emb_dict["min_interval_ms"]
     embedding = EmbeddingSettings(**emb_kwargs)
     envs = d.get("envs", {})
     return UserSettings(embedding=embedding, envs=envs)

{cocoindex_code-0.2.10 → cocoindex_code-0.2.11}/src/cocoindex_code/shared.py

@@ -21,6 +21,7 @@ from .settings import EmbeddingSettings
 logger = logging.getLogger(__name__)
 
 SBERT_PREFIX = "sbert/"
+DEFAULT_LITELLM_MIN_INTERVAL_MS = 5
 
 # Models that define a "query" prompt for asymmetric retrieval.
 _QUERY_PROMPT_MODELS = {"nomic-ai/nomic-embed-code", "nomic-ai/CodeRankEmbed"}
@@ -63,11 +64,23 @@ def create_embedder(settings: EmbeddingSettings) -> Embedder:
         )
         logger.info("Embedding model: %s | device: %s", settings.model, settings.device)
     else:
-        from cocoindex.ops.litellm import LiteLLMEmbedder
+        from .litellm_embedder import PacedLiteLLMEmbedder
 
-        instance = LiteLLMEmbedder(settings.model)
+        min_interval_ms = (
+            settings.min_interval_ms
+            if settings.min_interval_ms is not None
+            else DEFAULT_LITELLM_MIN_INTERVAL_MS
+        )
+        instance = PacedLiteLLMEmbedder(
+            settings.model,
+            min_interval_ms=min_interval_ms,
+        )
         query_prompt_name = None
-        logger.info("Embedding model (LiteLLM): %s", settings.model)
+        logger.info(
+            "Embedding model (LiteLLM): %s | min_interval_ms: %s",
+            settings.model,
+            min_interval_ms,
+        )
 
     embedder = instance
     return instance

cocoindex_code-0.2.10/src/cocoindex_code/_version.py

@@ -1,34 +0,0 @@
-# file generated by setuptools-scm
-# don't change, don't track in version control
-
-__all__ = [
-    "__version__",
-    "__version_tuple__",
-    "version",
-    "version_tuple",
-    "__commit_id__",
-    "commit_id",
-]
-
-TYPE_CHECKING = False
-if TYPE_CHECKING:
-    from typing import Tuple
-    from typing import Union
-
-    VERSION_TUPLE = Tuple[Union[int, str], ...]
-    COMMIT_ID = Union[str, None]
-else:
-    VERSION_TUPLE = object
-    COMMIT_ID = object
-
-version: str
-__version__: str
-__version_tuple__: VERSION_TUPLE
-version_tuple: VERSION_TUPLE
-commit_id: COMMIT_ID
-__commit_id__: COMMIT_ID
-
-__version__ = version = '0.2.10'
-__version_tuple__ = version_tuple = (0, 2, 10)
-
-__commit_id__ = commit_id = None