Custom Tool: Writing JSON Schema-Validated Data to Memory Blocks

Writing Schema-Validated Data to Memory Blocks

Use case: You want agents to store structured data (tables, datasets) in memory blocks, but need to enforce a specific JSON schema so your UI can reliably parse and display it.

Problem: The built-in memory tools (memory_insert, memory_replace, memory_rethink) don’t enforce schemas - agents can write any text format.

Solution: Create a custom tool that validates data against your schema before writing to the memory block.


Example: Storing Tabular Data

Let’s build a tool that stores datasets following the Frictionless Data Table Schema format.

Step 1: Define Your Custom Tool

from letta_client import Letta
import json

def store_table_data(
    block_label: str,
    table_name: str,
    columns: list[dict],
    rows: list[list]
) -> str:
    """
    Store tabular data in a memory block following Frictionless Table Schema format.
    
    Args:
        block_label: The label of the memory block to write to (e.g., "datasets")
        table_name: Human-readable name for this dataset
        columns: List of column definitions, e.g. [{"name": "age", "type": "integer"}, {"name": "name", "type": "string"}]
        rows: List of row data, e.g. [[25, "Alice"], [30, "Bob"]]
    
    Returns:
        Success message or error details
    """
    # Import inside function for sandboxed execution
    import json
    import os
    
    # Validate schema structure
    required_column_keys = {"name", "type"}
    for col in columns:
        if not all(key in col for key in required_column_keys):
            return f"Error: Each column must have 'name' and 'type' keys. Got: {col}"
    
    # Build Frictionless Table Schema
    table_schema = {
        "name": table_name,
        "schema": {
            "fields": columns
        },
        "data": rows
    }
    
    # Validate it's valid JSON
    try:
        schema_json = json.dumps(table_schema, indent=2)
    except Exception as e:
        return f"Error: Failed to serialize schema to JSON: {e}"
    
    # Write to memory block using Letta client
    # Note: This requires LETTA_API_KEY environment variable
    try:
        from letta_client import Letta
        
        api_key = os.getenv("LETTA_API_KEY")
        if not api_key:
            return "Error: LETTA_API_KEY environment variable not set"
        
        client = Letta(token=api_key)
        
        # Get agent ID from environment (passed by Letta during tool execution)
        agent_id = os.getenv("LETTA_AGENT_ID")
        if not agent_id:
            return "Error: Could not determine agent ID"
        
        # Find the memory block by label
        agent = client.agents.retrieve(agent_id=agent_id)
        target_block = None
        for block_id in agent.memory_block_ids:
            block = client.blocks.retrieve(block_id=block_id)
            if block.label == block_label:
                target_block = block
                break
        
        if not target_block:
            return f"Error: No memory block found with label '{block_label}'"
        
        # Update the block value
        client.blocks.update(
            block_id=target_block.id,
            value=schema_json
        )
        
        return f"Successfully stored table '{table_name}' with {len(columns)} columns and {len(rows)} rows in memory block '{block_label}'"
        
    except Exception as e:
        return f"Error writing to memory block: {e}"


# Register the tool
client = Letta(token="your_letta_api_key")

table_tool = client.tools.create(
    name="store_table_data",
    source_code=store_table_data.__code__.co_consts[0],  # Extract source
    source_type="python",
    tags=["memory", "schema", "tables"]
)

print(f"Created tool: {table_tool.id}")

Step 2: Create Agent with Schema-Validated Memory Block

# Create memory block for structured data
dataset_block = client.blocks.create(
    label="datasets",
    value="{}",  # Initialize empty
    description="Stores tabular datasets in Frictionless Table Schema format. "
                "Use store_table_data tool to write data here."
)

# Create agent with the custom tool
agent = client.agents.create(
    name="data_manager",
    system="You manage datasets. When users provide tabular data, use the store_table_data tool "
           "to validate and store it in the 'datasets' memory block.",
    memory_blocks=[dataset_block.id],
    tool_ids=[table_tool.id]
)

Step 3: Test It

# User provides unstructured data
response = client.agents.messages.create(
    agent_id=agent.id,
    messages=[{
        "role": "user",
        "content": "Store this customer data: Alice (age 25), Bob (age 30), Carol (age 28)"
    }]
)

# Agent calls store_table_data with structured format:
# {
#   "block_label": "datasets",
#   "table_name": "customers",
#   "columns": [
#     {"name": "name", "type": "string"},
#     {"name": "age", "type": "integer"}
#   ],
#   "rows": [
#     ["Alice", 25],
#     ["Bob", 30],
#     ["Carol", 28]
#   ]
# }

# Check memory block
updated_block = client.blocks.retrieve(block_id=dataset_block.id)
print(json.loads(updated_block.value))

Simpler Pattern: JSON Validation Only

If you don’t need full table schema support, here’s a minimal version:

def write_json_to_block(block_label: str, json_data: dict) -> str:
    """
    Write validated JSON to a memory block.
    
    Args:
        block_label: Memory block to write to
        json_data: Dictionary that will be serialized to JSON
    
    Returns:
        Success or error message
    """
    import json
    import os
    
    # Validate it's valid JSON-serializable
    try:
        json_str = json.dumps(json_data, indent=2)
    except Exception as e:
        return f"Error: Data is not valid JSON: {e}"
    
    # Write to block (simplified - requires environment setup)
    from letta_client import Letta
    client = Letta(token=os.getenv("LETTA_API_KEY"))
    agent_id = os.getenv("LETTA_AGENT_ID")
    
    agent = client.agents.retrieve(agent_id=agent_id)
    target_block = next(
        (b for b in agent.memory_blocks if b.label == block_label),
        None
    )
    
    if not target_block:
        return f"Error: Block '{block_label}' not found"
    
    client.blocks.update(block_id=target_block.id, value=json_str)
    return f"Successfully wrote JSON to block '{block_label}'"

Key Benefits

  1. Schema enforcement: Agent can’t write invalid data
  2. UI-friendly: Your frontend can reliably parse the memory block
  3. Error feedback: Agent sees validation errors and can retry
  4. Flexible: Adapt the validation logic to any schema you need

Important Notes

Environment Variables Required

Custom tools that modify memory blocks need:

  • LETTA_API_KEY: Your API key
  • LETTA_AGENT_ID: Auto-populated by Letta during tool execution

Set these when creating the tool:

client.tools.create(
    name="store_table_data",
    source_code=code,
    environment_variables=["LETTA_API_KEY", "LETTA_AGENT_ID"]
)

Alternative: Use memory_rethink with Validation

If you want to keep the native memory tools but add validation:

def validate_and_store(block_label: str, content: str) -> str:
    """Validate content is valid JSON, then write using memory_rethink"""
    import json
    
    try:
        # Validate it's valid JSON
        json.loads(content)
    except json.JSONDecodeError as e:
        return f"Error: Content must be valid JSON. Parse error: {e}"
    
    # If valid, agent can call memory_rethink separately
    return "Content validated - you may now call memory_rethink to store it"

Questions?

Reply with:

  • What schemas you’re trying to enforce
  • Issues with the implementation
  • Feature requests for schema validation

This pattern works for any structured format: JSON Schema, Pydantic models, CSV schemas, etc.