RAG/api_models.py at master · Augustus-118/RAG · GitHub

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
32
33
"""
API request and response models
"""
from pydantic import BaseModel, Field
from typing import List, Optional

class QueryRequest(BaseModel):
    question: str = Field(..., min_length=1, max_length=500)
    user_id: str = Field(default="anonymous")

    class Config:
        json_schema_extra = {
            "example": {
                "question": "What is RAG?",
                "user_id": "user_123"
            }
        }

class QueryResponse(BaseModel):
    answer: str
    sources: List[str]

class ErrorResponse(BaseModel):
    error: str
    detail: Optional[str] = None

class HealthResponse(BaseModel):
    status: str
    version: str

class StatsResponse(BaseModel):
    total_chunks: int
    total_documents: int