Spaces:
Runtime error
Runtime error
File size: 4,372 Bytes
37a6ee1 | 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 40 41 42 43 44 45 46 47 48 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 85 86 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 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 | import re
from pydantic import BaseModel, Field, field_validator
from typing import List, Any
# 1. Models for GET /datasets
class DatasetInfo(BaseModel):
db_name: str
display_name: str
description: str
columns: List[str]
row_count: int
is_sample: bool
class DatasetsResponse(BaseModel):
datasets: List[DatasetInfo]
total: int
# 2. Models for GET /datasets/{db_name}/schema
class ColumnDetail(BaseModel):
name: str
type: str
sample_values: List[Any]
class DatasetSchemaResponse(BaseModel):
db_name: str
table_name: str
columns: List[ColumnDetail]
row_count: int
# 3. Model for POST /upload
class UploadResponse(BaseModel):
success: bool
db_name: str
display_name: str
columns: List[str]
row_count: int
message: str
# 4. Model for DELETE /datasets/{db_name}
class DeleteResponse(BaseModel):
success: bool
message: str
# 5. Models for POST /query
class SQLResults(BaseModel):
columns: List[str]
rows: List[List[Any]]
row_count: int
execution_time_ms: float
class QueryRequest(BaseModel):
question: str = Field(..., min_length=3, max_length=300)
db_name: str
@field_validator("db_name")
@classmethod
def validate_db_name(cls, v: str) -> str:
# Enforce that db_name must only contain alphanumeric characters and underscores
if not re.match(r"^[a-zA-Z0-9_]+$", v):
raise ValueError("db_name must only contain alphanumeric characters and underscores")
return v
class QueryResponse(BaseModel):
question: str
sql: str
results: SQLResults
insight: str
retries: int
db_name: str
# 6. Model for GET /health
class HealthResponse(BaseModel):
status: str
groq_configured: bool
datasets_loaded: int
sample_datasets: List[str]
# Schemas for new Suggestions and History endpoints
class SuggestionsResponse(BaseModel):
db_name: str
suggestions: List[str]
class QueryHistoryItem(BaseModel):
question: str
sql: str
insight: str
timestamp: str
row_count: int
class QueryHistoryResponse(BaseModel):
db_name: str
history: List[QueryHistoryItem]
total: int
# 7. Models for Schema/Relational databases
class SchemaRelationship(BaseModel):
from_table: str
from_column: str
to_table: str
to_column: str
cardinality: Any = None
from_participation: Any = None
to_participation: Any = None
relationship_name: Any = None
class SchemaDatasetInfo(BaseModel):
db_name: str
display_name: str
description: str
mode: str = "schema"
is_sample: bool
total_tables: int
total_rows: int
tables: List[str]
relationships: List[SchemaRelationship]
class SchemaDatasetsResponse(BaseModel):
datasets: List[SchemaDatasetInfo]
total: int
class SchemaColumnDetail(BaseModel):
name: str
type: str
is_primary_key: bool
is_foreign_key: bool
references_table: Any = None
references_column: Any = None
sample_values: List[Any] = []
class SchemaTableDetail(BaseModel):
name: str
columns: List[SchemaColumnDetail]
row_count: int
class SchemaInfoResponse(BaseModel):
db_name: str
mode: str = "schema"
tables: List[SchemaTableDetail]
relationships: List[SchemaRelationship]
total_tables: int
total_rows: int
class TablePreviewItem(BaseModel):
table_name: str
columns: List[str]
rows: List[List[Any]]
total_rows: int
class SchemaPreviewsResponse(BaseModel):
db_name: str
previews: List[TablePreviewItem]
class SchemaQueryRequest(BaseModel):
question: str = Field(..., min_length=3, max_length=300)
db_name: str
@field_validator("db_name")
@classmethod
def validate_db_name(cls, v: str) -> str:
if not re.match(r"^[a-zA-Z0-9_]+$", v):
raise ValueError("db_name must only contain alphanumeric characters and underscores")
return v
class SchemaQueryResponse(BaseModel):
question: str
sql: str
mode: str = "schema"
results: SQLResults
insight: str
tables_used: List[str]
retries: int
db_name: str
class SchemaUploadResponse(BaseModel):
success: bool
db_name: str
mode: str = "schema"
total_tables: int
total_rows: int
tables: List[str]
relationships: List[SchemaRelationship]
message: str
|