feat: support jina rerank for text & image (#62)

* feat: support jina rerank for text & image

Signed-off-by: Keming <[email protected]>

* align the interface

Signed-off-by: Keming <[email protected]>

---------

Signed-off-by: Keming <[email protected]>

Author: kemingy

vechord/model/__init__.py

@@ -13,7 +13,12 @@
     SparseEmbedding,
     UMBRELAScore,
 )
-from vechord.model.jina import JinaEmbeddingRequest, JinaEmbeddingResponse
+from vechord.model.jina import (
+    JinaEmbeddingRequest,
+    JinaEmbeddingResponse,
+    JinaRerankRequest,
+    JinaRerankResponse,
+)
 from vechord.model.llamacloud import (
     LlamaCloudMimeType,
     LlamaCloudParseRequest,
@@ -44,6 +49,8 @@
     "InputType",
     "JinaEmbeddingRequest",
     "JinaEmbeddingResponse",
+    "JinaRerankRequest",
+    "JinaRerankResponse",
     "LlamaCloudMimeType",
     "LlamaCloudParseRequest",
     "LlamaCloudParseResponse",

vechord/model/jina.py

@@ -84,3 +84,61 @@ def get_emb(self) -> np.ndarray:
         if isinstance(emb, list):
             return np.array(emb, dtype=np.float32)
         return np.frombuffer(emb, dtype=np.float32)
+
+
+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]
+    return_documents: bool = False
+
+    @classmethod
+    def from_query_docs(
+        cls,
+        query: str,
+        documents: list[str],
+        model: Literal["jina-reranker-m0", "jina-reranker-v2-base-multilingual"],
+    ) -> Self:
+        if not query or not documents:
+            raise RequestError("Query and documents must be provided")
+
+        return JinaRerankRequest(
+            model=model,
+            query=query,
+            top_n=len(documents),
+            documents=[JinaInput(text=doc) for doc in documents]
+            if model == "jina-reranker-m0"
+            else documents,
+        )
+
+    @classmethod
+    def from_query_multimodal(
+        cls,
+        query: str,
+        documents: list[str],
+        doc_type: Literal["text", "image"],
+        model: Literal["jina-reranker-m0"] = "jina-reranker-m0",
+    ) -> Self:
+        docs = [
+            JinaInput(text=doc) if doc_type == "text" else JinaInput(image=doc)
+            for doc in documents
+        ]
+        return JinaRerankRequest(
+            model=model,
+            query=query,
+            top_n=len(docs),
+            documents=docs,
+        )
+
+
+class RerankObject(msgspec.Struct, kw_only=True):
+    index: int
+    relevance_score: float
+
+
+class JinaRerankResponse(msgspec.Struct, kw_only=True):
+    results: list[RerankObject]
+
+    def get_indices(self) -> list[int]:
+        return [result.index for result in self.results]

vechord/pipeline.py

@@ -37,7 +37,7 @@
     RunRequest,
     RunResponse,
 )
-from vechord.rerank import CohereReranker
+from vechord.rerank import CohereReranker, JinaReranker
 from vechord.spec import (
     AnyOf,
     DefaultDocument,
@@ -128,7 +128,7 @@ class _Relation(Table, kw_only=True):
         "jina": JinaMultiModalEmbedding,
     },
     "ocr": {"gemini": GeminiExtractor, "llamaparse": LlamaParseExtractor},
-    "rerank": {"cohere": CohereReranker},
+    "rerank": {"cohere": CohereReranker, "jina": JinaReranker},
     "graph": {"gemini": GeminiEntityRecognizer},
     "index": {"vectorchord": IndexOption},
     "search": {"vectorchord": SearchOption},
@@ -184,8 +184,6 @@ def __post_init__(self):
             raise RequestError("Vector index is required if `index` is specified")
         if self.search and not (self.text_emb or self.multimodal_emb):
             raise RequestError("Search requires at least one embedding provider")
-        if self.search and self.rerank and self.multimodal_emb:
-            raise RequestError("Rerank only supports text")
 
     @classmethod
     def from_steps(cls, steps: list[ResourceRequest]) -> Self:
@@ -340,21 +338,17 @@ async def run_index(self, request: RunRequest, vr: "VechordRegistry") -> RunAck:
                     rels.extend(conv_rels)
                 chunks.append(chunk)
 
-        async with (
-            vr.client.get_connection() as conn,
-            limit_to_transaction_buffer_conn(conn),
-        ):
-            await vr.insert(doc)
-            for chunk in chunks:
-                await vr.insert(chunk)
-            if self.index.graph:
-                if request.input_type is not InputType.TEXT:
-                    # insert the fake chunk for image/pdf
-                    await vr.insert(fake_chunk)
-                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)
+        await vr.insert(doc)
+        for chunk in chunks:
+            await vr.insert(chunk)
+        if self.index.graph:
+            if request.input_type is not InputType.TEXT:
+                # insert the fake chunk for image/pdf
+                await vr.insert(fake_chunk)
+            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)
 
     async def graph_insert(
         self,
@@ -440,7 +434,16 @@ class Relation(_Relation):
         if self.search.graph:
             resp.extend(await self.graph_search(query, Chunk, Entity, Relation, vr))
         if self.rerank:
-            indices = await self.rerank.rerank(query, [chunk.text for chunk in resp])
+            if self.multimodal_emb:
+                indices = await self.rerank.rerank_multimodal(
+                    query=query,
+                    chunks=[chunk.text for chunk in resp],
+                    doc_type=resp.chunk_type,
+                )
+            else:
+                indices = await self.rerank.rerank(
+                    query=query, chunks=[chunk.text for chunk in resp]
+                )
             resp.reorder(indices)
         if self.evaluate:
             resp.metrics = await self.evaluate.evaluate_with_estimation(

vechord/provider.py

@@ -13,6 +13,8 @@
     GeminiGenerateResponse,
     JinaEmbeddingRequest,
     JinaEmbeddingResponse,
+    JinaRerankRequest,
+    JinaRerankResponse,
     VoyageEmbeddingRequest,
     VoyageEmbeddingResponse,
     VoyageMultiModalEmbeddingRequest,
@@ -22,6 +24,7 @@
     GEMINI_EMBEDDING_RPS,
     GEMINI_GENERATE_RPS,
     JINA_EMBEDDING_RPS,
+    JINA_RERANK_RPS,
     VOYAGE_EMBEDDING_RPS,
     RateLimitTransport,
 )
@@ -70,7 +73,6 @@ def __init__(self, model: str = "gemini-2.5-flash"):
         )
 
     async def query(self, req: GeminiGenerateRequest) -> GeminiGenerateResponse:
-        """Query the Gemini model with a request."""
         response = await self.client.post(self.url, content=self.encoder.encode(req))
         if response.is_error:
             raise HTTPCallError(
@@ -107,7 +109,6 @@ def __init__(
         self.decoder = msgspec.json.Decoder(GeminiEmbeddingResponse)
 
     async def query(self, req: GeminiEmbeddingRequest) -> GeminiEmbeddingResponse:
-        """Query the Gemini embedding model with a request."""
         response = await self.client.post(self.url, content=self.encoder.encode(req))
         if response.is_error:
             raise HTTPCallError(
@@ -137,7 +138,6 @@ def __init__(self, model: str = "jina-embeddings-v4", dim: int = 2048):
         self.decoder = msgspec.json.Decoder(JinaEmbeddingResponse)
 
     async def query(self, req: JinaEmbeddingRequest) -> JinaEmbeddingResponse:
-        """Query the Jina embedding model with a request."""
         response = await self.client.post(self.url, content=self.encoder.encode(req))
         if response.is_error:
             raise HTTPCallError(
@@ -146,6 +146,34 @@ async def query(self, req: JinaEmbeddingRequest) -> JinaEmbeddingResponse:
         return self.decoder.decode(response.content)
 
 
+class JinaRerankProvider(BaseProvider):
+    """Jina Rerank Provider."""
+
+    PROVIDER_NAME = "JINA"
+
+    def __init__(self, model: str = "jina-reranker-m0"):
+        super().__init__(model)
+        self.client = httpx.AsyncClient(
+            headers={
+                "Content-Type": "application/json",
+                "Authorization": f"Bearer {self.api_key}",
+            },
+            timeout=httpx.Timeout(120.0, connect=10.0),
+            transport=RateLimitTransport(max_per_second=JINA_RERANK_RPS),
+        )
+        self.url = "https://api.jina.ai/v1/rerank"
+        self.encoder = msgspec.json.Encoder()
+        self.decoder = msgspec.json.Decoder(JinaRerankResponse)
+
+    async def query(self, req: JinaRerankRequest) -> JinaRerankResponse:
+        response = await self.client.post(self.url, content=self.encoder.encode(req))
+        if response.is_error:
+            raise HTTPCallError(
+                "Failed to query Jina rerank", response.status_code, response.text
+            )
+        return self.decoder.decode(response.content)
+
+
 class VoyageEmbeddingProvider(BaseProvider):
     """Voyage Embedding Provider."""
 
@@ -169,7 +197,6 @@ def __init__(self, model: str = "voyage-3.5", dim: int = 1024):
     async def query(
         self, req: VoyageEmbeddingRequest | VoyageMultiModalEmbeddingRequest
     ) -> VoyageEmbeddingResponse:
-        """Query the Voyage embedding model with a request."""
         response = await self.client.post(self.url, content=self.encoder.encode(req))
         if response.is_error:
             raise HTTPCallError(

vechord/rerank.py

@@ -4,6 +4,8 @@
 from os import environ
 from typing import TypeVar
 
+from vechord.model import JinaRerankRequest
+from vechord.provider import JinaRerankProvider
 from vechord.spec import Table
 
 T = TypeVar("T", bound=Table)
@@ -15,9 +17,19 @@ async def rerank(self, query: str, chunks: list[str]) -> list[int]:
         """Return the indices of the reranked chunks."""
         raise NotImplementedError
 
+    @abstractmethod
+    async def rerank_multimodal(
+        self, query: str, chunks: list[str], doc_type: str
+    ) -> list[int]:
+        """Return the indices of the reranked multimodal chunks."""
+        raise NotImplementedError
+
 
 class CohereReranker(BaseReranker):
-    """Rerank chunks using Cohere API (requires env `COHERE_API_KEY`)."""
+    """Rerank chunks using Cohere API (requires env `COHERE_API_KEY`).
+
+    Only supports rerank documents.
+    """
 
     def __init__(self, model: str = "rerank-v3.5"):
         self.api_key = environ.get("COHERE_API_KEY")
@@ -45,6 +57,38 @@ async def rerank(self, query: str, chunks: list[str]) -> list[int]:
         )
         return [item.index for item in resp.results]
 
+    async def rerank_multimodal(
+        self, query: str, chunks: list[str], doc_type: str
+    ) -> list[int]:
+        raise NotImplementedError("Cohere does not support multimodal reranking.")
+
+
+class JinaReranker(BaseReranker, JinaRerankProvider):
+    """Rerank chunks using Jina Rerank API (requires env `JINA_API_KEY`)."""
+
+    def __init__(self, model: str = "jina-reranker-m0"):
+        super().__init__(model)
+
+    async def rerank(self, query: str, chunks: list[str]) -> list[int]:
+        resp = await self.query(
+            JinaRerankRequest.from_query_docs(query=query, docs=chunks)
+        )
+        return resp.get_indices()
+
+    async def rerank_multimodal(
+        self, query: str, chunks: list[str], doc_type: str
+    ) -> list[int]:
+        """
+        Args:
+            doc_type: "text" or "image"
+        """
+        resp = await self.query(
+            JinaRerankRequest.from_query_multimodal(
+                query=query, documents=chunks, doc_type=doc_type
+            )
+        )
+        return resp.get_indices()
+
 
 class ReciprocalRankFusion:
     """Fuse chunks using reciprocal rank."""

vechord/utils.py

@@ -12,6 +12,7 @@
 VOYAGE_EMBEDDING_RPS = 33.33
 # https://jina.ai/api-dashboard/rate-limit
 JINA_EMBEDDING_RPS = 8.33
+JINA_RERANK_RPS = 8.33
 
 
 class RateLimitTransport(httpx.AsyncHTTPTransport):