How to Do Remote API Calls with Potentially Expiring Tokens#

python-icon Download Python Script

Python script/notebook for this guide.

MCP how-to script

Prerequisites This guide assumes familiarity with:

When buidling assistants with tools that reply on remote API calls, it is important to handle the authentication failures gracefully—especially those caused by expired access tokens. In this guide, you will build an assistant that calls a mock service requiring a valid token for authentication.

Setup#

To demonstrate the concept in a safe environment, we first set up a local mock API server (here, using Starlette). This simulates an endpoint that requires and validates an authentication token. If the token provided is:

  • a valid token (valid-token): the service responds with a success message.

  • an expired token (expired-token) or invalid token: the 401 Unauthorized error is returned with details.

 1from starlette.applications import Starlette
 2from starlette.responses import JSONResponse
 3from starlette.requests import Request
 4from starlette.routing import Route
 5from starlette.exceptions import HTTPException
 6from starlette.status import HTTP_401_UNAUTHORIZED
 7
 8async def protected_endpoint(request: Request):
 9    user = request.query_params.get("user")
10    if user is None:
11        return JSONResponse({"detail": "Missing 'user' query parameter."}, status_code=400)
12
13    authorization = request.headers.get("authorization")
14    if authorization is None or not authorization.startswith("Bearer "):
15        raise HTTPException(
16            status_code=HTTP_401_UNAUTHORIZED,
17            detail="Missing or malformed Authorization header.",
18            headers={"WWW-Authenticate": "Bearer"},
19        )
20
21    token = authorization.split(" ")[1]
22    if token == "valid-token":
23        return JSONResponse({"response": f"Success! You are authenticated, {user}."})
24    elif token == "expired-token":
25        raise HTTPException(
26            status_code=HTTP_401_UNAUTHORIZED,
27            detail="Token has expired.",
28            headers={"WWW-Authenticate": "Bearer error='invalid_token', error_description='The access token expired'"},
29        )
30    else:
31        raise HTTPException(
32            status_code=HTTP_401_UNAUTHORIZED,
33            detail="Invalid access token.",
34            headers={"WWW-Authenticate": "Bearer error='invalid_token'"},
35        )
36
37app = Starlette(debug=True, routes=[
38    Route("/protected", protected_endpoint)
39])
40
41# Start the server: Uncomment these lines
42# import uvicorn
43# uvicorn.run(app, host="localhost", port=8001)

Basic implementation#

In this example, you will build a simple Agent that includes a Flow with three steps:

  • A start step to get the user name

  • A step to trigger a client tool that collects a token from the user

  • A step to call a remote API given the user name and the token

This guide requires the use of an LLM. WayFlow supports several LLM API providers. Select an LLM from the options below:

from wayflowcore.models import OCIGenAIModel

if __name__ == "__main__":

    llm = OCIGenAIModel(
        model_id="provider.model-id",
        service_endpoint="https://url-to-service-endpoint.com",
        compartment_id="compartment-id",
        auth_type="API_KEY",
    )

Importing libraries#

First import what is needed for this tutorial:

 1from wayflowcore.property import StringProperty
 2from wayflowcore.tools import ClientTool
 3from wayflowcore.steps import (
 4    StartStep,
 5    CompleteStep,
 6    ApiCallStep,
 7    ToolExecutionStep
 8)
 9from wayflowcore.flow import Flow
10from wayflowcore.controlconnection import ControlFlowEdge
11from wayflowcore.dataconnection import DataFlowEdge

Creating the steps#

Define the variable names and steps.

1TOKEN = "token"
2USER = "user"

 1# 1. Start step
 2start_step = StartStep(
 3    name="start_step",
 4    input_descriptors=[StringProperty(name=USER)]
 5)
 6
 7# 2. Get token step
 8# A client tool to get token at client side
 9get_token_tool = ClientTool(
10    name="get_token_tool",
11    description="Get token from user",
12    input_descriptors=[],
13    output_descriptors=[StringProperty(name=TOKEN)]
14)
15
16# A step gets token by using the get_token_tool
17get_token_tool_step = ToolExecutionStep(
18    name="get_token_step",
19    tool=get_token_tool,
20)
21
22# 3. Call API step
23call_api_step = ApiCallStep(
24    name="call_api_step",
25    url="http://localhost:8003/protected",
26    allow_insecure_http=True,
27    method="GET",
28    headers={"Authorization": "Bearer {{ token }}"},
29    params={"user": "{{ user }}"},
30)
31
32# 4. End step
33end_step = CompleteStep(name="end_step")

In this simple example, we manually input the user name and the token in the code. For a more interactive approach, consider using InputMessageStep to prompt the user to enter these values during execution.

Creating the flow#

 1remote_call_flow = Flow(
 2    name="Remote Call Flow",
 3    description="Perform a call to a remote endpoint given the `user` parameter.",
 4    begin_step=start_step,
 5    control_flow_edges=[
 6        ControlFlowEdge(source_step=start_step, destination_step=get_token_tool_step),
 7        ControlFlowEdge(source_step=get_token_tool_step, destination_step=call_api_step),
 8        ControlFlowEdge(source_step=call_api_step, destination_step=end_step),
 9    ],
10    data_flow_edges=[
11        DataFlowEdge(
12            source_step=start_step,
13            source_output=USER,
14            destination_step=call_api_step,
15            destination_input=USER,
16        ),
17        DataFlowEdge(
18            source_step=get_token_tool_step,
19            source_output=TOKEN,
20            destination_step=call_api_step,
21            destination_input=TOKEN,
22        ),
23    ]
24)

This flow simply proceeds through three steps as defined in the control_flow_edges. The data_flow_edges connect the outputs of each step—the user name from start_step and the token from get_token_tool_step—to the inputs required by call_api_step.

Testing the flow#

 1from wayflowcore.executors.executionstatus import ToolRequestStatus
 2from wayflowcore.tools import ToolResult
 3
 4inputs = {"user": "alice"}
 5conversation = remote_call_flow.start_conversation(inputs=inputs)
 6
 7auth_token = "valid-token"
 8# auth_token = "expired-token" # This will raise error
 9
10status = conversation.execute()
11if isinstance(status, ToolRequestStatus): # Asking for token
12    tool_request_id = status.tool_requests[0].tool_request_id # Need to be adapted when using parallel tool calling (not the case here)
13    conversation.append_tool_result(ToolResult(content=auth_token, tool_request_id=tool_request_id))
14else:
15    print(
16        f"Invalid execution status, expected ToolRequestStatus, received {type(status)}"
17    )

To simulate a valid user, provide auth_token = "valid-token". To test expiry handling, use auth_token = "expired-token", which is expected to raise an error. The flow should pause at the token step, mimicking a credential input prompt, then proceed upon receiving input.

Creating an agent#

Now, create an agent that includes the defined flow:

1from wayflowcore.agent import Agent
2
3agent = Agent(
4    name="Agent",
5    flows=[remote_call_flow],
6    llm=llm,
7)

Testing the agent#

 1from wayflowcore.executors.executionstatus import ToolRequestStatus, UserMessageRequestStatus
 2from wayflowcore.tools import ToolResult
 3
 4conversation = agent.start_conversation()
 5conversation.append_user_message("Call the remote tool with user `alice`")
 6
 7auth_token = "valid-token"
 8# auth_token = "expired-token" # This will raise error
 9status = conversation.execute()
10
11if isinstance(status, ToolRequestStatus): # Asking for token
12    tool_request_id = status.tool_requests[0].tool_request_id # Needs to be adapted when using parallel tool calling (not the case here)
13    conversation.append_tool_result(ToolResult(content=auth_token, tool_request_id=tool_request_id))
14else:
15    print(
16        f"Invalid execution status, expected ToolRequestStatus, received {type(status)}"
17    )
18
19status = conversation.execute() # Resuming the conversation after the client provided the auth token
20if isinstance(status, UserMessageRequestStatus):
21    assistant_reply = conversation.get_last_message()
22    print(f"---\nAssistant >>> {assistant_reply.content}\n---")
23else:
24    print(
25        f"Invalid execution status, expected UserMessageRequestStatus, received {type(status)}"
26    )

The code block above demonstrates an interaction flow between a user and the agent, simulating how the assistant processes a remote, authenticated API call. During the first execution, the agent determines that a token is required and issues a tool request to the client. This is reflected by the status being an instance of ToolRequestStatus. After the client provides the required credential (the token), the second execution resumes the conversation. If authentication is successful, the agent proceeds to call the API, processes the response, and generates a user message as its reply. At this stage, the status should be UserMessageRequestStatus, which indicates that the agent has completed processing and is now ready to present a message to the user or wait for the next user prompt. Checking for UserMessageRequestStatus ensures that your code only tries to access the assistant’s reply when it is actually available.

Agent Spec Exporting/Loading#

You can export the assistant configuration to its Agent Spec configuration using the AgentSpecExporter.

from wayflowcore.agentspec import AgentSpecExporter

serialized_assistant = AgentSpecExporter().to_json(agent)

Here is what the Agent Spec representation will look like ↓

Click here to see the assistant configuration.
{
  "component_type": "ExtendedAgent",
  "id": "f3afdfe1-4b5b-469e-9505-33f72b3cd342",
  "name": "Agent",
  "description": "",
  "metadata": {
    "__metadata_info__": {
      "name": "Agent",
      "description": ""
    }
  },
  "inputs": [],
  "outputs": [],
  "llm_config": {
    "component_type": "VllmConfig",
    "id": "ab50c1e4-d6f3-493b-84b6-c570da3c7464",
    "name": "LLAMA_MODEL_ID",
    "description": null,
    "metadata": {
      "__metadata_info__": {}
    },
    "default_generation_parameters": null,
    "url": "LLAMA_API_URL",
    "model_id": "LLAMA_MODEL_ID"
  },
  "system_prompt": "",
  "tools": [],
  "toolboxes": [],
  "context_providers": null,
  "can_finish_conversation": false,
  "max_iterations": 10,
  "initial_message": "Hi! How can I help you?",
  "caller_input_mode": "always",
  "agents": [],
  "flows": [
    {
      "component_type": "Flow",
      "id": "8e8bc09a-6b3c-47c3-b8b8-ec64743beca8",
      "name": "Remote Call Flow",
      "description": "Perform a call to a remote endpoint given the `user` parameter.",
      "metadata": {
        "__metadata_info__": {}
      },
      "inputs": [
        {
          "type": "string",
          "title": "user"
        }
      ],
      "outputs": [
        {
          "type": "string",
          "title": "token"
        },
        {
          "description": "returned http status code",
          "type": "integer",
          "title": "http_status_code"
        }
      ],
      "start_node": {
        "$component_ref": "2b14ec5f-efe2-4b00-bfd9-1a7070b501b7"
      },
      "nodes": [
        {
          "$component_ref": "2b14ec5f-efe2-4b00-bfd9-1a7070b501b7"
        },
        {
          "$component_ref": "d4b896d7-fa3d-49ec-a45c-79021bf74e28"
        },
        {
          "$component_ref": "ed43def1-9d9b-4f15-84a3-cc9068143021"
        },
        {
          "$component_ref": "969d49b1-785b-4108-bd1e-6082b09160da"
        }
      ],
      "control_flow_connections": [
        {
          "component_type": "ControlFlowEdge",
          "id": "f44393b4-b9fb-4acd-9dc9-f4d496ccecae",
          "name": "start_step_to_get_token_step_control_flow_edge",
          "description": null,
          "metadata": {
            "__metadata_info__": {}
          },
          "from_node": {
            "$component_ref": "2b14ec5f-efe2-4b00-bfd9-1a7070b501b7"
          },
          "from_branch": null,
          "to_node": {
            "$component_ref": "d4b896d7-fa3d-49ec-a45c-79021bf74e28"
          }
        },
        {
          "component_type": "ControlFlowEdge",
          "id": "212e435b-059c-4842-9204-ed3b02f2bd17",
          "name": "get_token_step_to_call_api_step_control_flow_edge",
          "description": null,
          "metadata": {
            "__metadata_info__": {}
          },
          "from_node": {
            "$component_ref": "d4b896d7-fa3d-49ec-a45c-79021bf74e28"
          },
          "from_branch": null,
          "to_node": {
            "$component_ref": "ed43def1-9d9b-4f15-84a3-cc9068143021"
          }
        },
        {
          "component_type": "ControlFlowEdge",
          "id": "1b509ffa-24fe-4cda-a441-5456675ad64a",
          "name": "call_api_step_to_end_step_control_flow_edge",
          "description": null,
          "metadata": {
            "__metadata_info__": {}
          },
          "from_node": {
            "$component_ref": "ed43def1-9d9b-4f15-84a3-cc9068143021"
          },
          "from_branch": null,
          "to_node": {
            "$component_ref": "969d49b1-785b-4108-bd1e-6082b09160da"
          }
        }
      ],
      "data_flow_connections": [
        {
          "component_type": "DataFlowEdge",
          "id": "a2efbfb4-7046-4cba-b88f-010adaa77784",
          "name": "start_step_user_to_call_api_step_user_data_flow_edge",
          "description": null,
          "metadata": {
            "__metadata_info__": {}
          },
          "source_node": {
            "$component_ref": "2b14ec5f-efe2-4b00-bfd9-1a7070b501b7"
          },
          "source_output": "user",
          "destination_node": {
            "$component_ref": "ed43def1-9d9b-4f15-84a3-cc9068143021"
          },
          "destination_input": "user"
        },
        {
          "component_type": "DataFlowEdge",
          "id": "6d1baa26-611e-48d3-8021-c2ef39ec60d4",
          "name": "get_token_step_token_to_call_api_step_token_data_flow_edge",
          "description": null,
          "metadata": {
            "__metadata_info__": {}
          },
          "source_node": {
            "$component_ref": "d4b896d7-fa3d-49ec-a45c-79021bf74e28"
          },
          "source_output": "token",
          "destination_node": {
            "$component_ref": "ed43def1-9d9b-4f15-84a3-cc9068143021"
          },
          "destination_input": "token"
        },
        {
          "component_type": "DataFlowEdge",
          "id": "91b4c20c-93e6-4238-8dcd-28aea890788c",
          "name": "get_token_step_token_to_end_step_token_data_flow_edge",
          "description": null,
          "metadata": {
            "__metadata_info__": {}
          },
          "source_node": {
            "$component_ref": "d4b896d7-fa3d-49ec-a45c-79021bf74e28"
          },
          "source_output": "token",
          "destination_node": {
            "$component_ref": "969d49b1-785b-4108-bd1e-6082b09160da"
          },
          "destination_input": "token"
        },
        {
          "component_type": "DataFlowEdge",
          "id": "1c10dfaa-f6f9-4634-9b83-7e286996650c",
          "name": "call_api_step_http_status_code_to_end_step_http_status_code_data_flow_edge",
          "description": null,
          "metadata": {
            "__metadata_info__": {}
          },
          "source_node": {
            "$component_ref": "ed43def1-9d9b-4f15-84a3-cc9068143021"
          },
          "source_output": "http_status_code",
          "destination_node": {
            "$component_ref": "969d49b1-785b-4108-bd1e-6082b09160da"
          },
          "destination_input": "http_status_code"
        }
      ],
      "$referenced_components": {
        "ed43def1-9d9b-4f15-84a3-cc9068143021": {
          "component_type": "ApiNode",
          "id": "ed43def1-9d9b-4f15-84a3-cc9068143021",
          "name": "call_api_step",
          "description": "",
          "metadata": {
            "__metadata_info__": {}
          },
          "inputs": [
            {
              "description": "string template variable named user",
              "type": "string",
              "title": "user"
            },
            {
              "description": "string template variable named token",
              "type": "string",
              "title": "token"
            }
          ],
          "outputs": [
            {
              "description": "returned http status code",
              "type": "integer",
              "title": "http_status_code"
            }
          ],
          "branches": [
            "next"
          ],
          "url": "http://localhost:8003/protected",
          "http_method": "GET",
          "api_spec_uri": null,
          "data": {},
          "query_params": {
            "user": "{{ user }}"
          },
          "headers": {
            "Authorization": "Bearer {{ token }}"
          }
        },
        "2b14ec5f-efe2-4b00-bfd9-1a7070b501b7": {
          "component_type": "StartNode",
          "id": "2b14ec5f-efe2-4b00-bfd9-1a7070b501b7",
          "name": "start_step",
          "description": "",
          "metadata": {
            "__metadata_info__": {}
          },
          "inputs": [
            {
              "type": "string",
              "title": "user"
            }
          ],
          "outputs": [
            {
              "type": "string",
              "title": "user"
            }
          ],
          "branches": [
            "next"
          ]
        },
        "d4b896d7-fa3d-49ec-a45c-79021bf74e28": {
          "component_type": "ExtendedToolNode",
          "id": "d4b896d7-fa3d-49ec-a45c-79021bf74e28",
          "name": "get_token_step",
          "description": "",
          "metadata": {
            "__metadata_info__": {}
          },
          "inputs": [],
          "outputs": [
            {
              "type": "string",
              "title": "token"
            }
          ],
          "branches": [
            "next"
          ],
          "tool": {
            "component_type": "ClientTool",
            "id": "41399e99-d42a-45c1-adc5-d3d016bd5a1c",
            "name": "get_token_tool",
            "description": "Get token from user",
            "metadata": {
              "__metadata_info__": {}
            },
            "inputs": [],
            "outputs": [
              {
                "type": "string",
                "title": "token"
              }
            ]
          },
          "input_mapping": {},
          "output_mapping": {},
          "raise_exceptions": false,
          "component_plugin_name": "NodesPlugin",
          "component_plugin_version": "25.4.0.dev0"
        },
        "969d49b1-785b-4108-bd1e-6082b09160da": {
          "component_type": "EndNode",
          "id": "969d49b1-785b-4108-bd1e-6082b09160da",
          "name": "end_step",
          "description": null,
          "metadata": {
            "__metadata_info__": {}
          },
          "inputs": [
            {
              "type": "string",
              "title": "token"
            },
            {
              "description": "returned http status code",
              "type": "integer",
              "title": "http_status_code"
            }
          ],
          "outputs": [
            {
              "type": "string",
              "title": "token"
            },
            {
              "description": "returned http status code",
              "type": "integer",
              "title": "http_status_code"
            }
          ],
          "branches": [],
          "branch_name": "end_step"
        }
      }
    }
  ],
  "agent_template": {
    "component_type": "PluginPromptTemplate",
    "id": "95b539ce-0908-4a45-9ab0-cd6fa1b346d7",
    "name": "",
    "description": "",
    "metadata": {
      "__metadata_info__": {}
    },
    "messages": [
      {
        "role": "system",
        "contents": [
          {
            "type": "text",
            "content": "{% if custom_instruction %}{{custom_instruction}}{% endif %}"
          }
        ],
        "tool_requests": null,
        "tool_result": null,
        "display_only": false,
        "sender": null,
        "recipients": [],
        "time_created": "2025-09-02T15:52:22.014400+00:00",
        "time_updated": "2025-09-02T15:52:22.014401+00:00"
      },
      {
        "role": "user",
        "contents": [],
        "tool_requests": null,
        "tool_result": null,
        "display_only": false,
        "sender": null,
        "recipients": [],
        "time_created": "2025-09-02T15:52:22.008803+00:00",
        "time_updated": "2025-09-02T15:52:22.010218+00:00"
      },
      {
        "role": "system",
        "contents": [
          {
            "type": "text",
            "content": "{% if __PLAN__ %}The current plan you should follow is the following: \n{{__PLAN__}}{% endif %}"
          }
        ],
        "tool_requests": null,
        "tool_result": null,
        "display_only": false,
        "sender": null,
        "recipients": [],
        "time_created": "2025-09-02T15:52:22.014421+00:00",
        "time_updated": "2025-09-02T15:52:22.014421+00:00"
      }
    ],
    "output_parser": null,
    "inputs": [
      {
        "description": "\"custom_instruction\" input variable for the template",
        "type": "string",
        "title": "custom_instruction",
        "default": ""
      },
      {
        "description": "\"__PLAN__\" input variable for the template",
        "type": "string",
        "title": "__PLAN__",
        "default": ""
      },
      {
        "type": "array",
        "items": {},
        "title": "__CHAT_HISTORY__"
      }
    ],
    "pre_rendering_transforms": null,
    "post_rendering_transforms": [
      {
        "component_type": "PluginRemoveEmptyNonUserMessageTransform",
        "id": "372d6f16-b945-4b10-b1c8-adc143ddab9d",
        "name": "removeemptynonusermessage_messagetransform",
        "description": null,
        "metadata": {
          "__metadata_info__": {}
        },
        "component_plugin_name": "MessageTransformPlugin",
        "component_plugin_version": "25.4.0.dev0"
      }
    ],
    "tools": null,
    "native_tool_calling": true,
    "response_format": null,
    "native_structured_generation": true,
    "generation_config": null,
    "component_plugin_name": "PromptTemplatePlugin",
    "component_plugin_version": "25.4.0.dev0"
  },
  "component_plugin_name": "AgentPlugin",
  "component_plugin_version": "25.4.0.dev0",
  "agentspec_version": "25.4.1"
}

You can then load the configuration back to an assistant using the AgentSpecLoader.

from wayflowcore.agentspec import AgentSpecLoader

assistant: Agent = AgentSpecLoader().load_json(serialized_assistant)

Note

This guide uses the following extension/plugin Agent Spec components:

  • PluginPromptTemplate

  • PluginRemoveEmptyNonUserMessageTransform

  • ExtendedToolNode

  • ExtendedAgent

See the list of available Agent Spec extension/plugin components in the API Reference

Next steps#

In this guide, you learned how to define a simple flow that retrieves a token from the user and uses it to authenticate remote API calls. To continue learning, checkout:

Full code#

  1# Copyright © 2025 Oracle and/or its affiliates.
  2#
  3# This software is under the Universal Permissive License
  4# %%[markdown]
  5# Code Example - How to Do Remote API Calls with Potentially Expiring Tokens
  6# --------------------------------------------------------------------------
  7
  8# How to use:
  9# Create a new Python virtual environment and install the latest WayFlow version.
 10# ```bash
 11# python -m venv venv-wayflowcore
 12# source venv-wayflowcore/bin/activate
 13# pip install --upgrade pip
 14# pip install "wayflowcore==26.1" 
 15# ```
 16
 17# You can now run the script
 18# 1. As a Python file:
 19# ```bash
 20# python howto_remote_tool_expired_token.py
 21# ```
 22# 2. As a Notebook (in VSCode):
 23# When viewing the file,
 24#  - press the keys Ctrl + Enter to run the selected cell
 25#  - or Shift + Enter to run the selected cell and move to the cell below# (UPL) 1.0 (LICENSE-UPL or https://oss.oracle.com/licenses/upl) or Apache License
 26# 2.0 (LICENSE-APACHE or http://www.apache.org/licenses/LICENSE-2.0), at your option.
 27
 28
 29
 30
 31# %%[markdown]
 32## Mock server
 33
 34# %%
 35from starlette.applications import Starlette
 36from starlette.responses import JSONResponse
 37from starlette.requests import Request
 38from starlette.routing import Route
 39from starlette.exceptions import HTTPException
 40from starlette.status import HTTP_401_UNAUTHORIZED
 41
 42async def protected_endpoint(request: Request):
 43    user = request.query_params.get("user")
 44    if user is None:
 45        return JSONResponse({"detail": "Missing 'user' query parameter."}, status_code=400)
 46
 47    authorization = request.headers.get("authorization")
 48    if authorization is None or not authorization.startswith("Bearer "):
 49        raise HTTPException(
 50            status_code=HTTP_401_UNAUTHORIZED,
 51            detail="Missing or malformed Authorization header.",
 52            headers={"WWW-Authenticate": "Bearer"},
 53        )
 54
 55    token = authorization.split(" ")[1]
 56    if token == "valid-token":
 57        return JSONResponse({"response": f"Success! You are authenticated, {user}."})
 58    elif token == "expired-token":
 59        raise HTTPException(
 60            status_code=HTTP_401_UNAUTHORIZED,
 61            detail="Token has expired.",
 62            headers={"WWW-Authenticate": "Bearer error='invalid_token', error_description='The access token expired'"},
 63        )
 64    else:
 65        raise HTTPException(
 66            status_code=HTTP_401_UNAUTHORIZED,
 67            detail="Invalid access token.",
 68            headers={"WWW-Authenticate": "Bearer error='invalid_token'"},
 69        )
 70
 71app = Starlette(debug=True, routes=[
 72    Route("/protected", protected_endpoint)
 73])
 74
 75# Start the server: Uncomment these lines
 76# import uvicorn
 77# uvicorn.run(app, host="localhost", port=8001)
 78
 79# %%[markdown]
 80## Import libraries
 81
 82# %%
 83from wayflowcore.property import StringProperty
 84from wayflowcore.tools import ClientTool
 85from wayflowcore.steps import (
 86    StartStep,
 87    CompleteStep,
 88    ApiCallStep,
 89    ToolExecutionStep
 90)
 91from wayflowcore.flow import Flow
 92from wayflowcore.controlconnection import ControlFlowEdge
 93from wayflowcore.dataconnection import DataFlowEdge
 94
 95# %%[markdown]
 96## Configure your LLM
 97
 98# %%
 99from wayflowcore.models import VllmModel
100llm = VllmModel(
101    model_id="LLAMA_MODEL_ID",
102    host_port="LLAMA_API_URL",
103)
104
105# %%[markdown]
106## Variable names
107
108# %%
109TOKEN = "token"
110USER = "user"
111
112# %%[markdown]
113## Defining steps
114
115# %%
116# 1. Start step
117start_step = StartStep(
118    name="start_step",
119    input_descriptors=[StringProperty(name=USER)]
120)
121
122# 2. Get token step
123# A client tool to get token at client side
124get_token_tool = ClientTool(
125    name="get_token_tool",
126    description="Get token from user",
127    input_descriptors=[],
128    output_descriptors=[StringProperty(name=TOKEN)]
129)
130
131# A step gets token by using the get_token_tool
132get_token_tool_step = ToolExecutionStep(
133    name="get_token_step",
134    tool=get_token_tool,
135)
136
137# 3. Call API step
138call_api_step = ApiCallStep(
139    name="call_api_step",
140    url="http://localhost:8003/protected",
141    allow_insecure_http=True,
142    method="GET",
143    headers={"Authorization": "Bearer {{ token }}"},
144    params={"user": "{{ user }}"},
145)
146
147# 4. End step
148end_step = CompleteStep(name="end_step")
149
150# %%[markdown]
151## Defining flow
152
153# %%
154remote_call_flow = Flow(
155    name="Remote Call Flow",
156    description="Perform a call to a remote endpoint given the `user` parameter.",
157    begin_step=start_step,
158    control_flow_edges=[
159        ControlFlowEdge(source_step=start_step, destination_step=get_token_tool_step),
160        ControlFlowEdge(source_step=get_token_tool_step, destination_step=call_api_step),
161        ControlFlowEdge(source_step=call_api_step, destination_step=end_step),
162    ],
163    data_flow_edges=[
164        DataFlowEdge(
165            source_step=start_step,
166            source_output=USER,
167            destination_step=call_api_step,
168            destination_input=USER,
169        ),
170        DataFlowEdge(
171            source_step=get_token_tool_step,
172            source_output=TOKEN,
173            destination_step=call_api_step,
174            destination_input=TOKEN,
175        ),
176    ]
177)
178
179# %%[markdown]
180## Testing flow
181
182# %%
183from wayflowcore.executors.executionstatus import ToolRequestStatus
184from wayflowcore.tools import ToolResult
185
186inputs = {"user": "alice"}
187conversation = remote_call_flow.start_conversation(inputs=inputs)
188
189auth_token = "valid-token"
190# auth_token = "expired-token" # This will raise error
191
192status = conversation.execute()
193if isinstance(status, ToolRequestStatus): # Asking for token
194    tool_request_id = status.tool_requests[0].tool_request_id # Need to be adapted when using parallel tool calling (not the case here)
195    conversation.append_tool_result(ToolResult(content=auth_token, tool_request_id=tool_request_id))
196else:
197    print(
198        f"Invalid execution status, expected ToolRequestStatus, received {type(status)}"
199    )
200
201# %%[markdown]
202## Defining agent
203
204# %%
205from wayflowcore.agent import Agent
206
207agent = Agent(
208    name="Agent",
209    flows=[remote_call_flow],
210    llm=llm,
211)
212
213# %%[markdown]
214## Testing agent
215
216# %%
217from wayflowcore.executors.executionstatus import ToolRequestStatus, UserMessageRequestStatus
218from wayflowcore.tools import ToolResult
219
220conversation = agent.start_conversation()
221conversation.append_user_message("Call the remote tool with user `alice`")
222
223auth_token = "valid-token"
224# auth_token = "expired-token" # This will raise error
225status = conversation.execute()
226
227if isinstance(status, ToolRequestStatus): # Asking for token
228    tool_request_id = status.tool_requests[0].tool_request_id # Needs to be adapted when using parallel tool calling (not the case here)
229    conversation.append_tool_result(ToolResult(content=auth_token, tool_request_id=tool_request_id))
230else:
231    print(
232        f"Invalid execution status, expected ToolRequestStatus, received {type(status)}"
233    )
234
235status = conversation.execute() # Resuming the conversation after the client provided the auth token
236if isinstance(status, UserMessageRequestStatus):
237    assistant_reply = conversation.get_last_message()
238    print(f"---\nAssistant >>> {assistant_reply.content}\n---")
239else:
240    print(
241        f"Invalid execution status, expected UserMessageRequestStatus, received {type(status)}"
242    )
243
244# %%[markdown]
245## Export config to Agent Spec
246
247# %%
248from wayflowcore.agentspec import AgentSpecExporter
249
250serialized_assistant = AgentSpecExporter().to_json(agent)
251
252# %%[markdown]
253## Load Agent Spec config
254
255# %%
256from wayflowcore.agentspec import AgentSpecLoader
257
258assistant: Agent = AgentSpecLoader().load_json(serialized_assistant)