letta-nightly 0.8.15.dev20250720104313__py3-none-any.whl → 0.8.16.dev20250721070720__py3-none-any.whl
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- letta/__init__.py +1 -1
- letta/agent.py +27 -11
- letta/agents/helpers.py +1 -1
- letta/agents/letta_agent.py +518 -322
- letta/agents/letta_agent_batch.py +1 -2
- letta/agents/voice_agent.py +15 -17
- letta/client/client.py +3 -3
- letta/constants.py +5 -0
- letta/embeddings.py +0 -2
- letta/errors.py +8 -0
- letta/functions/function_sets/base.py +3 -3
- letta/functions/helpers.py +2 -3
- letta/groups/sleeptime_multi_agent.py +0 -1
- letta/helpers/composio_helpers.py +2 -2
- letta/helpers/converters.py +1 -1
- letta/helpers/pinecone_utils.py +8 -0
- letta/helpers/tool_rule_solver.py +13 -18
- letta/llm_api/aws_bedrock.py +16 -2
- letta/llm_api/cohere.py +1 -1
- letta/llm_api/openai_client.py +1 -1
- letta/local_llm/grammars/gbnf_grammar_generator.py +1 -1
- letta/local_llm/llm_chat_completion_wrappers/zephyr.py +14 -14
- letta/local_llm/utils.py +1 -2
- letta/orm/agent.py +3 -3
- letta/orm/block.py +4 -4
- letta/orm/files_agents.py +0 -1
- letta/orm/identity.py +2 -0
- letta/orm/mcp_server.py +0 -2
- letta/orm/message.py +140 -14
- letta/orm/organization.py +5 -5
- letta/orm/passage.py +4 -4
- letta/orm/source.py +1 -1
- letta/orm/sqlalchemy_base.py +61 -39
- letta/orm/step.py +2 -0
- letta/otel/db_pool_monitoring.py +308 -0
- letta/otel/metric_registry.py +94 -1
- letta/otel/sqlalchemy_instrumentation.py +548 -0
- letta/otel/sqlalchemy_instrumentation_integration.py +124 -0
- letta/otel/tracing.py +37 -1
- letta/schemas/agent.py +0 -3
- letta/schemas/agent_file.py +283 -0
- letta/schemas/block.py +0 -3
- letta/schemas/file.py +28 -26
- letta/schemas/letta_message.py +15 -4
- letta/schemas/memory.py +1 -1
- letta/schemas/message.py +31 -26
- letta/schemas/openai/chat_completion_response.py +0 -1
- letta/schemas/providers.py +20 -0
- letta/schemas/source.py +11 -13
- letta/schemas/step.py +12 -0
- letta/schemas/tool.py +0 -4
- letta/serialize_schemas/marshmallow_agent.py +14 -1
- letta/serialize_schemas/marshmallow_block.py +23 -1
- letta/serialize_schemas/marshmallow_message.py +1 -3
- letta/serialize_schemas/marshmallow_tool.py +23 -1
- letta/server/db.py +110 -6
- letta/server/rest_api/app.py +85 -73
- letta/server/rest_api/routers/v1/agents.py +68 -53
- letta/server/rest_api/routers/v1/blocks.py +2 -2
- letta/server/rest_api/routers/v1/jobs.py +3 -0
- letta/server/rest_api/routers/v1/organizations.py +2 -2
- letta/server/rest_api/routers/v1/sources.py +18 -2
- letta/server/rest_api/routers/v1/tools.py +11 -12
- letta/server/rest_api/routers/v1/users.py +1 -1
- letta/server/rest_api/streaming_response.py +13 -5
- letta/server/rest_api/utils.py +8 -25
- letta/server/server.py +11 -4
- letta/server/ws_api/server.py +2 -2
- letta/services/agent_file_manager.py +616 -0
- letta/services/agent_manager.py +133 -46
- letta/services/block_manager.py +38 -17
- letta/services/file_manager.py +106 -21
- letta/services/file_processor/file_processor.py +93 -0
- letta/services/files_agents_manager.py +28 -0
- letta/services/group_manager.py +4 -5
- letta/services/helpers/agent_manager_helper.py +57 -9
- letta/services/identity_manager.py +22 -0
- letta/services/job_manager.py +210 -91
- letta/services/llm_batch_manager.py +9 -6
- letta/services/mcp/stdio_client.py +1 -2
- letta/services/mcp_manager.py +0 -1
- letta/services/message_manager.py +49 -26
- letta/services/passage_manager.py +0 -1
- letta/services/provider_manager.py +1 -1
- letta/services/source_manager.py +114 -5
- letta/services/step_manager.py +36 -4
- letta/services/telemetry_manager.py +9 -2
- letta/services/tool_executor/builtin_tool_executor.py +5 -1
- letta/services/tool_executor/core_tool_executor.py +3 -3
- letta/services/tool_manager.py +95 -20
- letta/services/user_manager.py +4 -12
- letta/settings.py +23 -6
- letta/system.py +1 -1
- letta/utils.py +26 -2
- {letta_nightly-0.8.15.dev20250720104313.dist-info → letta_nightly-0.8.16.dev20250721070720.dist-info}/METADATA +3 -2
- {letta_nightly-0.8.15.dev20250720104313.dist-info → letta_nightly-0.8.16.dev20250721070720.dist-info}/RECORD +99 -94
- {letta_nightly-0.8.15.dev20250720104313.dist-info → letta_nightly-0.8.16.dev20250721070720.dist-info}/LICENSE +0 -0
- {letta_nightly-0.8.15.dev20250720104313.dist-info → letta_nightly-0.8.16.dev20250721070720.dist-info}/WHEEL +0 -0
- {letta_nightly-0.8.15.dev20250720104313.dist-info → letta_nightly-0.8.16.dev20250721070720.dist-info}/entry_points.txt +0 -0
@@ -0,0 +1,616 @@
|
|
1
|
+
from datetime import datetime, timezone
|
2
|
+
from typing import Dict, List
|
3
|
+
|
4
|
+
from letta.errors import AgentFileExportError, AgentFileImportError
|
5
|
+
from letta.log import get_logger
|
6
|
+
from letta.schemas.agent import AgentState, CreateAgent
|
7
|
+
from letta.schemas.agent_file import (
|
8
|
+
AgentFileSchema,
|
9
|
+
AgentSchema,
|
10
|
+
BlockSchema,
|
11
|
+
FileAgentSchema,
|
12
|
+
FileSchema,
|
13
|
+
GroupSchema,
|
14
|
+
ImportResult,
|
15
|
+
MessageSchema,
|
16
|
+
SourceSchema,
|
17
|
+
ToolSchema,
|
18
|
+
)
|
19
|
+
from letta.schemas.block import Block
|
20
|
+
from letta.schemas.file import FileMetadata
|
21
|
+
from letta.schemas.message import Message
|
22
|
+
from letta.schemas.source import Source
|
23
|
+
from letta.schemas.tool import Tool
|
24
|
+
from letta.schemas.user import User
|
25
|
+
from letta.services.agent_manager import AgentManager
|
26
|
+
from letta.services.block_manager import BlockManager
|
27
|
+
from letta.services.file_manager import FileManager
|
28
|
+
from letta.services.file_processor.embedder.base_embedder import BaseEmbedder
|
29
|
+
from letta.services.file_processor.file_processor import FileProcessor
|
30
|
+
from letta.services.file_processor.parser.mistral_parser import MistralFileParser
|
31
|
+
from letta.services.files_agents_manager import FileAgentManager
|
32
|
+
from letta.services.group_manager import GroupManager
|
33
|
+
from letta.services.mcp_manager import MCPManager
|
34
|
+
from letta.services.message_manager import MessageManager
|
35
|
+
from letta.services.source_manager import SourceManager
|
36
|
+
from letta.services.tool_manager import ToolManager
|
37
|
+
from letta.utils import get_latest_alembic_revision
|
38
|
+
|
39
|
+
logger = get_logger(__name__)
|
40
|
+
|
41
|
+
|
42
|
+
class AgentFileManager:
|
43
|
+
"""
|
44
|
+
Manages export and import of agent files between database and AgentFileSchema format.
|
45
|
+
|
46
|
+
Handles:
|
47
|
+
- ID mapping between database IDs and human-readable file IDs
|
48
|
+
- Coordination across multiple entity managers
|
49
|
+
- Transaction safety during imports
|
50
|
+
- Referential integrity validation
|
51
|
+
"""
|
52
|
+
|
53
|
+
def __init__(
|
54
|
+
self,
|
55
|
+
agent_manager: AgentManager,
|
56
|
+
tool_manager: ToolManager,
|
57
|
+
source_manager: SourceManager,
|
58
|
+
block_manager: BlockManager,
|
59
|
+
group_manager: GroupManager,
|
60
|
+
mcp_manager: MCPManager,
|
61
|
+
file_manager: FileManager,
|
62
|
+
file_agent_manager: FileAgentManager,
|
63
|
+
message_manager: MessageManager,
|
64
|
+
embedder: BaseEmbedder,
|
65
|
+
file_parser: MistralFileParser,
|
66
|
+
using_pinecone: bool = False,
|
67
|
+
):
|
68
|
+
self.agent_manager = agent_manager
|
69
|
+
self.tool_manager = tool_manager
|
70
|
+
self.source_manager = source_manager
|
71
|
+
self.block_manager = block_manager
|
72
|
+
self.group_manager = group_manager
|
73
|
+
self.mcp_manager = mcp_manager
|
74
|
+
self.file_manager = file_manager
|
75
|
+
self.file_agent_manager = file_agent_manager
|
76
|
+
self.message_manager = message_manager
|
77
|
+
self.embedder = embedder
|
78
|
+
self.file_parser = file_parser
|
79
|
+
self.using_pinecone = using_pinecone
|
80
|
+
|
81
|
+
# ID mapping state for export
|
82
|
+
self._db_to_file_ids: Dict[str, str] = {}
|
83
|
+
|
84
|
+
# Counters for generating Stripe-style IDs
|
85
|
+
self._id_counters: Dict[str, int] = {
|
86
|
+
AgentSchema.__id_prefix__: 0,
|
87
|
+
GroupSchema.__id_prefix__: 0,
|
88
|
+
BlockSchema.__id_prefix__: 0,
|
89
|
+
FileSchema.__id_prefix__: 0,
|
90
|
+
SourceSchema.__id_prefix__: 0,
|
91
|
+
ToolSchema.__id_prefix__: 0,
|
92
|
+
MessageSchema.__id_prefix__: 0,
|
93
|
+
FileAgentSchema.__id_prefix__: 0,
|
94
|
+
# MCPServerSchema.__id_prefix__: 0,
|
95
|
+
}
|
96
|
+
|
97
|
+
def _reset_state(self):
|
98
|
+
"""Reset internal state for a new operation"""
|
99
|
+
self._db_to_file_ids.clear()
|
100
|
+
for key in self._id_counters:
|
101
|
+
self._id_counters[key] = 0
|
102
|
+
|
103
|
+
def _generate_file_id(self, entity_type: str) -> str:
|
104
|
+
"""Generate a Stripe-style ID for the given entity type"""
|
105
|
+
counter = self._id_counters[entity_type]
|
106
|
+
file_id = f"{entity_type}-{counter}"
|
107
|
+
self._id_counters[entity_type] += 1
|
108
|
+
return file_id
|
109
|
+
|
110
|
+
def _map_db_to_file_id(self, db_id: str, entity_type: str, allow_new: bool = True) -> str:
|
111
|
+
"""Map a database UUID to a file ID, creating if needed (export only)"""
|
112
|
+
if db_id in self._db_to_file_ids:
|
113
|
+
return self._db_to_file_ids[db_id]
|
114
|
+
|
115
|
+
if not allow_new:
|
116
|
+
raise AgentFileExportError(
|
117
|
+
f"Unexpected new {entity_type} ID '{db_id}' encountered during conversion. "
|
118
|
+
f"All IDs should have been mapped during agent processing."
|
119
|
+
)
|
120
|
+
|
121
|
+
file_id = self._generate_file_id(entity_type)
|
122
|
+
self._db_to_file_ids[db_id] = file_id
|
123
|
+
return file_id
|
124
|
+
|
125
|
+
def _extract_unique_tools(self, agent_states: List[AgentState]) -> List:
|
126
|
+
"""Extract unique tools across all agent states by ID"""
|
127
|
+
all_tools = []
|
128
|
+
for agent_state in agent_states:
|
129
|
+
if agent_state.tools:
|
130
|
+
all_tools.extend(agent_state.tools)
|
131
|
+
|
132
|
+
unique_tools = {}
|
133
|
+
for tool in all_tools:
|
134
|
+
unique_tools[tool.id] = tool
|
135
|
+
|
136
|
+
return sorted(unique_tools.values(), key=lambda x: x.name)
|
137
|
+
|
138
|
+
def _extract_unique_blocks(self, agent_states: List[AgentState]) -> List:
|
139
|
+
"""Extract unique blocks across all agent states by ID"""
|
140
|
+
all_blocks = []
|
141
|
+
for agent_state in agent_states:
|
142
|
+
if agent_state.memory and agent_state.memory.blocks:
|
143
|
+
all_blocks.extend(agent_state.memory.blocks)
|
144
|
+
|
145
|
+
unique_blocks = {}
|
146
|
+
for block in all_blocks:
|
147
|
+
unique_blocks[block.id] = block
|
148
|
+
|
149
|
+
return sorted(unique_blocks.values(), key=lambda x: x.label)
|
150
|
+
|
151
|
+
async def _extract_unique_sources_and_files_from_agents(
|
152
|
+
self, agent_states: List[AgentState], actor: User
|
153
|
+
) -> tuple[List[Source], List[FileMetadata]]:
|
154
|
+
"""Extract unique sources and files from agent states using bulk operations"""
|
155
|
+
|
156
|
+
all_source_ids = set()
|
157
|
+
all_file_ids = set()
|
158
|
+
|
159
|
+
for agent_state in agent_states:
|
160
|
+
files_agents = await self.file_agent_manager.list_files_for_agent(
|
161
|
+
agent_id=agent_state.id, actor=actor, is_open_only=False, return_as_blocks=False
|
162
|
+
)
|
163
|
+
for file_agent in files_agents:
|
164
|
+
all_source_ids.add(file_agent.source_id)
|
165
|
+
all_file_ids.add(file_agent.file_id)
|
166
|
+
sources = await self.source_manager.get_sources_by_ids_async(list(all_source_ids), actor)
|
167
|
+
files = await self.file_manager.get_files_by_ids_async(list(all_file_ids), actor, include_content=True)
|
168
|
+
|
169
|
+
return sources, files
|
170
|
+
|
171
|
+
async def _convert_agent_state_to_schema(self, agent_state: AgentState, actor: User) -> AgentSchema:
|
172
|
+
"""Convert AgentState to AgentSchema with ID remapping"""
|
173
|
+
|
174
|
+
agent_file_id = self._map_db_to_file_id(agent_state.id, AgentSchema.__id_prefix__)
|
175
|
+
files_agents = await self.file_agent_manager.list_files_for_agent(
|
176
|
+
agent_id=agent_state.id, actor=actor, is_open_only=False, return_as_blocks=False
|
177
|
+
)
|
178
|
+
agent_schema = await AgentSchema.from_agent_state(
|
179
|
+
agent_state, message_manager=self.message_manager, files_agents=files_agents, actor=actor
|
180
|
+
)
|
181
|
+
agent_schema.id = agent_file_id
|
182
|
+
|
183
|
+
if agent_schema.messages:
|
184
|
+
for message in agent_schema.messages:
|
185
|
+
message_file_id = self._map_db_to_file_id(message.id, MessageSchema.__id_prefix__)
|
186
|
+
message.id = message_file_id
|
187
|
+
message.agent_id = agent_file_id
|
188
|
+
|
189
|
+
if agent_schema.in_context_message_ids:
|
190
|
+
agent_schema.in_context_message_ids = [
|
191
|
+
self._map_db_to_file_id(message_id, MessageSchema.__id_prefix__, allow_new=False)
|
192
|
+
for message_id in agent_schema.in_context_message_ids
|
193
|
+
]
|
194
|
+
|
195
|
+
if agent_schema.tool_ids:
|
196
|
+
agent_schema.tool_ids = [self._map_db_to_file_id(tool_id, ToolSchema.__id_prefix__) for tool_id in agent_schema.tool_ids]
|
197
|
+
|
198
|
+
if agent_schema.source_ids:
|
199
|
+
agent_schema.source_ids = [
|
200
|
+
self._map_db_to_file_id(source_id, SourceSchema.__id_prefix__) for source_id in agent_schema.source_ids
|
201
|
+
]
|
202
|
+
|
203
|
+
if agent_schema.block_ids:
|
204
|
+
agent_schema.block_ids = [self._map_db_to_file_id(block_id, BlockSchema.__id_prefix__) for block_id in agent_schema.block_ids]
|
205
|
+
|
206
|
+
if agent_schema.files_agents:
|
207
|
+
for file_agent in agent_schema.files_agents:
|
208
|
+
file_agent.file_id = self._map_db_to_file_id(file_agent.file_id, FileSchema.__id_prefix__)
|
209
|
+
file_agent.source_id = self._map_db_to_file_id(file_agent.source_id, SourceSchema.__id_prefix__)
|
210
|
+
file_agent.agent_id = agent_file_id
|
211
|
+
|
212
|
+
return agent_schema
|
213
|
+
|
214
|
+
def _convert_tool_to_schema(self, tool) -> ToolSchema:
|
215
|
+
"""Convert Tool to ToolSchema with ID remapping"""
|
216
|
+
tool_file_id = self._map_db_to_file_id(tool.id, ToolSchema.__id_prefix__, allow_new=False)
|
217
|
+
tool_schema = ToolSchema.from_tool(tool)
|
218
|
+
tool_schema.id = tool_file_id
|
219
|
+
return tool_schema
|
220
|
+
|
221
|
+
def _convert_block_to_schema(self, block) -> BlockSchema:
|
222
|
+
"""Convert Block to BlockSchema with ID remapping"""
|
223
|
+
block_file_id = self._map_db_to_file_id(block.id, BlockSchema.__id_prefix__, allow_new=False)
|
224
|
+
block_schema = BlockSchema.from_block(block)
|
225
|
+
block_schema.id = block_file_id
|
226
|
+
return block_schema
|
227
|
+
|
228
|
+
def _convert_source_to_schema(self, source) -> SourceSchema:
|
229
|
+
"""Convert Source to SourceSchema with ID remapping"""
|
230
|
+
source_file_id = self._map_db_to_file_id(source.id, SourceSchema.__id_prefix__, allow_new=False)
|
231
|
+
source_schema = SourceSchema.from_source(source)
|
232
|
+
source_schema.id = source_file_id
|
233
|
+
return source_schema
|
234
|
+
|
235
|
+
def _convert_file_to_schema(self, file_metadata) -> FileSchema:
|
236
|
+
"""Convert FileMetadata to FileSchema with ID remapping"""
|
237
|
+
file_file_id = self._map_db_to_file_id(file_metadata.id, FileSchema.__id_prefix__, allow_new=False)
|
238
|
+
file_schema = FileSchema.from_file_metadata(file_metadata)
|
239
|
+
file_schema.id = file_file_id
|
240
|
+
file_schema.source_id = self._map_db_to_file_id(file_metadata.source_id, SourceSchema.__id_prefix__, allow_new=False)
|
241
|
+
return file_schema
|
242
|
+
|
243
|
+
async def export(self, agent_ids: List[str], actor: User) -> AgentFileSchema:
|
244
|
+
"""
|
245
|
+
Export agents and their related entities to AgentFileSchema format.
|
246
|
+
|
247
|
+
Args:
|
248
|
+
agent_ids: List of agent UUIDs to export
|
249
|
+
|
250
|
+
Returns:
|
251
|
+
AgentFileSchema with all related entities
|
252
|
+
|
253
|
+
Raises:
|
254
|
+
AgentFileExportError: If export fails
|
255
|
+
"""
|
256
|
+
try:
|
257
|
+
self._reset_state()
|
258
|
+
|
259
|
+
agent_states = await self.agent_manager.get_agents_by_ids_async(agent_ids=agent_ids, actor=actor)
|
260
|
+
|
261
|
+
# Validate that all requested agents were found
|
262
|
+
if len(agent_states) != len(agent_ids):
|
263
|
+
found_ids = {agent.id for agent in agent_states}
|
264
|
+
missing_ids = [agent_id for agent_id in agent_ids if agent_id not in found_ids]
|
265
|
+
raise AgentFileExportError(f"The following agent IDs were not found: {missing_ids}")
|
266
|
+
|
267
|
+
# Extract unique entities across all agents
|
268
|
+
tool_set = self._extract_unique_tools(agent_states)
|
269
|
+
block_set = self._extract_unique_blocks(agent_states)
|
270
|
+
|
271
|
+
# Extract sources and files from agent states BEFORE conversion
|
272
|
+
source_set, file_set = await self._extract_unique_sources_and_files_from_agents(agent_states, actor)
|
273
|
+
|
274
|
+
# Convert to schemas with ID remapping
|
275
|
+
agent_schemas = [await self._convert_agent_state_to_schema(agent_state, actor=actor) for agent_state in agent_states]
|
276
|
+
tool_schemas = [self._convert_tool_to_schema(tool) for tool in tool_set]
|
277
|
+
block_schemas = [self._convert_block_to_schema(block) for block in block_set]
|
278
|
+
source_schemas = [self._convert_source_to_schema(source) for source in source_set]
|
279
|
+
file_schemas = [self._convert_file_to_schema(file_metadata) for file_metadata in file_set]
|
280
|
+
|
281
|
+
logger.info(f"Exporting {len(agent_ids)} agents to agent file format")
|
282
|
+
|
283
|
+
# Return AgentFileSchema with converted entities
|
284
|
+
return AgentFileSchema(
|
285
|
+
agents=agent_schemas,
|
286
|
+
groups=[], # TODO: Extract and convert groups
|
287
|
+
blocks=block_schemas,
|
288
|
+
files=file_schemas,
|
289
|
+
sources=source_schemas,
|
290
|
+
tools=tool_schemas,
|
291
|
+
# mcp_servers=[], # TODO: Extract and convert MCP servers
|
292
|
+
metadata={"revision_id": await get_latest_alembic_revision()},
|
293
|
+
created_at=datetime.now(timezone.utc),
|
294
|
+
)
|
295
|
+
|
296
|
+
except Exception as e:
|
297
|
+
logger.error(f"Failed to export agent file: {e}")
|
298
|
+
raise AgentFileExportError(f"Export failed: {e}") from e
|
299
|
+
|
300
|
+
async def import_file(self, schema: AgentFileSchema, actor: User, dry_run: bool = False) -> ImportResult:
|
301
|
+
"""
|
302
|
+
Import AgentFileSchema into the database.
|
303
|
+
|
304
|
+
Args:
|
305
|
+
schema: The agent file schema to import
|
306
|
+
dry_run: If True, validate but don't commit changes
|
307
|
+
|
308
|
+
Returns:
|
309
|
+
ImportResult with success status and details
|
310
|
+
|
311
|
+
Raises:
|
312
|
+
AgentFileImportError: If import fails
|
313
|
+
"""
|
314
|
+
try:
|
315
|
+
self._reset_state()
|
316
|
+
|
317
|
+
if dry_run:
|
318
|
+
logger.info("Starting dry run import validation")
|
319
|
+
else:
|
320
|
+
logger.info("Starting agent file import")
|
321
|
+
|
322
|
+
# Validate schema first
|
323
|
+
self._validate_schema(schema)
|
324
|
+
|
325
|
+
if dry_run:
|
326
|
+
return ImportResult(
|
327
|
+
success=True,
|
328
|
+
message="Dry run validation passed",
|
329
|
+
imported_count=0,
|
330
|
+
)
|
331
|
+
|
332
|
+
# Import in dependency order
|
333
|
+
imported_count = 0
|
334
|
+
file_to_db_ids = {} # Maps file IDs to new database IDs
|
335
|
+
|
336
|
+
# 1. Create tools first (no dependencies)
|
337
|
+
for tool_schema in schema.tools:
|
338
|
+
# Convert ToolSchema back to ToolCreate
|
339
|
+
created_tool = await self.tool_manager.create_or_update_tool_async(
|
340
|
+
pydantic_tool=Tool(**tool_schema.model_dump(exclude={"id"})), actor=actor
|
341
|
+
)
|
342
|
+
file_to_db_ids[tool_schema.id] = created_tool.id
|
343
|
+
imported_count += 1
|
344
|
+
|
345
|
+
# 2. Create blocks (no dependencies)
|
346
|
+
for block_schema in schema.blocks:
|
347
|
+
# Convert BlockSchema back to CreateBlock
|
348
|
+
block = Block(**block_schema.model_dump(exclude={"id"}))
|
349
|
+
created_block = await self.block_manager.create_or_update_block_async(block, actor)
|
350
|
+
file_to_db_ids[block_schema.id] = created_block.id
|
351
|
+
imported_count += 1
|
352
|
+
|
353
|
+
# 3. Create sources (no dependencies)
|
354
|
+
for source_schema in schema.sources:
|
355
|
+
# Convert SourceSchema back to Source
|
356
|
+
source_data = source_schema.model_dump(exclude={"id", "embedding", "embedding_chunk_size"})
|
357
|
+
source = Source(**source_data)
|
358
|
+
created_source = await self.source_manager.create_source(source, actor)
|
359
|
+
file_to_db_ids[source_schema.id] = created_source.id
|
360
|
+
imported_count += 1
|
361
|
+
|
362
|
+
# 4. Create files (depends on sources)
|
363
|
+
for file_schema in schema.files:
|
364
|
+
# Convert FileSchema back to FileMetadata
|
365
|
+
file_data = file_schema.model_dump(exclude={"id", "content"})
|
366
|
+
# Remap source_id from file ID to database ID
|
367
|
+
file_data["source_id"] = file_to_db_ids[file_schema.source_id]
|
368
|
+
file_metadata = FileMetadata(**file_data)
|
369
|
+
created_file = await self.file_manager.create_file(file_metadata, actor, text=file_schema.content)
|
370
|
+
file_to_db_ids[file_schema.id] = created_file.id
|
371
|
+
imported_count += 1
|
372
|
+
|
373
|
+
# 5. Process files for chunking/embedding (depends on files and sources)
|
374
|
+
file_processor = FileProcessor(
|
375
|
+
file_parser=self.file_parser,
|
376
|
+
embedder=self.embedder,
|
377
|
+
actor=actor,
|
378
|
+
using_pinecone=self.using_pinecone,
|
379
|
+
)
|
380
|
+
|
381
|
+
for file_schema in schema.files:
|
382
|
+
if file_schema.content: # Only process files with content
|
383
|
+
file_db_id = file_to_db_ids[file_schema.id]
|
384
|
+
source_db_id = file_to_db_ids[file_schema.source_id]
|
385
|
+
|
386
|
+
# Get the created file metadata
|
387
|
+
file_metadata = await self.file_manager.get_file_by_id(file_db_id, actor)
|
388
|
+
|
389
|
+
# Save the db call of fetching content again
|
390
|
+
file_metadata.content = file_schema.content
|
391
|
+
|
392
|
+
# Process the file for chunking/embedding
|
393
|
+
passages = await file_processor.process_imported_file(file_metadata=file_metadata, source_id=source_db_id)
|
394
|
+
imported_count += len(passages)
|
395
|
+
|
396
|
+
# 6. Create agents with empty message history
|
397
|
+
for agent_schema in schema.agents:
|
398
|
+
# Convert AgentSchema back to CreateAgent, remapping tool/block IDs
|
399
|
+
agent_data = agent_schema.model_dump(exclude={"id", "in_context_message_ids", "messages"})
|
400
|
+
|
401
|
+
# Remap tool_ids from file IDs to database IDs
|
402
|
+
if agent_data.get("tool_ids"):
|
403
|
+
agent_data["tool_ids"] = [file_to_db_ids[file_id] for file_id in agent_data["tool_ids"]]
|
404
|
+
|
405
|
+
# Remap block_ids from file IDs to database IDs
|
406
|
+
if agent_data.get("block_ids"):
|
407
|
+
agent_data["block_ids"] = [file_to_db_ids[file_id] for file_id in agent_data["block_ids"]]
|
408
|
+
|
409
|
+
agent_create = CreateAgent(**agent_data)
|
410
|
+
created_agent = await self.agent_manager.create_agent_async(agent_create, actor, _init_with_no_messages=True)
|
411
|
+
file_to_db_ids[agent_schema.id] = created_agent.id
|
412
|
+
imported_count += 1
|
413
|
+
|
414
|
+
# 7. Create messages and update agent message_ids
|
415
|
+
for agent_schema in schema.agents:
|
416
|
+
agent_db_id = file_to_db_ids[agent_schema.id]
|
417
|
+
message_file_to_db_ids = {}
|
418
|
+
|
419
|
+
# Create messages for this agent
|
420
|
+
messages = []
|
421
|
+
for message_schema in agent_schema.messages:
|
422
|
+
# Convert MessageSchema back to Message, setting agent_id to new DB ID
|
423
|
+
message_data = message_schema.model_dump(exclude={"id"})
|
424
|
+
message_data["agent_id"] = agent_db_id # Remap agent_id to new database ID
|
425
|
+
message_obj = Message(**message_data)
|
426
|
+
messages.append(message_obj)
|
427
|
+
# Map file ID to the generated database ID immediately
|
428
|
+
message_file_to_db_ids[message_schema.id] = message_obj.id
|
429
|
+
|
430
|
+
created_messages = await self.message_manager.create_many_messages_async(pydantic_msgs=messages, actor=actor)
|
431
|
+
imported_count += len(created_messages)
|
432
|
+
|
433
|
+
# Remap in_context_message_ids from file IDs to database IDs
|
434
|
+
in_context_db_ids = [message_file_to_db_ids[message_schema_id] for message_schema_id in agent_schema.in_context_message_ids]
|
435
|
+
|
436
|
+
# Update agent with the correct message_ids
|
437
|
+
await self.agent_manager.update_message_ids_async(agent_id=agent_db_id, message_ids=in_context_db_ids, actor=actor)
|
438
|
+
|
439
|
+
# 8. Create file-agent relationships (depends on agents and files)
|
440
|
+
for agent_schema in schema.agents:
|
441
|
+
if agent_schema.files_agents:
|
442
|
+
agent_db_id = file_to_db_ids[agent_schema.id]
|
443
|
+
|
444
|
+
# Prepare files for bulk attachment
|
445
|
+
files_for_agent = []
|
446
|
+
visible_content_map = {}
|
447
|
+
|
448
|
+
for file_agent_schema in agent_schema.files_agents:
|
449
|
+
file_db_id = file_to_db_ids[file_agent_schema.file_id]
|
450
|
+
file_metadata = await self.file_manager.get_file_by_id(file_db_id, actor)
|
451
|
+
files_for_agent.append(file_metadata)
|
452
|
+
|
453
|
+
if file_agent_schema.visible_content:
|
454
|
+
visible_content_map[file_db_id] = file_agent_schema.visible_content
|
455
|
+
|
456
|
+
# Bulk attach files to agent
|
457
|
+
await self.file_agent_manager.attach_files_bulk(
|
458
|
+
agent_id=agent_db_id, files_metadata=files_for_agent, visible_content_map=visible_content_map, actor=actor
|
459
|
+
)
|
460
|
+
imported_count += len(files_for_agent)
|
461
|
+
|
462
|
+
return ImportResult(
|
463
|
+
success=True,
|
464
|
+
message=f"Import completed successfully. Imported {imported_count} entities.",
|
465
|
+
imported_count=imported_count,
|
466
|
+
id_mappings=file_to_db_ids,
|
467
|
+
)
|
468
|
+
|
469
|
+
except Exception as e:
|
470
|
+
logger.exception(f"Failed to import agent file: {e}")
|
471
|
+
raise AgentFileImportError(f"Import failed: {e}") from e
|
472
|
+
|
473
|
+
def _validate_id_format(self, schema: AgentFileSchema) -> List[str]:
|
474
|
+
"""Validate that all IDs follow the expected format"""
|
475
|
+
errors = []
|
476
|
+
|
477
|
+
# Define entity types and their expected prefixes
|
478
|
+
entity_checks = [
|
479
|
+
(schema.agents, AgentSchema.__id_prefix__),
|
480
|
+
(schema.groups, GroupSchema.__id_prefix__),
|
481
|
+
(schema.blocks, BlockSchema.__id_prefix__),
|
482
|
+
(schema.files, FileSchema.__id_prefix__),
|
483
|
+
(schema.sources, SourceSchema.__id_prefix__),
|
484
|
+
(schema.tools, ToolSchema.__id_prefix__),
|
485
|
+
]
|
486
|
+
|
487
|
+
for entities, expected_prefix in entity_checks:
|
488
|
+
for entity in entities:
|
489
|
+
if not entity.id.startswith(f"{expected_prefix}-"):
|
490
|
+
errors.append(f"Invalid ID format: {entity.id} should start with '{expected_prefix}-'")
|
491
|
+
else:
|
492
|
+
# Check that the suffix is a valid integer
|
493
|
+
try:
|
494
|
+
suffix = entity.id[len(expected_prefix) + 1 :]
|
495
|
+
int(suffix)
|
496
|
+
except ValueError:
|
497
|
+
errors.append(f"Invalid ID format: {entity.id} should have integer suffix")
|
498
|
+
|
499
|
+
# Also check message IDs within agents
|
500
|
+
for agent in schema.agents:
|
501
|
+
for message in agent.messages:
|
502
|
+
if not message.id.startswith(f"{MessageSchema.__id_prefix__}-"):
|
503
|
+
errors.append(f"Invalid message ID format: {message.id} should start with '{MessageSchema.__id_prefix__}-'")
|
504
|
+
else:
|
505
|
+
# Check that the suffix is a valid integer
|
506
|
+
try:
|
507
|
+
suffix = message.id[len(MessageSchema.__id_prefix__) + 1 :]
|
508
|
+
int(suffix)
|
509
|
+
except ValueError:
|
510
|
+
errors.append(f"Invalid message ID format: {message.id} should have integer suffix")
|
511
|
+
|
512
|
+
return errors
|
513
|
+
|
514
|
+
def _validate_duplicate_ids(self, schema: AgentFileSchema) -> List[str]:
|
515
|
+
"""Validate that there are no duplicate IDs within or across entity types"""
|
516
|
+
errors = []
|
517
|
+
all_ids = set()
|
518
|
+
|
519
|
+
# Check each entity type for internal duplicates and collect all IDs
|
520
|
+
entity_collections = [
|
521
|
+
("agents", schema.agents),
|
522
|
+
("groups", schema.groups),
|
523
|
+
("blocks", schema.blocks),
|
524
|
+
("files", schema.files),
|
525
|
+
("sources", schema.sources),
|
526
|
+
("tools", schema.tools),
|
527
|
+
]
|
528
|
+
|
529
|
+
for entity_type, entities in entity_collections:
|
530
|
+
entity_ids = [entity.id for entity in entities]
|
531
|
+
|
532
|
+
# Check for duplicates within this entity type
|
533
|
+
if len(entity_ids) != len(set(entity_ids)):
|
534
|
+
duplicates = [id for id in entity_ids if entity_ids.count(id) > 1]
|
535
|
+
errors.append(f"Duplicate {entity_type} IDs found: {set(duplicates)}")
|
536
|
+
|
537
|
+
# Check for duplicates across all entity types
|
538
|
+
for entity_id in entity_ids:
|
539
|
+
if entity_id in all_ids:
|
540
|
+
errors.append(f"Duplicate ID across entity types: {entity_id}")
|
541
|
+
all_ids.add(entity_id)
|
542
|
+
|
543
|
+
# Also check message IDs within agents
|
544
|
+
for agent in schema.agents:
|
545
|
+
message_ids = [msg.id for msg in agent.messages]
|
546
|
+
|
547
|
+
# Check for duplicates within agent messages
|
548
|
+
if len(message_ids) != len(set(message_ids)):
|
549
|
+
duplicates = [id for id in message_ids if message_ids.count(id) > 1]
|
550
|
+
errors.append(f"Duplicate message IDs in agent {agent.id}: {set(duplicates)}")
|
551
|
+
|
552
|
+
# Check for duplicates across all entity types
|
553
|
+
for message_id in message_ids:
|
554
|
+
if message_id in all_ids:
|
555
|
+
errors.append(f"Duplicate ID across entity types: {message_id}")
|
556
|
+
all_ids.add(message_id)
|
557
|
+
|
558
|
+
return errors
|
559
|
+
|
560
|
+
def _validate_file_source_references(self, schema: AgentFileSchema) -> List[str]:
|
561
|
+
"""Validate that all file source_id references exist"""
|
562
|
+
errors = []
|
563
|
+
source_ids = {source.id for source in schema.sources}
|
564
|
+
|
565
|
+
for file in schema.files:
|
566
|
+
if file.source_id not in source_ids:
|
567
|
+
errors.append(f"File {file.id} references non-existent source {file.source_id}")
|
568
|
+
|
569
|
+
return errors
|
570
|
+
|
571
|
+
def _validate_file_agent_references(self, schema: AgentFileSchema) -> List[str]:
|
572
|
+
"""Validate that all file-agent relationships reference existing entities"""
|
573
|
+
errors = []
|
574
|
+
file_ids = {file.id for file in schema.files}
|
575
|
+
source_ids = {source.id for source in schema.sources}
|
576
|
+
{agent.id for agent in schema.agents}
|
577
|
+
|
578
|
+
for agent in schema.agents:
|
579
|
+
for file_agent in agent.files_agents:
|
580
|
+
if file_agent.file_id not in file_ids:
|
581
|
+
errors.append(f"File-agent relationship references non-existent file {file_agent.file_id}")
|
582
|
+
if file_agent.source_id not in source_ids:
|
583
|
+
errors.append(f"File-agent relationship references non-existent source {file_agent.source_id}")
|
584
|
+
if file_agent.agent_id != agent.id:
|
585
|
+
errors.append(f"File-agent relationship has mismatched agent_id {file_agent.agent_id} vs {agent.id}")
|
586
|
+
|
587
|
+
return errors
|
588
|
+
|
589
|
+
def _validate_schema(self, schema: AgentFileSchema):
|
590
|
+
"""
|
591
|
+
Validate the agent file schema for consistency and referential integrity.
|
592
|
+
|
593
|
+
Args:
|
594
|
+
schema: The schema to validate
|
595
|
+
|
596
|
+
Raises:
|
597
|
+
AgentFileImportError: If validation fails
|
598
|
+
"""
|
599
|
+
errors = []
|
600
|
+
|
601
|
+
# 1. ID Format Validation
|
602
|
+
errors.extend(self._validate_id_format(schema))
|
603
|
+
|
604
|
+
# 2. Duplicate ID Detection
|
605
|
+
errors.extend(self._validate_duplicate_ids(schema))
|
606
|
+
|
607
|
+
# 3. File Source Reference Validation
|
608
|
+
errors.extend(self._validate_file_source_references(schema))
|
609
|
+
|
610
|
+
# 4. File-Agent Reference Validation
|
611
|
+
errors.extend(self._validate_file_agent_references(schema))
|
612
|
+
|
613
|
+
if errors:
|
614
|
+
raise AgentFileImportError(f"Schema validation failed: {'; '.join(errors)}")
|
615
|
+
|
616
|
+
logger.info("Schema validation passed")
|