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 building 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 FastAPI). 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 fastapi import FastAPI, HTTPException, Request, status as fastapi_status
 2from fastapi.responses import JSONResponse
 3
 4async def protected_endpoint(request: Request):
 5    user = request.query_params.get("user")
 6    if user is None:
 7        return JSONResponse({"detail": "Missing 'user' query parameter."}, status_code=400)
 8
 9    authorization = request.headers.get("authorization")
10    if authorization is None or not authorization.startswith("Bearer "):
11        raise HTTPException(
12            status_code=fastapi_status.HTTP_401_UNAUTHORIZED,
13            detail="Missing or malformed Authorization header.",
14            headers={"WWW-Authenticate": "Bearer"},
15        )
16
17    token = authorization.split(" ")[1]
18    if token == "valid-token":
19        return JSONResponse({"response": f"Success! You are authenticated, {user}."})
20    elif token == "expired-token":
21        raise HTTPException(
22            status_code=fastapi_status.HTTP_401_UNAUTHORIZED,
23            detail="Token has expired.",
24            headers={"WWW-Authenticate": "Bearer error='invalid_token', error_description='The access token expired'"},
25        )
26    else:
27        raise HTTPException(
28            status_code=fastapi_status.HTTP_401_UNAUTHORIZED,
29            detail="Invalid access token.",
30            headers={"WWW-Authenticate": "Bearer error='invalid_token'"},
31        )
32
33app = FastAPI(debug=True, docs_url=None, redoc_url=None, openapi_url=None)
34app.add_api_route("/protected", protected_endpoint, methods=["GET"])
35
36# Start the server: Uncomment these lines
37# import uvicorn
38# uvicorn.run(app, host="localhost", port=8003)

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, OCIClientConfigWithApiKey

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

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