fix: fit into the latest cocoindex API (#25)

Author: georgeh0

src/cocoindex_code/embedder.py

@@ -6,7 +6,6 @@
 from typing import TYPE_CHECKING, Any
 
 import cocoindex as coco
-import cocoindex.asyncio as coco_aio
 import numpy as np
 from cocoindex.resources import schema as _schema
 from numpy.typing import NDArray
@@ -74,7 +73,7 @@ def _get_model(self) -> SentenceTransformer:
                     )
         return self._model
 
-    @coco_aio.function(batching=True, runner=coco.GPU, memo=True, max_batch_size=_config.batch_size)
+    @coco.fn.as_async(batching=True, runner=coco.GPU, memo=True, max_batch_size=_config.batch_size)
     def embed(self, texts: list[str]) -> list[NDArray[np.float32]]:
         """Embed a batch of texts into float32 vectors."""
         model = self._get_model()
@@ -85,7 +84,7 @@ def embed(self, texts: list[str]) -> list[NDArray[np.float32]]:
         )  # type: ignore[assignment]
         return list(embeddings)
 
-    @coco_aio.function(batching=True, runner=coco.GPU, memo=True, max_batch_size=_config.batch_size)
+    @coco.fn.as_async(batching=True, runner=coco.GPU, memo=True, max_batch_size=_config.batch_size)
     def embed_query(self, texts: list[str]) -> list[NDArray[np.float32]]:
         """Embed query texts, applying query_prompt_name if configured."""
         model = self._get_model()
@@ -97,7 +96,7 @@ def embed_query(self, texts: list[str]) -> list[NDArray[np.float32]]:
         )  # type: ignore[assignment]
         return list(embeddings)
 
-    @coco_aio.function(runner=coco.GPU, memo=True)
+    @coco.fn.as_async(runner=coco.GPU, memo=True)
     def __coco_vector_schema__(self) -> _schema.VectorSchema:
         """Return the vector schema (dimension + dtype) for this model."""
         model = self._get_model()

tests/conftest.py

@@ -5,7 +5,7 @@
 from collections.abc import AsyncIterator
 from pathlib import Path
 
-import cocoindex.asyncio as coco_aio
+import cocoindex as coco
 import pytest
 import pytest_asyncio
 
@@ -29,5 +29,5 @@ async def coco_runtime() -> AsyncIterator[None]:
     Uses session-scoped event loop to ensure CocoIndex environment
     persists across all tests.
     """
-    async with coco_aio.runtime():
+    async with coco.runtime():
         yield

tests/test_embedder.py

@@ -51,7 +51,7 @@ def test_embed_query_forwards_prompt_name(self) -> None:
             embedder._get_model()
             mock_model.encode.reset_mock()
             # Call the underlying method directly, bypassing the CocoIndex batching decorator.
-            LocalEmbedder.embed_query.__wrapped__(embedder, ["find functions that embed text"])
+            LocalEmbedder.embed_query.__wrapped__(embedder, ["find functions that embed text"])  # type: ignore
             _, kwargs = mock_model.encode.call_args
             assert kwargs.get("prompt_name") == "query"
 
@@ -61,7 +61,7 @@ def test_embed_query_no_prompt_when_unset(self) -> None:
             embedder = LocalEmbedder("some-model", device="cpu", query_prompt_name=None)
             embedder._get_model()
             mock_model.encode.reset_mock()
-            LocalEmbedder.embed_query.__wrapped__(embedder, ["some query"])
+            LocalEmbedder.embed_query.__wrapped__(embedder, ["some query"])  # type: ignore
             _, kwargs = mock_model.encode.call_args
             assert kwargs.get("prompt_name") is None