diff --git a/providers/common/ai/docs/operators/agent.rst b/providers/common/ai/docs/operators/agent.rst index b3805aa34b765..11e773c75e431 100644 --- a/providers/common/ai/docs/operators/agent.rst +++ b/providers/common/ai/docs/operators/agent.rst @@ -156,6 +156,26 @@ tasks can consume it. :end-before: [END howto_agent_chain] +.. _howto/operator:agent-dynamic-system-prompt: + +Dynamic System Prompt +---------------------- + +``system_prompt`` is a templated field, so instead of a static string it +can be a Jinja expression that reads a value an earlier task already +computed -- for example, tailoring the agent's instructions to a +classification produced upstream. + +.. exampleinclude:: /../../ai/src/airflow/providers/common/ai/example_dags/example_agent.py + :language: python + :start-after: [START howto_agent_dynamic_system_prompt] + :end-before: [END howto_agent_dynamic_system_prompt] + +Open the **Rendered Template** tab on the task instance to see the +substituted ``system_prompt`` after Jinja fills in ``classify``'s XCom +values. + + Multi-turn Sessions ------------------- diff --git a/providers/common/ai/docs/operators/llm.rst b/providers/common/ai/docs/operators/llm.rst index 1f5a375ca9003..28187d1c183f5 100644 --- a/providers/common/ai/docs/operators/llm.rst +++ b/providers/common/ai/docs/operators/llm.rst @@ -186,6 +186,11 @@ to process a list of items in parallel: :start-after: [START howto_decorator_llm_pipeline] :end-before: [END howto_decorator_llm_pipeline] +.. seealso:: + :ref:`Dynamic System Prompt ` -- + ``system_prompt`` is templated identically on ``@task.llm``, so the same + upstream-XCom pattern applies here. + Human-in-the-Loop Approval -------------------------- 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 787b0d6dce2e1..e294260ca3117 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 @@ -301,3 +301,41 @@ def save_history(session_id: str, transcript: str) -> None: # [END howto_agent_session] example_agent_session() + + +# --------------------------------------------------------------------------- +# 9. Dynamic system prompt: template system_prompt from an upstream task's XCom +# --------------------------------------------------------------------------- + + +# [START howto_agent_dynamic_system_prompt] +@dag(tags=["example"]) +def example_agent_dynamic_system_prompt(): + @task + def classify(ticket: str) -> dict: + category = "shipping" if "order" in ticket.lower() else "other" + return {"priority": "high", "category": category} + + @task.agent( + llm_conn_id="pydanticai_default", + # system_prompt is a templated field -- Jinja renders it at task-run + # time, pulling the classification an upstream task already computed. + system_prompt=( + "You are handling a {{ ti.xcom_pull(task_ids='classify')['priority'] }}-priority " + "'{{ ti.xcom_pull(task_ids='classify')['category'] }}' ticket. " + "Draft a concise, friendly reply." + ), + ) + def draft_reply(ticket: str, triage: dict) -> str: + # `triage` creates the task dependency; its content also flows into + # system_prompt via Jinja above. The returned string is the *prompt* + # sent to the agent -- the drafted reply is this task's XCom output. + return f"Draft a reply for: {ticket}" + + ticket = "Where is my order? It still hasn't shipped." + draft_reply(ticket, classify(ticket)) + + +# [END howto_agent_dynamic_system_prompt] + +example_agent_dynamic_system_prompt()