Skip to content

Support an optional user ID everywhere #26

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 3 commits into from
Jun 26, 2025
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions CLAUDE.md
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@ uv run ruff check # Run linting
uv run ruff format # Format code
uv run pytest # Run tests
uv run pytest tests/ # Run specific test directory
uv run pytest --run-api-tests # Run all tests, including API tests
uv add <dependency> # Add a dependency to pyproject.toml and update lock file
uv remove <dependency> # Remove a dependency from pyproject.toml and update lock file

Expand Down
4 changes: 2 additions & 2 deletions agent-memory-client/README.md
Original file line number Diff line number Diff line change
Expand Up @@ -206,8 +206,8 @@ await client.update_working_memory_data(

# Append messages
new_messages = [
MemoryMessage(role="user", content="What's the weather?"),
MemoryMessage(role="assistant", content="It's sunny today!")
{"role": "user", "content": "What's the weather?"},
{"role": "assistant", "content": "It's sunny today!"}
]

await client.append_messages_to_working_memory(
Expand Down
2 changes: 1 addition & 1 deletion agent-memory-client/agent_memory_client/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,7 @@
memory management capabilities for AI agents and applications.
"""

__version__ = "0.9.0b3"
__version__ = "0.9.0b4"

from .client import MemoryAPIClient, MemoryClientConfig, create_memory_client
from .exceptions import (
Expand Down
144 changes: 76 additions & 68 deletions agent-memory-client/agent_memory_client/client.py

Large diffs are not rendered by default.

4 changes: 2 additions & 2 deletions agent-memory-client/tests/test_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -527,8 +527,8 @@ async def test_append_messages_to_working_memory(self, enhanced_test_client):
)

new_messages = [
MemoryMessage(role="assistant", content="Second message"),
MemoryMessage(role="user", content="Third message"),
{"role": "assistant", "content": "Second message"},
{"role": "user", "content": "Third message"},
]

with (
Expand Down
2 changes: 1 addition & 1 deletion agent_memory_server/__init__.py
Original file line number Diff line number Diff line change
@@ -1,3 +1,3 @@
"""Redis Agent Memory Server - A memory system for conversational AI."""

__version__ = "0.9.0b3"
__version__ = "0.9.0b4"
23 changes: 16 additions & 7 deletions agent_memory_server/api.py
Original file line number Diff line number Diff line change
Expand Up @@ -188,7 +188,7 @@ async def list_sessions(
Get a list of session IDs, with optional pagination.

Args:
options: Query parameters (page, size, namespace)
options: Query parameters (limit, offset, namespace, user_id)

Returns:
List of session IDs
Expand All @@ -200,6 +200,7 @@ async def list_sessions(
limit=options.limit,
offset=options.offset,
namespace=options.namespace,
user_id=options.user_id,
)

return SessionListResponse(
Expand All @@ -211,8 +212,8 @@ async def list_sessions(
@router.get("/v1/working-memory/{session_id}", response_model=WorkingMemoryResponse)
async def get_working_memory(
session_id: str,
user_id: str | None = None,
namespace: str | None = None,
window_size: int = settings.window_size, # Deprecated: kept for backward compatibility
model_name: ModelNameLiteral | None = None,
context_window_max: int | None = None,
current_user: UserInfo = Depends(get_current_user),
Expand All @@ -225,8 +226,8 @@ async def get_working_memory(

Args:
session_id: The session ID
user_id: The user ID to retrieve working memory for
namespace: The namespace to use for the session
window_size: DEPRECATED - The number of messages to include (kept for backward compatibility)
model_name: The client's LLM model name (will determine context window size if provided)
context_window_max: Direct specification of the context window max tokens (overrides model_name)

Expand All @@ -240,6 +241,7 @@ async def get_working_memory(
session_id=session_id,
namespace=namespace,
redis_client=redis,
user_id=user_id,
)

if not working_mem:
Expand All @@ -249,6 +251,7 @@ async def get_working_memory(
memories=[],
session_id=session_id,
namespace=namespace,
user_id=user_id,
)

# Apply token-based truncation if we have messages and model info
Expand All @@ -266,17 +269,14 @@ async def get_working_memory(
break
working_mem.messages = truncated_messages

# Fallback to message-count truncation for backward compatibility
elif len(working_mem.messages) > window_size:
working_mem.messages = working_mem.messages[-window_size:]

return working_mem


@router.put("/v1/working-memory/{session_id}", response_model=WorkingMemoryResponse)
async def put_working_memory(
session_id: str,
memory: WorkingMemory,
user_id: str | None = None,
model_name: ModelNameLiteral | None = None,
context_window_max: int | None = None,
background_tasks=Depends(get_background_tasks),
Expand All @@ -291,6 +291,7 @@ async def put_working_memory(
Args:
session_id: The session ID
memory: Working memory to save
user_id: Optional user ID for the session (overrides user_id in memory object)
model_name: The client's LLM model name for context window determination
context_window_max: Direct specification of context window max tokens
background_tasks: DocketBackgroundTasks instance (injected automatically)
Expand All @@ -303,6 +304,10 @@ async def put_working_memory(
# Ensure session_id matches
memory.session_id = session_id

# Override user_id if provided as query parameter
if user_id is not None:
memory.user_id = user_id

# Validate that all structured memories have id (if any)
for mem in memory.memories:
if not mem.id:
Expand Down Expand Up @@ -359,6 +364,7 @@ async def put_working_memory(
@router.delete("/v1/working-memory/{session_id}", response_model=AckResponse)
async def delete_working_memory(
session_id: str,
user_id: str | None = None,
namespace: str | None = None,
current_user: UserInfo = Depends(get_current_user),
):
Expand All @@ -369,6 +375,7 @@ async def delete_working_memory(

Args:
session_id: The session ID
user_id: Optional user ID for the session
namespace: Optional namespace for the session

Returns:
Expand All @@ -379,6 +386,7 @@ async def delete_working_memory(
# Delete unified working memory
await working_memory.delete_working_memory(
session_id=session_id,
user_id=user_id,
namespace=namespace,
redis_client=redis,
)
Expand Down Expand Up @@ -558,6 +566,7 @@ async def memory_prompt(
working_mem = await working_memory.get_working_memory(
session_id=params.session.session_id,
namespace=params.session.namespace,
user_id=params.session.user_id,
redis_client=redis,
)

Expand Down
4 changes: 3 additions & 1 deletion agent_memory_server/config.py
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,6 @@ def load_yaml_settings():
class Settings(BaseSettings):
redis_url: str = "redis://localhost:6379"
long_term_memory: bool = True
window_size: int = 20
openai_api_key: str | None = None
anthropic_api_key: str | None = None
generation_model: str = "gpt-4o-mini"
Expand Down Expand Up @@ -66,6 +65,9 @@ class Settings(BaseSettings):
auth0_client_id: str | None = None
auth0_client_secret: str | None = None

# Working memory settings
window_size: int = 20 # Default number of recent messages to return

# Other Application settings
log_level: Literal["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"] = "INFO"

Expand Down
17 changes: 13 additions & 4 deletions agent_memory_server/dependencies.py
Original file line number Diff line number Diff line change
Expand Up @@ -27,10 +27,19 @@ async def add_task(
logger.info("Scheduling task through Docket")
# Get the Redis connection that's already configured (will use testcontainer in tests)
redis_conn = await get_redis_conn()
# Use the connection's URL instead of settings.redis_url directly
redis_url = redis_conn.connection_pool.connection_kwargs.get(
"url", settings.redis_url
)

# Extract Redis URL from the connection pool
connection_kwargs = redis_conn.connection_pool.connection_kwargs
if "host" in connection_kwargs and "port" in connection_kwargs:
redis_url = (
f"redis://{connection_kwargs['host']}:{connection_kwargs['port']}"
)
if "db" in connection_kwargs:
redis_url += f"/{connection_kwargs['db']}"
else:
# Fallback to settings if we can't extract from connection
redis_url = settings.redis_url

logger.info("redis_url: %s", redis_url)
logger.info("docket_name: %s", settings.docket_name)
async with Docket(
Expand Down
76 changes: 52 additions & 24 deletions agent_memory_server/long_term_memory.py
Original file line number Diff line number Diff line change
Expand Up @@ -46,9 +46,6 @@
)


DEFAULT_MEMORY_LIMIT = 1000
MEMORY_INDEX = "memory_idx"

# Prompt for extracting memories from messages in working memory context
WORKING_MEMORY_EXTRACTION_PROMPT = """
You are a memory extraction assistant. Your job is to analyze conversation
Expand Down Expand Up @@ -354,15 +351,27 @@ async def compact_long_term_memories(
# Find all memories with this hash
# Use FT.SEARCH to find the actual memories with this hash
# TODO: Use RedisVL index
search_query = (
f"FT.SEARCH {index_name} "
f"(@memory_hash:{{{memory_hash}}}) {' '.join(filters)} "
"RETURN 6 id_ text last_accessed created_at user_id session_id "
"SORTBY last_accessed ASC" # Oldest first
)
if filters:
# Combine hash query with filters using boolean AND
query_expr = f"(@memory_hash:{{{memory_hash}}}) ({' '.join(filters)})"
else:
query_expr = f"@memory_hash:{{{memory_hash}}}"

search_results = await redis_client.execute_command(
search_query
"FT.SEARCH",
index_name,
f"'{query_expr}'",
"RETURN",
"6",
"id_",
"text",
"last_accessed",
"created_at",
"user_id",
"session_id",
"SORTBY",
"last_accessed",
"ASC",
)

if search_results and search_results[0] > 1:
Expand Down Expand Up @@ -1209,15 +1218,24 @@ async def deduplicate_by_hash(

# Use FT.SEARCH to find memories with this hash
# TODO: Use RedisVL
search_query = (
f"FT.SEARCH {index_name} "
f"(@memory_hash:{{{memory_hash}}}) {filter_str} "
"RETURN 1 id_ "
"SORTBY last_accessed DESC" # Newest first
if filter_str:
# Combine hash query with filters using boolean AND
query_expr = f"(@memory_hash:{{{memory_hash}}}) ({filter_str})"
else:
query_expr = f"@memory_hash:{{{memory_hash}}}"

search_results = await redis_client.execute_command(
"FT.SEARCH",
index_name,
f"'{query_expr}'",
"RETURN",
"1",
"id_",
"SORTBY",
"last_accessed",
"DESC",
)

search_results = await redis_client.execute_command(search_query)

if search_results and search_results[0] > 0:
# Found existing memory with the same hash
logger.info(f"Found existing memory with hash {memory_hash}")
Expand Down Expand Up @@ -1285,15 +1303,25 @@ async def deduplicate_by_id(

# Use FT.SEARCH to find memories with this id
# TODO: Use RedisVL
search_query = (
f"FT.SEARCH {index_name} "
f"(@id:{{{memory.id}}}) {filter_str} "
"RETURN 2 id_ persisted_at "
"SORTBY last_accessed DESC" # Newest first
if filter_str:
# Combine the id query with filters - Redis FT.SEARCH uses implicit AND between terms
query_expr = f"@id:{{{memory.id}}} {filter_str}"
else:
query_expr = f"@id:{{{memory.id}}}"

search_results = await redis_client.execute_command(
"FT.SEARCH",
index_name,
f"'{query_expr}'",
"RETURN",
"2",
"id_",
"persisted_at",
"SORTBY",
"last_accessed",
"DESC",
)

search_results = await redis_client.execute_command(search_query)

if search_results and search_results[0] > 0:
# Found existing memory with the same id
logger.info(f"Found existing memory with id {memory.id}, will overwrite")
Expand Down
1 change: 1 addition & 0 deletions agent_memory_server/mcp.py
Original file line number Diff line number Diff line change
Expand Up @@ -562,6 +562,7 @@ async def memory_prompt(
session = WorkingMemoryRequest(
session_id=_session_id,
namespace=namespace.eq if namespace and namespace.eq else None,
user_id=user_id.eq if user_id and user_id.eq else None,
window_size=window_size,
model_name=model_name,
context_window_max=context_window_max,
Expand Down
2 changes: 2 additions & 0 deletions agent_memory_server/models.py
Original file line number Diff line number Diff line change
Expand Up @@ -216,6 +216,7 @@ class WorkingMemoryRequest(BaseModel):

session_id: str
namespace: str | None = None
user_id: str | None = None
window_size: int = settings.window_size
model_name: ModelNameLiteral | None = None
context_window_max: int | None = None
Expand Down Expand Up @@ -257,6 +258,7 @@ class GetSessionsQuery(BaseModel):
limit: int = Field(default=20, ge=1, le=100)
offset: int = Field(default=0, ge=0)
namespace: str | None = None
user_id: str | None = None


class HealthCheckResponse(BaseModel):
Expand Down
21 changes: 15 additions & 6 deletions agent_memory_server/utils/keys.py
Original file line number Diff line number Diff line change
Expand Up @@ -56,13 +56,22 @@ def metadata_key(session_id: str, namespace: str | None = None) -> str:
)

@staticmethod
def working_memory_key(session_id: str, namespace: str | None = None) -> str:
def working_memory_key(
session_id: str, user_id: str | None = None, namespace: str | None = None
) -> str:
"""Get the working memory key for a session."""
return (
f"working_memory:{namespace}:{session_id}"
if namespace
else f"working_memory:{session_id}"
)
# Build key components, filtering out None values
key_parts = ["working_memory"]

if namespace:
key_parts.append(namespace)

if user_id:
key_parts.append(user_id)

key_parts.append(session_id)

return ":".join(key_parts)

@staticmethod
def search_index_name() -> str:
Expand Down
Loading