feat: support eval in dynamic pipeline

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

Author: kemingy

examples/dynamic.py

@@ -67,7 +67,7 @@ async def main():
         dir = Path.home() / "Pictures"
         await ingest(dir.glob("*.jpg"))
         res = await search("cat")
-        for item in res:
+        for item in res.chunks:
             print("=>", file_uuids.get(item.doc_id, "Unknown file"))
 
 

vechord/evaluate.py

@@ -5,7 +5,7 @@
 import msgspec
 import pytrec_eval
 
-from vechord.errors import DecodeStructuredOutputError
+from vechord.errors import DecodeStructuredOutputError, RequestError
 from vechord.model import GeminiGenerateRequest, RetrievedChunk, UMBRELAScore
 from vechord.provider import GeminiGenerateProvider
 
@@ -115,10 +115,15 @@ class GeminiUMBRELAEvaluator(BaseEvaluator, GeminiGenerateProvider):
     - paper: https://arxiv.org/pdf/2406.06519
     """
 
-    def __init__(self, model: str = "gemini-2.5-flash", relevant_threshold: int = 2):
+    def __init__(
+        self,
+        model: str = "gemini-2.5-flash",
+        relevant_threshold: int = 2,
+        k_values: Sequence[int] = (3, 5, 10),
+    ):
         super().__init__(model)
         self.relevant_threshold = relevant_threshold
-        self.k_values = (3, 5, 10)
+        self.k_values = k_values
         self.score_schema = msgspec.json.schema(UMBRELAScore)
         self.prompt = """
 Given a query and a passage, you must provide a score on an
@@ -153,6 +158,8 @@ def name(self) -> str:
         return f"gemini_umbrela_{self.model}"
 
     async def estimate(self, query: str, passage: str) -> int:
+        if not passage:
+            return 0
         content = self.prompt.format(query=query, passage=passage)
         resp = await self.query(
             GeminiGenerateRequest.from_prompt_structure_response(
@@ -171,6 +178,8 @@ async def evaluate_with_estimation(
         self, query: str, passages: list[str]
     ) -> dict[str, float]:
         """Calculate the Precision@K and Mean Reciprocal Rank (MRR)."""
+        if not query or not passages or not all(not p.strip() for p in passages):
+            raise RequestError("Query and passages must be non-empty strings.")
         scores = [await self.estimate(query, p) for p in passages]
         is_relevant = [score >= self.relevant_threshold for score in scores]
         metric = defaultdict(float)

vechord/model/__init__.py

@@ -24,7 +24,13 @@
     VoyageEmbeddingResponse,
     VoyageMultiModalEmbeddingRequest,
 )
-from vechord.model.web import InputType, ResourceRequest, RunAck, RunRequest
+from vechord.model.web import (
+    InputType,
+    ResourceRequest,
+    RunAck,
+    RunRequest,
+    RunResponse,
+)
 
 __all__ = [
     "Document",
@@ -45,6 +51,7 @@
     "RetrievedChunk",
     "RunAck",
     "RunRequest",
+    "RunResponse",
     "SparseEmbedding",
     "UMBRELAScore",
     "VoyageEmbeddingRequest",

vechord/model/web.py

@@ -35,6 +35,40 @@ class RunRequest(msgspec.Struct, kw_only=True, frozen=True):
 
 
 class RunAck(msgspec.Struct, kw_only=True, frozen=True):
+    """Acknowledgment of an index request."""
+
     name: str
     msg: str
     uid: UUID
+
+
+class SearchResponse(msgspec.Struct, kw_only=True):
+    uid: UUID
+    doc_id: UUID
+    text: str
+
+
+class RunResponse(msgspec.Struct, kw_only=True, omit_defaults=True):
+    """Response to a search request.
+
+    metrics:
+    - MRR
+    - precision@k
+    - average precision@k
+    """
+
+    chunks: list[SearchResponse] = msgspec.field(default_factory=list)
+    metrics: dict[str, float] = msgspec.field(default_factory=dict)
+
+    def extend(self, chunks: list):
+        for chunk in chunks:
+            self.chunks.append(
+                SearchResponse(
+                    uid=chunk.uid,
+                    doc_id=chunk.doc_id,
+                    text=chunk.text,
+                )
+            )
+
+    def reorder(self, indices: list[int]):
+        self.chunks = [self.chunks[i] for i in indices]

vechord/pipeline.py

@@ -21,6 +21,7 @@
 )
 from vechord.entity import GeminiEntityRecognizer
 from vechord.errors import RequestError
+from vechord.evaluate import GeminiUMBRELAEvaluator
 from vechord.extract import GeminiExtractor, LlamaParseExtractor
 from vechord.model import (
     GraphEntity,
@@ -29,6 +30,7 @@
     ResourceRequest,
     RunAck,
     RunRequest,
+    RunResponse,
 )
 from vechord.rerank import CohereReranker
 from vechord.spec import (
@@ -122,6 +124,7 @@ class _Relation(Table, kw_only=True):
     "graph": {"gemini": GeminiEntityRecognizer},
     "index": {"vectorchord": IndexOption},
     "search": {"vectorchord": SearchOption},
+    "evaluate": {"gemini": GeminiUMBRELAEvaluator},
 }
 
 
@@ -160,6 +163,7 @@ class DynamicPipeline(msgspec.Struct, kw_only=True):
     index: Optional[IndexOption] = None
     search: Optional[SearchOption] = None
     graph: Optional[GeminiEntityRecognizer] = None
+    evaluate: Optional[GeminiUMBRELAEvaluator] = None
 
     def __post_init__(self):
         if not (self.text_emb or self.multimodal_emb):
@@ -195,7 +199,9 @@ def from_steps(cls, steps: list[ResourceRequest]) -> Self:
                 calls[(step.kind).replace("-", "_")] = provider(**args)
         return msgspec.convert(calls, DynamicPipeline)
 
-    async def run(self, request: RunRequest, vr: "VechordRegistry"):
+    async def run(
+        self, request: RunRequest, vr: "VechordRegistry"
+    ) -> RunAck | RunResponse:
         """Run the dynamic pipeline with the given request."""
         if self.index:
             return await self.run_index(request, vr)
@@ -374,7 +380,9 @@ async def graph_insert(
             rel.vec = await self.text_emb.vectorize_chunk(f"{rel.description}")
             await vr.insert(rel)
 
-    async def run_search(self, request: RunRequest, vr: "VechordRegistry"):
+    async def run_search(
+        self, request: RunRequest, vr: "VechordRegistry"
+    ) -> RunResponse:
         query = request.data.decode("utf-8")
 
         # for type hint and compatibility
@@ -387,32 +395,34 @@ class Entity(_Entity):
         class Relation(_Relation):
             pass
 
-        retrieved: list[Chunk] = []
+        resp = RunResponse()
         if self.search.vector:
             vec = (
                 await self.text_emb.vectorize_query(query)
                 if self.text_emb
                 else await self.multimodal_emb.vectorize_multimodal_query(text=query)
             )
-            retrieved.extend(
+            resp.extend(
                 await vr.search_by_vector(
                     Chunk, vec, self.search.vector.topk, probe=self.search.vector.probe
                 )
             )
         if self.search.keyword:
-            retrieved.extend(
+            resp.extend(
                 await vr.search_by_keyword(Chunk, query, self.search.keyword.topk)
             )
         if self.search.graph:
-            retrieved.extend(
-                await self.graph_search(query, Chunk, Entity, Relation, vr)
-            )
+            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 retrieved]
+            indices = await self.rerank.rerank(query, [chunk.text for chunk in resp])
+            resp.reorder(indices)
+
+        if self.evaluate:
+            resp.metrics = await self.evaluate.evaluate_with_estimation(
+                query, [chunk.text for chunk in resp.chunks]
             )
-            retrieved = [retrieved[i] for i in indices]
-        return retrieved
+
+        return resp
 
     async def graph_search(
         self,