Tool Calling with Stained Glass Proxy (SGP) and Output Protected vLLM.

Large Language Models (LLMs) increasingly exhibit agentic behavior through their ability to use external tools. Instruction-tuned LLMs can interleave natural language chat with structured tool calls, enabling capabilities like web search, retrieval, and code execution.

This notebook demonstrates:

• Creating a custom tool schema to provide in-context to the target model hosted with Protopia's output protected vLLM.
• Stained Glass Transform (SGT) embedding protection of the conversation messages and tools.
• Attempted reconstruction failure: Given the SGT protection, a similarity search on the base model embedding layer cannot recover tool definition and outputs from the protected embeddings.
• Preserved utility: the target model retains tool-calling capability despite embedding protection.

Flow Diagram

sequenceDiagram
    autonumber
    participant Client
    participant Tokenizer as Chat Template & Tokenization
    participant SGT as Stained Glass Transform
    participant LLM as Output Protected vLLM

    Client->>Tokenizer: Prompt + Tools Definitions (JSON)
    Tokenizer->>SGT: Tokens from Template Formatted String
    SGT->>LLM: SGT Protected Embeddings
    LLM->>Client: Tool Calls (`tool_calls`)
    Client->>Client: Execute Tool Call(s)
        Note over Client: Run tools and append results to messages
    Client->>Tokenizer: Messages with Tool Call Outputs
    Tokenizer->>SGT: Tokens from Template Formatted String
    SGT->>LLM: SGT Protected Embeddings
    LLM->>Client: Final response

Pre-requisites

• A live instance of vLLM (>=v0.9.1) OpenAI-Compatible Server, with:
  • prompt embeddings enabled and
  • tool calling enabled (--enable-auto-tool-choice, --tool-call-parser)
  • recommended chat template (--chat-template)
• A live instance of SGT Proxy (Please refer to the deployment instructions) with tool calling enabled.

Enabling Tool Calling on SGP

Tool parsers detect and extract function calls from the model outputs. Set SGP_TOOL_PARSER to a valid parser (e.g., llama3_json for Llama-3.1-8B-Instruct). Find available parsers with vllm serve -h | grep tool-call-parser.

Setup SGT Proxy Access

In [1]:

import json

import markdown
import openai
from IPython.display import Markdown

import json import markdown import openai from IPython.display import Markdown

Configuration Required

Update these parameters for your specific setup:

• PROXY_URL: Your proxy server endpoint
• MODEL_NAME: The base model you want to test
• API_KEY: Your authentication key

In [2]:

# Set proxy access parameters.
PROXY_URL = "http://127.0.0.1:8601/v1"
MODEL_NAME = "meta-llama/Llama-3.1-8B-Instruct"
API_KEY = "<overwrite-with-your-api-key>"

# Set proxy access parameters. PROXY_URL = "http://127.0.0.1:8601/v1" MODEL_NAME = "meta-llama/Llama-3.1-8B-Instruct" API_KEY = "<overwrite-with-your-api-key>"

In [3]:

# Verify that the model is accessible through the proxy.
client = openai.OpenAI(
    base_url=PROXY_URL,
    api_key=API_KEY,
)
assert MODEL_NAME in [model.id for model in client.models.list()], (
    f"{MODEL_NAME=} was not found at {PROXY_URL=}/v1/models"
)

# Quick health check
response = client.chat.completions.create(
    model=MODEL_NAME,
    messages=[{"role": "user", "content": "Hello, world!"}],
    max_tokens=10,
)

# Hint: if this is not working check your API_KEY, if not API_KEY is needed set API_KEY="".
assert response.choices[0].message.content.strip() != "", (
    "Empty response from model"
)

"✅ SGT Proxy is accessible and responding correctly."

# Verify that the model is accessible through the proxy. client = openai.OpenAI( base_url=PROXY_URL, api_key=API_KEY, ) assert MODEL_NAME in [model.id for model in client.models.list()], ( f"{MODEL_NAME=} was not found at {PROXY_URL=}/v1/models" ) # Quick health check response = client.chat.completions.create( model=MODEL_NAME, messages=[{"role": "user", "content": "Hello, world!"}], max_tokens=10, ) # Hint: if this is not working check your API_KEY, if not API_KEY is needed set API_KEY="". assert response.choices[0].message.content.strip() != "", ( "Empty response from model" ) "✅ SGT Proxy is accessible and responding correctly."

Out[3]:

'✅ SGT Proxy is accessible and responding correctly.'

Tool Definitions

The chat completion API's tools parameter accepts an array of tool definitions. This demo uses a custom JSON tool to simulate scenarios where tool information privacy matters. The model infers when to invoke tools based on conversation context, then responds with the selected tools and parameters.

In [4]:

# Task
messages = [
    {
        "role": "system",
        "content": "You are a helpful tax assistant that can retrieve user information to answer tax-related questions.",
    },
    {
        "role": "user",
        "content": "What are important tax recommendations for user with ID 'user_123' based on their financial information?",
    },
]

# Tool (Function)
get_user_info = {
    "type": "function",
    "function": {
        "name": "get_user",
        "description": "Retrieve information about a user given their user ID.",
        "parameters": {
            "properties": {
                "user_id": {
                    "type": "string",
                    "description": "The unique identifier of the user to retrieve information for",
                }
            },
            "required": ["user_id"],
        },
    },
}

# Task messages = [ { "role": "system", "content": "You are a helpful tax assistant that can retrieve user information to answer tax-related questions.", }, { "role": "user", "content": "What are important tax recommendations for user with ID 'user_123' based on their financial information?", }, ] # Tool (Function) get_user_info = { "type": "function", "function": { "name": "get_user", "description": "Retrieve information about a user given their user ID.", "parameters": { "properties": { "user_id": { "type": "string", "description": "The unique identifier of the user to retrieve information for", } }, "required": ["user_id"], }, }, }

Our demo must execute the tool call generated by the model and return the results.

In [5]:

def get_user(user_id: str) -> dict:
    """Simulate function to get user information."""
    # In a real implementation, this would query a database or external service.
    return {
        "user_id": user_id,
        "name": "John Doe",
        "age": 25,
        "location": "New York, USA",
        "income": "$60,000",
        "filing_status": "Single",
    }

def get_user(user_id: str) -> dict: """Simulate function to get user information.""" # In a real implementation, this would query a database or external service. return { "user_id": user_id, "name": "John Doe", "age": 25, "location": "New York, USA", "income": "$60,000", "filing_status": "Single", }

In [6]:

skip-execution

response = client.chat.completions.create(
    model=MODEL_NAME,
    messages=messages,
    tools=[get_user_info],
    # when tool_choice='required' is set, the model is guaranteed to generate one or more tool calls.
    tool_choice="required",
    max_tokens=32,
)
assistant_message = response.choices[0].message
# Append assistant tool call message to the conversation history
messages.append(assistant_message)
if assistant_message.tool_calls:
    for tool_call in assistant_message.tool_calls:
        assert tool_call.type == "function"
        name = tool_call.function.name
        arguments = json.loads(tool_call.function.arguments)
        print(f"🛠️ Tool {name=} call with arguments {arguments=}")

        if name == "get_user":
            user_info = get_user(arguments["user_id"])
            # Append tool result to the conversation history
            messages.append(
                {
                    # ipython: A role introduced in Llama 3.1. Semantically, this role means "tool call output".
                    # This role is used to mark messages with the output of a tool call.
                    # https://www.llama.com/docs/model-cards-and-prompt-formats/llama3_1/#-supported-roles-
                    # both 'tool' and 'ipython' roles can be used to send function call results back to the model.
                    # https://github.com/vllm-project/vllm/blob/main/examples/tool_chat_template_llama3.1_json.jinja#L104
                    "role": "tool",
                    "content": json.dumps(user_info),
                    "tool_call_id": tool_call.id,
                }
            )
    # Get final response from the model after tool execution
    final_response = (
        client.chat.completions.create(
            model=MODEL_NAME,
            messages=messages,
            max_tokens=1024,
        )
        .choices[0]
        .message.content
    )
    display(Markdown(markdown.markdown(final_response)))
else:
    print(f"No tool calls were made by the model, {assistant_message=}")

response = client.chat.completions.create( model=MODEL_NAME, messages=messages, tools=[get_user_info], # when tool_choice='required' is set, the model is guaranteed to generate one or more tool calls. tool_choice="required", max_tokens=32, ) assistant_message = response.choices[0].message # Append assistant tool call message to the conversation history messages.append(assistant_message) if assistant_message.tool_calls: for tool_call in assistant_message.tool_calls: assert tool_call.type == "function" name = tool_call.function.name arguments = json.loads(tool_call.function.arguments) print(f"🛠️ Tool {name=} call with arguments {arguments=}") if name == "get_user": user_info = get_user(arguments["user_id"]) # Append tool result to the conversation history messages.append( { # ipython: A role introduced in Llama 3.1. Semantically, this role means "tool call output". # This role is used to mark messages with the output of a tool call. # https://www.llama.com/docs/model-cards-and-prompt-formats/llama3_1/#-supported-roles- # both 'tool' and 'ipython' roles can be used to send function call results back to the model. # https://github.com/vllm-project/vllm/blob/main/examples/tool_chat_template_llama3.1_json.jinja#L104 "role": "tool", "content": json.dumps(user_info), "tool_call_id": tool_call.id, } ) # Get final response from the model after tool execution final_response = ( client.chat.completions.create( model=MODEL_NAME, messages=messages, max_tokens=1024, ) .choices[0] .message.content ) display(Markdown(markdown.markdown(final_response))) else: print(f"No tool calls were made by the model, {assistant_message=}")

🛠️ Tool name='get_user' call with arguments arguments={'user_id': 'user_123'}

Based on the user's financial information, here are some important tax recommendations:

1. Take advantage of the standard deduction: As a single filer with an income of $60,000, the user may not have enough itemized deductions to exceed the standard deduction of $12,950 (for the 2022 tax year). Therefore, it's likely more beneficial to take the standard deduction.
2. Contribute to a retirement account: With a relatively high income, the user may be eligible to contribute to a retirement account such as a 401(k) or an IRA. Contributing to a retirement account can help reduce taxable income and provide a tax benefit.
3. Consider itemizing deductions: Although the user's income may not be high enough to exceed the standard deduction, they may still have itemized deductions that can reduce their taxable income. For example, they may have mortgage interest, property taxes, or charitable donations that can be itemized.
4. Be aware of the tax implications of living in New York: New York has a relatively high state income tax rate, with a top rate of 8.82%. The user should be aware of this when making tax-related decisions.
5. Consider consulting a tax professional: Given the user's income and location, it may be beneficial to consult a tax professional to ensure they are taking advantage of all available tax credits and deductions.

These are general recommendations and may not be applicable to the user's specific situation. It's always best to consult a tax professional for personalized advice.

SGT Proxy's Protection Mechanisms

The /stainedglass endpoint offers insights into the SGT Proxy's protection mechanisms by providing access to:

• Plain (un-transformed) LLM embeddings.
• Transformed LLM embeddings.
• Reconstructed text from protected embeddings.
• Obfuscation scores

In [7]:

import requests
import torch

headers = {
    "Content-Type": "application/json",
    "Authorization": f"Bearer {API_KEY}",
}
# prepare message for /stainedglass request body
messages = [
    msg.model_dump(mode="json") if hasattr(msg, "model_dump") else msg
    for msg in messages
]
payload = {"messages": messages}

import requests import torch headers = { "Content-Type": "application/json", "Authorization": f"Bearer {API_KEY}", } # prepare message for /stainedglass request body messages = [ msg.model_dump(mode="json") if hasattr(msg, "model_dump") else msg for msg in messages ] payload = {"messages": messages}

Plain Text Embeddings vs Transformed Embeddings

In [8]:

response = requests.post(
    f"{PROXY_URL}/stainedglass",
    headers=headers,
    json=payload
    | {
        "return_plain_text_embeddings": True,
        "return_transformed_embeddings": True,
        "return_reconstructed_prompt": False,
    },
    stream=False,
    timeout=30,
)
response.raise_for_status()

response = requests.post( f"{PROXY_URL}/stainedglass", headers=headers, json=payload | { "return_plain_text_embeddings": True, "return_transformed_embeddings": True, "return_reconstructed_prompt": False, }, stream=False, timeout=30, ) response.raise_for_status()

In [9]:

# check plain and transformed embeddings shapes
(
    torch.tensor(response.json()["plain_text_embeddings"]).shape,
    torch.tensor(response.json()["transformed_embeddings"]).shape,
)

# check plain and transformed embeddings shapes ( torch.tensor(response.json()["plain_text_embeddings"]).shape, torch.tensor(response.json()["transformed_embeddings"]).shape, )

Out[9]:

(torch.Size([85, 4096]), torch.Size([85, 4096]))

In [10]:

# these are plain text embeddings
torch.tensor(response.json()["plain_text_embeddings"])

# these are plain text embeddings torch.tensor(response.json()["plain_text_embeddings"])

Out[10]:

tensor([[-4.4556e-03,  9.5367e-04, -6.5308e-03,  ...,  1.1597e-02,
          3.1128e-03, -5.7602e-04],
        [ 3.6011e-03,  4.3869e-04,  1.1292e-03,  ...,  2.1973e-03,
          6.4087e-04,  8.2397e-03],
        [-7.4768e-04,  2.0885e-04, -1.0071e-03,  ..., -1.0986e-02,
         -3.9673e-03, -1.3733e-04],
        ...,
        [-1.1536e-02,  5.7068e-03, -8.0490e-04,  ...,  9.6436e-03,
          8.6594e-04,  3.2806e-03],
        [ 2.0142e-03,  6.9275e-03,  1.1353e-02,  ..., -8.7280e-03,
          3.5553e-03,  8.1635e-04],
        [ 5.9814e-03,  1.8234e-03,  7.2937e-03,  ...,  1.3809e-03,
          4.8876e-05, -8.4305e-04]])

In [11]:

# check plain text conversation
display(
    Markdown(
        markdown.markdown(
            f"# Conversation Plain Text\n\n{''.join(response.json()['tokenized_plain_text'])}"
        )
    )
)

# check plain text conversation display( Markdown( markdown.markdown( f"# Conversation Plain Text\n\n{''.join(response.json()['tokenized_plain_text'])}" ) ) )

Conversation Plain Text

You are a helpful tax assistant that can retrieve user information to answer tax-related questions.What are important tax recommendations for user with ID 'user_123' based on their financial information?{\"user_id\": \"user_123\", \"name\": \"John Doe\", \"age\": 25, \"location\": \"New York, USA\", \"income\": \"$60,000\", \"filing_status\": \"Single\"}

In [12]:

# these are the SGT transformed embeddings
torch.tensor(response.json()["transformed_embeddings"])

# these are the SGT transformed embeddings torch.tensor(response.json()["transformed_embeddings"])

Out[12]:

tensor([[-0.0182,  0.0271,  0.0454,  ..., -0.1001, -0.0361,  0.0364],
        [-0.0227, -0.0161,  0.0138,  ..., -0.0312,  0.0046, -0.0432],
        [ 0.0136, -0.0075,  0.0123,  ...,  0.0559,  0.0064,  0.0086],
        ...,
        [ 0.0398,  0.0025,  0.0403,  ..., -0.0571, -0.0240, -0.0601],
        [-0.0045,  0.0031, -0.0070,  ..., -0.0239,  0.0322, -0.0079],
        [ 0.0064,  0.0073,  0.0161,  ..., -0.0161,  0.0047, -0.0248]])

Reconstructed Prompt and Obfuscation Score

In [13]:

response = requests.post(
    f"{PROXY_URL}/stainedglass",
    headers=headers,
    json=payload
    | {
        "return_reconstructed_prompt": True,
        "return_obfuscation_score": True,
        "return_transformed_embeddings": True,
    },
    stream=False,
    timeout=30,
)
response.raise_for_status()

response = requests.post( f"{PROXY_URL}/stainedglass", headers=headers, json=payload | { "return_reconstructed_prompt": True, "return_obfuscation_score": True, "return_transformed_embeddings": True, }, stream=False, timeout=30, ) response.raise_for_status()

In [14]:

display(
    Markdown(
        markdown.markdown(
            f"# Reconstructed Text\n\n{response.json()['reconstructed_prompt']}"
        )
    )
)

display( Markdown( markdown.markdown( f"# Reconstructed Text\n\n{response.json()['reconstructed_prompt']}" ) ) )

Reconstructed Text

rumpetalyaquotelevãesteџџческихalardan.** диза€€€€€€€€€€€€€€€€erusformární ।”

();

ČeskoslovenuseRalativeImagePath mücadel дозволя },

џџџџџџџџџџџџџџџџџџџџџџџџџџџџџџџџџџџџџџџџџџџџџџџџџџџџџџџџџџџџџџџџ -->

€€€€€€€€uyordu ।”

useRalativeImagePath';

_ComCallableWrapperџџџџџџџџџџџџџџџџџџџџџџџџџџџџџџџџuseRalativeImagePath PodsDummyerusform_ComCallableWrapper€€€€€€€€€€€€€€€€useRalativeImagePath€€€€€€€€rabilir ।”

iyesi использованиextracomment_ComCallableWrapper 네이트온iyesi uvědom ÜNİVERS uvědom найкра диза Doeiyesi zvlátalya uvědomuyordu swingerclubuyordu zvlá;

muştur найкра диза дизаuyordu переваiyesi использованиãeste uvědomextracomment(InitializedTypeInfouyordu },

найкра использовани ।”

---

।”

uvědom найкра jednoduchselectorMethod№№№№

Obfuscations Score

An obfuscation score is a percentage metric that indicates how many of the plain text tokens are different from the transformed tokens. A higher score indicates a higher level of obfuscation and data privacy.

In [15]:

display(
    Markdown(
        markdown.markdown(
            f"# Obfuscation Score\n\n{round(response.json()['obfuscation_score'], 2)}"
        )
    )
)

display( Markdown( markdown.markdown( f"# Obfuscation Score\n\n{round(response.json()['obfuscation_score'], 2)}" ) ) )

Obfuscation Score

0.99