feat(client): Allow customizing memory tool descriptions with ToolSchema wrapper classes

[fcenedes]

Summary

This PR introduces ToolSchema and ToolSchemaCollection wrapper classes that provide a fluent API for customizing tool descriptions, names, and parameter descriptions before passing them to LLMs. This allows users of the agent-memory-client to tune tool descriptions for their specific use cases and LLM providers.

Motivation

Different LLMs and use cases may benefit from customized tool descriptions. For example:

• Adjusting descriptions to match the tone/style of a specific application
• Adding domain-specific context to tool descriptions
• Renaming tools to avoid conflicts with other tools in the same LLM context
• Customizing parameter descriptions for better LLM understanding

Previously, users had to manually modify the returned dictionaries, which was error-prone and didn't provide a clean API.
Attempt to fix #96

Changes

New Classes

ToolSchema - Wrapper for individual tool schema dictionaries:

• set_description(text) - Set custom tool description
• set_name(name) - Set custom tool name
• set_parameter_description(param, text) - Customize parameter descriptions
• get_description() / get_name() / get_parameter_description(param) - Getters
• to_dict() - Convert to dict for LLM consumption (returns deep copy)
• copy() - Create independent copy of the schema
• Supports both OpenAI and Anthropic formats via format property
• Backwards compatible with dict-like access (schema["function"]["name"])

ToolSchemaCollection - Wrapper for lists of tool schemas:

• get_by_name(name) - Get specific tool by name
• set_description(name, text) - Set description for a tool by name
• set_name(old_name, new_name) - Rename a tool
• names() - Get list of all tool names
• to_list() - Convert to list of dicts for LLM consumption
• copy() - Create independent copy of the collection
• Iterable and indexable

Updated Methods

All 20 public schema methods now return ToolSchema or ToolSchemaCollection:

Method	Returns
get_memory_search_tool_schema()	ToolSchema
get_working_memory_tool_schema()	ToolSchema
get_add_memory_tool_schema()	ToolSchema
get_update_memory_data_tool_schema()	ToolSchema
get_long_term_memory_tool_schema()	ToolSchema
edit_long_term_memory_tool_schema()	ToolSchema
create_long_term_memory_tool_schema()	ToolSchema
delete_long_term_memories_tool_schema()	ToolSchema
get_current_datetime_tool_schema()	ToolSchema
get_all_memory_tool_schemas()	ToolSchemaCollection
*_anthropic() variants	Same pattern

Usage Examples

from agent_memory_client import MemoryAPIClient, ToolSchema, ToolSchemaCollection

# Customize a single tool
schema = MemoryAPIClient.get_memory_search_tool_schema()
schema.set_description("Search through the user's personal knowledge base")
schema.set_name("search_knowledge_base")
schema.set_parameter_description("query", "Natural language search query")

# Method chaining
schema = (MemoryAPIClient.get_memory_search_tool_schema()
    .set_description("Custom description")
    .set_name("custom_search"))

# Customize all tools
all_tools = MemoryAPIClient.get_all_memory_tool_schemas()
all_tools.set_description("search_memory", "Find relevant memories")
all_tools.set_name("search_memory", "find_memories")

# Get specific tool from collection
search_tool = all_tools.get_by_name("find_memories")

# Convert to dict/list for LLM consumption
tool_dict = schema.to_dict()
tools_list = all_tools.to_list()

# Backwards compatible - existing code still works
tools = MemoryAPIClient.get_all_memory_tool_schemas()
response = await openai_client.chat.completions.create(
    model="gpt-4o",
    messages=messages,
    tools=tools.to_list()
)

Backwards Compatibility

• ToolSchema implements __getitem__ and __setitem__ for dict-like access
• ToolSchemaCollection is iterable and indexable
• Existing code that accesses schema["function"]["name"] continues to work
• to_dict() and to_list() return deep copies to prevent accidental mutations

Files Changed

File	Change
agent-memory-client/agent_memory_client/tool_schema.py	New - ToolSchema and ToolSchemaCollection classes
agent-memory-client/agent_memory_client/client.py	Updated all schema methods to return wrapper types
agent-memory-client/agent_memory_client/__init__.py	Added exports for new classes
agent-memory-client/tests/test_tool_schemas.py	Added 16 new tests for customization API
docs/python-sdk.md	Added documentation for tool customization

Testing

• ✅ 78 tests pass in agent-memory-client/tests/
• ✅ 30 tests pass in tests/test_client_tool_calls.py
• ✅ Linting passes with ruff check

Checklist

• New classes with fluent API for customization
• All 20 schema methods updated
• Backwards compatible with existing code
• Deep copy on to_dict()/to_list() to prevent mutations
• Comprehensive tests added
• Documentation updated

[Copilot]

Pull request overview

This PR introduces ToolSchema and ToolSchemaCollection wrapper classes that provide a fluent API for customizing tool descriptions, names, and parameter descriptions before passing them to LLMs. This enables users to tailor tool schemas for specific use cases, LLM providers, or application domains without manually modifying dictionaries. The implementation maintains full backwards compatibility through dict-like access patterns.

Key Changes:

• New ToolSchema class wraps individual tool schema dictionaries with methods for customization (set_description, set_name, set_parameter_description) and supports both OpenAI and Anthropic formats
• New ToolSchemaCollection class wraps lists of tool schemas with bulk customization operations (get_by_name, set_description, set_name)
• All 20 tool schema methods in MemoryAPIClient now return wrapper types instead of plain dicts, with backwards compatibility maintained via __getitem__ implementation

Reviewed changes

Copilot reviewed 5 out of 5 changed files in this pull request and generated 6 comments.

Show a summary per file

File	Description
agent-memory-client/agent_memory_client/tool_schema.py	New module defining ToolSchema and ToolSchemaCollection classes with fluent API for customization
agent-memory-client/agent_memory_client/client.py	Updated all 20 schema methods to return ToolSchema/ToolSchemaCollection wrappers; modified _convert_openai_to_anthropic_schema to handle both wrapper and dict inputs
agent-memory-client/agent_memory_client/__init__.py	Added exports for ToolSchema and ToolSchemaCollection classes
agent-memory-client/tests/test_tool_schemas.py	Added 16 new tests covering customization methods, backwards compatibility, and collection operations
docs/python-sdk.md	Added comprehensive documentation section with examples for basic customization, method chaining, bulk operations, and API reference tables

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[abrookins]

Looks good, my dude! 👍 LGTM

[abrookins]

All tests passing locally with API keys.