/ src / python / txtai / api / routers / reranker.py
reranker.py
 1  """
 2  Defines API paths for reranking endpoints.
 3  """
 4  
 5  from typing import List, Optional
 6  
 7  from fastapi import APIRouter, Body
 8  
 9  from .. import application
10  from ..route import EncodingAPIRoute
11  
12  router = APIRouter(route_class=EncodingAPIRoute)
13  
14  
15  @router.get("/rerank")
16  def rerank(query: str, limit: Optional[int] = 3, factor: Optional[int] = 10):
17      """
18      Queries an embeddings database and reranks the results with a similarity pipeline.
19  
20      Args:
21          query: query text
22          limit: maximum results
23          factor: factor to multiply limit by for the initial embeddings search
24  
25      Returns:
26          query results
27      """
28  
29      return application.get().pipeline("reranker", (query, limit, factor))
30  
31  
32  @router.post("/batchrerank")
33  def batchrerank(queries: List[str] = Body(...), limit: Optional[int] = Body(default=3), factor: Optional[int] = Body(default=10)):
34      """
35      Queries an embeddings database and reranks the results with a similarity pipeline.
36  
37      Args:
38          queries: list of queries
39          limit: maximum results
40          factor: factor to multiply limit by for the initial embeddings search
41  
42      Returns:
43          query results
44      """
45  
46      return application.get().pipeline("reranker", (queries, limit, factor))