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

agent/app/workflow.py
232 строки10 KB
pv_hum
Initial commit
21 май 2026, 22:25
21 май 2026, 22:250de9fa8
"""Controlled iterative agent workflow: generation → review → decision."""

from __future__ import annotations

import json
import logging
import os
from typing import Any

from .gigachat_provider import GigaChatError, chat_completion, estimate_tokens, is_configured
from .models import (
    AgentWorkflowRequest,
    AgentWorkflowResponse,
    ChangedFlags,
    IterationSummary,
    SchemaModel,
    SchemaResponse,
    TokenBudgetSummary,
    WorkflowFailure,
)
from .payload_builder import build_agent_input_payload
from .reviewer import run_schema_review
from .services import explain_changes, normalize_schema, schema_to_sql, validate_schema

logger = logging.getLogger(__name__)

MAX_ITERATIONS = int(os.getenv('AGENT_MAX_ITERATIONS', '3'))
MAX_TOTAL_TOKENS = int(os.getenv('AGENT_MAX_TOTAL_TOKENS', '6000'))


def _programmatic_decision(review_decision: str, validation_valid: bool, critical_count: int) -> str:
    if review_decision == 'reject':
        return 'fail'
    if critical_count > 0 or not validation_valid:
        return 'revise'
    if review_decision == 'accept' and validation_valid:
        return 'accept'
    if review_decision == 'revise':
        return 'revise'
    return 'accept' if validation_valid else 'revise'


def _critical_issues(review_issues: list) -> int:
    return sum(1 for i in review_issues if i.severity == 'critical')


async def run_agent_workflow(request: AgentWorkflowRequest) -> AgentWorkflowResponse | WorkflowFailure:
    if not is_configured():
        return WorkflowFailure(
            success=False,
            stopReason='no_gigachat_connection',
            message='GigaChat не настроен. Задайте GIGACHAT_CREDENTIALS и GIGACHAT_MODEL.',
        )

    if request.useDatabaseContext and not request.databaseContextSummary:
        return WorkflowFailure(
            success=False,
            stopReason='db_context_unavailable',
            message='Запрошен контекст БД, но сводка метаданных недоступна.',
        )

    total_tokens = 0
    iterations: list[IterationSummary] = []
    previous_issues: list[dict[str, Any]] | None = None
    current_schema: SchemaModel | None = request.erSchema
    dialect = request.sqlDialect or 'postgresql'
    flags = request.changedFlags or ChangedFlags()

    for iteration in range(1, MAX_ITERATIONS + 1):
        payload = build_agent_input_payload(
            mode=request.mode,
            active_view=request.activeView,
            text_prompt=request.textPrompt,
            sql=request.sql,
            er_schema=current_schema or request.erSchema,
            revision_prompt=request.revisionPrompt,
            sql_dialect=dialect,
            use_database_context=request.useDatabaseContext,
            database_context_summary=request.databaseContextSummary,
            changed_flags=flags,
            iteration=iteration,
            previous_issues=previous_issues,
        )

        gen_estimate = estimate_tokens(payload.systemPrompt) + estimate_tokens(payload.userPrompt)
        if total_tokens + gen_estimate > MAX_TOTAL_TOKENS:
            logger.warning('token budget exceeded before generation iteration=%s', iteration)
            return WorkflowFailure(
                success=False,
                stopReason='token_budget_exceeded',
                message=f'Превышен бюджет {MAX_TOTAL_TOKENS} токенов до итерации {iteration}.',
                iterationCount=iteration - 1,
                tokenBudget=TokenBudgetSummary(estimatedTotal=total_tokens, maxAllowed=MAX_TOTAL_TOKENS),
            )

        logger.info('iteration %s/%s generation started', iteration, MAX_ITERATIONS)
        try:
            raw, gen_tokens = await chat_completion(payload.systemPrompt, payload.userPrompt)
            total_tokens += gen_tokens
        except GigaChatError as exc:
            return WorkflowFailure(
                success=False,
                stopReason=exc.code.lower(),
                message=exc.message,
                iterationCount=iteration - 1,
                tokenBudget=TokenBudgetSummary(estimatedTotal=total_tokens, maxAllowed=MAX_TOTAL_TOKENS),
            )

        try:
            schema = normalize_schema(SchemaModel.model_validate(raw), source_mode=request.mode, dialect=dialect)
        except Exception as exc:
            return WorkflowFailure(
                success=False,
                stopReason='malformed_model_output',
                message=f'Не удалось нормализовать ответ модели: {exc}',
                iterationCount=iteration,
                tokenBudget=TokenBudgetSummary(estimatedTotal=total_tokens, maxAllowed=MAX_TOTAL_TOKENS),
            )

        validation = validate_schema(schema)
        review_estimate = estimate_tokens(json.dumps(schema.model_dump())) + 400
        if total_tokens + review_estimate > MAX_TOTAL_TOKENS:
            logger.warning('token budget exceeded before review iteration=%s', iteration)
            return WorkflowFailure(
                success=False,
                stopReason='token_budget_exceeded',
                message=f'Бюджет токенов исчерпан перед ревью на итерации {iteration}.',
                iterationCount=iteration,
                tokenBudget=TokenBudgetSummary(estimatedTotal=total_tokens, maxAllowed=MAX_TOTAL_TOKENS),
            )

        try:
            review, review_tokens = await run_schema_review(
                schema,
                mode=request.mode,
                dialect=dialect,
                db_summary=request.databaseContextSummary if request.useDatabaseContext else None,
                iteration=iteration,
            )
            total_tokens += review_tokens
        except GigaChatError as exc:
            return WorkflowFailure(
                success=False,
                stopReason=exc.code.lower(),
                message=exc.message,
                iterationCount=iteration,
                tokenBudget=TokenBudgetSummary(estimatedTotal=total_tokens, maxAllowed=MAX_TOTAL_TOKENS),
            )

        critical = _critical_issues(review.issues)
        decision = _programmatic_decision(review.decision, validation.valid, critical)
        iterations.append(
            IterationSummary(
                iteration=iteration,
                generationTokens=gen_tokens,
                reviewTokens=review_tokens,
                decision=decision,
                reviewDecision=review.decision,
                issueCount=len(review.issues),
            )
        )
        logger.info(
            'iteration %s/%s completed decision=%s tokens_total=%s/%s',
            iteration, MAX_ITERATIONS, decision, total_tokens, MAX_TOTAL_TOKENS,
        )

        if decision == 'accept':
            old = request.erSchema
            change_summary = explain_changes(old, schema) if old and old.entities else None
            return AgentWorkflowResponse(
                success=True,
                schema=schema,
                sql=schema_to_sql(schema, dialect),
                validation=validation,
                review=review,
                changeSummary=change_summary,
                iterationSummary=iterations,
                tokenBudget=TokenBudgetSummary(estimatedTotal=total_tokens, maxAllowed=MAX_TOTAL_TOKENS),
                stopReason='accepted',
            )

        if decision == 'fail':
            return WorkflowFailure(
                success=False,
                stopReason='validation_failed_after_review',
                message=review.summary or 'Ревьюер отклонил схему.',
                iterationCount=iteration,
                tokenBudget=TokenBudgetSummary(estimatedTotal=total_tokens, maxAllowed=MAX_TOTAL_TOKENS),
                lastSchema=schema,
                lastValidation=validation,
                lastReview=review,
            )

        current_schema = schema
        previous_issues = [i.model_dump() for i in review.issues]
        flags = ChangedFlags(promptChanged=False, sqlChanged=False, erChanged=True)

        next_estimate = estimate_tokens(json.dumps(schema.model_dump())) + 800
        if iteration >= MAX_ITERATIONS:
            return WorkflowFailure(
                success=False,
                stopReason='validation_failed_after_max_retries',
                message='Достигнут лимит итераций, критические замечания не устранены.',
                iterationCount=iteration,
                tokenBudget=TokenBudgetSummary(estimatedTotal=total_tokens, maxAllowed=MAX_TOTAL_TOKENS),
                lastSchema=schema,
                lastValidation=validation,
                lastReview=review,
            )
        if total_tokens + next_estimate > MAX_TOTAL_TOKENS:
            return WorkflowFailure(
                success=False,
                stopReason='token_budget_exceeded',
                message='Недостаточно бюджета токенов для следующей итерации.',
                iterationCount=iteration,
                tokenBudget=TokenBudgetSummary(estimatedTotal=total_tokens, maxAllowed=MAX_TOTAL_TOKENS),
                lastSchema=schema,
                lastValidation=validation,
                lastReview=review,
            )

    return WorkflowFailure(
        success=False,
        stopReason='validation_failed_after_max_retries',
        message='Исчерпан лимит итераций.',
        iterationCount=MAX_ITERATIONS,
        tokenBudget=TokenBudgetSummary(estimatedTotal=total_tokens, maxAllowed=MAX_TOTAL_TOKENS),
    )


async def workflow_generate_or_revise(request: AgentWorkflowRequest) -> AgentWorkflowResponse | WorkflowFailure:
    return await run_agent_workflow(request)