genai-toolbox/sdks/langchain/README.md

GenAI Toolbox SDK

This SDK allows you to seamlessly integrate the functionalities of Toolbox into your LLM applications, enabling advanced orchestration and interaction with GenAI models.

Table of Contents

• Installation
• Usage
• Load a toolset
• Load a single tool
• Use with LangChain
• Use with LangGraph
  • Represent Tools as Nodes
  • Connect Tools with LLM
• Manual usage
• Authenticating Tools
  • Supported Authentication Mechanisms
  • Configuring Tools for Authentication
  • Configure SDK for Authentication
  • Complete Example

Installation

Important

This SDK is not yet available on PyPI. For now, install it from source by following these installation instructions.

You can install the Toolbox SDK for LangChain using pip.

pip install toolbox-langchain-sdk

Usage

Import and initialize the toolbox client.

from toolbox_langchain_sdk import ToolboxClient

# Replace with your Toolbox service's URL
toolbox = ToolboxClient("http://127.0.0.1:5000")

Important

The toolbox client requires an asynchronous environment. For guidance on running asynchronous Python programs, see running an async program in python.

Tip

You can also pass your own ClientSession so that the ToolboxClient can reuse the same session.

async with ClientSession() as session:
  toolbox = ToolboxClient("http://localhost:5000", session)

Load a toolset

You can load a toolset, a collection of related tools.

# Load all tools
tools = await toolbox.load_toolset()

# Load a specific toolset
tools = await toolbox.load_toolset("my-toolset")

Load a single tool

You can also load a single tool.

tool = await toolbox.load_tool("my-tool")

Use with LangChain

LangChain's agents can dynamically choose and execute tools based on the user input. The user can include the tools loaded from the Toolbox SDK in the agent's toolkit.

from langchain_google_vertexai import ChatVertexAI

model = ChatVertexAI(model="gemini-1.5-pro-002")

# Initialize agent with tools
agent = model.bind_tools(tools)

# Run the agent
result = agent.invoke("Do something with the tools")

Use with LangGraph

The Toolbox SDK can be seamlessly integrated with LangGraph to enable the use of Toolbox service tools within a graph-based workflow. Using this SDK, we can follow the official guide with minimal changes.

Represent Tools as Nodes

Each tool generated by the SDK can be represented as a LangGraph node. The node's functionality would encapsulate the execution of the corresponding tool.

from toolbox_langchain_sdk import ToolboxClient
from langgraph.graph import StateGraph, MessagesState
from langgraph.prebuilt import ToolNode

# Define the function that calls the model
def call_model(state: MessagesState):
    messages = state['messages']
    response = model.invoke(messages)
    # We return a list, because this will get added to the existing list
    return {"messages": [response]}

model = ChatVertexAI(model="gemini-1.5-pro-002")
builder = StateGraph(MessagesState)
tool_node = ToolNode(tools)

builder.add_node("agent", call_model)
builder.add_node("tools", tool_node)

Connect Tools with LLM

Now we can connect the tool nodes with LLM nodes. The LLM can decide which tool to use based on the user input or the context of the conversation. The output from a tool can then be fed back into the LLM for further processing or decision-making.

from typing import Literal
from langgraph.graph import END, START
from langchain_core.messages import HumanMessage

# Define the function that determines whether to continue or not
def should_continue(state: MessagesState) -> Literal["tools", END]:
    messages = state['messages']
    last_message = messages[-1]
    # If the LLM makes a tool call, then we route to the "tools" node
    if last_message.tool_calls:
        return "tools"
    # Otherwise, we stop (reply to the user)
    return END

builder.add_edge(START, "agent")
builder.add_conditional_edges(
    "agent",
    should_continue,
)
builder.add_edge("tools", 'agent')

graph = builder.compile()

graph.invoke(
    {"messages": [HumanMessage(content="Do something with the tools")]},
)

Manual usage

You can also execute a tool manually using the arun method.

result = await tools[0].arun({ "name": "Alice", "age": 30 })

Authenticating Tools

Warning

Always use HTTPS to connect your application with the Toolbox service, especially when using tools with authentication configured. Using HTTP exposes your application to serious security risks, including unauthorized access to user information and man-in-the-middle attacks, where sensitive data can be intercepted.

Some tools in your Toolbox configuration might require user authentication to access sensitive data. This section guides you on how to configure tools for authentication and use them with the SDK.

Supported Authentication Mechanisms

The Toolbox SDK currently supports authentication using OIDC protocol. Specifically, it uses ID tokens and not access tokens for Google OAuth 2.0.

Configuring Tools for Authentication

Refer to these instructions on configuring tools for authenticated parameters.

Configure SDK for Authentication

Provide the auth_tokens parameter to the load_tool or load_toolset calls with a dictionary. The keys of this dictionary should match the names of the authentication sources configured in your tools file (e.g., my_auth_service), and the values should be callable functions (e.g., lambdas or regular functions) that return the ID token of the logged-in user.

Here's an example:

def get_auth_token():
    # ... Logic to retrieve ID token (e.g., from local storage, OAuth flow)
    # This example just returns a placeholder. Replace with your actual token retrieval.
    return "YOUR_ID_TOKEN"

toolbox = ToolboxClient("http://localhost:5000")

tools = toolbox.load_toolset(auth_tokens={ "my_auth_service": get_auth_token })

# OR

tool = toolbox.load_tool("my_tool", auth_tokens={ "my_auth_service": get_auth_token })

Alternatively, you can call the add_auth_token method to configure authentication separately.

toolbox.add_auth_token("my_auth_service", get_auth_token)

Note

Authentication tokens added via load_tool, load_toolset, or add_auth_token apply to all subsequent tool invocations, regardless of when the tool was loaded. This ensures a consistent authentication context.

Complete Example

import asyncio
from toolbox_langchain_sdk import ToolboxClient

async def get_auth_token():
    # Replace with your actual ID token retrieval logic.
    # For example, using a library like google-auth
    # from google.oauth2 import id_token
    # from google.auth.transport import requests
    # request = requests.Request()
    # id_token_string = id_token.fetch_id_token(request, "YOUR_AUDIENCE")# Replace with your audience
    # return id_token_string
    return "YOUR_ACTUAL_ID_TOKEN" # placeholder

async def main():
    toolbox = ToolboxClient("http://localhost:5000")
    toolbox.add_auth_token("my_auth_service", get_auth_token)
    tools = await toolbox.load_toolset()
    result = await tools[0].arun({"input": "some input"})
    print(result)

if __name__ == "__main__":
    asyncio.run(main())