diff --git a/providers/common/ai/docs/index.rst b/providers/common/ai/docs/index.rst index a3be5daf4d7e4..9fce0fd08e070 100644 --- a/providers/common/ai/docs/index.rst +++ b/providers/common/ai/docs/index.rst @@ -134,6 +134,7 @@ See the Optional dependencies table below for the exact package each extra insta Operators Examples Retry Policies + Self-hosted models HITL Review Observability diff --git a/providers/common/ai/docs/retry_policies.rst b/providers/common/ai/docs/retry_policies.rst index 8456b1262aa3f..a24e5df448143 100644 --- a/providers/common/ai/docs/retry_policies.rst +++ b/providers/common/ai/docs/retry_policies.rst @@ -161,7 +161,8 @@ Local LLM support ----------------- For environments where exception data must not leave the infrastructure, point -to a local model via Ollama or vLLM: +to a local model via Ollama or vLLM -- see :ref:`howto/self_hosted_models` for +general self-hosted connection setup: .. code-block:: python diff --git a/providers/common/ai/docs/self_hosted_models.rst b/providers/common/ai/docs/self_hosted_models.rst new file mode 100644 index 0000000000000..8919051683b80 --- /dev/null +++ b/providers/common/ai/docs/self_hosted_models.rst @@ -0,0 +1,378 @@ + .. 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/self_hosted_models: + +Self-hosted models +=================== + +This guide serves an open model on your own infrastructure, points +:class:`~airflow.providers.common.ai.hooks.pydantic_ai.PydanticAIHook` at it +instead of a cloud provider, runs a ``@task.llm`` task and an +``AgentOperator`` with tool calling against it, and then routes the same +setup through a bearer-token AI gateway instead of talking to the server +directly. Running a model locally keeps prompt and response data inside your +own infrastructure, avoids per-call cost, and works offline. + +Before you start +------------------ + +This guide assumes a working :doc:`apache-airflow:installation/index` +(Airflow 3.0+) already exists. Its job stops at wiring Airflow to a server +that's already running -- it doesn't cover installing or operating the +model-serving stack itself. + +Prerequisites +^^^^^^^^^^^^^ + +You need an OpenAI-compatible model server reachable from where Airflow +runs, and hardware able to run the model you pick. The provider is built on +pydantic-ai, which reaches these servers by speaking the OpenAI-compatible +wire protocol -- that's why the server needs to expose an OpenAI-compatible +API. Any server that does works the same way; see +pydantic-ai's `OpenAI-compatible models docs +`__ +for the full list of what qualifies. Ollama, vLLM, and LM Studio are three +common choices, used as the examples throughout this guide: + +- `Ollama `__ -- see its + `hardware support docs `__ for GPU/CPU + requirements. +- `vLLM `__ -- a higher-throughput alternative, + commonly GPU-hosted; see its + `GPU installation requirements `__. +- `LM Studio `__ -- a desktop app with a built-in + OpenAI-compatible server; see its + `system requirements `__. + +This page uses these three as running examples rather than trying to cover +every option; for how to point Airflow at any of the others, see +:ref:`self_hosted_models_other_servers` below. + +1. Start an OpenAI-compatible server +-------------------------------------- + +Ollama is the reference server this guide proves end to end. The vLLM and LM +Studio tabs show equivalent wiring for readers already running one of those +instead, but this guide hasn't exercised them itself. + +.. tab-set:: + + .. tab-item:: Ollama + :sync: ollama + + `Install it `__ and pull a model: + + .. code-block:: bash + + ollama pull gemma4 + + Installing Ollama sets up a background service (the macOS/Windows app, + or a systemd service on Linux), so the server is already running -- + ``ollama pull`` only downloads the model. Ollama then serves an + OpenAI-compatible API at ``http://localhost:11434``. On a headless host + with no service, start it yourself with ``ollama serve``. + + .. tab-item:: vLLM + :sync: vllm + + Unlike Ollama and LM Studio, vLLM serves a model by its full + HuggingFace identifier (``organization/model-name``) rather than a + short tag. Serve one, aliasing the served name to ``gemma4`` to match + the connection below: + + .. code-block:: bash + + vllm serve Qwen/Qwen2.5-1.5B-Instruct --served-model-name gemma4 + + See `vLLM's installation docs `__ + for install and GPU requirements. vLLM then serves an + OpenAI-compatible API at ``http://localhost:8000``. + + .. tab-item:: LM Studio + :sync: lmstudio + + Download a model, load it, then start the server: + + .. code-block:: bash + + lms get gemma4 + lms load gemma4 + lms server start + + See `LM Studio's CLI docs `__ for the + full ``lms`` reference (you can also download and load a model from + the LM Studio UI instead). LM Studio then serves an OpenAI-compatible + API at ``http://localhost:1234``. + +2. Create the connection +-------------------------- + +Every LLM call goes through a Pydantic AI connection (``conn_type`` +``pydanticai``, default connection id ``pydanticai_default``). + +.. tab-set:: + + .. tab-item:: Ollama + :sync: ollama + + For a local, keyless server: + + .. code-block:: bash + + export AIRFLOW_CONN_PYDANTICAI_DEFAULT='{"conn_type": "pydanticai", "host": "http://localhost:11434/v1", "extra": {"model": "ollama:gemma4"}}' + + .. tab-item:: vLLM + :sync: vllm + + .. code-block:: bash + + export AIRFLOW_CONN_PYDANTICAI_DEFAULT='{"conn_type": "pydanticai", "host": "http://localhost:8000/v1", "extra": {"model": "openai:gemma4"}}' + + .. tab-item:: LM Studio + :sync: lmstudio + + .. code-block:: bash + + export AIRFLOW_CONN_PYDANTICAI_DEFAULT='{"conn_type": "pydanticai", "host": "http://localhost:1234/v1", "extra": {"model": "openai:gemma4"}}' + + ``http://localhost:1234`` is the default server port shown in `LM + Studio's OpenAI-compatibility docs + `__. + +- ``host`` -- the server's base URL. +- ``password`` -- left empty for a keyless local server. That triggers two + separate things: the Airflow hook forwards no key at all when ``password`` + is empty (only ``base_url`` goes to the provider constructor); pydantic-ai's + provider class then fills in its own placeholder API key and does send + that placeholder in the actual HTTP request. A server with no auth check + doesn't care, so the request still succeeds either way. +- ``extra["model"]`` -- ``ollama:`` picks up Ollama-specific tuning; + vLLM and LM Studio have no dedicated pydantic-ai provider class, so they + use ``openai:`` instead. See + :ref:`Model identifier format ` below for + details. + +The examples above all use ``gemma4`` to illustrate the format -- the actual +value must match a model your server is really serving: + +- Ollama -- ``ollama list`` shows the tags you've pulled. +- vLLM -- defaults to its ``--model`` startup argument; the + ``--served-model-name`` flag can set an alias explicitly (e.g. to + ``gemma4``) if you want the served name to differ from ``--model``. +- LM Studio -- the model identifier shown for whatever you loaded in its UI. +- Any of the three (and any other OpenAI-compatible server) -- + ``curl /v1/models`` lists what the server actually reports as + available. + +See :ref:`howto/connection:pydanticai` for the full connection field +reference. + +3. Run your first Dag with ``@task.llm`` +------------------------------------------- + +The ``@task.llm`` decorator turns a function that returns a prompt string +into a task that sends that prompt to the connection above and returns the +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``); the +``summarize`` task completes using the local model, not a cloud provider. + +4. Run an agent with tool calling +------------------------------------ + +Tool calling needs a model that supports pydantic-ai's ``tools`` capability +-- not every small local model does. + +.. tab-set:: + + .. tab-item:: Ollama + :sync: ollama + + The connection from the previous steps already points at a model + that supports it, so this step runs against the same connection + with no changes. + + .. tab-item:: vLLM + :sync: vllm + + vLLM needs extra server startup flags to enable tool calling for a + given model (``--enable-auto-tool-choice`` plus a matching + ``--tool-call-parser``) -- see its + `Tool Calling docs `__ + for how to pick the right ones for your model. + + .. tab-item:: LM Studio + :sync: lmstudio + + LM Studio's built-in server supports tool calling out of the box + through ``/v1/chat/completions`` -- see its + `Tool Use docs `__. + Its docs note that smaller models, or ones not trained for tool use, + may return improperly formatted tool calls. + +``AgentOperator`` can wrap an existing Airflow hook as a tool via +``HookToolset``. This example wraps an HTTP connection named ``my_api`` -- +point it at any reachable HTTP endpoint that returns JSON: + +.. code-block:: bash + + export AIRFLOW_CONN_MY_API='{"conn_type": "http", "host": ""}' + +.. exampleinclude:: /../../ai/src/airflow/providers/common/ai/example_dags/example_agent.py + :language: python + :start-after: [START howto_agent_self_hosted] + :end-before: [END howto_agent_self_hosted] + +The hook above is constructed with ``method="GET"`` because the target +endpoint -- the OpenAI-compatible model list (``/v1/models``) every one of +these servers exposes -- is GET-only, while ``HttpHook`` +defaults to ``method="POST"`` +(:class:`~airflow.providers.http.hooks.http.HttpHook`) -- fixed once the +hook is constructed, not something the model can change per call. + +The prompt itself just asks a natural question -- it's the ``system_prompt`` +that tells the agent where to find the answer (``GET /v1/models``), so the +model doesn't have to guess a path. That matters because ``HttpHook`` raises +on any response outside the 2xx/3xx range: if the model guessed wrong, the +task would fail outright rather than the agent recovering gracefully. + +Routing through an AI gateway +-------------------------------- + +The same wiring routes through a bearer-token AI gateway instead of talking +to the server directly -- swap three connection fields and nothing else +changes: + +- ``host`` -- the gateway's base URL, with no ``/v1`` suffix; pydantic-ai's + ``openai`` provider appends the API path itself. +- ``password`` -- the bearer token the gateway expects. +- ``extra["model"]`` -- ``openai:``, where ```` is the route + name the gateway exposes for the model you want. + +Only ``api_key`` and ``base_url`` reach the provider class -- there is no +extension point for custom headers. A gateway that authenticates with a +bearer token works through this connection shape; one that requires +additional custom headers does not. + +You can validate this pattern locally with any OpenAI-compatible proxy that +speaks bearer-token auth -- for example, LiteLLM's proxy mode running in +front of a model server. That's a different role than the ``litellm:`` model +prefix covered in the reference section below: here LiteLLM stands in for +the gateway itself, reached through the generic ``openai:`` prefix, +not addressed via its own model-prefix. + +Reference +---------- + +The rest of this page is reference material for the connection shapes +above. + +.. _self_hosted_models_model_id: + +Model identifier format +^^^^^^^^^^^^^^^^^^^^^^^^ + +The ``extra`` JSON's ``model`` value keeps the ``provider:model`` format the +hook always expects, but which ``provider`` prefix to use depends on the +endpoint (verified against ``pydantic-ai-slim`` 2.5.0, the version this +guide was developed against, and 2.10.0, the latest at the time of +writing): + +- **vLLM** has no dedicated provider class in pydantic-ai -- ``openai:`` + is the only option. pydantic-ai's ``openai`` provider class talks to + whatever ``host`` points at, so it works against vLLM's OpenAI-compatible + API even though the weights being served aren't OpenAI's. +- **Ollama** has two valid options: the dedicated ``ollama:`` prefix + (recommended) applies Ollama-specific model-profile tuning based on the + model name (``llama``, ``gemma``, ``qwen``, ``deepseek``, ``mistral``, + ``command``, and ``gpt-oss`` prefixes get tuned defaults); the generic + ``openai:`` prefix also works, since Ollama exposes an + OpenAI-compatible API too, but skips that tuning. + +Both prefixes need the same ``openai`` extra +(``pip install "apache-airflow-providers-common-ai[openai]"``) -- pydantic-ai's +Ollama provider is built on the same underlying ``openai`` Python package as +its OpenAI provider, and there's no separate ``ollama`` extra. + +.. note:: + + ``LiteLLMProvider`` is an exception: its constructor takes ``api_base`` + instead of ``base_url``, so passing ``host`` raises a ``TypeError`` that + the hook catches and silently falls back to environment-variable auth, + ignoring ``host`` entirely. Point a LiteLLM proxy's OpenAI-compatible + endpoint via ``openai:`` instead -- this is the same proxy role + used in the gateway section above, just addressed through the working + prefix rather than the rejected ``litellm:`` one. + +.. _self_hosted_models_other_servers: + +Other OpenAI-compatible servers +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +The same ``openai:`` plus ``host`` setup works against any other +server that exposes an OpenAI-compatible API -- for example +`LM Studio `__, +`llama.cpp `__'s ``llama-server``, or +`LocalAI `__. Point ``host`` at that server's +OpenAI-compatible endpoint the same way as the vLLM example above. + +Cross-hook naming differences +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +:class:`~airflow.providers.common.ai.hooks.langchain.LangChainHook` and +:class:`~airflow.providers.common.ai.hooks.llamaindex.LlamaIndexHook` also +support self-hosted endpoints through the same ``host`` connection field, but +the underlying constructor keyword each hook passes it to is not uniform: + +.. list-table:: + :header-rows: 1 + + * - Hook + - Connection field + - Constructor keyword + * - ``PydanticAIHook`` + - ``host`` + - ``base_url`` + * - ``LangChainHook`` + - ``host`` + - ``base_url`` + * - ``LlamaIndexHook`` + - ``host`` + - ``api_base`` + +Where to go next +------------------- + +- :ref:`howto/connection:pydanticai` -- the full connection field reference. +- :doc:`retry_policies` -- the "Local LLM support" section covers pointing + ``LLMRetryPolicy`` at a self-hosted endpoint. +- :doc:`examples` -- more runnable Dags against the ``pydanticai`` + connection type; point one at any of the connections on this page. +- pydantic-ai isn't limited to Ollama and vLLM -- any OpenAI-compatible + endpoint works the same way. See the + `pydantic-ai install docs `__ + for the full list of providers it supports. +- `pydantic-ai's OpenAI-compatible models docs + `__ + -- background on the ``openai`` provider class this guide's + ``openai:`` examples use. diff --git a/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_agent.py b/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_agent.py index 79a0c5fe333b2..16bca0700eb4e 100644 --- a/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_agent.py +++ b/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_agent.py @@ -110,6 +110,42 @@ def example_agent_operator_hook(): example_agent_operator_hook() +# --------------------------------------------------------------------------- +# 2b. Hook-based tools against a GET-only endpoint (self-hosted models tutorial) +# --------------------------------------------------------------------------- + + +# [START howto_agent_self_hosted] +@dag(tags=["example"]) +def example_agent_self_hosted(): + from airflow.providers.http.hooks.http import HttpHook + + # The OpenAI-compatible model list (GET /v1/models) is GET-only; HttpHook defaults to POST. + http_hook = HttpHook(http_conn_id="my_api", method="GET") + + AgentOperator( + task_id="list_models", + prompt="Which models are available?", + llm_conn_id="pydanticai_default", + system_prompt=( + "You are an API assistant. Use the tools to answer questions; " + "the server's model list is served at GET /v1/models." + ), + toolsets=[ + HookToolset( + http_hook, + allowed_methods=["run"], + tool_name_prefix="http_", + ) + ], + ) + + +# [END howto_agent_self_hosted] + +example_agent_self_hosted() + + # --------------------------------------------------------------------------- # 3. @task.agent decorator with dynamic prompt # ---------------------------------------------------------------------------