286 lines
12 KiB
Python
286 lines
12 KiB
Python
from datetime import datetime
|
|
from typing import Dict, List, Optional
|
|
|
|
from pydantic import BaseModel, Field
|
|
|
|
from letta.schemas.agent import AgentState, CreateAgent
|
|
from letta.schemas.block import Block, CreateBlock
|
|
from letta.schemas.enums import MessageRole
|
|
from letta.schemas.file import FileAgent, FileAgentBase, FileMetadata, FileMetadataBase
|
|
from letta.schemas.group import GroupCreate
|
|
from letta.schemas.message import Message, MessageCreate
|
|
from letta.schemas.source import Source, SourceCreate
|
|
from letta.schemas.tool import Tool
|
|
from letta.schemas.user import User
|
|
from letta.services.message_manager import MessageManager
|
|
|
|
|
|
class ImportResult:
|
|
"""Result of an agent file import operation"""
|
|
|
|
def __init__(
|
|
self,
|
|
success: bool,
|
|
message: str = "",
|
|
imported_count: int = 0,
|
|
errors: Optional[List[str]] = None,
|
|
id_mappings: Optional[Dict[str, str]] = None,
|
|
):
|
|
self.success = success
|
|
self.message = message
|
|
self.imported_count = imported_count
|
|
self.errors = errors or []
|
|
self.id_mappings = id_mappings or {}
|
|
|
|
|
|
class MessageSchema(MessageCreate):
|
|
"""Message with human-readable ID for agent file"""
|
|
|
|
__id_prefix__ = "message"
|
|
id: str = Field(..., description="Human-readable identifier for this message in the file")
|
|
|
|
# Override the role field to accept all message roles, not just user/system/assistant
|
|
role: MessageRole = Field(..., description="The role of the participant.")
|
|
model: Optional[str] = Field(None, description="The model used to make the function call")
|
|
agent_id: Optional[str] = Field(None, description="The unique identifier of the agent")
|
|
|
|
@classmethod
|
|
def from_message(cls, message: Message) -> "MessageSchema":
|
|
"""Convert Message to MessageSchema"""
|
|
|
|
# Create MessageSchema directly without going through MessageCreate
|
|
# to avoid role validation issues
|
|
return cls(
|
|
id=message.id,
|
|
role=message.role,
|
|
content=message.content,
|
|
name=message.name,
|
|
otid=None, # TODO
|
|
sender_id=None, # TODO
|
|
batch_item_id=message.batch_item_id,
|
|
group_id=message.group_id,
|
|
model=message.model,
|
|
agent_id=message.agent_id,
|
|
)
|
|
|
|
|
|
class FileAgentSchema(FileAgentBase):
|
|
"""File-Agent relationship with human-readable ID for agent file"""
|
|
|
|
__id_prefix__ = "file_agent"
|
|
id: str = Field(..., description="Human-readable identifier for this file-agent relationship in the file")
|
|
|
|
@classmethod
|
|
def from_file_agent(cls, file_agent: FileAgent) -> "FileAgentSchema":
|
|
"""Convert FileAgent to FileAgentSchema"""
|
|
|
|
create_file_agent = FileAgentBase(
|
|
agent_id=file_agent.agent_id,
|
|
file_id=file_agent.file_id,
|
|
source_id=file_agent.source_id,
|
|
file_name=file_agent.file_name,
|
|
is_open=file_agent.is_open,
|
|
visible_content=file_agent.visible_content,
|
|
last_accessed_at=file_agent.last_accessed_at,
|
|
)
|
|
|
|
# Create FileAgentSchema with the file_agent's ID (will be remapped later)
|
|
return cls(id=file_agent.id, **create_file_agent.model_dump())
|
|
|
|
|
|
class AgentSchema(CreateAgent):
|
|
"""Agent with human-readable ID for agent file"""
|
|
|
|
__id_prefix__ = "agent"
|
|
id: str = Field(..., description="Human-readable identifier for this agent in the file")
|
|
in_context_message_ids: List[str] = Field(
|
|
default_factory=list, description="List of message IDs that are currently in the agent's context"
|
|
)
|
|
messages: List[MessageSchema] = Field(default_factory=list, description="List of messages in the agent's conversation history")
|
|
files_agents: List[FileAgentSchema] = Field(default_factory=list, description="List of file-agent relationships for this agent")
|
|
|
|
@classmethod
|
|
async def from_agent_state(
|
|
cls, agent_state: AgentState, message_manager: MessageManager, files_agents: List[FileAgent], actor: User
|
|
) -> "AgentSchema":
|
|
"""Convert AgentState to AgentSchema"""
|
|
|
|
create_agent = CreateAgent(
|
|
name=agent_state.name,
|
|
memory_blocks=[], # TODO: Convert from agent_state.memory if needed
|
|
tools=[],
|
|
tool_ids=[tool.id for tool in agent_state.tools] if agent_state.tools else [],
|
|
source_ids=[], # [source.id for source in agent_state.sources] if agent_state.sources else [],
|
|
block_ids=[block.id for block in agent_state.memory.blocks],
|
|
tool_rules=agent_state.tool_rules,
|
|
tags=agent_state.tags,
|
|
system=agent_state.system,
|
|
agent_type=agent_state.agent_type,
|
|
llm_config=agent_state.llm_config,
|
|
embedding_config=agent_state.embedding_config,
|
|
initial_message_sequence=None,
|
|
include_base_tools=False,
|
|
include_multi_agent_tools=False,
|
|
include_base_tool_rules=False,
|
|
include_default_source=False,
|
|
description=agent_state.description,
|
|
metadata=agent_state.metadata,
|
|
model=None,
|
|
embedding=None,
|
|
context_window_limit=None,
|
|
embedding_chunk_size=None,
|
|
max_tokens=None,
|
|
max_reasoning_tokens=None,
|
|
enable_reasoner=False,
|
|
from_template=None, # TODO: Need to get passed in
|
|
template=False, # TODO: Need to get passed in
|
|
project=None, # TODO: Need to get passed in
|
|
tool_exec_environment_variables=agent_state.get_agent_env_vars_as_dict(),
|
|
memory_variables=None, # TODO: Need to get passed in
|
|
project_id=None, # TODO: Need to get passed in
|
|
template_id=None, # TODO: Need to get passed in
|
|
base_template_id=None, # TODO: Need to get passed in
|
|
identity_ids=None, # TODO: Need to get passed in
|
|
message_buffer_autoclear=agent_state.message_buffer_autoclear,
|
|
enable_sleeptime=False, # TODO: Need to figure out how to patch this
|
|
response_format=agent_state.response_format,
|
|
timezone=agent_state.timezone or "UTC",
|
|
max_files_open=agent_state.max_files_open,
|
|
per_file_view_window_char_limit=agent_state.per_file_view_window_char_limit,
|
|
)
|
|
|
|
messages = await message_manager.list_messages_for_agent_async(
|
|
agent_id=agent_state.id, actor=actor, limit=50
|
|
) # TODO: Expand to get more messages
|
|
|
|
# Convert messages to MessageSchema objects
|
|
message_schemas = [MessageSchema.from_message(msg) for msg in messages]
|
|
|
|
# Create AgentSchema with agent state ID (remapped later)
|
|
return cls(
|
|
id=agent_state.id,
|
|
in_context_message_ids=agent_state.message_ids or [],
|
|
messages=message_schemas, # Messages will be populated separately by the manager
|
|
files_agents=[FileAgentSchema.from_file_agent(f) for f in files_agents],
|
|
**create_agent.model_dump(),
|
|
)
|
|
|
|
|
|
class GroupSchema(GroupCreate):
|
|
"""Group with human-readable ID for agent file"""
|
|
|
|
__id_prefix__ = "group"
|
|
id: str = Field(..., description="Human-readable identifier for this group in the file")
|
|
|
|
|
|
class BlockSchema(CreateBlock):
|
|
"""Block with human-readable ID for agent file"""
|
|
|
|
__id_prefix__ = "block"
|
|
id: str = Field(..., description="Human-readable identifier for this block in the file")
|
|
|
|
@classmethod
|
|
def from_block(cls, block: Block) -> "BlockSchema":
|
|
"""Convert Block to BlockSchema"""
|
|
|
|
create_block = CreateBlock(
|
|
value=block.value,
|
|
limit=block.limit,
|
|
template_name=block.template_name,
|
|
is_template=block.is_template,
|
|
preserve_on_migration=block.preserve_on_migration,
|
|
label=block.label,
|
|
read_only=block.read_only,
|
|
description=block.description,
|
|
metadata=block.metadata or {},
|
|
)
|
|
|
|
# Create BlockSchema with the block's ID (will be remapped later)
|
|
return cls(id=block.id, **create_block.model_dump())
|
|
|
|
|
|
class FileSchema(FileMetadataBase):
|
|
"""File with human-readable ID for agent file"""
|
|
|
|
__id_prefix__ = "file"
|
|
id: str = Field(..., description="Human-readable identifier for this file in the file")
|
|
|
|
@classmethod
|
|
def from_file_metadata(cls, file_metadata: FileMetadata) -> "FileSchema":
|
|
"""Convert FileMetadata to FileSchema"""
|
|
|
|
create_file = FileMetadataBase(
|
|
source_id=file_metadata.source_id,
|
|
file_name=file_metadata.file_name,
|
|
original_file_name=file_metadata.original_file_name,
|
|
file_path=file_metadata.file_path,
|
|
file_type=file_metadata.file_type,
|
|
file_size=file_metadata.file_size,
|
|
file_creation_date=file_metadata.file_creation_date,
|
|
file_last_modified_date=file_metadata.file_last_modified_date,
|
|
processing_status=file_metadata.processing_status,
|
|
error_message=file_metadata.error_message,
|
|
total_chunks=file_metadata.total_chunks,
|
|
chunks_embedded=file_metadata.chunks_embedded,
|
|
content=file_metadata.content,
|
|
)
|
|
|
|
# Create FileSchema with the file's ID (will be remapped later)
|
|
return cls(id=file_metadata.id, **create_file.model_dump())
|
|
|
|
|
|
class SourceSchema(SourceCreate):
|
|
"""Source with human-readable ID for agent file"""
|
|
|
|
__id_prefix__ = "source"
|
|
id: str = Field(..., description="Human-readable identifier for this source in the file")
|
|
|
|
@classmethod
|
|
def from_source(cls, source: Source) -> "SourceSchema":
|
|
"""Convert Block to BlockSchema"""
|
|
|
|
create_block = SourceCreate(
|
|
name=source.name,
|
|
description=source.description,
|
|
instructions=source.instructions,
|
|
metadata=source.metadata,
|
|
embedding_config=source.embedding_config,
|
|
)
|
|
|
|
# Create SourceSchema with the block's ID (will be remapped later)
|
|
return cls(id=source.id, **create_block.model_dump())
|
|
|
|
|
|
# TODO: This one is quite thin, just a wrapper over Tool
|
|
class ToolSchema(Tool):
|
|
"""Tool with human-readable ID for agent file"""
|
|
|
|
__id_prefix__ = "tool"
|
|
id: str = Field(..., description="Human-readable identifier for this tool in the file")
|
|
|
|
@classmethod
|
|
def from_tool(cls, tool: Tool) -> "ToolSchema":
|
|
"""Convert Tool to ToolSchema"""
|
|
return cls(**tool.model_dump())
|
|
|
|
|
|
# class MCPServerSchema(RegisterMCPServer):
|
|
# """MCP Server with human-readable ID for agent file"""
|
|
# id: str = Field(..., description="Human-readable identifier for this MCP server in the file")
|
|
|
|
|
|
class AgentFileSchema(BaseModel):
|
|
"""Schema for serialized agent file that can be exported to JSON and imported into agent server."""
|
|
|
|
agents: List[AgentSchema] = Field(..., description="List of agents in this agent file")
|
|
groups: List[GroupSchema] = Field(..., description="List of groups in this agent file")
|
|
blocks: List[BlockSchema] = Field(..., description="List of memory blocks in this agent file")
|
|
files: List[FileSchema] = Field(..., description="List of files in this agent file")
|
|
sources: List[SourceSchema] = Field(..., description="List of sources in this agent file")
|
|
tools: List[ToolSchema] = Field(..., description="List of tools in this agent file")
|
|
# mcp_servers: List[MCPServerSchema] = Field(..., description="List of MCP servers in this agent file")
|
|
metadata: Dict[str, str] = Field(
|
|
default_factory=dict, description="Metadata for this agent file, including revision_id and other export information."
|
|
)
|
|
created_at: Optional[datetime] = Field(default=None, description="The timestamp when the object was created.")
|