feat: Add conversation_id filtering to message endpoints (#8324)

* feat: Add conversation_id filtering to message list and search endpoints

Add optional conversation_id parameter to filter messages by conversation:
- client.agents.messages.list
- client.messages.list
- client.messages.search

Changes:
- Added conversation_id field to MessageSearchRequest and SearchAllMessagesRequest schemas
- Added conversation_id filtering to list_messages in message_manager.py
- Updated get_agent_recall_async and get_all_messages_recall_async in server.py
- Added conversation_id query parameter to router endpoints
- Updated Turbopuffer client to support conversation_id filtering in searches

Fixes #8320

🤖 Generated with [Letta Code](https://letta.com)

Co-Authored-By: Charles Packer <cpacker@users.noreply.github.com>

* add conversation_id to message and tpuf

* default messages filter for backward compatibility

* add test and auto gen

* fix integration test

* fix test

* update test

---------

Co-authored-by: letta-code <248085862+letta-code@users.noreply.github.com>
Co-authored-by: Charles Packer <cpacker@users.noreply.github.com>
Co-authored-by: christinatong01 <christina@letta.com>
This commit is contained in:
Charles Packer
2026-01-07 13:45:52 -08:00
committed by Caren Thomas
parent 737d6e2550
commit ed6284cedb
9 changed files with 334 additions and 24 deletions

View File

@@ -400,6 +400,7 @@ class TurbopufferClient:
created_ats: List[datetime],
project_id: Optional[str] = None,
template_id: Optional[str] = None,
conversation_ids: Optional[List[Optional[str]]] = None,
) -> bool:
"""Insert messages into Turbopuffer.
@@ -413,6 +414,7 @@ class TurbopufferClient:
created_ats: List of creation timestamps for each message
project_id: Optional project ID for all messages
template_id: Optional template ID for all messages
conversation_ids: Optional list of conversation IDs (one per message, must match 1:1 with message_texts)
Returns:
True if successful
@@ -441,22 +443,26 @@ class TurbopufferClient:
raise ValueError(f"message_ids length ({len(message_ids)}) must match roles length ({len(roles)})")
if len(message_ids) != len(created_ats):
raise ValueError(f"message_ids length ({len(message_ids)}) must match created_ats length ({len(created_ats)})")
if conversation_ids is not None and len(conversation_ids) != len(message_ids):
raise ValueError(f"conversation_ids length ({len(conversation_ids)}) must match message_ids length ({len(message_ids)})")
# prepare column-based data for turbopuffer - optimized for batch insert
ids = []
vectors = []
texts = []
organization_ids = []
agent_ids = []
organization_ids_list = []
agent_ids_list = []
message_roles = []
created_at_timestamps = []
project_ids = []
template_ids = []
project_ids_list = []
template_ids_list = []
conversation_ids_list = []
for (original_idx, text), embedding in zip(filtered_messages, embeddings):
message_id = message_ids[original_idx]
role = roles[original_idx]
created_at = created_ats[original_idx]
conversation_id = conversation_ids[original_idx] if conversation_ids else None
# ensure the provided timestamp is timezone-aware and in UTC
if created_at.tzinfo is None:
@@ -470,31 +476,36 @@ class TurbopufferClient:
ids.append(message_id)
vectors.append(embedding)
texts.append(text)
organization_ids.append(organization_id)
agent_ids.append(agent_id)
organization_ids_list.append(organization_id)
agent_ids_list.append(agent_id)
message_roles.append(role.value)
created_at_timestamps.append(timestamp)
project_ids.append(project_id)
template_ids.append(template_id)
project_ids_list.append(project_id)
template_ids_list.append(template_id)
conversation_ids_list.append(conversation_id)
# build column-based upsert data
upsert_columns = {
"id": ids,
"vector": vectors,
"text": texts,
"organization_id": organization_ids,
"agent_id": agent_ids,
"organization_id": organization_ids_list,
"agent_id": agent_ids_list,
"role": message_roles,
"created_at": created_at_timestamps,
}
# only include conversation_id if it's provided
if conversation_ids is not None:
upsert_columns["conversation_id"] = conversation_ids_list
# only include project_id if it's provided
if project_id is not None:
upsert_columns["project_id"] = project_ids
upsert_columns["project_id"] = project_ids_list
# only include template_id if it's provided
if template_id is not None:
upsert_columns["template_id"] = template_ids
upsert_columns["template_id"] = template_ids_list
try:
# use global semaphore to limit concurrent Turbopuffer writes
@@ -506,7 +517,10 @@ class TurbopufferClient:
await namespace.write(
upsert_columns=upsert_columns,
distance_metric="cosine_distance",
schema={"text": {"type": "string", "full_text_search": True}},
schema={
"text": {"type": "string", "full_text_search": True},
"conversation_id": {"type": "string"},
},
)
logger.info(f"Successfully inserted {len(ids)} messages to Turbopuffer for agent {agent_id}")
return True
@@ -792,6 +806,7 @@ class TurbopufferClient:
roles: Optional[List[MessageRole]] = None,
project_id: Optional[str] = None,
template_id: Optional[str] = None,
conversation_id: Optional[str] = None,
vector_weight: float = 0.5,
fts_weight: float = 0.5,
start_date: Optional[datetime] = None,
@@ -809,6 +824,7 @@ class TurbopufferClient:
roles: Optional list of message roles to filter by
project_id: Optional project ID to filter messages by
template_id: Optional template ID to filter messages by
conversation_id: Optional conversation ID to filter messages by (use "default" for NULL)
vector_weight: Weight for vector search results in hybrid mode (default: 0.5)
fts_weight: Weight for FTS results in hybrid mode (default: 0.5)
start_date: Optional datetime to filter messages created after this date
@@ -875,6 +891,19 @@ class TurbopufferClient:
if template_id:
template_filter = ("template_id", "Eq", template_id)
# build conversation_id filter if provided
# three cases:
# 1. conversation_id=None (omitted) -> return all messages (no filter)
# 2. conversation_id="default" -> return only default messages (conversation_id is none), for backward compatibility
# 3. conversation_id="xyz" -> return only messages in that conversation
conversation_filter = None
if conversation_id == "default":
# "default" is reserved for default messages only (conversation_id is none)
conversation_filter = ("conversation_id", "Eq", None)
elif conversation_id is not None:
# Specific conversation
conversation_filter = ("conversation_id", "Eq", conversation_id)
# combine all filters
all_filters = [agent_filter] # always include agent_id filter
if role_filter:
@@ -883,6 +912,8 @@ class TurbopufferClient:
all_filters.append(project_filter)
if template_filter:
all_filters.append(template_filter)
if conversation_filter:
all_filters.append(conversation_filter)
if date_filters:
all_filters.extend(date_filters)
@@ -901,7 +932,7 @@ class TurbopufferClient:
query_embedding=query_embedding,
query_text=query_text,
top_k=top_k,
include_attributes=["text", "organization_id", "agent_id", "role", "created_at"],
include_attributes=["text", "organization_id", "agent_id", "role", "created_at", "conversation_id"],
filters=final_filter,
vector_weight=vector_weight,
fts_weight=fts_weight,
@@ -952,6 +983,7 @@ class TurbopufferClient:
agent_id: Optional[str] = None,
project_id: Optional[str] = None,
template_id: Optional[str] = None,
conversation_id: Optional[str] = None,
vector_weight: float = 0.5,
fts_weight: float = 0.5,
start_date: Optional[datetime] = None,
@@ -969,6 +1001,10 @@ class TurbopufferClient:
agent_id: Optional agent ID to filter messages by
project_id: Optional project ID to filter messages by
template_id: Optional template ID to filter messages by
conversation_id: Optional conversation ID to filter messages by. Special values:
- None (omitted): Return all messages
- "default": Return only default messages (conversation_id IS NULL)
- Any other value: Return messages in that specific conversation
vector_weight: Weight for vector search results in hybrid mode (default: 0.5)
fts_weight: Weight for FTS results in hybrid mode (default: 0.5)
start_date: Optional datetime to filter messages created after this date
@@ -1017,6 +1053,18 @@ class TurbopufferClient:
if template_id:
all_filters.append(("template_id", "Eq", template_id))
# conversation filter
# three cases:
# 1. conversation_id=None (omitted) -> return all messages (no filter)
# 2. conversation_id="default" -> return only default messages (conversation_id is none), for backward compatibility
# 3. conversation_id="xyz" -> return only messages in that conversation
if conversation_id == "default":
# "default" is reserved for default messages only (conversation_id is none)
all_filters.append(("conversation_id", "Eq", None))
elif conversation_id is not None:
# Specific conversation
all_filters.append(("conversation_id", "Eq", conversation_id))
# date filters
if start_date:
# Convert to UTC to match stored timestamps
@@ -1049,7 +1097,7 @@ class TurbopufferClient:
query_embedding=query_embedding,
query_text=query_text,
top_k=top_k,
include_attributes=["text", "organization_id", "agent_id", "role", "created_at"],
include_attributes=["text", "organization_id", "agent_id", "role", "created_at", "conversation_id"],
filters=final_filter,
vector_weight=vector_weight,
fts_weight=fts_weight,
@@ -1134,6 +1182,7 @@ class TurbopufferClient:
"agent_id": getattr(row, "agent_id", None),
"role": getattr(row, "role", None),
"created_at": getattr(row, "created_at", None),
"conversation_id": getattr(row, "conversation_id", None),
}
messages.append(message_dict)