chore: upgrade to cocoindex 1.0.0a24 and use its embedder (#36)

Author: georgeh0

pyproject.toml

@@ -23,7 +23,7 @@ classifiers = [
 
 dependencies = [
     "mcp>=1.0.0",
-    "cocoindex[litellm]==1.0.0a23",
+    "cocoindex[litellm]==1.0.0a24",
     "sentence-transformers>=2.2.0",
     "sqlite-vec>=0.1.0",
     "pydantic>=2.0.0",
@@ -82,10 +82,6 @@ python_version = "3.11"
 strict = true
 ignore_missing_imports = true
 
-[[tool.mypy.overrides]]
-module = "cocoindex_code.embedder"
-warn_unused_ignores = false
-
 [tool.pytest.ini_options]
 testpaths = ["tests"]
 python_files = ["test_*.py"]

src/cocoindex_code/config.py

@@ -6,7 +6,6 @@
 from dataclasses import dataclass
 from pathlib import Path
 
-_SBERT_PREFIX = "sbert/"
 _DEFAULT_MODEL = "sbert/sentence-transformers/all-MiniLM-L6-v2"
 
 
@@ -65,7 +64,6 @@ class Config:
     index_dir: Path
     device: str
     trust_remote_code: bool
-    batch_size: int
     extra_extensions: dict[str, str | None]
 
     @classmethod
@@ -101,19 +99,6 @@ def from_env(cls) -> Config:
             "yes",
         )
 
-        # Batch size for local embedding model
-        _raw_batch_size = os.environ.get("COCOINDEX_CODE_BATCH_SIZE", "16")
-        try:
-            batch_size = int(_raw_batch_size)
-        except ValueError:
-            raise ValueError(
-                f"COCOINDEX_CODE_BATCH_SIZE must be a positive integer, got: {_raw_batch_size!r}"
-            ) from None
-        if batch_size <= 0:
-            raise ValueError(
-                f"COCOINDEX_CODE_BATCH_SIZE must be a positive integer, got: {batch_size}"
-            )
-
         # Extra file extensions (format: "inc:php,yaml,toml" — optional lang after colon)
         raw_extra = os.environ.get("COCOINDEX_CODE_EXTRA_EXTENSIONS", "")
         extra_extensions: dict[str, str | None] = {}
@@ -133,7 +118,6 @@ def from_env(cls) -> Config:
             index_dir=index_dir,
             device=device,
             trust_remote_code=trust_remote_code,
-            batch_size=batch_size,
             extra_extensions=extra_extensions,
         )
 

src/cocoindex_code/embedder.py

@@ -14,8 +14,7 @@
 
 if TYPE_CHECKING:
     from cocoindex.ops.litellm import LiteLLMEmbedder
-
-    from .embedder import LocalEmbedder
+    from cocoindex.ops.sentence_transformers import SentenceTransformerEmbedder
 
 from .config import config
 
@@ -24,9 +23,9 @@
 SBERT_PREFIX = "sbert/"
 
 # Initialize embedder at module level based on model prefix
-embedder: LocalEmbedder | LiteLLMEmbedder
+embedder: SentenceTransformerEmbedder | LiteLLMEmbedder
 if config.embedding_model.startswith(SBERT_PREFIX):
-    from .embedder import LocalEmbedder
+    from cocoindex.ops.sentence_transformers import SentenceTransformerEmbedder
 
     _model_name = config.embedding_model[len(SBERT_PREFIX) :]
     # Models that define a "query" prompt for asymmetric retrieval.
@@ -35,7 +34,7 @@
     # Models whose custom remote code is known-compatible with transformers 5.x.
     _KNOWN_REMOTE_CODE_MODELS = {"nomic-ai/CodeRankEmbed"}
     _trust = config.trust_remote_code or _model_name in _KNOWN_REMOTE_CODE_MODELS
-    embedder = LocalEmbedder(
+    embedder = SentenceTransformerEmbedder(
         _model_name,
         device=config.device,
         trust_remote_code=_trust,

src/cocoindex_code/shared.py

@@ -6,8 +6,6 @@
 from pathlib import Path
 from unittest.mock import patch
 
-import pytest
-
 from cocoindex_code.config import Config, _detect_device
 
 
@@ -83,63 +81,6 @@ def test_default_model_is_minilm(self, tmp_path: Path) -> None:
             assert "all-MiniLM-L6-v2" in config.embedding_model
 
 
-class TestConfigBatchSize:
-    """Tests for COCOINDEX_CODE_BATCH_SIZE env var."""
-
-    def test_default_batch_size_is_16(self, tmp_path: Path) -> None:
-        with patch.dict(
-            os.environ,
-            {"COCOINDEX_CODE_ROOT_PATH": str(tmp_path)},
-        ):
-            os.environ.pop("COCOINDEX_CODE_BATCH_SIZE", None)
-            config = Config.from_env()
-            assert config.batch_size == 16
-
-    def test_batch_size_reads_env_var(self, tmp_path: Path) -> None:
-        with patch.dict(
-            os.environ,
-            {
-                "COCOINDEX_CODE_ROOT_PATH": str(tmp_path),
-                "COCOINDEX_CODE_BATCH_SIZE": "32",
-            },
-        ):
-            config = Config.from_env()
-            assert config.batch_size == 32
-
-    def test_batch_size_raises_on_non_integer(self, tmp_path: Path) -> None:
-        with patch.dict(
-            os.environ,
-            {
-                "COCOINDEX_CODE_ROOT_PATH": str(tmp_path),
-                "COCOINDEX_CODE_BATCH_SIZE": "notanint",
-            },
-        ):
-            with pytest.raises(ValueError, match="COCOINDEX_CODE_BATCH_SIZE"):
-                Config.from_env()
-
-    def test_batch_size_raises_on_zero(self, tmp_path: Path) -> None:
-        with patch.dict(
-            os.environ,
-            {
-                "COCOINDEX_CODE_ROOT_PATH": str(tmp_path),
-                "COCOINDEX_CODE_BATCH_SIZE": "0",
-            },
-        ):
-            with pytest.raises(ValueError, match="COCOINDEX_CODE_BATCH_SIZE"):
-                Config.from_env()
-
-    def test_batch_size_raises_on_negative(self, tmp_path: Path) -> None:
-        with patch.dict(
-            os.environ,
-            {
-                "COCOINDEX_CODE_ROOT_PATH": str(tmp_path),
-                "COCOINDEX_CODE_BATCH_SIZE": "-1",
-            },
-        ):
-            with pytest.raises(ValueError, match="COCOINDEX_CODE_BATCH_SIZE"):
-                Config.from_env()
-
-
 class TestExtraExtensions:
     """Tests for COCOINDEX_CODE_EXTRA_EXTENSIONS env var."""