File size: 639 Bytes
21b82de
 
cc246c6
2cfed75
21b82de
 
 
 
 
 
e6a77f3
 
2cfed75
3dd609a
cc246c6
e6a77f3
 
 
 
 
 
 
 
21b82de
 
 
 
 
e6a77f3
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
"""Pydantic models for API request and response bodies."""

from typing import List, Literal, Optional

from pydantic import BaseModel


class ChatRequest(BaseModel):
    """Schema for chat input message sent by the user."""

    query: str
    top_k: int = 5
    mode: Literal["qa", "summarize"] = "qa"
    session_id: str = "default"
    doc_ids: Optional[List[str]]


class Citation(BaseModel):
    """Schema for citations in LLM's response."""

    page_start: int
    page_end: int
    snippet: str


class ChatResponse(BaseModel):
    """Schema for LLM-generated assistant response."""

    answer: str
    citations: list[Citation]