agent/app/models.py - master - pv_hum/schemaforge

agent/app/models.py
207 строк6 KB
pv_hum
Initial commit
21 май 2026, 22:25
21 май 2026, 22:250de9fa8
from typing import Any, Literal
from pydantic import BaseModel, Field, model_validator

SqlDialect = Literal['postgresql', 'mysql', 'unknown']
RelationshipType = Literal['one-to-one', 'one-to-many', 'many-to-one', 'many-to-many']
ProductMode = Literal['prompt', 'sql', 'er']
ActiveView = Literal['er', 'sql']
ReviewDecision = Literal['accept', 'revise', 'reject']
IssueSeverity = Literal['critical', 'warning', 'info']

class AttributeReference(BaseModel):
    entity: str
    attribute: str

class Attribute(BaseModel):
    name: str
    type: str
    nullable: bool = True
    primaryKey: bool = False
    unique: bool = False
    foreignKey: bool = False
    references: AttributeReference | None = None

class Entity(BaseModel):
    name: str
    displayName: str | None = None
    description: str | None = None
    attributes: list[Attribute] = Field(default_factory=list)
    x: float | None = None
    y: float | None = None

class Relationship(BaseModel):
    name: str | None = None
    fromEntity: str
    toEntity: str
    type: RelationshipType
    fromField: str | None = None
    toField: str | None = None
    label: str | None = None
    junctionEntity: str | None = None

class SchemaModel(BaseModel):
    entities: list[Entity]
    relationships: list[Relationship] = Field(default_factory=list)
    metadata: dict[str, Any] | None = None
    sourceMode: Literal['prompt', 'sql', 'er'] | None = None
    sqlDialect: SqlDialect | None = 'unknown'

class ValidationIssue(BaseModel):
    level: Literal['info', 'warning', 'error']
    code: str
    message: str
    entity: str | None = None
    attribute: str | None = None

class ValidationResult(BaseModel):
    valid: bool
    issues: list[ValidationIssue] = Field(default_factory=list)

class ChangeItem(BaseModel):
    type: Literal['added', 'removed', 'updated']
    target: str
    message: str

class ChangeSummary(BaseModel):
    items: list[ChangeItem] = Field(default_factory=list)

class SchemaResponse(BaseModel):
    schema: SchemaModel
    sql: str
    validation: ValidationResult
    changeSummary: ChangeSummary | None = None

class ChangedFlags(BaseModel):
    sqlChanged: bool = False
    erChanged: bool = False
    promptChanged: bool = False

class AgentInputPayload(BaseModel):
    systemPrompt: str
    userPrompt: str
    mode: str
    activeView: str
    includedContext: dict[str, bool] = Field(default_factory=dict)

class ReviewIssue(BaseModel):
    severity: IssueSeverity
    issueCode: str
    message: str
    entity: str | None = None
    attribute: str | None = None
    recommendation: str | None = None

class ReviewResult(BaseModel):
    decision: ReviewDecision
    issues: list[ReviewIssue] = Field(default_factory=list)
    summary: str = ''

class IterationSummary(BaseModel):
    iteration: int
    generationTokens: int = 0
    reviewTokens: int = 0
    decision: str = ''
    reviewDecision: str = ''
    issueCount: int = 0

class TokenBudgetSummary(BaseModel):
    estimatedTotal: int
    maxAllowed: int

class AgentWorkflowRequest(BaseModel):
    mode: ProductMode
    activeView: ActiveView
    textPrompt: str | None = None
    sql: str | None = None
    erSchema: SchemaModel | None = None
    revisionPrompt: str | None = None
    sqlDialect: SqlDialect = 'postgresql'
    useDatabaseContext: bool = False
    databaseContextSummary: str | None = None
    changedFlags: ChangedFlags | None = None

class AgentWorkflowResponse(BaseModel):
    success: bool = True
    schema: SchemaModel
    sql: str
    validation: ValidationResult
    review: ReviewResult | None = None
    changeSummary: ChangeSummary | None = None
    iterationSummary: list[IterationSummary] = Field(default_factory=list)
    tokenBudget: TokenBudgetSummary | None = None
    stopReason: str = 'accepted'

class WorkflowFailure(BaseModel):
    success: bool = False
    stopReason: str
    message: str
    iterationCount: int = 0
    tokenBudget: TokenBudgetSummary | None = None
    lastSchema: SchemaModel | None = None
    lastValidation: ValidationResult | None = None
    lastReview: ReviewResult | None = None

class GenerateSchemaRequest(BaseModel):
    mode: ProductMode = 'prompt'
    activeView: ActiveView = 'er'
    text: str | None = None
    textPrompt: str | None = None
    sql: str | None = None
    erSchema: SchemaModel | None = None
    revisionPrompt: str | None = None
    dialect: SqlDialect = 'postgresql'
    sqlDialect: SqlDialect | None = None
    useDatabaseContext: bool = False
    databaseContextSummary: str | None = None
    changedFlags: ChangedFlags | None = None

    def resolved_prompt(self) -> str:
        return (self.textPrompt or self.text or '').strip()

    def resolved_dialect(self) -> SqlDialect:
        return self.sqlDialect or self.dialect

class ReviseSchemaRequest(BaseModel):
    mode: ProductMode = 'prompt'
    activeView: ActiveView = 'er'
    currentSchema: SchemaModel
    currentSql: str | None = None
    sql: str | None = None
    editPrompt: str | None = None
    revisionPrompt: str | None = None
    dialect: SqlDialect = 'postgresql'
    sqlDialect: SqlDialect | None = None
    useDatabaseContext: bool = False
    databaseContextSummary: str | None = None
    changedFlags: ChangedFlags | None = None

    def resolved_revision(self) -> str:
        return (self.revisionPrompt or self.editPrompt or '').strip()

    def resolved_dialect(self) -> SqlDialect:
        return self.sqlDialect or self.dialect

class ValidateSchemaRequest(BaseModel):
    schema: SchemaModel

class GenerateSqlRequest(BaseModel):
    schema: SchemaModel
    dialect: SqlDialect = 'postgresql'

class ExplainChangesRequest(BaseModel):
    oldSchema: SchemaModel
    newSchema: SchemaModel

class ParseSqlRequest(BaseModel):
    sql: str
    dialect: SqlDialect = 'postgresql'

    @model_validator(mode='after')
    def ensure_sql(self):
        if not self.sql.strip():
            raise ValueError('SQL must not be empty')
        return self

class FromErRequest(BaseModel):
    erSchema: SchemaModel
    dialect: SqlDialect = 'postgresql'