mirror of
https://github.com/langgenius/dify.git
synced 2026-11-11 05:41:09 +08:00
60 lines
2.4 KiB
Python
60 lines
2.4 KiB
Python
"""Authenticated, idempotent ingestion of Agent observations."""
|
|
|
|
from typing import Literal
|
|
from uuid import UUID
|
|
|
|
from dify_agent.protocol.tracing import AgentTraceBatch
|
|
from flask import request
|
|
from flask_restx import Resource
|
|
from pydantic import ValidationError
|
|
from werkzeug.exceptions import BadRequest, Conflict, NotFound, RequestEntityTooLarge
|
|
|
|
from configs import dify_config
|
|
from controllers.common.schema import register_response_schema_models, register_schema_models
|
|
from controllers.inner_api import inner_api_ns
|
|
from controllers.inner_api.wraps import agent_inner_api_only
|
|
from extensions.ext_application_services import application_services
|
|
from fields.base import ResponseModel
|
|
from services.tracing.errors import (
|
|
TraceCollectionConflictError,
|
|
TraceCollectionError,
|
|
TraceCollectionLimitError,
|
|
TraceCollectionNotFoundError,
|
|
)
|
|
|
|
|
|
class AgentTraceBatchResponse(ResponseModel):
|
|
result: Literal["success"]
|
|
batch_id: UUID
|
|
accepted_spans: int
|
|
|
|
|
|
register_schema_models(inner_api_ns, AgentTraceBatch)
|
|
register_response_schema_models(inner_api_ns, AgentTraceBatchResponse)
|
|
|
|
|
|
@inner_api_ns.route("/agent/tracing/batches")
|
|
class AgentTraceBatchesApi(Resource):
|
|
@agent_inner_api_only
|
|
@inner_api_ns.doc("inner_agent_trace_batches")
|
|
@inner_api_ns.expect(inner_api_ns.models[AgentTraceBatch.__name__])
|
|
@inner_api_ns.response(200, "Durably accepted", inner_api_ns.models[AgentTraceBatchResponse.__name__])
|
|
def post(self):
|
|
request.max_content_length = dify_config.OPS_AGENT_TRACE_BATCH_MAX_BYTES
|
|
try:
|
|
payload = AgentTraceBatch.model_validate(inner_api_ns.payload or {})
|
|
accepted = application_services().tracing.agent_traces.ingest(payload)
|
|
except ValidationError as error:
|
|
raise BadRequest("Invalid Agent trace batch") from error
|
|
except TraceCollectionNotFoundError as error:
|
|
raise NotFound(error.code) from error
|
|
except TraceCollectionConflictError as error:
|
|
raise Conflict(error.code) from error
|
|
except TraceCollectionLimitError as error:
|
|
raise RequestEntityTooLarge(error.code) from error
|
|
except TraceCollectionError as error:
|
|
raise BadRequest(error.code) from error
|
|
return AgentTraceBatchResponse(result="success", batch_id=payload.batch_id, accepted_spans=accepted).model_dump(
|
|
mode="json"
|
|
)
|