Apache Airflow에서 Common AI Provider로 LLM Agentic 워크플로우 구성하기

[choo121600]

LLM을 파이프라인에 넣는 방법은 지금까지 대개 두 가지였던 것 같습니다.

하나는 애플리케이션 코드에서 벤더 SDK를 직접 부르는 것, 다른 하나는 벤더가 관리하는 배치나 에이전트 서비스에 일을 맡겨두고 결과를 기다리는 것. Airflow 3의 common-ai provider(apache-airflow-providers-common-ai)는 여기에 세 번째를 더합니다. 모델 호출도 에이전트 루프도 전부 Airflow 워커 안에서 돌기 때문에, AI 스텝이 다른 태스크와 똑같이 다뤄집니다. 저는 개인적으로 이 대목, 그러니까 'AI가 한 일도 다른 태스크처럼 들여다보고 되돌릴 수 있다’는 점을 꽤 크게 봅니다. 왜 그렇게까지 보는지는 맨 뒤에서 잠깐 더 얘기하겠습니다.

이 글에서 Common AI Provider의 기능별로 한번 알아보도록 하겠습니다.

설치와 커넥션

provider는 pydantic-ai 위에 있어서 실제 모델 SDK는 extra로 골라 답니다. Anthropic을 쓴다면 이렇게 설치합니다.

pip install 'apache-airflow-providers-common-ai[anthropic]'

openai, google, bedrock 같은 extra 이름은 pydantic-ai 쪽 그룹과 같고, Airflow는 3.x 이상이면 됩니다.

커넥션은 pydanticai 타입 하나입니다. 모델은 provider:model 형식으로 적고(extra.model에 저장됩니다), API 키는 비밀번호 칸에 넣습니다.

{
  "conn_type": "pydanticai",
  "password": "sk-ant-...",
  "extra": "{\"model\": \"anthropic:claude-sonnet-4-5\"}"
}

모델 ID는 pydantic-ai의 provider:model 형식이라 anthropic:..., openai:..., google:...처럼 씁니다(현재 이름은 각 벤더 모델 문서에서 확인하시면 됩니다). Hook 기본값이 pydanticai_default라, 커넥션 ID를 따로 안 주면 이걸 씁니다. OpenAI로 가려면 openai:... 모델과 OpenAI 키로 커넥션만 바꾸면 되고, Ollama나 vLLM 같은 로컬 엔드포인트는 host에 베이스 URL을 넣습니다. 코드는 그대로입니다.

Azure OpenAI, Bedrock, Vertex는 인증 방식이 달라서 전용 커넥션 타입(pydanticai-azure, pydanticai-bedrock, pydanticai-vertex)이 따로 있고, MCP 서버, LangChain, LlamaIndex도 각자 커넥션 타입이 있습니다.
따라서 꼭 각각의 인증 방식에 대해서 먼저 확인해보고 적용해주세요.

@task.llm: 한 번 부르고 결과를 받는다

가장 단순한 것부터 봅시다. 프롬프트를 주면 모델이 답하고, 그 답이 XCom으로 넘어갑니다. 요약이나 분류, 추출처럼 도구 없이 한 번에 끝나는 일에 씁니다. @task.llm으로 감싼 함수는 프롬프트 문자열을 반환하고, 나머지(커넥션, 시스템 프롬프트)는 데코레이터 인자로 줍니다.

from airflow.sdk import dag, task

@dag
def summarize_ticket():
    @task.llm(
        llm_conn_id="pydanticai_default",
        system_prompt="You summarize support tickets in one sentence.",
    )
    def summarize(body: str) -> str:
        return f"Summarize this ticket:\n{body}"

    summarize("Checkout returns a 500 whenever I apply a coupon, since this morning.")

summarize_ticket()

문자열 대신 구조를 강제하고 싶으면 output_type에 Pydantic 모델을 줍니다. 그러면 모델이 그 스키마대로 답하고, 다음 태스크는 타입 힌트로 그대로 받습니다. 파싱이나 검증 코드를 따로 안 짜도 됩니다.

from typing import Literal
from pydantic import BaseModel
from airflow.sdk import dag, task

# XCom 직렬화 때문에 출력 모델은 모듈 최상단에 둡니다.
class Ticket(BaseModel):
    category: Literal["billing", "auth", "bug", "other"]
    urgency: Literal["low", "medium", "high"]
    summary: str

@dag
def classify_ticket():
    @task.llm(
        llm_conn_id="pydanticai_default",
        system_prompt="Classify the support ticket.",
        output_type=Ticket,
    )
    def classify(body: str) -> str:
        return f"Classify this ticket:\n{body}"

    @task
    def route(t: Ticket):          # 타입 그대로 받아 속성 접근
        print(t.category, t.urgency)

    route(classify("I was charged twice for my subscription this month."))

classify_ticket()

output_type으로 준 클래스는 워커가 Dag를 로드할 때 XCom 역직렬화용으로 등록되니까, 같은 Dag 안에서는 따로 설정할 게 없습니다. 모듈 최상단 정의라는 조건만 지키면 됩니다.

비용이 걱정되면 usage_limits로 상한을 겁니다. 요청 수나 입출력 토큰을 넘기면 태스크가 실패합니다. 모델 하이퍼파라미터는 agent_params로 넘깁니다(예: model_settings의 temperature).

from pydantic_ai.usage import UsageLimits

@task.llm(
    llm_conn_id="pydanticai_default",
    system_prompt="You are a concise reviewer.",
    usage_limits=UsageLimits(request_limit=5, input_tokens_limit=4000, output_tokens_limit=1000),
    agent_params={"model_settings": {"temperature": 0}},
)
def review(doc: str) -> str:
    return f"Review this design doc in three bullets:\n{doc}"

@task.llm_branch: 모델 판단으로 흐름을 가른다

'이 다음에 뭘 할지’는 보통 우리가 코드로 정합니다. 그걸 모델한테 맡기고 싶을 때가 있죠. 이때 씁니다. 다른 오퍼레이터와 달리 반환값을 파싱하지 않습니다. Dag 토폴로지에서 다운스트림 태스크들을 스스로 찾아서, 그 태스크 ID를 pydantic-ai의 구조화 출력으로 모델에게 '고를 수 있는 목록’으로 넘깁니다. 모델은 목록 밖을 답할 수 없습니다.

@dag
def route_ticket():
    @task.llm_branch(
        llm_conn_id="pydanticai_default",
        system_prompt="Route the support ticket to the right team.",
    )
    def route(message: str) -> str:
        return f"Route this ticket: {message}"

    @task
    def billing(): ...
    @task
    def auth(): ...
    @task
    def general(): ...

    route("My password reset email never arrived.") >> [billing(), auth(), general()]

route_ticket()

모델이 auth를 고르면 그 태스크만 돌고 나머지는 스킵됩니다. 한 입력이 여러 갈래에 걸리면 allow_multiple_branches=True를 줘서 해당하는 걸 다 고르게 합니다(나머지는 스킵).

사람이 한 번 확인하고 넘어가야 하는 분기라면 require_approval=True를 줍니다. 모델이 고른 뒤 태스크가 멈춰 리뷰를 기다리고, 리뷰 화면에 모델의 선택과 유효한 다운스트림 태스크 ID가 뜹니다. allow_modifications=True면 사람이 선택을 바꿀 수도 있습니다.

@task.agent : 도구를 쓰는 에이전트 루프

한 번 부르고 끝이 아니라, 모델이 도구를 부르고 결과를 보고 다시 판단하는 여러 턴짜리 루프가 필요할 때가 있습니다. @task.agent(= AgentOperator)는 그 루프를 Airflow 워커 안에서 돌립니다. 도구는 '툴셋’으로 줍니다.

가장 흔한 건 SQL 에이전트입니다. 데이터베이스를 주고, 질문에 SQL을 직접 써서 답하게 합니다.

from airflow.providers.common.ai.toolsets.sql import SQLToolset
from airflow.sdk import dag, task

@dag
def sales_qa():
    @task.agent(
        llm_conn_id="pydanticai_default",
        system_prompt="You are a SQL analyst. Use the tools to answer with data.",
        toolsets=[
            SQLToolset(
                db_conn_id="postgres_default",
                allowed_tables=["customers", "orders"],
                max_rows=20,
            )
        ],
    )
    def ask(question: str) -> str:
        return question

    ask("Which five customers placed the most orders last month?")

sales_qa()

SQLToolset은 네 가지 도구를 에이전트에 줍니다. list_tables(테이블 목록), get_schema(컬럼), query(질의 실행), check_query(문법 검사). 에이전트는 보통 스키마부터 훑고 나서 질의를 씁니다.

여기서 한 가지 짚고 싶은 게 있습니다. allowed_tables는 에이전트가 만질 테이블을 좁혀 주지만, 문서가 분명히 적어 둔 대로 이건 보안 경계가 아니라 애플리케이션 레벨 가드레일입니다. sqlglot으로 파싱해서 목록 밖 테이블을 막긴 하는데, 파서가 못 보는 함수를 통해 우회될 수 있습니다(프롬프트 주입이 그런 경우죠). 진짜 경계는 커넥션이 가리키는 최소 권한 DB 역할입니다. 그러니 allowed_tables로는 의도를 좁히고, 그 커넥션은 딱 그 테이블만 SELECT할 수 있는 역할에 물리는 게 맞다고 봅니다. 무엇을 막아주고 무엇은 못 막는지를 문서가 솔직하게 적어둔 점이 저는 좋았습니다.

좁은 목록은 비용 관점에서도 낫습니다. 에이전트는 도구 루프가 길어지기 쉬워서, 런당 상한이 필요하면 usage_limits의 tool_calls_limit으로 못을 박습니다.

에이전트 출력도 output_type으로 구조를 받을 수 있습니다.

class Answer(BaseModel):
    summary: str
    top_customers: list[str]

@task.agent(
    llm_conn_id="pydanticai_default",
    system_prompt="Answer with data. Return a structured result.",
    output_type=Answer,
    toolsets=[SQLToolset(db_conn_id="postgres_default", allowed_tables=["orders"])],
)
def ask(question: str) -> str:
    return question

옵션이 몇 가지 더 있습니다. durable=True를 주고 태스크에 재시도가 설정돼 있으면, 실패했을 때 이미 끝낸 LLM, 도구 스텝을 캐시에서 재생하고 안 끝난 데서 이어갑니다(성공하면 캐시는 지워집니다). message_history로 이전 대화를 넘기면 여러 런에 걸친 세션이 되고, system_prompt는 Jinja 템플릿이라 앞 태스크가 만든 값을 끼워 넣을 수 있습니다. pydantic-ai Agent 생성자 인자는 agent_params로 전달합니다(Thinking, WebSearch 같은 capability도 이 길로 줍니다).

툴셋: 에이전트에 줄 도구

에이전트 도구는 툴셋으로 묶습니다. 전부 pydantic-ai의 AbstractToolset이라 pydantic-ai 에이전트 어디에나 넘길 수 있고, provider에는 이만큼 들어 있습니다.

• SQLToolset: 위에서 쓴 그 툴셋. DB에 SQL로 질의하는 도구 4개, allowed_tables, max_rows(기본 50), max_result_bytes. [sql] extra.
• DataFusionToolset: S3나 로컬, Iceberg 같은 오브젝트 스토어의 Parquet, CSV를 SQL로 질의(Apache DataFusion). common-sql[datafusion].
• HookToolset: 아무 Airflow Hook의 메서드나 도구로 만듭니다. 자동 발견은 없고 allowed_methods로 직접 명시합니다.
• MCPToolset: Airflow 커넥션으로 설정한 MCP 서버에 연결. [mcp] extra.
• AgentSkillsToolset: Agent Skills(SKILL.md 번들)를 필요할 때만 로드합니다. 이름과 설명 카탈로그만 프롬프트에 두고 모델이 고를 때 본문을 읽어서, 스킬이 많아도 토큰이 덜 듭니다. [skills] extra.
• SandboxToolset: 워커 밖 격리 샌드박스에서 셸과 파일시스템을 다룹니다(run_command, read_file, write_file, list_directory).
• LoggingToolset: 도구 호출을 실시간으로 로깅하는 래퍼입니다. AgentOperator가 자동으로 씌웁니다.

HookToolset은 이미 있는 Hook을 그대로 도구로 만듭니다. HTTP API 하나를 에이전트에 쥐여 주려면 이렇게 합니다.

from airflow.providers.http.hooks.http import HttpHook
from airflow.providers.common.ai.toolsets.hook import HookToolset

@task.agent(
    llm_conn_id="pydanticai_default",
    system_prompt="Use the tools to call the API and answer.",
    toolsets=[
        HookToolset(HttpHook(http_conn_id="my_api"),
                    allowed_methods=["run"], tool_name_prefix="http_"),
    ],
)
def explore(question: str) -> str:
    return question

Airflow provider Hook이 350개가 넘습니다. 자격증명 관리와 타입이 이미 붙어 있는 그 Hook들을 에이전트 도구로 그대로 쓸 수 있다는 거죠. query 결과는 컬럼명 한 번에 위치 기반 행으로 반환되고 max_rows, max_result_bytes로 크기가 묶입니다. 도구 결과는 이후 모든 모델 호출에 다시 실려서 비용이 재부과되니까, 이 상한이 곧 비용입니다.

Human In The Loop(HITL)

AI가 낸 걸 사람이 한 번 보고 넘기고 싶을 때가 있습니다. common-ai의 HITL은 두 갈래인데, 워커 슬롯을 다루는 방식이 다릅니다.

하나는 승인 게이트입니다. @task.llm이나 @task.llm_branch에 require_approval=True를 주면, 모델이 출력을 내거나(llm) 분기를 고른 뒤(llm_branch) 태스크가 멈춰서 사람의 승인이나 거부를 기다립니다. allow_modifications=True면 사람이 출력이나 분기 선택을 고쳐서 승인할 수 있고, approval_timeout으로 마감을 겁니다. 거부하거나 타임아웃이 나면 태스크가 실패하고 다운스트림은 upstream_failed가 됩니다. 이쪽은 태스크가 defer돼서, 기다리는 동안 워커 슬롯을 놓습니다.

@task.llm(
    llm_conn_id="pydanticai_default",
    system_prompt="Draft a one-sentence reply to the support ticket. A human approves it before it is sent.",
    require_approval=True,
    allow_modifications=True,
    approval_timeout=timedelta(hours=24),
)
def draft_reply(body: str) -> str:
    return f"Draft a reply to this ticket:\n{body}"

다른 하나는 반복 리뷰입니다. @task.agent에 enable_hitl_review=True를 주면 첫 생성 뒤에 리뷰 루프가 돕니다. 사람이 채팅 UI에서 승인, 거부뿐 아니라 수정 요청까지 할 수 있고, 요청하면 에이전트가 피드백을 반영해 다시 생성합니다. max_hitl_iterations(기본 5)까지 돌고, hitl_timeout으로 전체 대기 시간을 겁니다. 이건 hitl_review 플러그인(API 서버에 /hitl-review로 붙는 FastAPI)과 Airflow 3.1 이상이 필요하고, 상태는 태스크 인스턴스의 XCom에 저장됩니다.

from datetime import timedelta
from airflow.providers.common.ai.operators.agent import AgentOperator

AgentOperator(
    task_id="summarize_with_review",
    prompt="Summarize the Q4 sales report in 3 bullet points.",
    llm_conn_id="pydanticai_default",
    system_prompt="You are a concise business analyst.",
    enable_hitl_review=True,
    max_hitl_iterations=5,
    hitl_timeout=timedelta(minutes=30),
    hitl_poll_interval=10.0,
)

여기서 주의할 게 하나 있습니다. 반복 리뷰는 승인 게이트와 달리 defer하지 않고 XCom을 폴링해서, 리뷰가 끝날 때까지 워커 슬롯을 붙잡습니다. 30분짜리 리뷰라면 그동안 워커 하나가 묶여 있는 겁니다. 리뷰 화면은 태스크 인스턴스 페이지의 ‘HITL Review’ 링크로 엽니다.

어느 쪽이든 승인, 거부, 타임아웃이 Airflow 태스크의 상태로 남습니다. AI 출력을 사람이 게이트하는 일을 파이프라인 밖 별도 시스템이 아니라 태스크 하나로 끝낸다는 것. 저는 이게 HITL을 파이프라인 안에 두는 이유라고 봅니다.

그 밖의 오퍼레이터

핵심 셋 말고도 같은 커넥션, 같은 워커에서 도는 오퍼레이터가 더 있습니다.

• @task.llm_file_analysis(LLMFileAnalysisOperator): 오브젝트 스토리지나 로컬 파일, 이미지, PDF를 프롬프트로 읽어 분석합니다(읽기 전용).
• @task.llm_sql(LLMSQLQueryOperator): 자연어를 SQL로 만들어 XCom으로 돌려줍니다(실행은 안 하고, 다운스트림이 실행합니다).
• @task.llm_schema_compare(LLMSchemaCompareOperator): 시스템 간 스키마를 비교해서, 단순 비교로는 놓치는 드리프트를 잡습니다.
• DocumentLoaderOperator → LlamaIndexEmbeddingOperator → LlamaIndexRetrievalOperator: 파일을 구조화 텍스트로 파싱하고 임베딩한 뒤 검색하는 RAG 파이프라인입니다(데코레이터 없이 오퍼레이터로).

모델을 바꾸는 건 커넥션을 바꾸는 것

지금까지 예제의 코드는 어떤 모델을 쓰는지 모릅니다. llm_conn_id가 가리키는 커넥션의 extra.model만 바뀌면 됩니다.

{"conn_type": "pydanticai", "extra": "{\"model\": \"openai:gpt-4o\"}"}
{"conn_type": "pydanticai", "extra": "{\"model\": \"google:gemini-2.0-flash\"}"}

pydantic-ai가 지원하는 모델은 이보다 많고, 태스크 코드는 어느 쪽이든 그대로입니다. 기존 LangChain, LlamaIndex 자산도 잠기지 않습니다. pydantic-ai가 LangChain 도구를 감싸는 툴셋을 제공하고, provider에도 Airflow 툴셋을 LangChain 도구로 바꾸는 브리지가 있어서 양방향으로 붙습니다.

운영에서 챙기는 것들

• 비용 상한: @task.llm, @task.agent에 usage_limits(요청 수, 토큰, 도구 호출)로 런당 예산을 걸어 초과 시 실패시킵니다.
• 가드레일: pydantic-ai-shields의 입력, 출력, 도구 가드(InputGuard, OutputGuard, ToolGuard)와 비용 추적이 같은 에이전트 루프에 붙습니다([shields] extra).
• code-mode: 에이전트의 여러 도구 호출을 샌드박스 안의 코드 한 번으로 묶어 실행하는 옵션입니다([code-mode] extra).
• durable 실행: 위의 durable=True. 재시도는 Airflow 표준 정책 그대로 얹힙니다.
• 관측: 모델 호출, 도구 호출, 토큰이 태스크 로그와 관측 훅에 남습니다.
• 셀프호스팅 모델: Ollama나 vLLM 같은 OpenAI 호환 엔드포인트는 커넥션 host만 바꾸면 됩니다.

마무리하며

앞에서 미뤄뒀던 에이전트에 대한 제 개인적인 생각을 하나 꺼내고 싶습니다. 저는 에이전트와 관련해서 결과보다 과정에 대한 신뢰가 중요하다고 생각합니다.
무언가 권위 있는 존재(탑티어 모델)이기 때문에 믿는 것과, 내가 언제든 들여다볼 수 있고, 필요하면 수정할 수 있으며, 원한다면 되돌릴 수도 있기 때문에 믿는 것은 전혀 다른 종류의 신뢰입니다. 후자는 맹신이 아니라 검증 가능한 신뢰입니다. 관측할 수 있어야 하고, 제어할 수 있어야 하고, 되돌릴 수 있어야 합니다.

흥미롭게도 common-ai가 AI를 다루는 방식이 정확히 이 지점에 있습니다. 에이전트 하네스는 판단 과정을 루프 안쪽에 숨겨두지만, common-ai는 그 호출을 Airflow 태스크로 꺼냅니다. 공식 블로그의 표현을 빌리면 Observable, Retryable, Auditable by Design인데, 이 글 내내 보여드린 게 정확히 그거였습니다. classify 로그의 요청 수와 토큰, SQL 에이전트가 부른 도구 한 줄 한 줄, durable 재생, 태스크 상태로 박히는 HITL 승인, 거부. 에이전트 하네스가 안쪽에 숨겨두는 정보를, Airflow는 태스크로 표면에 드러냅니다.

그렇다고 Airflow가 에이전트가 되려는 건 아닙니다. 에이전트는 판단하고, 생성하고, 추론합니다. Airflow가 하는 일은 오케스트레이션, 스케줄링, 재시도, 관측, 운영 제어죠. common-ai는 LLM 호출과 에이전트 루프를 Airflow가 부르는 방법을 표준화한 것이지, Airflow가 LLM이 된 게 아닙니다. Airflow는 AI를 만드는 게 아니라, AI를 포함한 워크플로우를 오케스트레이션합니다.

정리하면, common-ai가 하는 일은 LLM 호출과 에이전트 루프를 그냥 Airflow 태스크로 만드는 것입니다. AI를 파이프라인에 넣겠다고 새 오케스트레이터를 하나 더 들이는 게 아니라, Airflow에 태스크 종류가 하나 늘어난 것뿐입니다.
관심있으신 분들은 써보시면 좋을 것 같네요.

Reference

• apache-airflow-providers-common-ai 개요
• Common AI Operators (전체 오퍼레이터 표)
• Pydantic AI connection
• LLMOperator / @task.llm
• LLMBranchOperator / @task.llm_branch
• AgentOperator / @task.agent (durable, message_history, capabilities)
• Toolsets: SQL, DataFusion, Hook, MCP, AgentSkills, Sandbox, Logging
• HITL Review: 에이전트 출력을 사람이 리뷰
• pydantic-ai
• Agentic Workloads on Airflow: Observable, Retryable, and Auditable by Design