fix: some type check and argument alignment (#64)

* fix: some type check and argument alignment

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

* fix lint

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

---------

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

Author: kemingy

pyproject.toml

@@ -15,7 +15,7 @@ classifiers = [
 ]
 dependencies = [
     "defspec>=0.4.0",
-    "falcon>=4.0.2",
+    "falcon>=4.1.0",
     "httpx>=0.28.1",
     "msgspec>=0.19.0",
     "numpy>=2.0.2",
@@ -106,6 +106,7 @@ python_version = "3.10"
 warn_redundant_casts = true
 warn_unreachable = true
 pretty = true
+exclude = ["docs/build/html/_downloads/.*"]
 
 [[tool.mypy.overrides]]
 module = [

uv.lock

@@ -158,7 +158,7 @@ async def segment(self, text: str) -> list[str]:
             return chunks
 
         all_chunks = []
-        for chunk in self.regex_chunker.segment(text):
+        for chunk in await self.regex_chunker.segment(text):
             chunks = await self.structure_query(
                 self.prompt.format(size=self.size, document=chunk)
             )

vechord/chunk.py

@@ -96,21 +96,29 @@ def __init__(self, model: str = "en_core_web_sm"):
 
     def recognize(self, text) -> list[GraphEntity]:
         doc = self.nlp(text)
-        return [GraphEntity(text=ent.text, label=ent.label_) for ent in doc.ents]
+        return [
+            GraphEntity(text=ent.text, label=ent.label_, description="")
+            for ent in doc.ents
+        ]
 
     def recognize_with_relations(
         self, text
     ) -> tuple[list[GraphEntity], list[GraphRelation]]:
         doc = self.nlp(text)
-        ents = [GraphEntity(text=ent.text, label=ent.label_) for ent in doc.ents]
+        ents = [
+            GraphEntity(text=ent.text, label=ent.label_, description="")
+            for ent in doc.ents
+        ]
         relations: list[GraphRelation] = []
         matches = self.matcher(doc)
         for _, start, end in matches:
             span = doc[start:end]
             ent0 = ent1 = None
             for token in span:
                 if token.ent_type_:
-                    ent = GraphEntity(text=token.text, label=token.ent_type_)
+                    ent = GraphEntity(
+                        text=token.text, label=token.ent_type_, description=""
+                    )
                     if ent0 is None:
                         ent0 = ent
                     else:
@@ -119,9 +127,13 @@ def recognize_with_relations(
             relations.append(
                 GraphRelation(
                     source=ent0
-                    or GraphEntity(text=span[0].text, label=span[0].ent_type_),
+                    or GraphEntity(
+                        text=span[0].text, label=span[0].ent_type_, description=""
+                    ),
                     target=ent1
-                    or GraphEntity(text=span[-1].text, label=span[-1].ent_type_),
+                    or GraphEntity(
+                        text=span[-1].text, label=span[-1].ent_type_, description=""
+                    ),
                     description=" ".join(token.text for token in span),
                 )
             )

vechord/graph.py

@@ -32,9 +32,9 @@
 from vechord.model.web import (
     InputType,
     ResourceRequest,
-    RunAck,
+    RunIngestAck,
     RunRequest,
-    RunResponse,
+    RunSearchResponse,
 )
 
 __all__ = [
@@ -56,9 +56,9 @@
     "LlamaCloudParseResponse",
     "ResourceRequest",
     "RetrievedChunk",
-    "RunAck",
+    "RunIngestAck",
     "RunRequest",
-    "RunResponse",
+    "RunSearchResponse",
     "SparseEmbedding",
     "UMBRELAScore",
     "VoyageEmbeddingRequest",

vechord/model/__init__.py

@@ -90,7 +90,7 @@ class JinaRerankRequest(msgspec.Struct, kw_only=True):
     model: Literal["jina-reranker-v2-base-multilingual", "jina-reranker-m0"]
     query: str
     top_n: int
-    documents: list[str | JinaInput]
+    documents: list[str] | list[JinaInput]
     return_documents: bool = False
 
     @classmethod

vechord/model/jina.py

@@ -41,7 +41,7 @@ class RunRequest(msgspec.Struct, kw_only=True, frozen=True):
     steps: list[ResourceRequest] = msgspec.field(default_factory=list)
 
 
-class RunAck(msgspec.Struct, kw_only=True, frozen=True):
+class RunIngestAck(msgspec.Struct, kw_only=True, frozen=True, tag="ingest"):
     """Acknowledgment of an index request."""
 
     name: str
@@ -55,7 +55,7 @@ class SearchResponse(msgspec.Struct, kw_only=True, omit_defaults=True):
     text: Optional[str] = None
 
 
-class RunResponse(msgspec.Struct, kw_only=True, omit_defaults=True):
+class RunSearchResponse(msgspec.Struct, kw_only=True, omit_defaults=True, tag="search"):
     """Response to a search request.
 
     metrics:
@@ -87,5 +87,13 @@ def cleanup(self):
             for chunk in self.chunks:
                 chunk.text = None
 
+    def deduplicate(self):
+        """Deduplicate chunks while maintain the order."""
+        unique_chunks = {}
+        for chunk in self.chunks:
+            if chunk.uid not in unique_chunks:
+                unique_chunks[chunk.uid] = chunk
+        self.chunks = list(unique_chunks.values())
+
     def reorder(self, indices: list[int]):
         self.chunks = [self.chunks[i] for i in indices]

vechord/model/web.py

@@ -33,9 +33,9 @@
     GraphRelation,
     InputType,
     ResourceRequest,
-    RunAck,
+    RunIngestAck,
     RunRequest,
-    RunResponse,
+    RunSearchResponse,
 )
 from vechord.rerank import BaseReranker, CohereReranker, JinaReranker
 from vechord.spec import (
@@ -209,7 +209,7 @@ def from_steps(cls, steps: list[ResourceRequest]) -> Self:
 
     async def run(
         self, request: RunRequest, vr: "VechordRegistry"
-    ) -> RunAck | RunResponse:
+    ) -> RunIngestAck | RunSearchResponse:
         """Run the dynamic pipeline with the given request."""
         async with set_namespace(request.name):
             resp = (
@@ -248,7 +248,9 @@ def _convert_from_extracted_graph(
         ]
         return converted_ents, converted_rels
 
-    async def run_index(self, request: RunRequest, vr: "VechordRegistry") -> RunAck:  # noqa: PLR0912
+    async def run_index(  # noqa: PLR0912
+        self, request: RunRequest, vr: "VechordRegistry"
+    ) -> RunIngestAck:
         dim = (
             self.text_emb.get_dim() if self.text_emb else self.multimodal_emb.get_dim()
         )
@@ -335,7 +337,7 @@ async def run_index(self, request: RunRequest, vr: "VechordRegistry") -> RunAck:
             await self.graph_insert(
                 ents=ents, rels=rels, ent_cls=Entity, rel_cls=Relation, vr=vr
             )
-        return RunAck(name=request.name, msg="succeed", uid=doc.uid)
+        return RunIngestAck(name=request.name, msg="succeed", uid=doc.uid)
 
     async def graph_insert(
         self,
@@ -392,7 +394,7 @@ async def graph_insert(
 
     async def run_search(
         self, request: RunRequest, vr: "VechordRegistry"
-    ) -> RunResponse:
+    ) -> RunSearchResponse:
         query = request.data.decode("utf-8")
 
         # for type hint and compatibility
@@ -405,7 +407,7 @@ class Entity(_Entity):
         class Relation(_Relation):
             pass
 
-        resp = RunResponse()
+        resp = RunSearchResponse()
         if self.search.vector:
             vec = (
                 await self.text_emb.vectorize_query(query)
@@ -423,6 +425,7 @@ class Relation(_Relation):
             )
         if self.search.graph:
             resp.extend(await self.graph_search(query, Chunk, Entity, Relation, vr))
+        resp.deduplicate()
         if self.rerank:
             if self.multimodal_emb:
                 indices = await self.rerank.rerank_multimodal(
@@ -432,7 +435,7 @@ class Relation(_Relation):
                 )
             else:
                 indices = await self.rerank.rerank(
-                    query=query, chunks=[chunk.text for chunk in resp]
+                    query=query, chunks=[chunk.text for chunk in resp.chunks]
                 )
             resp.reorder(indices)
         if self.evaluate:

vechord/pipeline.py

@@ -1,6 +1,7 @@
 import asyncio
 from os import environ
 from typing import Literal
+from uuid import UUID
 
 import httpx
 import msgspec
@@ -236,7 +237,7 @@ async def parse(self, req: LlamaCloudParseRequest) -> LlamaCloudParseResponse:
             )
         return self.decoder.decode(response.content)
 
-    async def get_text(self, job_id: str) -> str:
+    async def get_text(self, job_id: UUID) -> str:
         """Get the text result from a LlamaCloud job."""
         loop = asyncio.get_running_loop()
         deadline = loop.time() + EXTRACT_MAX_POLLING_TIME

vechord/provider.py

@@ -71,7 +71,9 @@ def __init__(self, model: str = "jina-reranker-m0"):
 
     async def rerank(self, query: str, chunks: list[str]) -> list[int]:
         resp = await self.query(
-            JinaRerankRequest.from_query_docs(query=query, docs=chunks)
+            JinaRerankRequest.from_query_docs(
+                query=query, documents=chunks, model=self.model
+            )
         )
         return resp.get_indices()
 
@@ -84,7 +86,7 @@ async def rerank_multimodal(
         """
         resp = await self.query(
             JinaRerankRequest.from_query_multimodal(
-                query=query, documents=chunks, doc_type=doc_type
+                query=query, documents=chunks, doc_type=doc_type, model=self.model
             )
         )
         return resp.get_indices()