fix: avoid cyclic import (#65)

* fix: avoid cyclic import

Signed-off-by: Keming <kemingyang@tensorchord.ai>

* fix comments

Signed-off-by: Keming <kemingyang@tensorchord.ai>

---------

Signed-off-by: Keming <kemingyang@tensorchord.ai>

Author: kemingy

vechord/__init__.py

@@ -1,18 +1,5 @@
-from vechord.augment import GeminiAugmenter
-from vechord.chunk import GeminiChunker, RegexChunker, SpacyChunker
 from vechord.client import VechordClient
-from vechord.embedding import (
-    GeminiDenseEmbedding,
-    OpenAIDenseEmbedding,
-    SpacyDenseEmbedding,
-)
-from vechord.evaluate import GeminiEvaluator
-from vechord.extract import GeminiExtractor, SimpleExtractor
-from vechord.load import LocalLoader
-from vechord.model import Document
 from vechord.registry import VechordPipeline, VechordRegistry
-from vechord.rerank import CohereReranker
-from vechord.service import create_web_app
 from vechord.spec import (
     DefaultDocument,
     ForeignKey,
@@ -29,33 +16,19 @@
 )
 
 __all__ = [
-    "CohereReranker",
     "DefaultDocument",
-    "Document",
     "ForeignKey",
-    "GeminiAugmenter",
-    "GeminiChunker",
-    "GeminiDenseEmbedding",
-    "GeminiEvaluator",
-    "GeminiExtractor",
     "IndexColumn",
     "Keyword",
     "KeywordIndex",
-    "LocalLoader",
     "MultiVectorIndex",
-    "OpenAIDenseEmbedding",
     "PrimaryKeyAutoIncrease",
     "PrimaryKeyUUID",
-    "RegexChunker",
-    "SimpleExtractor",
-    "SpacyChunker",
-    "SpacyDenseEmbedding",
     "Table",
     "VechordClient",
     "VechordPipeline",
     "VechordRegistry",
     "Vector",
     "VectorIndex",
     "create_chunk_with_dim",
-    "create_web_app",
 ]

vechord/pipeline.py

@@ -1,17 +1,15 @@
 import base64
 import itertools
-from collections.abc import Callable, Iterable
+from collections.abc import Iterable
 from contextlib import contextmanager
 from os import environ
-from typing import TYPE_CHECKING, Annotated, Any, Optional
+from typing import Annotated, Any, Optional
 from uuid import UUID
 
 import msgspec
 
 from vechord.chunk import BaseChunker, GeminiChunker, RegexChunker
 from vechord.client import (
-    VechordClient,
-    limit_to_transaction_buffer_conn,
     set_namespace,
 )
 from vechord.embedding import (
@@ -37,6 +35,7 @@
     RunRequest,
     RunSearchResponse,
 )
+from vechord.registry import VechordRegistry
 from vechord.rerank import BaseReranker, CohereReranker, JinaReranker
 from vechord.spec import (
     AnyOf,
@@ -51,9 +50,6 @@
 )
 from vechord.typing import Self
 
-if TYPE_CHECKING:
-    from vechord.registry import VechordRegistry
-
 
 class GraphIndex(msgspec.Struct):
     """Graph index for entities and relations extracted from the text/image."""
@@ -208,7 +204,7 @@ def from_steps(cls, steps: list[ResourceRequest]) -> Self:
         return msgspec.convert(calls, DynamicPipeline)
 
     async def run(
-        self, request: RunRequest, vr: "VechordRegistry"
+        self, request: RunRequest, vr: VechordRegistry
     ) -> RunIngestAck | RunSearchResponse:
         """Run the dynamic pipeline with the given request."""
         async with set_namespace(request.name):
@@ -249,7 +245,7 @@ def _convert_from_extracted_graph(
         return converted_ents, converted_rels
 
     async def run_index(  # noqa: PLR0912
-        self, request: RunRequest, vr: "VechordRegistry"
+        self, request: RunRequest, vr: VechordRegistry
     ) -> RunIngestAck:
         dim = (
             self.text_emb.get_dim() if self.text_emb else self.multimodal_emb.get_dim()
@@ -345,7 +341,7 @@ async def graph_insert(
         rels: list[_Relation],
         ent_cls: type[Table],
         rel_cls: type[Table],
-        vr: "VechordRegistry",
+        vr: VechordRegistry,
     ):
         """Insert entities and relations into the graph index."""
         ent_map: dict[str, _Entity] = {}
@@ -393,7 +389,7 @@ async def graph_insert(
             await vr.insert(rel)
 
     async def run_search(
-        self, request: RunRequest, vr: "VechordRegistry"
+        self, request: RunRequest, vr: VechordRegistry
     ) -> RunSearchResponse:
         query = request.data.decode("utf-8")
 
@@ -451,7 +447,7 @@ async def graph_search(
         chunk_cls: type[Table],
         ent_cls: type[Table],
         rel_cls: type[Table],
-        vr: "VechordRegistry",
+        vr: VechordRegistry,
     ):
         ents, rels = await self.graph.recognize_with_relations(query)
         emb_func = (
@@ -499,40 +495,3 @@ def deduplicate_uid(uuids: Iterable[UUID], limit: Optional[int] = None) -> list[
     """Maintain the order of the occurrence of UUIDs and deduplicate them."""
     uuids = {uid: None for uid in uuids}
     return list(uuids.keys())[:limit]
-
-
-class VechordPipeline:
-    """Set up the pipeline to run multiple functions in a transaction.
-
-    Args:
-        client: :class:`VectorChordClient` to be used for the transaction.
-        steps: a list of functions to be run in the pipeline. The first function
-            will be used to accept the input, and the last function will be used
-            to return the output. The rest of the functions will be used to
-            process the data in between. The functions will be run in the order
-            they are defined in the list.
-    """
-
-    def __init__(self, client: VechordClient, steps: list[Callable]):
-        self.client = client
-        self.steps = steps
-
-    async def run(self, *args, **kwargs) -> Any:
-        """Execute the pipeline in a transactional manner.
-
-        All the `args` and `kwargs` will be passed to the first function in the
-        pipeline. The pipeline will run in *one* transaction, and all the `inject`
-        can only see the data inserted in this transaction (to guarantee only the
-        new inserted data will be processed in this pipeline).
-
-        This will also return the final result of the last function in the pipeline.
-        """
-        async with (
-            self.client.get_connection() as conn,
-            limit_to_transaction_buffer_conn(conn),
-        ):
-            # only the 1st one can accept input (could be empty)
-            await self.steps[0](*args, **kwargs)
-            for func in self.steps[1:-1]:
-                await func()
-            return await self.steps[-1]()

vechord/registry.py

@@ -3,6 +3,7 @@
 from functools import wraps
 from inspect import isasyncgenfunction, iscoroutinefunction
 from typing import (
+    Any,
     Callable,
     Optional,
     Sequence,
@@ -13,9 +14,12 @@
 
 import numpy as np
 
-from vechord.client import VechordClient, select_transaction_buffer_conn
+from vechord.client import (
+    VechordClient,
+    limit_to_transaction_buffer_conn,
+    select_transaction_buffer_conn,
+)
 from vechord.log import logger
-from vechord.pipeline import VechordPipeline
 from vechord.spec import Table, Vector
 
 T = TypeVar("T", bound=Table)
@@ -34,6 +38,43 @@ def get_iterator_type(typ) -> type:
     return get_iterator_type(typ.__args__[0])
 
 
+class VechordPipeline:
+    """Set up the pipeline to run multiple functions in a transaction.
+
+    Args:
+        client: :class:`VechordClient` to be used for the transaction.
+        steps: a list of functions to be run in the pipeline. The first function
+            will be used to accept the input, and the last function will be used
+            to return the output. The rest of the functions will be used to
+            process the data in between. The functions will be run in the order
+            they are defined in the list.
+    """
+
+    def __init__(self, client: VechordClient, steps: list[Callable]):
+        self.client = client
+        self.steps = steps
+
+    async def run(self, *args, **kwargs) -> Any:
+        """Execute the pipeline in a transactional manner.
+
+        All the `args` and `kwargs` will be passed to the first function in the
+        pipeline. The pipeline will run in *one* transaction, and all the `inject`
+        can only see the data inserted in this transaction (to guarantee only the
+        new inserted data will be processed in this pipeline).
+
+        This will also return the final result of the last function in the pipeline.
+        """
+        async with (
+            self.client.get_connection() as conn,
+            limit_to_transaction_buffer_conn(conn),
+        ):
+            # only the 1st one can accept input (could be empty)
+            await self.steps[0](*args, **kwargs)
+            for func in self.steps[1:-1]:
+                await func()
+            return await self.steps[-1]()
+
+
 class VechordRegistry:
     """Create a registry for the given namespace and PostgreSQL URL.