621 lines
24 KiB
Python
621 lines
24 KiB
Python
from typing import TYPE_CHECKING, Dict, List, Optional
|
|
|
|
if TYPE_CHECKING:
|
|
pass
|
|
|
|
# Import AgentState outside TYPE_CHECKING for @enforce_types decorator
|
|
from sqlalchemy import delete, func, select
|
|
|
|
from letta.errors import LettaInvalidArgumentError
|
|
from letta.orm.agent import Agent as AgentModel
|
|
from letta.orm.block import Block as BlockModel
|
|
from letta.orm.blocks_conversations import BlocksConversations
|
|
from letta.orm.conversation import Conversation as ConversationModel
|
|
from letta.orm.conversation_messages import ConversationMessage as ConversationMessageModel
|
|
from letta.orm.message import Message as MessageModel
|
|
from letta.otel.tracing import trace_method
|
|
from letta.schemas.agent import AgentState
|
|
from letta.schemas.block import Block as PydanticBlock
|
|
from letta.schemas.conversation import Conversation as PydanticConversation, CreateConversation, UpdateConversation
|
|
from letta.schemas.letta_message import LettaMessage
|
|
from letta.schemas.message import Message as PydanticMessage
|
|
from letta.schemas.user import User as PydanticUser
|
|
from letta.server.db import db_registry
|
|
from letta.services.helpers.agent_manager_helper import validate_agent_exists_async
|
|
from letta.utils import enforce_types
|
|
|
|
|
|
class ConversationManager:
|
|
"""Manager class to handle business logic related to Conversations."""
|
|
|
|
@enforce_types
|
|
@trace_method
|
|
async def create_conversation(
|
|
self,
|
|
agent_id: str,
|
|
conversation_create: CreateConversation,
|
|
actor: PydanticUser,
|
|
) -> PydanticConversation:
|
|
"""Create a new conversation for an agent.
|
|
|
|
Args:
|
|
agent_id: The ID of the agent this conversation belongs to
|
|
conversation_create: The conversation creation request, optionally including
|
|
isolated_block_labels for conversation-specific memory blocks
|
|
actor: The user performing the action
|
|
|
|
Returns:
|
|
The created conversation with isolated_block_ids if any were created
|
|
"""
|
|
async with db_registry.async_session() as session:
|
|
# Validate that the agent exists before creating the conversation
|
|
await validate_agent_exists_async(session, agent_id, actor)
|
|
conversation = ConversationModel(
|
|
agent_id=agent_id,
|
|
summary=conversation_create.summary,
|
|
organization_id=actor.organization_id,
|
|
)
|
|
await conversation.create_async(session, actor=actor)
|
|
|
|
# Handle isolated blocks if requested
|
|
isolated_block_ids = []
|
|
if conversation_create.isolated_block_labels:
|
|
isolated_block_ids = await self._create_isolated_blocks(
|
|
session=session,
|
|
conversation=conversation,
|
|
agent_id=agent_id,
|
|
isolated_block_labels=conversation_create.isolated_block_labels,
|
|
actor=actor,
|
|
)
|
|
|
|
pydantic_conversation = conversation.to_pydantic()
|
|
pydantic_conversation.isolated_block_ids = isolated_block_ids
|
|
return pydantic_conversation
|
|
|
|
@enforce_types
|
|
@trace_method
|
|
async def get_conversation_by_id(
|
|
self,
|
|
conversation_id: str,
|
|
actor: PydanticUser,
|
|
) -> PydanticConversation:
|
|
"""Retrieve a conversation by its ID, including in-context message IDs."""
|
|
async with db_registry.async_session() as session:
|
|
conversation = await ConversationModel.read_async(
|
|
db_session=session,
|
|
identifier=conversation_id,
|
|
actor=actor,
|
|
check_is_deleted=True,
|
|
)
|
|
|
|
# Get the in-context message IDs for this conversation
|
|
message_ids = await self.get_message_ids_for_conversation(
|
|
conversation_id=conversation_id,
|
|
actor=actor,
|
|
)
|
|
|
|
# Build the pydantic model with in_context_message_ids
|
|
pydantic_conversation = conversation.to_pydantic()
|
|
pydantic_conversation.in_context_message_ids = message_ids
|
|
return pydantic_conversation
|
|
|
|
@enforce_types
|
|
@trace_method
|
|
async def list_conversations(
|
|
self,
|
|
agent_id: str,
|
|
actor: PydanticUser,
|
|
limit: int = 50,
|
|
after: Optional[str] = None,
|
|
summary_search: Optional[str] = None,
|
|
) -> List[PydanticConversation]:
|
|
"""List conversations for an agent with cursor-based pagination.
|
|
|
|
Args:
|
|
agent_id: The agent ID to list conversations for
|
|
actor: The user performing the action
|
|
limit: Maximum number of conversations to return
|
|
after: Cursor for pagination (conversation ID)
|
|
summary_search: Optional text to search for within the summary field
|
|
|
|
Returns:
|
|
List of conversations matching the criteria
|
|
"""
|
|
async with db_registry.async_session() as session:
|
|
# If summary search is provided, use custom query
|
|
if summary_search:
|
|
from sqlalchemy import and_
|
|
|
|
stmt = (
|
|
select(ConversationModel)
|
|
.where(
|
|
and_(
|
|
ConversationModel.agent_id == agent_id,
|
|
ConversationModel.organization_id == actor.organization_id,
|
|
ConversationModel.summary.isnot(None),
|
|
ConversationModel.summary.contains(summary_search),
|
|
)
|
|
)
|
|
.order_by(ConversationModel.created_at.desc())
|
|
.limit(limit)
|
|
)
|
|
|
|
if after:
|
|
# Add cursor filtering
|
|
after_conv = await ConversationModel.read_async(
|
|
db_session=session,
|
|
identifier=after,
|
|
actor=actor,
|
|
)
|
|
stmt = stmt.where(ConversationModel.created_at < after_conv.created_at)
|
|
|
|
result = await session.execute(stmt)
|
|
conversations = result.scalars().all()
|
|
return [conv.to_pydantic() for conv in conversations]
|
|
|
|
# Use default list logic
|
|
conversations = await ConversationModel.list_async(
|
|
db_session=session,
|
|
actor=actor,
|
|
agent_id=agent_id,
|
|
limit=limit,
|
|
after=after,
|
|
ascending=False,
|
|
)
|
|
return [conv.to_pydantic() for conv in conversations]
|
|
|
|
@enforce_types
|
|
@trace_method
|
|
async def update_conversation(
|
|
self,
|
|
conversation_id: str,
|
|
conversation_update: UpdateConversation,
|
|
actor: PydanticUser,
|
|
) -> PydanticConversation:
|
|
"""Update a conversation."""
|
|
async with db_registry.async_session() as session:
|
|
conversation = await ConversationModel.read_async(
|
|
db_session=session,
|
|
identifier=conversation_id,
|
|
actor=actor,
|
|
)
|
|
|
|
# Set attributes on the model
|
|
update_data = conversation_update.model_dump(exclude_none=True)
|
|
for key, value in update_data.items():
|
|
setattr(conversation, key, value)
|
|
|
|
# Commit the update
|
|
updated_conversation = await conversation.update_async(
|
|
db_session=session,
|
|
actor=actor,
|
|
)
|
|
return updated_conversation.to_pydantic()
|
|
|
|
@enforce_types
|
|
@trace_method
|
|
async def delete_conversation(
|
|
self,
|
|
conversation_id: str,
|
|
actor: PydanticUser,
|
|
) -> None:
|
|
"""Soft delete a conversation and hard-delete its isolated blocks."""
|
|
async with db_registry.async_session() as session:
|
|
conversation = await ConversationModel.read_async(
|
|
db_session=session,
|
|
identifier=conversation_id,
|
|
actor=actor,
|
|
)
|
|
|
|
# Get isolated blocks before modifying conversation
|
|
isolated_blocks = list(conversation.isolated_blocks)
|
|
|
|
# Soft delete the conversation first
|
|
conversation.is_deleted = True
|
|
await conversation.update_async(db_session=session, actor=actor)
|
|
|
|
# Hard-delete isolated blocks (Block model doesn't support soft-delete)
|
|
# Following same pattern as block_manager.delete_block_async
|
|
for block in isolated_blocks:
|
|
# Delete junction table entry first
|
|
await session.execute(delete(BlocksConversations).where(BlocksConversations.block_id == block.id))
|
|
await session.flush()
|
|
# Then hard-delete the block
|
|
await block.hard_delete_async(db_session=session, actor=actor)
|
|
|
|
# ==================== Message Management Methods ====================
|
|
|
|
@enforce_types
|
|
@trace_method
|
|
async def get_message_ids_for_conversation(
|
|
self,
|
|
conversation_id: str,
|
|
actor: PydanticUser,
|
|
) -> List[str]:
|
|
"""
|
|
Get ordered message IDs for a conversation.
|
|
|
|
Returns message IDs ordered by position in the conversation.
|
|
Only returns messages that are currently in_context.
|
|
"""
|
|
async with db_registry.async_session() as session:
|
|
query = (
|
|
select(ConversationMessageModel.message_id)
|
|
.where(
|
|
ConversationMessageModel.conversation_id == conversation_id,
|
|
ConversationMessageModel.organization_id == actor.organization_id,
|
|
ConversationMessageModel.in_context == True,
|
|
ConversationMessageModel.is_deleted == False,
|
|
)
|
|
.order_by(ConversationMessageModel.position)
|
|
)
|
|
result = await session.execute(query)
|
|
return list(result.scalars().all())
|
|
|
|
@enforce_types
|
|
@trace_method
|
|
async def get_messages_for_conversation(
|
|
self,
|
|
conversation_id: str,
|
|
actor: PydanticUser,
|
|
) -> List[PydanticMessage]:
|
|
"""
|
|
Get ordered Message objects for a conversation.
|
|
|
|
Returns full Message objects ordered by position in the conversation.
|
|
Only returns messages that are currently in_context.
|
|
"""
|
|
async with db_registry.async_session() as session:
|
|
query = (
|
|
select(MessageModel)
|
|
.join(
|
|
ConversationMessageModel,
|
|
MessageModel.id == ConversationMessageModel.message_id,
|
|
)
|
|
.where(
|
|
ConversationMessageModel.conversation_id == conversation_id,
|
|
ConversationMessageModel.organization_id == actor.organization_id,
|
|
ConversationMessageModel.in_context == True,
|
|
ConversationMessageModel.is_deleted == False,
|
|
)
|
|
.order_by(ConversationMessageModel.position)
|
|
)
|
|
result = await session.execute(query)
|
|
return [msg.to_pydantic() for msg in result.scalars().all()]
|
|
|
|
@enforce_types
|
|
@trace_method
|
|
async def add_messages_to_conversation(
|
|
self,
|
|
conversation_id: str,
|
|
agent_id: str,
|
|
message_ids: List[str],
|
|
actor: PydanticUser,
|
|
starting_position: Optional[int] = None,
|
|
) -> None:
|
|
"""
|
|
Add messages to a conversation's tracking table.
|
|
|
|
Creates ConversationMessage entries with auto-incrementing positions.
|
|
|
|
Args:
|
|
conversation_id: The conversation to add messages to
|
|
agent_id: The agent ID
|
|
message_ids: List of message IDs to add
|
|
actor: The user performing the action
|
|
starting_position: Optional starting position (defaults to next available)
|
|
"""
|
|
if not message_ids:
|
|
return
|
|
|
|
async with db_registry.async_session() as session:
|
|
# Get starting position if not provided
|
|
if starting_position is None:
|
|
query = select(func.coalesce(func.max(ConversationMessageModel.position), -1)).where(
|
|
ConversationMessageModel.conversation_id == conversation_id,
|
|
ConversationMessageModel.organization_id == actor.organization_id,
|
|
)
|
|
result = await session.execute(query)
|
|
max_position = result.scalar()
|
|
# Use explicit None check instead of `or` to handle position=0 correctly
|
|
if max_position is None:
|
|
max_position = -1
|
|
starting_position = max_position + 1
|
|
|
|
# Create ConversationMessage entries
|
|
for i, message_id in enumerate(message_ids):
|
|
conv_msg = ConversationMessageModel(
|
|
conversation_id=conversation_id,
|
|
agent_id=agent_id,
|
|
message_id=message_id,
|
|
position=starting_position + i,
|
|
in_context=True,
|
|
organization_id=actor.organization_id,
|
|
)
|
|
session.add(conv_msg)
|
|
|
|
await session.commit()
|
|
|
|
@enforce_types
|
|
@trace_method
|
|
async def update_in_context_messages(
|
|
self,
|
|
conversation_id: str,
|
|
in_context_message_ids: List[str],
|
|
actor: PydanticUser,
|
|
) -> None:
|
|
"""
|
|
Update which messages are in context for a conversation.
|
|
|
|
Sets in_context=True for messages in the list, False for others.
|
|
Also updates positions to preserve the order specified in in_context_message_ids.
|
|
|
|
This is critical for correctness: when summarization inserts a summary message
|
|
that needs to appear before an approval request, the positions must reflect
|
|
the intended order, not the insertion order.
|
|
|
|
Args:
|
|
conversation_id: The conversation to update
|
|
in_context_message_ids: List of message IDs in the desired order
|
|
actor: The user performing the action
|
|
"""
|
|
async with db_registry.async_session() as session:
|
|
# Get all conversation messages for this conversation
|
|
query = select(ConversationMessageModel).where(
|
|
ConversationMessageModel.conversation_id == conversation_id,
|
|
ConversationMessageModel.organization_id == actor.organization_id,
|
|
ConversationMessageModel.is_deleted == False,
|
|
)
|
|
result = await session.execute(query)
|
|
conv_messages = result.scalars().all()
|
|
|
|
# Build lookup dict
|
|
conv_msg_dict = {cm.message_id: cm for cm in conv_messages}
|
|
|
|
# Update in_context status AND positions
|
|
in_context_set = set(in_context_message_ids)
|
|
for conv_msg in conv_messages:
|
|
conv_msg.in_context = conv_msg.message_id in in_context_set
|
|
|
|
# Update positions to match the order in in_context_message_ids
|
|
# This ensures ORDER BY position returns messages in the correct order
|
|
for position, message_id in enumerate(in_context_message_ids):
|
|
if message_id in conv_msg_dict:
|
|
conv_msg_dict[message_id].position = position
|
|
|
|
await session.commit()
|
|
|
|
@enforce_types
|
|
@trace_method
|
|
async def list_conversation_messages(
|
|
self,
|
|
conversation_id: str,
|
|
actor: PydanticUser,
|
|
limit: Optional[int] = 100,
|
|
before: Optional[str] = None,
|
|
after: Optional[str] = None,
|
|
reverse: bool = False,
|
|
group_id: Optional[str] = None,
|
|
include_err: Optional[bool] = None,
|
|
) -> List[LettaMessage]:
|
|
"""
|
|
List all messages in a conversation with pagination support.
|
|
|
|
Unlike get_messages_for_conversation, this returns ALL messages
|
|
(not just in_context) and supports cursor-based pagination.
|
|
|
|
Args:
|
|
conversation_id: The conversation to list messages for
|
|
actor: The user performing the action
|
|
limit: Maximum number of messages to return
|
|
before: Return messages before this message ID
|
|
after: Return messages after this message ID
|
|
reverse: If True, return messages in descending order (newest first)
|
|
group_id: Optional group ID to filter messages by
|
|
include_err: Optional boolean to include error messages and error statuses
|
|
|
|
Returns:
|
|
List of LettaMessage objects
|
|
"""
|
|
async with db_registry.async_session() as session:
|
|
# Build base query joining Message with ConversationMessage
|
|
query = (
|
|
select(MessageModel)
|
|
.join(
|
|
ConversationMessageModel,
|
|
MessageModel.id == ConversationMessageModel.message_id,
|
|
)
|
|
.where(
|
|
ConversationMessageModel.conversation_id == conversation_id,
|
|
ConversationMessageModel.organization_id == actor.organization_id,
|
|
ConversationMessageModel.is_deleted == False,
|
|
)
|
|
)
|
|
|
|
# Filter by group_id if provided
|
|
if group_id:
|
|
query = query.where(MessageModel.group_id == group_id)
|
|
|
|
# Handle cursor-based pagination
|
|
if before:
|
|
# Get the position of the cursor message
|
|
cursor_query = select(ConversationMessageModel.position).where(
|
|
ConversationMessageModel.conversation_id == conversation_id,
|
|
ConversationMessageModel.message_id == before,
|
|
)
|
|
cursor_result = await session.execute(cursor_query)
|
|
cursor_position = cursor_result.scalar_one_or_none()
|
|
if cursor_position is not None:
|
|
query = query.where(ConversationMessageModel.position < cursor_position)
|
|
|
|
if after:
|
|
# Get the position of the cursor message
|
|
cursor_query = select(ConversationMessageModel.position).where(
|
|
ConversationMessageModel.conversation_id == conversation_id,
|
|
ConversationMessageModel.message_id == after,
|
|
)
|
|
cursor_result = await session.execute(cursor_query)
|
|
cursor_position = cursor_result.scalar_one_or_none()
|
|
if cursor_position is not None:
|
|
query = query.where(ConversationMessageModel.position > cursor_position)
|
|
|
|
# Order by position
|
|
if reverse:
|
|
query = query.order_by(ConversationMessageModel.position.desc())
|
|
else:
|
|
query = query.order_by(ConversationMessageModel.position.asc())
|
|
|
|
# Apply limit
|
|
if limit is not None:
|
|
query = query.limit(limit)
|
|
|
|
result = await session.execute(query)
|
|
messages = [msg.to_pydantic() for msg in result.scalars().all()]
|
|
|
|
# Convert to LettaMessages (reverse=False keeps sub-messages in natural order)
|
|
return PydanticMessage.to_letta_messages_from_list(
|
|
messages, reverse=False, include_err=include_err, text_is_assistant_message=True
|
|
)
|
|
|
|
# ==================== Isolated Blocks Methods ====================
|
|
|
|
async def _create_isolated_blocks(
|
|
self,
|
|
session,
|
|
conversation: ConversationModel,
|
|
agent_id: str,
|
|
isolated_block_labels: List[str],
|
|
actor: PydanticUser,
|
|
) -> List[str]:
|
|
"""Create conversation-specific copies of blocks for isolated labels.
|
|
|
|
Args:
|
|
session: The database session
|
|
conversation: The conversation model (must be created but not yet committed)
|
|
agent_id: The agent ID to get source blocks from
|
|
isolated_block_labels: List of block labels to isolate
|
|
actor: The user performing the action
|
|
|
|
Returns:
|
|
List of created block IDs
|
|
|
|
Raises:
|
|
LettaInvalidArgumentError: If a block label is not found on the agent
|
|
"""
|
|
# Get the agent with its blocks
|
|
agent = await AgentModel.read_async(db_session=session, identifier=agent_id, actor=actor)
|
|
|
|
# Map of label -> agent block
|
|
agent_blocks_by_label = {block.label: block for block in agent.core_memory}
|
|
|
|
created_block_ids = []
|
|
for label in isolated_block_labels:
|
|
if label not in agent_blocks_by_label:
|
|
raise LettaInvalidArgumentError(
|
|
f"Block with label '{label}' not found on agent '{agent_id}'",
|
|
argument_name="isolated_block_labels",
|
|
)
|
|
|
|
source_block = agent_blocks_by_label[label]
|
|
|
|
# Create a copy of the block with a new ID using Pydantic schema (which auto-generates ID)
|
|
new_block_pydantic = PydanticBlock(
|
|
label=source_block.label,
|
|
description=source_block.description,
|
|
value=source_block.value,
|
|
limit=source_block.limit,
|
|
metadata=source_block.metadata_,
|
|
read_only=source_block.read_only,
|
|
)
|
|
|
|
# Convert to ORM model
|
|
block_data = new_block_pydantic.model_dump(to_orm=True, exclude_none=True)
|
|
new_block = BlockModel(**block_data, organization_id=actor.organization_id)
|
|
await new_block.create_async(session, actor=actor)
|
|
|
|
# Create the junction table entry
|
|
blocks_conv = BlocksConversations(
|
|
conversation_id=conversation.id,
|
|
block_id=new_block.id,
|
|
block_label=label,
|
|
)
|
|
session.add(blocks_conv)
|
|
created_block_ids.append(new_block.id)
|
|
|
|
return created_block_ids
|
|
|
|
@enforce_types
|
|
@trace_method
|
|
async def get_isolated_blocks_for_conversation(
|
|
self,
|
|
conversation_id: str,
|
|
actor: PydanticUser,
|
|
) -> Dict[str, PydanticBlock]:
|
|
"""Get isolated blocks for a conversation, keyed by label.
|
|
|
|
Args:
|
|
conversation_id: The conversation ID
|
|
actor: The user performing the action
|
|
|
|
Returns:
|
|
Dictionary mapping block labels to their conversation-specific blocks
|
|
"""
|
|
async with db_registry.async_session() as session:
|
|
conversation = await ConversationModel.read_async(
|
|
db_session=session,
|
|
identifier=conversation_id,
|
|
actor=actor,
|
|
check_is_deleted=True,
|
|
)
|
|
return {block.label: block.to_pydantic() for block in conversation.isolated_blocks}
|
|
|
|
@enforce_types
|
|
@trace_method
|
|
async def apply_isolated_blocks_to_agent_state(
|
|
self,
|
|
agent_state: "AgentState",
|
|
conversation_id: str,
|
|
actor: PydanticUser,
|
|
) -> "AgentState":
|
|
"""Apply conversation-specific block overrides to an agent state.
|
|
|
|
This method modifies the agent_state.memory to replace blocks that have
|
|
conversation-specific isolated versions.
|
|
|
|
Args:
|
|
agent_state: The agent state to modify (will be modified in place)
|
|
conversation_id: The conversation ID to get isolated blocks from
|
|
actor: The user performing the action
|
|
|
|
Returns:
|
|
The modified agent state (same object, modified in place)
|
|
"""
|
|
from letta.schemas.memory import Memory
|
|
|
|
# Get conversation's isolated blocks
|
|
isolated_blocks = await self.get_isolated_blocks_for_conversation(
|
|
conversation_id=conversation_id,
|
|
actor=actor,
|
|
)
|
|
|
|
if not isolated_blocks:
|
|
return agent_state
|
|
|
|
# Override agent's blocks with conversation-specific blocks
|
|
memory_blocks = []
|
|
for block in agent_state.memory.blocks:
|
|
if block.label in isolated_blocks:
|
|
memory_blocks.append(isolated_blocks[block.label])
|
|
else:
|
|
memory_blocks.append(block)
|
|
|
|
# Create new Memory with overridden blocks
|
|
agent_state.memory = Memory(
|
|
blocks=memory_blocks,
|
|
file_blocks=agent_state.memory.file_blocks,
|
|
agent_type=agent_state.memory.agent_type,
|
|
git_enabled=agent_state.memory.git_enabled,
|
|
)
|
|
|
|
return agent_state
|