parent
d5d204c642
commit
e40c0af2f7
@ -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…
Reference in new issue