Add models.

main
KKlochko 1 month ago
parent d5d204c642
commit e40c0af2f7
Signed by: KKlochko
GPG Key ID: 572ECCD219BBA91B

@ -0,0 +1,119 @@
from pydantic import BaseModel, Field, field_validator
from typing import List, Optional, Union, Dict, Any, Tuple
import os
# ==================== Request Models ====================
class GeneralRequest(BaseModel):
text: Optional[str] = None
texts: Optional[List[str]] = None
entities: List[str]
prompt: Optional[str] = None
threshold: Optional[float] = Field(default=0.5, ge=0.0, le=1.0)
@field_validator('entities')
@classmethod
def validate_entities(cls, v):
if not v or len(v) == 0:
raise ValueError("entities list cannot be empty")
return v
@field_validator('text', 'texts')
@classmethod
def validate_text_length(cls, v):
max_length = int(os.getenv("MAX_TEXT_LENGTH", "10000"))
if isinstance(v, str) and len(v) > max_length:
raise ValueError(f"text exceeds maximum length of {max_length}")
if isinstance(v, list):
for text in v:
if len(text) > max_length:
raise ValueError(f"text exceeds maximum length of {max_length}")
return v
class RelationRule(BaseModel):
relation: str
pairs_filter: List[Tuple[str, str]]
distance: int
class RelationExtractionRequest(BaseModel):
text: Optional[str] = None
texts: Optional[List[str]] = None
relations: List[str]
entities: Optional[List[str]] = None
rules: Optional[List[RelationRule]] = None
prompt: Optional[str] = None
threshold: Optional[float] = Field(default=0.5, ge=0.0, le=1.0)
@field_validator('relations')
@classmethod
def validate_relations(cls, v):
if not v or len(v) == 0:
raise ValueError("relations list cannot be empty")
return v
@field_validator('text', 'texts')
@classmethod
def validate_text_length(cls, v):
max_length = int(os.getenv("MAX_TEXT_LENGTH", "10000"))
if isinstance(v, str) and len(v) > max_length:
raise ValueError(f"text exceeds maximum length of {max_length}")
if isinstance(v, list):
for text in v:
if len(text) > max_length:
raise ValueError(f"text exceeds maximum length of {max_length}")
return v
class SummarizationRequest(BaseModel):
text: Optional[str] = None
texts: Optional[List[str]] = None
prompt: Optional[str] = None
threshold: Optional[float] = Field(default=0.5, ge=0.0, le=1.0)
@field_validator('text', 'texts')
@classmethod
def validate_text_length(cls, v):
max_length = int(os.getenv("MAX_TEXT_LENGTH", "10000"))
if isinstance(v, str) and len(v) > max_length:
raise ValueError(f"text exceeds maximum length of {max_length}")
if isinstance(v, list):
for text in v:
if len(text) > max_length:
raise ValueError(f"text exceeds maximum length of {max_length}")
return v
# ==================== Response Models ====================
class EntityOutput(BaseModel):
entity: str
span: str
start: int
end: int
score: float
class RelationOutput(BaseModel):
source: EntityOutput
relation: str
target: EntityOutput
score: float
class SummaryOutput(BaseModel):
text: str
start: int
end: int
score: float
# ==================== Error Response Model ====================
class ErrorDetail(BaseModel):
code: str
message: str
details: Dict[str, Any] = Field(default_factory=dict)
class ErrorResponse(BaseModel):
error: ErrorDetail
Loading…
Cancel
Save