File size: 804 Bytes
3ef53b4
 
570f7bd
c1bc4eb
570f7bd
 
370553a
 
 
 
 
570f7bd
c1bc4eb
570f7bd
 
76df10c
570f7bd
 
 
 
 
c1bc4eb
570f7bd
3ef53b4
a337fad
570f7bd
3ef53b4
 
570f7bd
c1bc4eb
570f7bd
 
 
 
c1bc4eb
570f7bd
 
 
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
34
35
36
37
38
39
from pydantic import BaseModel
from typing import List, Optional, Any, Dict


class NL2SQLRequest(BaseModel):
    query: str
    db_id: Optional[str] = None
    schema_preview: Optional[str] = None

    class Config:
        extra = "ignore"


class TraceModel(BaseModel):
    stage: str
    duration_ms: int
    token_in: int | None = 0
    token_out: int | None = 0
    cost_usd: float | None = 0
    notes: Dict[str, Any] | None = None


class NL2SQLResponse(BaseModel):
    ambiguous: bool
    sql: Optional[str] = None
    rationale: Optional[str] = None
    traces: List[Dict[str, Any]] = []
    result: Dict[str, Any] = {}


class ClarifyResponse(BaseModel):
    ambiguous: bool = True
    questions: List[str]


class ErrorResponse(BaseModel):
    error: str
    details: List[str] | None = None