Skip to content

rerank

RerankInput

Bases: BaseModel

Rerank Input

Source code in api/routes/rerank.py 
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
class RerankInput(BaseModel):
    """Rerank Input"""

    model_config = {
        "title": "Rerank Input",
        "strict": True,
    }
    query: str = Field(
        ...,
        title="Query",
        description="Query",
        examples=["Hoang Sa and Truong Sa belong to Vietnam"],
    )
    documents: list[Document] = Field(
        ...,
        title="Documents",
        description="List of documents",
        examples=[
            [
                {
                    "page_content": "Document 1",
                    "metadata": {
                        "reference_id": "1",
                    },
                },
                {
                    "page_content": "Document 2",
                    "metadata": {
                        "reference_id": "2",
                    },
                },
            ]
        ],
    )

RerankOutput

Bases: BaseModel

Rerank Output

Source code in api/routes/rerank.py 
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
class RerankOutput(BaseModel):
    """Rerank Output"""

    model_config = {
        "title": "Rerank Output",
        "strict": True,
    }
    query: str = Field(
        ...,
        title="Query",
        description="Query",
        examples=["Hoang Sa and Truong Sa belong to Vietnam"],
    )
    documents: list[DocumentWithScore] = Field(
        ...,
        title="Documents",
        description="List of documents",
        examples=[
            [
                {
                    "page_content": "Document 1",
                    "metadata": {
                        "reference_id": "1",
                    },
                    "score": 0.9,
                },
                {
                    "page_content": "Document 2",
                    "metadata": {
                        "reference_id": "2",
                    },
                    "score": 0.8,
                },
            ]
        ],
    )

rerank_documents(rerank_input, rerank_model)

Rerank the documents based on the query

Parameters:

Name Type Description Default
rerank_input RerankInput

Rerank input

 required
rerank_model RerankModel

Rerank model

 required

Returns:

Name Type Description
RerankOutput RerankOutput

Rerank output
Source code in api/routes/rerank.py 
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
@router.post(
    "",
    description="Rerank the documents based on the query",
    summary="Rerank the documents",
    response_description="List of documents with ranked score and sorted by score",
)
def rerank_documents(
    rerank_input: RerankInput,
    rerank_model: Annotated[RerankModel, Depends(get_rerank_model)],
) -> RerankOutput:
    """
    Rerank the documents based on the query

    Args:
        rerank_input (RerankInput): Rerank input
        rerank_model (RerankModel): Rerank model

    Returns:
        RerankOutput: Rerank output
    """
    rerank = Rerank(model=rerank_model)
    reranked_documents = rerank.rerank_documents(
        rerank_input.query, rerank_input.documents
    )
    return RerankOutput(query=rerank_input.query, documents=reranked_documents)