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.
Files changed (99) hide show
  1. letta/__init__.py +1 -1
  2. letta/agent.py +27 -11
  3. letta/agents/helpers.py +1 -1
  4. letta/agents/letta_agent.py +518 -322
  5. letta/agents/letta_agent_batch.py +1 -2
  6. letta/agents/voice_agent.py +15 -17
  7. letta/client/client.py +3 -3
  8. letta/constants.py +5 -0
  9. letta/embeddings.py +0 -2
  10. letta/errors.py +8 -0
  11. letta/functions/function_sets/base.py +3 -3
  12. letta/functions/helpers.py +2 -3
  13. letta/groups/sleeptime_multi_agent.py +0 -1
  14. letta/helpers/composio_helpers.py +2 -2
  15. letta/helpers/converters.py +1 -1
  16. letta/helpers/pinecone_utils.py +8 -0
  17. letta/helpers/tool_rule_solver.py +13 -18
  18. letta/llm_api/aws_bedrock.py +16 -2
  19. letta/llm_api/cohere.py +1 -1
  20. letta/llm_api/openai_client.py +1 -1
  21. letta/local_llm/grammars/gbnf_grammar_generator.py +1 -1
  22. letta/local_llm/llm_chat_completion_wrappers/zephyr.py +14 -14
  23. letta/local_llm/utils.py +1 -2
  24. letta/orm/agent.py +3 -3
  25. letta/orm/block.py +4 -4
  26. letta/orm/files_agents.py +0 -1
  27. letta/orm/identity.py +2 -0
  28. letta/orm/mcp_server.py +0 -2
  29. letta/orm/message.py +140 -14
  30. letta/orm/organization.py +5 -5
  31. letta/orm/passage.py +4 -4
  32. letta/orm/source.py +1 -1
  33. letta/orm/sqlalchemy_base.py +61 -39
  34. letta/orm/step.py +2 -0
  35. letta/otel/db_pool_monitoring.py +308 -0
  36. letta/otel/metric_registry.py +94 -1
  37. letta/otel/sqlalchemy_instrumentation.py +548 -0
  38. letta/otel/sqlalchemy_instrumentation_integration.py +124 -0
  39. letta/otel/tracing.py +37 -1
  40. letta/schemas/agent.py +0 -3
  41. letta/schemas/agent_file.py +283 -0
  42. letta/schemas/block.py +0 -3
  43. letta/schemas/file.py +28 -26
  44. letta/schemas/letta_message.py +15 -4
  45. letta/schemas/memory.py +1 -1
  46. letta/schemas/message.py +31 -26
  47. letta/schemas/openai/chat_completion_response.py +0 -1
  48. letta/schemas/providers.py +20 -0
  49. letta/schemas/source.py +11 -13
  50. letta/schemas/step.py +12 -0
  51. letta/schemas/tool.py +0 -4
  52. letta/serialize_schemas/marshmallow_agent.py +14 -1
  53. letta/serialize_schemas/marshmallow_block.py +23 -1
  54. letta/serialize_schemas/marshmallow_message.py +1 -3
  55. letta/serialize_schemas/marshmallow_tool.py +23 -1
  56. letta/server/db.py +110 -6
  57. letta/server/rest_api/app.py +85 -73
  58. letta/server/rest_api/routers/v1/agents.py +68 -53
  59. letta/server/rest_api/routers/v1/blocks.py +2 -2
  60. letta/server/rest_api/routers/v1/jobs.py +3 -0
  61. letta/server/rest_api/routers/v1/organizations.py +2 -2
  62. letta/server/rest_api/routers/v1/sources.py +18 -2
  63. letta/server/rest_api/routers/v1/tools.py +11 -12
  64. letta/server/rest_api/routers/v1/users.py +1 -1
  65. letta/server/rest_api/streaming_response.py +13 -5
  66. letta/server/rest_api/utils.py +8 -25
  67. letta/server/server.py +11 -4
  68. letta/server/ws_api/server.py +2 -2
  69. letta/services/agent_file_manager.py +616 -0
  70. letta/services/agent_manager.py +133 -46
  71. letta/services/block_manager.py +38 -17
  72. letta/services/file_manager.py +106 -21
  73. letta/services/file_processor/file_processor.py +93 -0
  74. letta/services/files_agents_manager.py +28 -0
  75. letta/services/group_manager.py +4 -5
  76. letta/services/helpers/agent_manager_helper.py +57 -9
  77. letta/services/identity_manager.py +22 -0
  78. letta/services/job_manager.py +210 -91
  79. letta/services/llm_batch_manager.py +9 -6
  80. letta/services/mcp/stdio_client.py +1 -2
  81. letta/services/mcp_manager.py +0 -1
  82. letta/services/message_manager.py +49 -26
  83. letta/services/passage_manager.py +0 -1
  84. letta/services/provider_manager.py +1 -1
  85. letta/services/source_manager.py +114 -5
  86. letta/services/step_manager.py +36 -4
  87. letta/services/telemetry_manager.py +9 -2
  88. letta/services/tool_executor/builtin_tool_executor.py +5 -1
  89. letta/services/tool_executor/core_tool_executor.py +3 -3
  90. letta/services/tool_manager.py +95 -20
  91. letta/services/user_manager.py +4 -12
  92. letta/settings.py +23 -6
  93. letta/system.py +1 -1
  94. letta/utils.py +26 -2
  95. {letta_nightly-0.8.15.dev20250720104313.dist-info → letta_nightly-0.8.16.dev20250721070720.dist-info}/METADATA +3 -2
  96. {letta_nightly-0.8.15.dev20250720104313.dist-info → letta_nightly-0.8.16.dev20250721070720.dist-info}/RECORD +99 -94
  97. {letta_nightly-0.8.15.dev20250720104313.dist-info → letta_nightly-0.8.16.dev20250721070720.dist-info}/LICENSE +0 -0
  98. {letta_nightly-0.8.15.dev20250720104313.dist-info → letta_nightly-0.8.16.dev20250721070720.dist-info}/WHEEL +0 -0
  99. {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")