Common AI

apache-airflow-providers-common-ai

AI/LLM hooks and operators for Airflow pipelines using pydantic-ai.

Works with Airflow 2.11+
137K
Monthly Downloads
38
Total Modules
View Docs Source
Install:
pip install apache-airflow-providers-common-ai==1.0.0

Airflow

2.11+

Python

>=3.11

Dependencies (5)

Show all Hide apache-airflow>=2.11.0 apache-airflow-providers-common-compat>=1.21.0 apache-airflow-providers-standard>=1.20.0 pydantic-ai-slim>=2.36.0 structlog>=24.2.0

Connections (7)

Connection type External services (examples)
pydanticai Alibaba Cloud Model Studio, Anthropic, AWS Bedrock, AWS Bedrock Mantle, Azure OpenAI, Cerebras, Cohere, Crusoe, DeepSeek, Fireworks AI, GitHub Copilot, Google Gemini, Google Vertex AI, Groq, Heroku, Hugging Face, LiteLLM, Mistral AI, Moonshot AI, Nebius, Ollama, OpenAI, OpenRouter, OVHcloud, Pydantic AI Gateway, SambaNova, Snowflake Cortex, Together AI, TypeSafe, Vercel AI Gateway, vLLM, xAI, Z.AI
pydanticai_azure Azure OpenAI
pydanticai_bedrock AWS Bedrock
pydanticai_vertex Google Vertex AI
mcp Any MCP server (user-supplied endpoint)
langchain OpenAI, Anthropic, Groq, Mistral AI, DeepSeek, Ollama, vLLM
llamaindex OpenAI

Modules

Durable and deferrable operators across all providers

O

AgentOperator Durable

Run a pydantic-ai Agent with tools and multi-turn reasoning.

airflow.providers.common.ai.operators.agent.AgentOperator
View Docs Guide Source
O

LLMOperator

Call an LLM with a prompt and return the output.

airflow.providers.common.ai.operators.llm.LLMOperator
View Docs Guide Source
O

LLMBatchOperator Deferrable

Submit prompts as a provider batch job and land results on object storage.

airflow.providers.common.ai.operators.llm_batch.LLMBatchOperator
View Docs Guide Source
O

LLMFileAnalysisOperator

Analyze files from object storage or local storage using a single LLM call.

airflow.providers.common.ai.operators.llm_file_analysis.LLMFileAnalysisOperator
View Docs Guide Source
O

LLMBranchOperator

Ask an LLM to choose which downstream task(s) to execute.

airflow.providers.common.ai.operators.llm_branch.LLMBranchOperator
View Docs Guide Source
O

LLMSQLQueryOperator

Generate SQL queries from natural language using an LLM.

airflow.providers.common.ai.operators.llm_sql.LLMSQLQueryOperator
View Docs Guide Source
O

LLMSchemaCompareOperator

Compare schemas across different database systems and detect drift using LLM reasoning.

airflow.providers.common.ai.operators.llm_schema_compare.LLMSchemaCompareOperator
View Docs Guide Source
O

DocumentLoaderOperator

Parse files into ``list[dict(text, metadata)]`` for downstream embedding.

airflow.providers.common.ai.operators.document_loader.DocumentLoaderOperator
View Docs Guide Source
O

LlamaIndexEmbeddingOperator

Chunk documents and produce embedding vectors using LlamaIndex.

airflow.providers.common.ai.operators.llamaindex_embedding.LlamaIndexEmbeddingOperator
View Docs Guide Source
O

LlamaIndexRetrievalOperator

Retrieve relevant document chunks from a persisted LlamaIndex index.

airflow.providers.common.ai.operators.llamaindex_retrieval.LlamaIndexRetrievalOperator
View Docs Guide Source
H

PydanticAIAzureHook

Hook for Azure OpenAI via pydantic-ai.

airflow.providers.common.ai.hooks.pydantic_ai.PydanticAIAzureHook
View Docs Source
H

PydanticAIBedrockHook

Hook for AWS Bedrock via pydantic-ai.

airflow.providers.common.ai.hooks.pydantic_ai.PydanticAIBedrockHook
View Docs Source
H

PydanticAIHook

Hook for LLM access via pydantic-ai.

airflow.providers.common.ai.hooks.pydantic_ai.PydanticAIHook
View Docs Guide Source
H

PydanticAIVertexHook

Hook for Google Vertex AI (or Generative Language API) via pydantic-ai.

airflow.providers.common.ai.hooks.pydantic_ai.PydanticAIVertexHook
View Docs Source
H

MCPHook

Hook for connecting to MCP (Model Context Protocol) servers.

airflow.providers.common.ai.hooks.mcp.MCPHook
View Docs Guide Source
H

LangChainHook

Bridge an Airflow connection to LangChain chat and embedding models.

airflow.providers.common.ai.hooks.langchain.LangChainHook
View Docs Guide Source
H

LlamaIndexHook

Bridge an Airflow connection to LlamaIndex chat and embedding models.

airflow.providers.common.ai.hooks.llamaindex.LlamaIndexHook
View Docs Guide Source
T

LLMBatchTrigger

Poll a batch adapter until the batch reaches a terminal state.

airflow.providers.common.ai.triggers.llm_batch.LLMBatchTrigger
View Docs Source
U

HookToolset

Expose selected methods of an Airflow Hook as pydantic-ai tools.

airflow.providers.common.ai.toolsets.hook.HookToolset
View Docs Guide Source
U

SQLToolset

Curated toolset that gives an LLM agent safe access to a SQL database.

airflow.providers.common.ai.toolsets.sql.SQLToolset
View Docs Guide Source
U

DataFusionToolset

Curated toolset that gives an LLM agent SQL access to object-storage data via Apache DataFusion.

airflow.providers.common.ai.toolsets.datafusion.DataFusionToolset
View Docs Guide Source
U

LoggingToolset

Wrap a toolset to log each tool call with timing.

airflow.providers.common.ai.toolsets.logging.LoggingToolset
View Docs Guide Source
U

MCPToolset

Toolset that connects to an MCP server configured via an Airflow connection.

airflow.providers.common.ai.toolsets.mcp.MCPToolset
View Docs Guide Source
U

ObjectStorageToolset

Give an agent read-only access to the files under one object-storage path.

airflow.providers.common.ai.toolsets.object_storage.ObjectStorageToolset
View Docs Guide Source
U

SandboxToolset

Give an agent shell and file access inside a disposable sandbox, off the Airflow worker.

airflow.providers.common.ai.toolsets.sandbox.SandboxToolset
View Docs Source
U

AgentSkillsToolset

A pydantic-ai toolset that loads Agent Skills, with Git credentials from Airflow connections.

airflow.providers.common.ai.toolsets.skills.AgentSkillsToolset
View Docs Guide Source
U

ManagedAgentToolset

Expose any managed-agent client as one tool.

airflow.providers.common.ai.toolsets.managed_agent.ManagedAgentToolset
View Docs Guide Source
R

ClassifierRetryPolicy

Retry policy where the model names the kind of failure and the author's table decides.

airflow.providers.common.ai.policies.retry.ClassifierRetryPolicy
View Docs Source
R

LLMRetryPolicy

Retry policy that uses an LLM to classify errors and decide retry behaviour.

airflow.providers.common.ai.policies.retry.LLMRetryPolicy
View Docs Source
P

HITLReviewPlugin

Register the HITL Review REST API + chat UI on the Airflow API server.

airflow.providers.common.ai.plugins.hitl_review.HITLReviewPlugin
View Docs Source
P

ModelPanelPlugin

Register the 'Model' tab showing the resolved LLM model name and token usage.

airflow.providers.common.ai.plugins.model_panel.ModelPanelPlugin
View Docs Source
@

@task.agent Durable

Wrap a function that returns a prompt into an agentic LLM task.

airflow.providers.common.ai.decorators.agent.agent_task
View Docs Guide Source
@

@task.llm

Wrap a function that returns a prompt into a general-purpose LLM task.

airflow.providers.common.ai.decorators.llm.llm_task
View Docs Guide Source
@

@task.llm_batch Deferrable

Wrap a function that returns a batch's inputs into an ``@task.llm_batch`` task.

airflow.providers.common.ai.decorators.llm_batch.llm_batch_task
View Docs Guide Source
@

@task.llm_file_analysis

Wrap a callable that returns a prompt into an LLM-backed file-analysis task.

airflow.providers.common.ai.decorators.llm_file_analysis.llm_file_analysis_task
View Docs Guide Source
@

@task.llm_branch

Wrap a function that returns a prompt into an LLM-driven branching task.

airflow.providers.common.ai.decorators.llm_branch.llm_branch_task
View Docs Guide Source
@

@task.llm_sql

Wrap a function that returns a natural language prompt into an LLM SQL task.

airflow.providers.common.ai.decorators.llm_sql.llm_sql_task
View Docs Guide Source
@

@task.llm_schema_compare

Wrap a function that returns a prompt into an LLM schema comparison task.

airflow.providers.common.ai.decorators.llm_schema_compare.llm_schema_compare_task
View Docs Guide Source