Research agent with human review¶

Someone has a question that needs a knowledge base, a dataset and a web search to answer, and the right sequence of lookups depends on the question. This Dag runs a LangChain ReAct agent that decides for itself which tools to call and in what order, then hands the raw findings to a separate formatting step and to a reviewer. Airflow puts a person in front of the agent to edit the question, exposes the findings and tool calls as XCom, makes formatting its own retryable task, and holds the report for approval.

This is the shape for teams with existing LangChain tools. The agent loop is LangChain’s; the connection, the formatting call and the review gates are Airflow’s.

What this demonstrates¶

LangChain models: LangChainHook – LangChainHook supplies the chat model and embeddings from one Airflow connection.
Single prompts: LLMOperator and @task.llm – LLMOperator formats the findings with a Jinja template that reads the agent’s XCom.
Human in the Loop (HITL) Operators – HITLEntryOperator to edit the question, ApprovalOperator to sign off the report.
LangChain tools in both directions – the reverse direction, giving a LangChain agent tools built from Airflow connections.

Run it¶

Install the LangChain extra and the packages the tools use:

pip install "apache-airflow-providers-common-ai[langchain]" \
    langchain-openai langchain-text-splitters langchain-community faiss-cpu

Create a langchain connection named langchain_default with your API key in the password field (see LangChain connection).
Optionally put documents under DOCS_PATH and the survey CSV at SURVEY_CSV_PATH; without them the Dag writes sample pages so the tools have something to search.

Trigger the Dag:

airflow dags test example_langchain_tool_agent

The run pauses at prompt_review and report_approval; answer from Required Actions in the UI (see the note on What you can build). The run_research_agent log shows every tool call, and its XCom keeps the findings and call list.

The Dag¶

airflow/providers/common/ai/example_dags/example_langchain_tool_agent.py[source]

@dag(tags=["example"])
def example_langchain_tool_agent():
    """
    Research agent with LangChain tools and human review.

    Task graph::

        prompt_review (HITLEntryOperator)
            -> prepare_tools (@task)
            -> run_research_agent (@task)
            -> format_report (LLMOperator)
            -> report_approval (ApprovalOperator)

    The agent uses LangChain's ``create_agent`` with a ReAct reasoning
    loop.  It autonomously decides which tools to call -- knowledge base
    search, survey data query, web search, or current-time lookup --
    based on the user's question.  The number and sequence of tool calls
    is determined by the LLM at runtime.

    The surrounding Airflow DAG provides what the agent cannot:
    human review of the question (HITLEntryOperator), formatted report
    generation (LLMOperator), and human approval of the final output
    (ApprovalOperator).
    """

    prompt_review = HITLEntryOperator(
        task_id="prompt_review",
        subject="Review the research question",
        params={
            "question": Param(
                DEFAULT_QUESTION,
                type="string",
                description="The research question for the agent to investigate",
            ),
        },
        response_timeout=datetime.timedelta(hours=1),
    )

    @task
    def prepare_tools(hitl_response: dict) -> dict:
        """Build the FAISS knowledge base index and resolve tool config."""
        from airflow.providers.common.ai.hooks.langchain import LangChainHook

        hook = LangChainHook(
            llm_conn_id=LLM_CONN_ID,
            llm_model=LLM_MODEL,
            embed_model=EMBEDDING_MODEL,
        )
        index_dir = _ensure_knowledge_base(hook)

        question = hitl_response["params_input"]["question"]
        return {
            "question": question,
            "index_dir": index_dir,
            "survey_csv_path": SURVEY_CSV_PATH,
        }

    @task
    def run_research_agent(config: dict) -> dict:
        """Run a LangChain ReAct agent that autonomously researches the question.

        The agent decides which tools to call and in what order.  The number
        of tool calls depends on the complexity of the question.  All
        reasoning steps, tool calls, and observations are logged.
        """
        from langchain.agents import create_agent

        from airflow.providers.common.ai.hooks.langchain import LangChainHook

        hook = LangChainHook(
            llm_conn_id=LLM_CONN_ID,
            llm_model=LLM_MODEL,
            embed_model=EMBEDDING_MODEL,
        )
        model = hook.get_chat_model()
        tools = _build_tools(hook, config["index_dir"], config["survey_csv_path"])

        agent = create_agent(
            model,
            tools=tools,
            system_prompt=(
                "You are a thorough research assistant for Apache Airflow. "
                "You have access to tools for searching a knowledge base, "
                "querying survey data, searching the web, and checking the "
                "current UTC time. "
                "Use the appropriate tools to fully answer the question. "
                "Combine information from multiple sources when relevant. "
                "Always cite which tool provided each piece of information."
            ),
        )

        question = config["question"]
        print(f"Research question: {question}")
        print("Agent starting research...")

        tool_calls_log = []
        final_answer = ""

        for step in agent.stream(
            {"messages": [{"role": "user", "content": question}]},
            stream_mode="values",
        ):
            msg = step["messages"][-1]
            if hasattr(msg, "tool_calls") and msg.tool_calls:
                for tc in msg.tool_calls:
                    tool_calls_log.append(
                        {
                            "tool": tc["name"],
                            "args": str(tc.get("args", {}))[:200],
                        }
                    )
                    print(f"  Tool call: {tc['name']}({tc.get('args', {})})")
            elif hasattr(msg, "content") and msg.content:
                final_answer = msg.content

        print(f"Agent completed. Tool calls made: {len(tool_calls_log)}")

        return {
            "question": question,
            "findings": final_answer,
            "tool_calls": tool_calls_log,
            "tool_call_count": len(tool_calls_log),
        }

    tools_config = prepare_tools(prompt_review.output)
    research_result = run_research_agent(tools_config)

    format_report = LLMOperator(
        task_id="format_report",
        llm_conn_id=LLM_CONN_ID,
        system_prompt=REPORT_SYSTEM_PROMPT,
        prompt="""\
Format the following research findings into a clear report.

{% set result = ti.xcom_pull(task_ids='run_research_agent') -%}
Question: {{ result['question'] }}

Raw findings:
{{ result['findings'] }}

Tools used: {{ result['tool_call_count'] }} calls
{% for tc in result['tool_calls'] -%}
  - {{ tc['tool'] }}: {{ tc['args'][:100] }}
{% endfor -%}""",
    )
    research_result >> format_report

    report_approval = ApprovalOperator(  # noqa: F841
        task_id="report_approval",
        subject="Review the research report",
        body=format_report.output,
        response_timeout=datetime.timedelta(hours=1),
    )

The tools (knowledge-base search, survey query, a stubbed web search, a clock) are ordinary LangChain @tool functions built in _build_tools.

Adapting it¶

Replace _build_tools with your own LangChain tools. Anything from the LangChain ecosystem works unchanged.
To run the same shape on pydantic-ai instead, use Agents with tools: AgentOperator and @task.agent with Toolsets; the surrounding review and formatting tasks stay as they are.
For routine questions, drop prompt_review and take the question from a Dag param; keep report_approval.