from dataclasses import dataclass

from pydantic_ai import Agent, RunContext, UsageLimits
from pydantic_ai.models.test import TestModel


@dataclass
class AppDeps:
    org_id: str
    actor_id: str


analysis_agent = Agent(
    TestModel(custom_output_text="A-001 계정의 처리 대기 항목은 3건입니다."),
    deps_type=AppDeps,
    instructions="분석/처리 대기 항목 질문만 처리합니다.",
)

router_agent = Agent(
    TestModel(
        call_tools=["ask_analysis"],
        custom_output_text="분석 agent 결과를 바탕으로 답변했습니다.",
    ),
    deps_type=AppDeps,
    instructions="필요하면 전문 agent tool을 호출한 뒤 최종 답변합니다.",
)


@router_agent.tool
async def ask_analysis(ctx: RunContext[AppDeps], question: str) -> str:
    result = await analysis_agent.run(
        question,
        deps=ctx.deps,
        usage=ctx.usage,
    )
    return result.output


result = router_agent.run_sync(
    "A-001 계정 처리 대기 항목 알려줘",
    deps=AppDeps(org_id="ORG-001", actor_id="ACTOR-007"),
    usage_limits=UsageLimits(request_limit=5, tool_calls_limit=2),
)
print(result.output)
print(result.usage)

from typing import Literal

from pydantic import BaseModel
from pydantic_ai import Agent


class RouteDecision(BaseModel):
    domain: Literal["profile", "analysis", "catalog", "general"]
    reason: str


router_agent = Agent(
    "openai:gpt-5-mini",
    output_type=RouteDecision,
    instructions="사용자 요청을 application workflow 도메인으로 분류합니다.",
    defer_model_check=True,
)

profile_agent = Agent("openai:gpt-5-mini", instructions="프로필 업무를 처리합니다.", defer_model_check=True)
analysis_agent = Agent("openai:gpt-5-mini", instructions="분석 업무를 처리합니다.", defer_model_check=True)
catalog_agent = Agent("openai:gpt-5-mini", instructions="리소스 업무를 처리합니다.", defer_model_check=True)
general_agent = Agent("openai:gpt-5-mini", instructions="일반 업무 안내를 처리합니다.", defer_model_check=True)


async def answer_by_route(user_message: str) -> str:
    route = await router_agent.run(user_message)

    agents = {
        "profile": profile_agent,
        "analysis": analysis_agent,
        "catalog": catalog_agent,
        "general": general_agent,
    }
    selected_agent = agents[route.output.domain]
    result = await selected_agent.run(user_message)
    return result.output

from dataclasses import dataclass


@dataclass
class SplitDecisionInput:
    distinct_toolsets: bool
    distinct_permissions: bool
    needs_different_model: bool
    requires_separate_trace: bool
    simple_read_only_request: bool


def should_split_agent(case: SplitDecisionInput) -> bool:
    if case.simple_read_only_request and not case.distinct_permissions:
        return False

    split_signals = [
        case.distinct_toolsets,
        case.distinct_permissions,
        case.needs_different_model,
        case.requires_separate_trace,
    ]
    return sum(split_signals) >= 2


assert should_split_agent(
    SplitDecisionInput(
        distinct_toolsets=True,
        distinct_permissions=True,
        needs_different_model=False,
        requires_separate_trace=True,
        simple_read_only_request=False,
    )
)

import asyncio
from dataclasses import dataclass, field
from typing import Literal

from pydantic import BaseModel
from pydantic_graph import GraphBuilder, StepContext, TypeExpression


Domain = Literal["profile", "analysis", "catalog", "general"]


class AgentAnswer(BaseModel):
    domain: Domain
    text: str
    requires_confirmation: bool = False


@dataclass
class WorkflowGraphState:
    org_id: str
    actor_id: str
    route: Domain | None = None
    trace_events: list[str] = field(default_factory=list)


g = GraphBuilder(
    state_type=WorkflowGraphState,
    input_type=str,
    output_type=AgentAnswer,
)


@g.step
async def route_request(ctx: StepContext[WorkflowGraphState, None, str]) -> Domain:
    message = ctx.inputs
    if "레코드" in message or "처리 대기" in message:
        route: Domain = "analysis"
    elif "리소스" in message:
        route = "catalog"
    elif "계정" in message:
        route = "profile"
    else:
        route = "general"

    ctx.state.route = route
    ctx.state.trace_events.append(f"route={route}")
    return route


@g.step
async def profile_step(ctx: StepContext[WorkflowGraphState, None, Domain]) -> AgentAnswer:
    ctx.state.trace_events.append("profile_agent")
    return AgentAnswer(domain="profile", text="계정 정보를 조회했습니다.")


@g.step
async def analysis_step(ctx: StepContext[WorkflowGraphState, None, Domain]) -> AgentAnswer:
    ctx.state.trace_events.append("analysis_agent")
    return AgentAnswer(domain="analysis", text="처리 대기 항목 요약을 조회했습니다.")


@g.step
async def catalog_step(ctx: StepContext[WorkflowGraphState, None, Domain]) -> AgentAnswer:
    ctx.state.trace_events.append("catalog_agent")
    return AgentAnswer(domain="catalog", text="리소스 현황을 조회했습니다.")


@g.step
async def general_step(ctx: StepContext[WorkflowGraphState, None, Domain]) -> AgentAnswer:
    ctx.state.trace_events.append("general_agent")
    return AgentAnswer(domain="general", text="처리 가능한 application workflow를 안내했습니다.")


@g.step
async def policy_validator(ctx: StepContext[WorkflowGraphState, None, AgentAnswer]) -> AgentAnswer:
    if "확인" in ctx.inputs.text or "수정" in ctx.inputs.text:
        ctx.inputs.requires_confirmation = True
    ctx.state.trace_events.append("policy_validated")
    return ctx.inputs


g.add(
    g.edge_from(g.start_node).to(route_request),
    g.edge_from(route_request).to(
        g.decision()
        .branch(g.match(TypeExpression[Literal["profile"]]).to(profile_step))
        .branch(g.match(TypeExpression[Literal["analysis"]]).to(analysis_step))
        .branch(g.match(TypeExpression[Literal["catalog"]]).to(catalog_step))
        .branch(g.match(TypeExpression[Literal["general"]]).to(general_step))
    ),
    g.edge_from(profile_step, analysis_step, catalog_step, general_step).to(policy_validator),
    g.edge_from(policy_validator).to(g.end_node),
)

workflow_graph = g.build()


async def main():
    state = WorkflowGraphState(org_id="ORG-001", actor_id="ACTOR-007")
    answer = await workflow_graph.run(state=state, inputs="A-001 계정 처리 대기 항목 요약")
    print(answer)
    print(state.trace_events)


asyncio.run(main())

from dataclasses import dataclass, field

from pydantic_ai import Agent, UsageLimits


@dataclass
class OrchestrationTrace:
    request_id: str
    events: list[str] = field(default_factory=list)


analysis_agent = Agent(
    "openai:gpt-5-mini",
    instructions="분석 조회 요청만 처리합니다.",
    defer_model_check=True,
)


async def run_analysis_with_fallback(
    user_message: str,
    trace: OrchestrationTrace,
) -> str:
    trace.events.append("analysis_agent:start")

    try:
        result = await analysis_agent.run(
            user_message,
            usage_limits=UsageLimits(request_limit=3, total_tokens_limit=2_000),
        )
    except TimeoutError:
        trace.events.append("analysis_agent:timeout")
        return "분석 조회가 지연되고 있습니다. 잠시 후 다시 시도해 주세요."
    except Exception as exc:
        trace.events.append(f"analysis_agent:error:{type(exc).__name__}")
        return "분석 조회 중 오류가 발생했습니다. 요청 ID로 관리자에게 문의해 주세요."

    trace.events.append("analysis_agent:success")
    return result.output

from opentelemetry import trace


tracer = trace.get_tracer("sample.ai.agent-service")


async def handle_agent_request(
    *,
    org_id: str,
    actor_id: str,
    request_id: str,
    message: str,
) -> str:
    with tracer.start_as_current_span("sample.agent.request") as span:
        span.set_attribute("sample.org_id", org_id)
        span.set_attribute("sample.actor_id", actor_id)
        span.set_attribute("sample.request_id", request_id)
        span.set_attribute("sample.module", "conversation")

        # 실제 구현에서는 여기서 agent.run(...)을 호출합니다.
        return f"처리 대상 메시지: {message}"

import os

import logfire


os.environ["OTEL_EXPORTER_OTLP_ENDPOINT"] = "http://otel-collector:4318"

logfire.configure(
    service_name="sample-agent-service",
    send_to_logfire=False,
)
logfire.instrument_pydantic_ai()

from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
from opentelemetry.sdk.trace import TracerProvider
from opentelemetry.sdk.trace.export import BatchSpanProcessor
from opentelemetry.trace import set_tracer_provider

from pydantic_ai import Agent
from pydantic_ai.models.instrumented import InstrumentationSettings


tracer_provider = TracerProvider()
tracer_provider.add_span_processor(BatchSpanProcessor(OTLPSpanExporter()))
set_tracer_provider(tracer_provider)

Agent.instrument_all(
    InstrumentationSettings(
        version=2,
        use_aggregated_usage_attribute_names=True,
    )
)

from pydantic_ai import Agent
from pydantic_ai.capabilities import Instrumentation
from pydantic_ai.models.instrumented import InstrumentationSettings


instrumented_agent = Agent(
    "openai:gpt-5-mini",
    name="catalog_agent",
    instructions="리소스 조회 요청만 처리합니다.",
    defer_model_check=True,
    capabilities=[
        Instrumentation(
            settings=InstrumentationSettings(
                include_content=False,
                include_binary_content=False,
                use_aggregated_usage_attribute_names=True,
            )
        )
    ],
)

from pydantic_ai import ModelRequest
from pydantic_ai.direct import model_request_sync

response = model_request_sync(
    "anthropic:claude-haiku-4-5",
    [ModelRequest.user_text_prompt("대기 중인 요청 요약 문구를 한 문장으로 작성해줘")],
    instrument=True,
)

from pydantic_ai import ModelRequest
from pydantic_ai.direct import model_request


async def rewrite_summary(summary: str) -> str:
    response = await model_request(
        "openai:gpt-5-mini",
        [
            ModelRequest.user_text_prompt(
                f"다음 업무 요약을 간결하고 전문적인 문체로 다듬어줘: {summary}"
            )
        ],
        instrument=True,
    )
    return response.parts[0].content

from pydantic_ai import ModelRequest
from pydantic_ai.direct import model_request_stream


async def stream_rewrite(summary: str) -> list[str]:
    chunks: list[str] = []
    messages = [ModelRequest.user_text_prompt(summary)]

    async with model_request_stream(
        "openai:gpt-5-mini",
        messages,
        instrument=True,
    ) as stream:
        async for event in stream:
            chunks.append(type(event).__name__)

    return chunks

from dataclasses import dataclass

from fastapi import FastAPI
from starlette.requests import Request
from starlette.responses import Response

from pydantic_ai import Agent
from pydantic_ai.ui.ag_ui import AGUIAdapter


@dataclass
class AppDeps:
    org_id: str
    actor_id: str


agent = Agent(
    "openai:gpt-5-mini",
    deps_type=AppDeps,
    name="sample_ui_agent",
    instructions="업무 애플리케이션 채팅 UI 요청을 처리합니다.",
    defer_model_check=True,
)

app = FastAPI()


@app.post("/ag-ui")
async def ag_ui_chat(request: Request) -> Response:
    org_id = request.headers["x-org-id"]
    actor_id = request.headers["x-actor-id"]

    return await AGUIAdapter.dispatch_request(
        request,
        agent=agent,
        deps=AppDeps(org_id=org_id, actor_id=actor_id),
    )

from fastapi import FastAPI
from starlette.requests import Request
from starlette.responses import Response

from pydantic_ai import Agent
from pydantic_ai.ui.vercel_ai import VercelAIAdapter


agent = Agent(
    "openai:gpt-5-mini",
    name="sample_web_ui_agent",
    instructions="Vercel AI SDK 기반 채팅 요청을 처리합니다.",
    defer_model_check=True,
)

app = FastAPI()


@app.post("/chat")
async def chat(request: Request) -> Response:
    return await VercelAIAdapter.dispatch_request(request, agent=agent)

from pydantic_ai import Agent


agent = Agent(
    "openai:gpt-5-mini",
    name="sample_cli_smoke_agent",
    instructions="업무 애플리케이션 agent 로컬 smoke 테스트용입니다.",
    defer_model_check=True,
)

from temporalio import workflow

from pydantic_ai import Agent
from pydantic_ai.durable_exec.temporal import PydanticAIWorkflow, TemporalAgent


agent = Agent(
    "openai:gpt-5-mini",
    name="workflow_confirmation_agent",
    instructions="긴 확인 검토 workflow에서 확인 문서를 분석합니다.",
    defer_model_check=True,
)

temporal_agent = TemporalAgent(agent)


@workflow.defn
class ConfirmationReviewWorkflow(PydanticAIWorkflow):
    __pydantic_ai_agents__ = [temporal_agent]

    @workflow.run
    async def run(self, request_id: str, user_message: str) -> str:
        result = await temporal_agent.run(
            user_message,
            metadata={"request_id": request_id},
        )
        return result.output

from temporalio.client import Client

from pydantic_ai.durable_exec.temporal import LogfirePlugin, PydanticAIPlugin


async def connect_temporal() -> Client:
    return await Client.connect(
        "localhost:7233",
        plugins=[
            PydanticAIPlugin(),
            LogfirePlugin(),
        ],
    )

from dbos import DBOS, DBOSConfig

from pydantic_ai import Agent
from pydantic_ai.durable_exec.dbos import DBOSAgent


dbos_config: DBOSConfig = {
    "name": "sample_ai_agent_service",
    "system_database_url": "sqlite:///dbos.sqlite",
}
DBOS(config=dbos_config)

agent = Agent(
    "openai:gpt-5-mini",
    name="workflow_dbos_agent",
    instructions="재시도 가능한 workflow agent workflow를 처리합니다.",
    defer_model_check=True,
)

durable_agent = DBOSAgent(agent)


async def run_durable_message(user_message: str) -> str:
    DBOS.launch()
    result = await durable_agent.run(user_message)
    return result.output

from dataclasses import dataclass
from uuid import uuid4

import logfire
from pydantic_ai import Agent, UsageLimits


@dataclass
class ConversationTrace:
    request_id: str
    org_id: str
    actor_id: str
    agent_name: str
    input_tokens: int
    output_tokens: int
    requests: int
    tool_calls: int


agent = Agent(
    "openai:gpt-5-mini",
    name="workflow_agent",
    instructions="업무 애플리케이션 사용자를 지원합니다.",
    defer_model_check=True,
)


async def traceed_run(
    *,
    org_id: str,
    actor_id: str,
    user_message: str,
) -> tuple[str, ConversationTrace]:
    request_id = str(uuid4())

    with logfire.span("workflow.agent.run", request_id=request_id, org_id=org_id):
        result = await agent.run(
            user_message,
            usage_limits=UsageLimits(request_limit=5, total_tokens_limit=4_000),
            metadata={
                "request_id": request_id,
                "org_id": org_id,
                "actor_id": actor_id,
            },
        )

    usage = result.usage
    trace = ConversationTrace(
        request_id=request_id,
        org_id=org_id,
        actor_id=actor_id,
        agent_name="workflow_agent",
        input_tokens=usage.input_tokens or 0,
        output_tokens=usage.output_tokens or 0,
        requests=usage.requests or 0,
        tool_calls=usage.tool_calls or 0,
    )
    return result.output, trace

import re


PHONE_PATTERN = re.compile(r"\b01[016789]-\d{3,4}-\d{4}\b")
ACCOUNT_PATTERN = re.compile(r"\b\d{2,6}-\d{2,6}-\d{2,8}\b")


def redact_message(value: str) -> str:
    value = PHONE_PATTERN.sub("[PHONE]", value)
    value = ACCOUNT_PATTERN.sub("[ACCOUNT]", value)
    return value