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

agent/app/main.py
158 строк6 KB
pv_hum
Initial commit
21 май 2026, 22:25
21 май 2026, 22:250de9fa8
import logging
import os
from pathlib import Path

from dotenv import load_dotenv

load_dotenv(Path(__file__).resolve().parent.parent / '.env')
from fastapi import FastAPI, HTTPException
from fastapi.responses import JSONResponse

from .gigachat_provider import GigaChatError, get_model, is_configured
from .models import *
from .services import explain_changes, parse_sql_to_schema, schema_to_sql, validate_schema
from .workflow import workflow_generate_or_revise

logging.basicConfig(level=logging.INFO, format='%(asctime)s %(levelname)s %(name)s %(message)s')
logger = logging.getLogger(__name__)

app = FastAPI(title='SchemaForge Agent', version='1.0.0')


@app.get('/health')
async def health():
    from .gigachat_provider import check_connection

    gigachat = await check_connection()
    return {
        'status': 'ok',
        'provider': 'gigachat',
        'configured': is_configured(),
        'model': os.getenv('GIGACHAT_MODEL', '') if is_configured() else None,
        'gigachat': gigachat,
        'maxIterations': int(os.getenv('AGENT_MAX_ITERATIONS', '3')),
        'maxTotalTokens': int(os.getenv('AGENT_MAX_TOTAL_TOKENS', '6000')),
    }


def _to_schema_response(result: AgentWorkflowResponse) -> SchemaResponse:
    return SchemaResponse(
        schema=result.schema,
        sql=result.sql,
        validation=result.validation,
        changeSummary=result.changeSummary,
    )


def _workflow_request_from_generate(req: GenerateSchemaRequest) -> AgentWorkflowRequest:
    flags = req.changedFlags or ChangedFlags(
        promptChanged=bool(req.resolved_prompt()),
        sqlChanged=bool(req.sql),
        erChanged=bool(req.erSchema and req.erSchema.entities),
    )
    return AgentWorkflowRequest(
        mode=req.mode,
        activeView=req.activeView,
        textPrompt=req.resolved_prompt() or None,
        sql=req.sql,
        erSchema=req.erSchema,
        revisionPrompt=None,
        sqlDialect=req.resolved_dialect(),
        useDatabaseContext=req.useDatabaseContext,
        databaseContextSummary=req.databaseContextSummary,
        changedFlags=flags,
    )


def _workflow_request_from_revise(req: ReviseSchemaRequest) -> AgentWorkflowRequest:
    flags = req.changedFlags or ChangedFlags(
        promptChanged=False,
        sqlChanged=bool(req.currentSql or req.sql),
        erChanged=True,
    )
    return AgentWorkflowRequest(
        mode=req.mode,
        activeView=req.activeView,
        textPrompt=None,
        sql=req.sql or req.currentSql,
        erSchema=req.currentSchema,
        revisionPrompt=req.resolved_revision(),
        sqlDialect=req.resolved_dialect(),
        useDatabaseContext=req.useDatabaseContext,
        databaseContextSummary=req.databaseContextSummary,
        changedFlags=flags,
    )


@app.post('/agent/workflow')
async def agent_workflow(request: AgentWorkflowRequest):
    result = await workflow_generate_or_revise(request)
    if isinstance(result, WorkflowFailure):
        return JSONResponse(status_code=422, content=result.model_dump())
    return result


@app.post('/agent/generate')
async def generate(request: GenerateSchemaRequest):
    wf = _workflow_request_from_generate(request)
    if not wf.textPrompt and not wf.sql and not (wf.erSchema and wf.erSchema.entities):
        raise HTTPException(status_code=400, detail='Не задан контекст для генерации.')
    try:
        result = await workflow_generate_or_revise(wf)
    except GigaChatError as exc:
        raise HTTPException(status_code=503, detail={'code': exc.code, 'message': exc.message}) from exc
    if isinstance(result, WorkflowFailure):
        return JSONResponse(status_code=422, content=result.model_dump())
    return _to_schema_response(result)


@app.post('/agent/revise')
async def revise(request: ReviseSchemaRequest):
    if not request.resolved_revision():
        raise HTTPException(status_code=400, detail='Не задан промпт доработки.')
    wf = _workflow_request_from_revise(request)
    try:
        result = await workflow_generate_or_revise(wf)
    except GigaChatError as exc:
        raise HTTPException(status_code=503, detail={'code': exc.code, 'message': exc.message}) from exc
    if isinstance(result, WorkflowFailure):
        return JSONResponse(status_code=422, content=result.model_dump())
    return _to_schema_response(result)


@app.post('/agent/validate', response_model=ValidationResult)
async def validate(request: ValidateSchemaRequest):
    return validate_schema(request.schema)


@app.post('/agent/sql')
async def sql(request: GenerateSqlRequest):
    return {'sql': schema_to_sql(request.schema, request.dialect)}


@app.post('/agent/explain-changes', response_model=ChangeSummary)
async def changes(request: ExplainChangesRequest):
    return explain_changes(request.oldSchema, request.newSchema)


@app.post('/agent/parse-sql', response_model=SchemaResponse)
async def parse_sql(request: ParseSqlRequest):
    schema = parse_sql_to_schema(request.sql, request.dialect)
    validation = validate_schema(schema)
    return SchemaResponse(
        schema=schema,
        sql=schema_to_sql(schema, request.dialect),
        validation=validation,
        changeSummary=ChangeSummary(items=[ChangeItem(type='added', target='schema', message='SQL разобран во внутреннюю модель схемы.')]),
    )


@app.post('/agent/from-er', response_model=SchemaResponse)
async def from_er(request: FromErRequest):
    validation = validate_schema(request.erSchema)
    return SchemaResponse(
        schema=request.erSchema,
        sql=schema_to_sql(request.erSchema, request.dialect),
        validation=validation,
        changeSummary=ChangeSummary(items=[ChangeItem(type='updated', target='schema', message='SQL сгенерирован из ER-схемы.')]),
    )