diff --git a/providers/common/ai/docs/index.rst b/providers/common/ai/docs/index.rst index 7398a86ce73fc..7de5f5a8ac431 100644 --- a/providers/common/ai/docs/index.rst +++ b/providers/common/ai/docs/index.rst @@ -73,6 +73,12 @@ a service the vendor runs for you, which no vendor-neutral operator wraps: Managed Agents sessions where the agent loop runs on Anthropic's infrastructure rather than in the Airflow worker. * :doc:`apache-airflow-providers-cohere:index` — Cohere's own Embed API. +* :doc:`apache-airflow-providers-google:index` — Vertex AI's Batch Prediction jobs + (``CreateBatchPredictionJobOperator``), a managed batch service like OpenAI's Batch API. +* :doc:`apache-airflow-providers-amazon:index` — Bedrock's Batch Inference + (``BedrockBatchInferenceOperator``), and Bedrock AgentCore's managed agent runtime + (``BedrockCreateAgentRuntimeOperator`` / ``BedrockInvokeAgentRuntimeOperator``), where the + agent loop runs on AWS's infrastructure rather than in the Airflow worker. As a rule of thumb: if Airflow should *run* the AI step (and the model should stay swappable), use ``common.ai``; if the Dag *submits work to* a vendor-managed service and @@ -86,6 +92,26 @@ OpenAI, Anthropic, or other pydantic-ai-supported connection: :start-after: [START howto_operator_llm_basic] :end-before: [END howto_operator_llm_basic] +Choosing extras +---------------- + +The provider's extras split into a few groups: + +* **Model providers** — ``openai``, ``anthropic``, ``google``, ``bedrock``: pick the one + matching your ``llm_conn_id`` connection. Each extra name mirrors the identically named + ``pydantic-ai-slim`` optional dependency group; pydantic-ai supports more model providers + than these four, each under its own extra name, so check the + `pydantic-ai install docs `__ for the full list. +* **Agent tooling** — ``mcp``, ``skills``, ``code-mode``: MCP servers, Agent Skills, and + code-mode tool execution. +* **Document loading** — ``pdf``, ``docx``, ``avro``, ``parquet``: file formats for + document pipelines. +* **Retrieval / SQL** — ``sql``, ``common.sql``, ``langchain``, ``llamaindex``: RAG and + SQL-schema tooling. +* **Git-backed content** — ``git``: pulling Agent Skills or documents from a git connection. + +See the Optional dependencies table below for the exact package each extra installs. + .. toctree:: :hidden: :maxdepth: 1 @@ -100,6 +126,7 @@ OpenAI, Anthropic, or other pydantic-ai-supported connection: :maxdepth: 1 :caption: Guides + Quick start Connection types MCP connection Hooks diff --git a/providers/common/ai/docs/quickstart.rst b/providers/common/ai/docs/quickstart.rst new file mode 100644 index 0000000000000..ec3d07941ebe3 --- /dev/null +++ b/providers/common/ai/docs/quickstart.rst @@ -0,0 +1,92 @@ + .. Licensed to the Apache Software Foundation (ASF) under one + or more contributor license agreements. See the NOTICE file + distributed with this work for additional information + regarding copyright ownership. The ASF licenses this file + to you under the Apache License, Version 2.0 (the + "License"); you may not use this file except in compliance + with the License. You may obtain a copy of the License at + + .. http://www.apache.org/licenses/LICENSE-2.0 + + .. Unless required by applicable law or agreed to in writing, + software distributed under the License is distributed on an + "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + KIND, either express or implied. See the License for the + specific language governing permissions and limitations + under the License. + +.. _howto/quickstart: + +Quick start +=========== + +This guide installs the provider, configures a connection, and runs a first +LLM task. + +Before you start: this assumes a working :doc:`apache-airflow:installation/index` +(Airflow 3.0+) already exists, you have an API key for the LLM provider you +plan to use, and step 3 below makes a real, billed API call to that provider. + +1. Install +---------- + +Install the provider together with the extra matching the model SDK you plan +to use — ``openai``, ``anthropic``, ``google``, or ``bedrock`` (see +:doc:`index` for the full list of available extras). Replace ```` +below with the one you need: + +.. code-block:: bash + + pip install "apache-airflow-providers-common-ai[]" + +2. Configure the connection +---------------------------- + +Every LLM call goes through a Pydantic AI connection (``conn_type`` ``pydanticai``, +default connection id ``pydanticai_default``). The model is set in ``provider:model`` +format and the API key goes in the password field. See :ref:`howto/connection:pydanticai` +for the full reference, including providers that +don't need an API key (Bedrock, Vertex AI). + +The quickest way to set one up is an environment variable. Replace +``openai:gpt-5.3`` with a model you have access to and ``sk-...`` with your +actual API key: + +.. code-block:: bash + + export AIRFLOW_CONN_PYDANTICAI_DEFAULT='{"conn_type": "pydanticai", "password": "sk-...", "extra": {"model": "openai:gpt-5.3"}}' + +Or add it through the Airflow UI (``Admin > Connections``) or the CLI (``airflow connections add``). + +3. Write your first Dag +------------------------ + +The ``@task.llm`` decorator turns a function that returns a prompt string into +a task that sends that prompt to the LLM and returns its response: + +.. exampleinclude:: /../../ai/src/airflow/providers/common/ai/example_dags/example_quickstart.py + :language: python + :start-after: [START howto_quickstart_llm] + :end-before: [END howto_quickstart_llm] + +Run it like any other Dag (``airflow dags test quickstart_llm``) and the +``summarize`` task pushes the LLM's response to XCom. + +Structured output +^^^^^^^^^^^^^^^^^^ + +Need typed data instead of a string? Set ``output_type`` to a Pydantic +``BaseModel`` and the model instance is pushed to XCom unchanged. See the +"Structured Output" section of the :ref:`howto/operator:llm` guide for the +full example and its XCom-deserialization requirements. + +Where to go next +----------------- + +- :doc:`operators/index` — the full set of operators and ``@task`` decorators + (file analysis, SQL, branching, schema comparison). +- :doc:`toolsets` — give an agent tools built from Airflow hooks, SQL + databases, or MCP servers. +- :ref:`howto/operator:agent` — run a multi-turn agent that reasons and calls + tools instead of a single prompt-response call. +- :doc:`observability` — trace LLM and tool calls with OpenTelemetry. diff --git a/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_quickstart.py b/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_quickstart.py new file mode 100644 index 0000000000000..478cb332e68f4 --- /dev/null +++ b/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_quickstart.py @@ -0,0 +1,37 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +"""Quickstart example: a first @task.llm Dag.""" + +from __future__ import annotations + +# [START howto_quickstart_llm] +from airflow.sdk import dag, task + + +@dag(schedule=None, tags=["example"]) +def quickstart_llm(): + @task.llm(llm_conn_id="pydanticai_default", system_prompt="You are a helpful assistant. Be concise.") + def summarize(text: str): + return f"Summarize this article: {text}" + + summarize( + "Apache Airflow is a platform for programmatically authoring, scheduling, and monitoring workflows." + ) + + +quickstart_llm() +# [END howto_quickstart_llm]