Fimeg/letta-server
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

feat: add agent_id to search results (#6867)

Browse Source
This commit is contained in:
Sarah Wooders
2025-12-14 16:54:00 -08:00
committed by Caren Thomas
parent 4b9485a484
commit a721a00899
6 changed files with 354 additions and 46 deletions
Download Patch File Download Diff File Expand all files Collapse all files

211
fern/openapi.json
View File

@@ -15632,53 +15632,25 @@
"items": {
"oneOf": [
{
"$ref": "#/components/schemas/SystemMessage"
"$ref": "#/components/schemas/SystemMessageListResult"
},
{
"$ref": "#/components/schemas/UserMessage"
"$ref": "#/components/schemas/UserMessageListResult"
},
{
"$ref": "#/components/schemas/ReasoningMessage"
"$ref": "#/components/schemas/ReasoningMessageListResult"
},
{
"$ref": "#/components/schemas/HiddenReasoningMessage"
},
{
"$ref": "#/components/schemas/ToolCallMessage"
},
{
"$ref": "#/components/schemas/ToolReturnMessage"
},
{
"$ref": "#/components/schemas/AssistantMessage"
},
{
"$ref": "#/components/schemas/ApprovalRequestMessage"
},
{
"$ref": "#/components/schemas/ApprovalResponseMessage"
},
{
"$ref": "#/components/schemas/SummaryMessage"
},
{
"$ref": "#/components/schemas/EventMessage"
"$ref": "#/components/schemas/AssistantMessageListResult"
}
],
"discriminator": {
"propertyName": "message_type",
"mapping": {
"system_message": "#/components/schemas/SystemMessage",
"user_message": "#/components/schemas/UserMessage",
"reasoning_message": "#/components/schemas/ReasoningMessage",
"hidden_reasoning_message": "#/components/schemas/HiddenReasoningMessage",
"tool_call_message": "#/components/schemas/ToolCallMessage",
"tool_return_message": "#/components/schemas/ToolReturnMessage",
"assistant_message": "#/components/schemas/AssistantMessage",
"approval_request_message": "#/components/schemas/ApprovalRequestMessage",
"approval_response_message": "#/components/schemas/ApprovalResponseMessage",
"summary": "#/components/schemas/SummaryMessage",
"event": "#/components/schemas/EventMessage"
"system_message": "#/components/schemas/SystemMessageListResult",
"user_message": "#/components/schemas/UserMessageListResult",
"reasoning_message": "#/components/schemas/ReasoningMessageListResult",
"assistant_message": "#/components/schemas/AssistantMessageListResult"
}
}
},
@@ -21957,6 +21929,53 @@
"title": "AssistantMessage",
"description": "A message sent by the LLM in response to user input. Used in the LLM context.\n\nArgs:\n id (str): The ID of the message\n date (datetime): The date the message was created in ISO format\n name (Optional[str]): The name of the sender of the message\n content (Union[str, List[LettaAssistantMessageContentUnion]]): The message content sent by the agent (can be a string or an array of content parts)"
},
"AssistantMessageListResult": {
"properties": {
"message_type": {
"type": "string",
"const": "assistant_message",
"title": "Message Type",
"default": "assistant_message"
},
"content": {
"anyOf": [
{
"items": {
"$ref": "#/components/schemas/LettaAssistantMessageContentUnion"
},
"type": "array"
},
{
"type": "string"
}
],
"title": "Content",
"description": "The message content sent by the assistant (can be a string or an array of content parts)"
},
"agent_id": {
"anyOf": [
{
"type": "string"
},
{
"type": "null"
}
],
"title": "Agent Id",
"description": "The unique identifier of the agent that owns the message."
},
"created_at": {
"type": "string",
"format": "date-time",
"title": "Created At",
"description": "The time the message was created in ISO format."
}
},
"type": "object",
"required": ["content", "created_at"],
"title": "AssistantMessageListResult",
"description": "Assistant message list result with agent context.\n\nShape is identical to UpdateAssistantMessage but includes the owning agent_id."
},
"Audio": {
"properties": {
"id": {
@@ -34195,6 +34214,42 @@
"title": "ReasoningMessage",
"description": "Representation of an agent's internal reasoning.\n\nArgs:\n id (str): The ID of the message\n date (datetime): The date the message was created in ISO format\n name (Optional[str]): The name of the sender of the message\n source (Literal[\"reasoner_model\", \"non_reasoner_model\"]): Whether the reasoning\n content was generated natively by a reasoner model or derived via prompting\n reasoning (str): The internal reasoning of the agent\n signature (Optional[str]): The model-generated signature of the reasoning step"
},
"ReasoningMessageListResult": {
"properties": {
"reasoning": {
"type": "string",
"title": "Reasoning"
},
"message_type": {
"type": "string",
"const": "reasoning_message",
"title": "Message Type",
"default": "reasoning_message"
},
"agent_id": {
"anyOf": [
{
"type": "string"
},
{
"type": "null"
}
],
"title": "Agent Id",
"description": "The unique identifier of the agent that owns the message."
},
"created_at": {
"type": "string",
"format": "date-time",
"title": "Created At",
"description": "The time the message was created in ISO format."
}
},
"type": "object",
"required": ["reasoning", "created_at"],
"title": "ReasoningMessageListResult",
"description": "Reasoning message list result with agent context.\n\nShape is identical to UpdateReasoningMessage but includes the owning agent_id."
},
"RedactedReasoningContent": {
"properties": {
"type": {
@@ -36652,6 +36707,43 @@
"title": "SystemMessage",
"description": "A message generated by the system. Never streamed back on a response, only used for cursor pagination.\n\nArgs:\n id (str): The ID of the message\n date (datetime): The date the message was created in ISO format\n name (Optional[str]): The name of the sender of the message\n content (str): The message content sent by the system"
},
"SystemMessageListResult": {
"properties": {
"message_type": {
"type": "string",
"const": "system_message",
"title": "Message Type",
"default": "system_message"
},
"content": {
"type": "string",
"title": "Content",
"description": "The message content sent by the system (can be a string or an array of multi-modal content parts)"
},
"agent_id": {
"anyOf": [
{
"type": "string"
},
{
"type": "null"
}
],
"title": "Agent Id",
"description": "The unique identifier of the agent that owns the message."
},
"created_at": {
"type": "string",
"format": "date-time",
"title": "Created At",
"description": "The time the message was created in ISO format."
}
},
"type": "object",
"required": ["content", "created_at"],
"title": "SystemMessageListResult",
"description": "System message list result with agent context.\n\nShape is identical to UpdateSystemMessage but includes the owning agent_id."
},
"TagSchema": {
"properties": {
"tag": {
@@ -39276,6 +39368,53 @@
"title": "UserMessage",
"description": "A message sent by the user. Never streamed back on a response, only used for cursor pagination.\n\nArgs:\n id (str): The ID of the message\n date (datetime): The date the message was created in ISO format\n name (Optional[str]): The name of the sender of the message\n content (Union[str, List[LettaUserMessageContentUnion]]): The message content sent by the user (can be a string or an array of multi-modal content parts)"
},
"UserMessageListResult": {
"properties": {
"message_type": {
"type": "string",
"const": "user_message",
"title": "Message Type",
"default": "user_message"
},
"content": {
"anyOf": [
{
"items": {
"$ref": "#/components/schemas/LettaUserMessageContentUnion"
},
"type": "array"
},
{
"type": "string"
}
],
"title": "Content",
"description": "The message content sent by the user (can be a string or an array of multi-modal content parts)"
},
"agent_id": {
"anyOf": [
{
"type": "string"
},
{
"type": "null"
}
],
"title": "Agent Id",
"description": "The unique identifier of the agent that owns the message."
},
"created_at": {
"type": "string",
"format": "date-time",
"title": "Created At",
"description": "The time the message was created in ISO format."
}
},
"type": "object",
"required": ["content", "created_at"],
"title": "UserMessageListResult",
"description": "User message list result with agent context.\n\nShape is identical to UpdateUserMessage but includes the owning agent_id."
},
"UserUpdate": {
"properties": {
"id": {

72
letta/schemas/letta_message.py
View File

@@ -555,6 +555,78 @@ LettaMessageUpdateUnion = Annotated[
]
# ------------------------------
# Message Search Result Schemas
# ------------------------------
class SystemMessageListResult(UpdateSystemMessage):
"""System message list result with agent context.
Shape is identical to UpdateSystemMessage but includes the owning agent_id.
"""
agent_id: str | None = Field(
default=None,
description="The unique identifier of the agent that owns the message.",
)
created_at: datetime = Field(..., description="The time the message was created in ISO format.")
class UserMessageListResult(UpdateUserMessage):
"""User message list result with agent context.
Shape is identical to UpdateUserMessage but includes the owning agent_id.
"""
agent_id: str | None = Field(
default=None,
description="The unique identifier of the agent that owns the message.",
)
created_at: datetime = Field(..., description="The time the message was created in ISO format.")
class ReasoningMessageListResult(UpdateReasoningMessage):
"""Reasoning message list result with agent context.
Shape is identical to UpdateReasoningMessage but includes the owning agent_id.
"""
agent_id: str | None = Field(
default=None,
description="The unique identifier of the agent that owns the message.",
)
created_at: datetime = Field(..., description="The time the message was created in ISO format.")
class AssistantMessageListResult(UpdateAssistantMessage):
"""Assistant message list result with agent context.
Shape is identical to UpdateAssistantMessage but includes the owning agent_id.
"""
agent_id: str | None = Field(
default=None,
description="The unique identifier of the agent that owns the message.",
)
created_at: datetime = Field(..., description="The time the message was created in ISO format.")
LettaMessageSearchResult = Annotated[
Union[
SystemMessageListResult,
UserMessageListResult,
ReasoningMessageListResult,
AssistantMessageListResult,
],
Field(discriminator="message_type"),
]
# --------------------------
# Deprecated Message Schemas
# --------------------------

79
letta/schemas/message.py
View File

@@ -29,17 +29,22 @@ from letta.schemas.letta_message import (
ApprovalResponseMessage,
ApprovalReturn,
AssistantMessage,
AssistantMessageListResult,
HiddenReasoningMessage,
LettaMessage,
LettaMessageReturnUnion,
LettaMessageSearchResult,
MessageType,
ReasoningMessage,
ReasoningMessageListResult,
SystemMessage,
SystemMessageListResult,
ToolCall,
ToolCallMessage,
ToolReturn as LettaToolReturn,
ToolReturnMessage,
UserMessage,
UserMessageListResult,
)
from letta.schemas.letta_message_content import (
ImageContent,
@@ -319,6 +324,80 @@ class Message(BaseMessage):
)
]
@staticmethod
@trace_method
def to_letta_search_results_from_list(
search_results: List["MessageSearchResult"],
use_assistant_message: bool = True,
assistant_message_tool_name: str = DEFAULT_MESSAGE_TOOL,
assistant_message_tool_kwarg: str = DEFAULT_MESSAGE_TOOL_KWARG,
reverse: bool = True,
include_err: Optional[bool] = None,
text_is_assistant_message: bool = False,
) -> List[LettaMessageSearchResult]:
"""Convert MessageSearchResult objects into LettaMessageSearchResult objects.
This mirrors the behavior of to_letta_messages_from_list, but preserves the
originating Message.agent_id on each search result variant.
"""
letta_search_results: List[LettaMessageSearchResult] = []
for result in search_results:
message = result.message
# Convert the underlying Message into LettaMessage variants
letta_messages = message.to_letta_messages(
use_assistant_message=use_assistant_message,
assistant_message_tool_name=assistant_message_tool_name,
assistant_message_tool_kwarg=assistant_message_tool_kwarg,
reverse=reverse,
include_err=include_err,
text_is_assistant_message=text_is_assistant_message,
)
for lm in letta_messages:
if isinstance(lm, SystemMessage):
letta_search_results.append(
SystemMessageListResult(
message_type=lm.message_type,
content=lm.content,
agent_id=message.agent_id,
created_at=message.created_at,
)
)
elif isinstance(lm, UserMessage):
letta_search_results.append(
UserMessageListResult(
message_type=lm.message_type,
content=lm.content,
agent_id=message.agent_id,
created_at=message.created_at,
)
)
elif isinstance(lm, ReasoningMessage):
letta_search_results.append(
ReasoningMessageListResult(
message_type=lm.message_type,
reasoning=lm.reasoning,
agent_id=message.agent_id,
created_at=message.created_at,
)
)
elif isinstance(lm, AssistantMessage):
letta_search_results.append(
AssistantMessageListResult(
message_type=lm.message_type,
content=lm.content,
agent_id=message.agent_id,
created_at=message.created_at,
)
)
# Other LettaMessage variants (tool, approval, etc.) are not part of
# LettaMessageSearchResult and are intentionally skipped here.
return letta_search_results
def to_letta_messages(
self,
use_assistant_message: bool = False,

6
letta/server/rest_api/routers/v1/messages.py
View File

@@ -8,7 +8,7 @@ from letta.agents.letta_agent_batch import LettaAgentBatch
from letta.errors import LettaInvalidArgumentError
from letta.log import get_logger
from letta.schemas.job import BatchJob, JobStatus, JobType, JobUpdate
from letta.schemas.letta_message import LettaMessageUnion
from letta.schemas.letta_message import LettaMessageSearchResult, LettaMessageUnion
from letta.schemas.letta_request import CreateBatch
from letta.schemas.letta_response import LettaBatchMessages
from letta.schemas.message import Message, MessageSearchRequest, MessageSearchResult, SearchAllMessagesRequest
@@ -55,7 +55,7 @@ async def list_all_messages(
)
@router.post("/search", response_model=List[LettaMessageUnion], operation_id="search_all_messages")
@router.post("/search", response_model=List[LettaMessageSearchResult], operation_id="search_all_messages")
async def search_all_messages(
request: SearchAllMessagesRequest = Body(...),
server: SyncServer = Depends(get_letta_server),
@@ -77,7 +77,7 @@ async def search_all_messages(
start_date=request.start_date,
end_date=request.end_date,
)
return Message.to_letta_messages_from_list(messages=[result.message for result in results], text_is_assistant_message=True)
return Message.to_letta_search_results_from_list(search_results=results, text_is_assistant_message=True)
@router.post(

18
tests/sdk/search_test.py
View File

@@ -15,7 +15,6 @@ from letta_client import Letta
from letta_client.types import CreateBlockParam, MessageCreateParam
from letta.config import LettaConfig
from letta.schemas.message import MessageSearchResult
from letta.schemas.tool import ToolSearchResult
from letta.server.rest_api.routers.v1.passages import PassageSearchResult
from letta.server.server import SyncServer
@@ -387,10 +386,12 @@ def test_message_search_basic(client: Letta, enable_message_embedding):
time.sleep(10)
# Test FTS search for messages
# Note: The endpoint returns LettaMessageUnion, not MessageSearchResult
# Note: The endpoint returns LettaMessageSearchResult (API schema)
# and we treat the response as generic dicts here to avoid tight
# coupling to internal server-side models.
results = client.post(
"/v1/messages/search",
cast_to=list[MessageSearchResult],
cast_to=list[dict[str, Any]],
body={
"query": "capital Saudi Arabia",
"search_mode": "fts",
@@ -403,6 +404,17 @@ def test_message_search_basic(client: Letta, enable_message_embedding):
print(f"First result type: {type(results[0])}")
print(f"First result keys: {results[0].keys() if isinstance(results[0], dict) else 'N/A'}")
for result in results:
assert "agent_id" in result, "Result should have agent_id field"
# created_at should always be present and parseable
assert "created_at" in result, "Result should have created_at field"
assert result["created_at"], "created_at should be set"
created_at = result["created_at"]
if isinstance(created_at, str):
# Handle both "+00:00" and "Z" suffixes
datetime.fromisoformat(created_at.replace("Z", "+00:00"))
assert len(results) > 0, f"Should find at least one message. Got {len(results)} results."
finally:

14
tests/test_client.py
View File

@@ -58,8 +58,8 @@ def agent(client: Letta):
agent_state = client.agents.create(
name="test_client",
memory_blocks=[{"label": "human", "value": ""}, {"label": "persona", "value": ""}],
model="letta/letta-free",
embedding="letta/letta-free",
model="anthropic/claude-haiku-4-5-20251001",
embedding="openai/text-embedding-3-small",
)
yield agent_state
@@ -74,8 +74,8 @@ def search_agent_one(client: Letta):
agent_state = client.agents.create(
name="Search Agent One",
memory_blocks=[{"label": "human", "value": ""}, {"label": "persona", "value": ""}],
model="letta/letta-free",
embedding="letta/letta-free",
model="anthropic/claude-haiku-4-5-20251001",
embedding="openai/text-embedding-3-small",
)
yield agent_state
@@ -460,6 +460,12 @@ def test_messages(client: Letta, agent: AgentState):
messages_response = client.agents.messages.list(agent_id=agent.id, limit=1).items
assert len(messages_response) > 0, "Retrieving messages failed"
search_response = list(client.messages.search(query="test"))
assert len(search_response) > 0, "Searching messages failed"
for result in search_response:
assert result.agent_id == agent.id
assert result.created_at
# TODO: Add back when new agent loop hits
# @pytest.mark.asyncio