Parameter Injection in Legion

Parameter injection allows you to securely provide sensitive data to tools without exposing it to the agent or including it in the prompt. This feature is essential for handling API keys, credentials, and other sensitive information.

What is Parameter Injection?

Parameter injection is a technique that:

  • Allows secure passing of sensitive data to tools
  • Keeps credentials out of prompts and agent memory
  • Supports per-message parameter customization
  • Provides default values for convenience
  • Works with both standalone and agent-specific tools

Creating Tools with Injectable Parameters

Here's how to create and use tools with injectable parameters:

from typing import Annotated
from pydantic import Field
from legion import agent, tool

# Define a tool with injectable parameters
@tool(
    inject=["api_key", "endpoint", "custom_header"],  # Mark these as injectable
    description="Process a query using an external API with configurable parameters",
    defaults={  # Default values for injectable parameters
        "api_key": "sk_test_default_key",
        "endpoint": "https://api.example.com/v1",
        "custom_header": "default-header"
    }
)
def process_api_query(
    query: Annotated[str, Field(description="The query to process")],
    api_key: Annotated[str, Field(description="API key for the service", default=None)],
    endpoint: Annotated[str, Field(description="API endpoint", default=None)],
    custom_header: Annotated[str, Field(
        description="Custom header for the request",
        default=None
    )]
) -> str:
    """Process a query using an external API with injected credentials and headers."""
    # In a real application, you would make an actual API call here
    return (
        f"Processed '{query}' using API at {endpoint}
"
        f"Auth: {api_key[:4]}...
"
        f"Header: {custom_header}"
    )

@agent(
    model="openai:gpt-4o-mini",
    temperature=0.2,
    tools=[process_api_query]
)
class APIAgent:
    """An agent that demonstrates using tools with injected parameters."""
    pass

Using Parameter Injection

When processing messages, you can inject parameters:

# Create an instance of the agent
agent = APIAgent()

# Example 1: Using default values (no injection)
response = await agent.aprocess(
    "Process the query 'hello world' using our API."
)

# Example 2: Override API credentials for this message only
response = await agent.aprocess(
    "Process the query 'test message' with production credentials.",
    injected_parameters=[
        {
            "tool": process_api_query,
            "parameters": {
                "api_key": "sk_prod_key_123",
                "endpoint": "https://api.prod.example.com/v1"
            }
        }
    ]
)

# Example 3: Inject just one parameter
response = await agent.aprocess(
    "Process the query 'another test'.",
    injected_parameters=[
        {
            "tool": process_api_query,
            "parameters": {
                "custom_header": "header-for-special-request"
            }
        }
    ]
)

Security Benefits

  1. Credential Protection: Sensitive data never appears in prompts
  2. Isolation: The agent can't access or leak injected parameters
  3. Flexibility: Change credentials without modifying agent code
  4. Defaults: Provide fallback values for convenience

Best Practices

  1. Always mark sensitive parameters as injectable
  2. Provide meaningful default values when possible
  3. Use per-message injection for context-specific credentials
  4. Keep injected parameters focused on security-sensitive data
  5. Use descriptive parameter names and documentation