letta-nightly 0.7.13.dev20250512104305__py3-none-any.whl → 0.7.14.dev20250513020711__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 (43) hide show
  1. letta/__init__.py +1 -1
  2. letta/agent.py +14 -17
  3. letta/agents/base_agent.py +112 -1
  4. letta/agents/letta_agent.py +35 -55
  5. letta/agents/letta_agent_batch.py +22 -45
  6. letta/agents/voice_agent.py +10 -42
  7. letta/functions/schema_generator.py +7 -3
  8. letta/llm_api/anthropic.py +4 -2
  9. letta/llm_api/openai.py +4 -2
  10. letta/orm/agents_tags.py +5 -2
  11. letta/orm/blocks_agents.py +3 -1
  12. letta/orm/sqlalchemy_base.py +91 -1
  13. letta/schemas/message.py +1 -1
  14. letta/serialize_schemas/marshmallow_agent.py +4 -4
  15. letta/server/db.py +180 -88
  16. letta/server/rest_api/app.py +6 -3
  17. letta/server/rest_api/chat_completions_interface.py +1 -0
  18. letta/server/rest_api/interface.py +54 -16
  19. letta/server/rest_api/routers/v1/sources.py +1 -0
  20. letta/server/server.py +1 -2
  21. letta/services/agent_manager.py +40 -31
  22. letta/services/block_manager.py +61 -34
  23. letta/services/group_manager.py +11 -15
  24. letta/services/identity_manager.py +9 -13
  25. letta/services/job_manager.py +12 -17
  26. letta/services/llm_batch_manager.py +17 -21
  27. letta/services/message_manager.py +53 -31
  28. letta/services/organization_manager.py +7 -14
  29. letta/services/passage_manager.py +6 -10
  30. letta/services/provider_manager.py +5 -9
  31. letta/services/sandbox_config_manager.py +13 -17
  32. letta/services/source_manager.py +13 -17
  33. letta/services/step_manager.py +5 -9
  34. letta/services/tool_manager.py +9 -14
  35. letta/services/user_manager.py +7 -12
  36. letta/settings.py +2 -0
  37. letta/streaming_interface.py +2 -0
  38. letta/utils.py +1 -1
  39. {letta_nightly-0.7.13.dev20250512104305.dist-info → letta_nightly-0.7.14.dev20250513020711.dist-info}/METADATA +2 -1
  40. {letta_nightly-0.7.13.dev20250512104305.dist-info → letta_nightly-0.7.14.dev20250513020711.dist-info}/RECORD +43 -43
  41. {letta_nightly-0.7.13.dev20250512104305.dist-info → letta_nightly-0.7.14.dev20250513020711.dist-info}/LICENSE +0 -0
  42. {letta_nightly-0.7.13.dev20250512104305.dist-info → letta_nightly-0.7.14.dev20250513020711.dist-info}/WHEEL +0 -0
  43. {letta_nightly-0.7.13.dev20250512104305.dist-info → letta_nightly-0.7.14.dev20250513020711.dist-info}/entry_points.txt +0 -0
@@ -472,6 +472,7 @@ class StreamingServerInterface(AgentChunkStreamingInterface):
472
472
  expect_reasoning_content: bool = False,
473
473
  name: Optional[str] = None,
474
474
  message_index: int = 0,
475
+ prev_message_type: Optional[str] = None,
475
476
  ) -> Optional[Union[ReasoningMessage, ToolCallMessage, AssistantMessage]]:
476
477
  """
477
478
  Example data from non-streaming response looks like:
@@ -488,7 +489,6 @@ class StreamingServerInterface(AgentChunkStreamingInterface):
488
489
 
489
490
  choice = chunk.choices[0]
490
491
  message_delta = choice.delta
491
- otid = Message.generate_otid_from_id(message_id, message_index)
492
492
 
493
493
  if (
494
494
  message_delta.content is None
@@ -503,6 +503,8 @@ class StreamingServerInterface(AgentChunkStreamingInterface):
503
503
 
504
504
  # inner thoughts
505
505
  if expect_reasoning_content and message_delta.reasoning_content is not None:
506
+ if prev_message_type and prev_message_type != "reasoning_message":
507
+ message_index += 1
506
508
  processed_chunk = ReasoningMessage(
507
509
  id=message_id,
508
510
  date=message_date,
@@ -510,16 +512,18 @@ class StreamingServerInterface(AgentChunkStreamingInterface):
510
512
  signature=message_delta.reasoning_content_signature,
511
513
  source="reasoner_model" if message_delta.reasoning_content else "non_reasoner_model",
512
514
  name=name,
513
- otid=otid,
515
+ otid=Message.generate_otid_from_id(message_id, message_index),
514
516
  )
515
517
  elif expect_reasoning_content and message_delta.redacted_reasoning_content is not None:
518
+ if prev_message_type and prev_message_type != "hidden_reasoning_message":
519
+ message_index += 1
516
520
  processed_chunk = HiddenReasoningMessage(
517
521
  id=message_id,
518
522
  date=message_date,
519
523
  hidden_reasoning=message_delta.redacted_reasoning_content,
520
524
  state="redacted",
521
525
  name=name,
522
- otid=otid,
526
+ otid=Message.generate_otid_from_id(message_id, message_index),
523
527
  )
524
528
  elif expect_reasoning_content and message_delta.content is not None:
525
529
  # "ignore" content if we expect reasoning content
@@ -537,6 +541,8 @@ class StreamingServerInterface(AgentChunkStreamingInterface):
537
541
  # NOTE: this is hardcoded for our DeepSeek API integration
538
542
  json_reasoning_content = parse_json(self.expect_reasoning_content_buffer)
539
543
 
544
+ if prev_message_type and prev_message_type != "tool_call_message":
545
+ message_index += 1
540
546
  processed_chunk = ToolCallMessage(
541
547
  id=message_id,
542
548
  date=message_date,
@@ -546,7 +552,7 @@ class StreamingServerInterface(AgentChunkStreamingInterface):
546
552
  tool_call_id=None,
547
553
  ),
548
554
  name=name,
549
- otid=otid,
555
+ otid=Message.generate_otid_from_id(message_id, message_index),
550
556
  )
551
557
 
552
558
  except json.JSONDecodeError as e:
@@ -576,12 +582,14 @@ class StreamingServerInterface(AgentChunkStreamingInterface):
576
582
  # print(f"Hiding content delta stream: '{message_delta.content}'")
577
583
  # return None
578
584
  elif message_delta.content is not None:
585
+ if prev_message_type and prev_message_type != "reasoning_message":
586
+ message_index += 1
579
587
  processed_chunk = ReasoningMessage(
580
588
  id=message_id,
581
589
  date=message_date,
582
590
  reasoning=message_delta.content,
583
591
  name=name,
584
- otid=otid,
592
+ otid=Message.generate_otid_from_id(message_id, message_index),
585
593
  )
586
594
 
587
595
  # tool calls
@@ -629,7 +637,15 @@ class StreamingServerInterface(AgentChunkStreamingInterface):
629
637
  # TODO: Assumes consistent state and that prev_content is subset of new_content
630
638
  diff = new_content.replace(prev_content, "", 1)
631
639
  self.current_json_parse_result = parsed_args
632
- processed_chunk = AssistantMessage(id=message_id, date=message_date, content=diff, name=name, otid=otid)
640
+ if prev_message_type and prev_message_type != "assistant_message":
641
+ message_index += 1
642
+ processed_chunk = AssistantMessage(
643
+ id=message_id,
644
+ date=message_date,
645
+ content=diff,
646
+ name=name,
647
+ otid=Message.generate_otid_from_id(message_id, message_index),
648
+ )
633
649
  else:
634
650
  return None
635
651
 
@@ -653,6 +669,8 @@ class StreamingServerInterface(AgentChunkStreamingInterface):
653
669
  processed_chunk = None
654
670
  print("skipping empty chunk...")
655
671
  else:
672
+ if prev_message_type and prev_message_type != "tool_call_message":
673
+ message_index += 1
656
674
  processed_chunk = ToolCallMessage(
657
675
  id=message_id,
658
676
  date=message_date,
@@ -662,7 +680,7 @@ class StreamingServerInterface(AgentChunkStreamingInterface):
662
680
  tool_call_id=tool_call_delta.get("id"),
663
681
  ),
664
682
  name=name,
665
- otid=otid,
683
+ otid=Message.generate_otid_from_id(message_id, message_index),
666
684
  )
667
685
 
668
686
  elif self.inner_thoughts_in_kwargs and tool_call.function:
@@ -694,12 +712,14 @@ class StreamingServerInterface(AgentChunkStreamingInterface):
694
712
 
695
713
  # If we have inner thoughts, we should output them as a chunk
696
714
  if updates_inner_thoughts:
715
+ if prev_message_type and prev_message_type != "reasoning_message":
716
+ message_index += 1
697
717
  processed_chunk = ReasoningMessage(
698
718
  id=message_id,
699
719
  date=message_date,
700
720
  reasoning=updates_inner_thoughts,
701
721
  name=name,
702
- otid=otid,
722
+ otid=Message.generate_otid_from_id(message_id, message_index),
703
723
  )
704
724
  # Additionally inner thoughts may stream back with a chunk of main JSON
705
725
  # In that case, since we can only return a chunk at a time, we should buffer it
@@ -727,6 +747,8 @@ class StreamingServerInterface(AgentChunkStreamingInterface):
727
747
  self.prev_assistant_message_id = self.function_id_buffer
728
748
 
729
749
  else:
750
+ if prev_message_type and prev_message_type != "tool_call_message":
751
+ message_index += 1
730
752
  processed_chunk = ToolCallMessage(
731
753
  id=message_id,
732
754
  date=message_date,
@@ -736,7 +758,7 @@ class StreamingServerInterface(AgentChunkStreamingInterface):
736
758
  tool_call_id=self.function_id_buffer,
737
759
  ),
738
760
  name=name,
739
- otid=otid,
761
+ otid=Message.generate_otid_from_id(message_id, message_index),
740
762
  )
741
763
 
742
764
  # Record what the last function name we flushed was
@@ -789,12 +811,14 @@ class StreamingServerInterface(AgentChunkStreamingInterface):
789
811
  # In this case, we should release the buffer + new data at once
790
812
  combined_chunk = self.function_args_buffer + updates_main_json
791
813
 
814
+ if prev_message_type and prev_message_type != "assistant_message":
815
+ message_index += 1
792
816
  processed_chunk = AssistantMessage(
793
817
  id=message_id,
794
818
  date=message_date,
795
819
  content=combined_chunk,
796
820
  name=name,
797
- otid=otid,
821
+ otid=Message.generate_otid_from_id(message_id, message_index),
798
822
  )
799
823
  # Store the ID of the tool call so allow skipping the corresponding response
800
824
  if self.function_id_buffer:
@@ -818,8 +842,14 @@ class StreamingServerInterface(AgentChunkStreamingInterface):
818
842
  # TODO: Assumes consistent state and that prev_content is subset of new_content
819
843
  diff = new_content.replace(prev_content, "", 1)
820
844
  self.current_json_parse_result = parsed_args
845
+ if prev_message_type and prev_message_type != "assistant_message":
846
+ message_index += 1
821
847
  processed_chunk = AssistantMessage(
822
- id=message_id, date=message_date, content=diff, name=name, otid=otid
848
+ id=message_id,
849
+ date=message_date,
850
+ content=diff,
851
+ name=name,
852
+ otid=Message.generate_otid_from_id(message_id, message_index),
823
853
  )
824
854
  else:
825
855
  return None
@@ -836,6 +866,8 @@ class StreamingServerInterface(AgentChunkStreamingInterface):
836
866
  if self.function_args_buffer:
837
867
  # In this case, we should release the buffer + new data at once
838
868
  combined_chunk = self.function_args_buffer + updates_main_json
869
+ if prev_message_type and prev_message_type != "tool_call_message":
870
+ message_index += 1
839
871
  processed_chunk = ToolCallMessage(
840
872
  id=message_id,
841
873
  date=message_date,
@@ -845,13 +877,15 @@ class StreamingServerInterface(AgentChunkStreamingInterface):
845
877
  tool_call_id=self.function_id_buffer,
846
878
  ),
847
879
  name=name,
848
- otid=otid,
880
+ otid=Message.generate_otid_from_id(message_id, message_index),
849
881
  )
850
882
  # clear buffer
851
883
  self.function_args_buffer = None
852
884
  self.function_id_buffer = None
853
885
  else:
854
886
  # If there's no buffer to clear, just output a new chunk with new data
887
+ if prev_message_type and prev_message_type != "tool_call_message":
888
+ message_index += 1
855
889
  processed_chunk = ToolCallMessage(
856
890
  id=message_id,
857
891
  date=message_date,
@@ -861,7 +895,7 @@ class StreamingServerInterface(AgentChunkStreamingInterface):
861
895
  tool_call_id=self.function_id_buffer,
862
896
  ),
863
897
  name=name,
864
- otid=otid,
898
+ otid=Message.generate_otid_from_id(message_id, message_index),
865
899
  )
866
900
  self.function_id_buffer = None
867
901
 
@@ -982,6 +1016,8 @@ class StreamingServerInterface(AgentChunkStreamingInterface):
982
1016
  processed_chunk = None
983
1017
  print("skipping empty chunk...")
984
1018
  else:
1019
+ if prev_message_type and prev_message_type != "tool_call_message":
1020
+ message_index += 1
985
1021
  processed_chunk = ToolCallMessage(
986
1022
  id=message_id,
987
1023
  date=message_date,
@@ -991,7 +1027,7 @@ class StreamingServerInterface(AgentChunkStreamingInterface):
991
1027
  tool_call_id=tool_call_delta.get("id"),
992
1028
  ),
993
1029
  name=name,
994
- otid=otid,
1030
+ otid=Message.generate_otid_from_id(message_id, message_index),
995
1031
  )
996
1032
 
997
1033
  elif choice.finish_reason is not None:
@@ -1074,6 +1110,7 @@ class StreamingServerInterface(AgentChunkStreamingInterface):
1074
1110
  expect_reasoning_content: bool = False,
1075
1111
  name: Optional[str] = None,
1076
1112
  message_index: int = 0,
1113
+ prev_message_type: Optional[str] = None,
1077
1114
  ):
1078
1115
  """Process a streaming chunk from an OpenAI-compatible server.
1079
1116
 
@@ -1101,6 +1138,7 @@ class StreamingServerInterface(AgentChunkStreamingInterface):
1101
1138
  expect_reasoning_content=expect_reasoning_content,
1102
1139
  name=name,
1103
1140
  message_index=message_index,
1141
+ prev_message_type=prev_message_type,
1104
1142
  )
1105
1143
  if processed_chunk is None:
1106
1144
  return
@@ -1303,7 +1341,7 @@ class StreamingServerInterface(AgentChunkStreamingInterface):
1303
1341
  stdout=msg_obj.tool_returns[0].stdout if msg_obj.tool_returns else None,
1304
1342
  stderr=msg_obj.tool_returns[0].stderr if msg_obj.tool_returns else None,
1305
1343
  name=msg_obj.name,
1306
- otid=Message.generate_otid_from_id(msg_obj.id, chunk_index),
1344
+ otid=Message.generate_otid_from_id(msg_obj.id, chunk_index) if chunk_index is not None else None,
1307
1345
  )
1308
1346
 
1309
1347
  elif msg.startswith("Error: "):
@@ -1319,7 +1357,7 @@ class StreamingServerInterface(AgentChunkStreamingInterface):
1319
1357
  stdout=msg_obj.tool_returns[0].stdout if msg_obj.tool_returns else None,
1320
1358
  stderr=msg_obj.tool_returns[0].stderr if msg_obj.tool_returns else None,
1321
1359
  name=msg_obj.name,
1322
- otid=Message.generate_otid_from_id(msg_obj.id, chunk_index),
1360
+ otid=Message.generate_otid_from_id(msg_obj.id, chunk_index) if chunk_index is not None else None,
1323
1361
  )
1324
1362
 
1325
1363
  else:
@@ -112,6 +112,7 @@ def create_source(
112
112
  name=source_create.name,
113
113
  embedding_config=source_create.embedding_config,
114
114
  description=source_create.description,
115
+ instructions=source_create.instructions,
115
116
  metadata=source_create.metadata,
116
117
  )
117
118
  return server.source_manager.create_source(source=source, actor=actor)
letta/server/server.py CHANGED
@@ -223,7 +223,6 @@ class SyncServer(Server):
223
223
  if init_with_default_org_and_user:
224
224
  self.default_org = self.organization_manager.create_default_organization()
225
225
  self.default_user = self.user_manager.create_default_user()
226
- self.block_manager.add_default_blocks(actor=self.default_user)
227
226
  self.tool_manager.upsert_base_tools(actor=self.default_user)
228
227
 
229
228
  # Add composio keys to the tool sandbox env vars of the org
@@ -1115,7 +1114,7 @@ class SyncServer(Server):
1115
1114
  ),
1116
1115
  CreateBlock(
1117
1116
  label="instructions",
1118
- value=source.description,
1117
+ value=source.instructions,
1119
1118
  ),
1120
1119
  ],
1121
1120
  llm_config=main_agent.llm_config,
@@ -56,6 +56,7 @@ from letta.serialize_schemas import MarshmallowAgentSchema
56
56
  from letta.serialize_schemas.marshmallow_message import SerializedMessageSchema
57
57
  from letta.serialize_schemas.marshmallow_tool import SerializedToolSchema
58
58
  from letta.serialize_schemas.pydantic_agent_schema import AgentSchema
59
+ from letta.server.db import db_registry
59
60
  from letta.services.block_manager import BlockManager
60
61
  from letta.services.helpers.agent_manager_helper import (
61
62
  _apply_filters,
@@ -85,9 +86,6 @@ class AgentManager:
85
86
  """Manager class to handle business logic related to Agents."""
86
87
 
87
88
  def __init__(self):
88
- from letta.server.db import db_context
89
-
90
- self.session_maker = db_context
91
89
  self.block_manager = BlockManager()
92
90
  self.tool_manager = ToolManager()
93
91
  self.source_manager = SourceManager()
@@ -200,7 +198,7 @@ class AgentManager:
200
198
  identity_ids = agent_create.identity_ids or []
201
199
  tag_values = agent_create.tags or []
202
200
 
203
- with self.session_maker() as session:
201
+ with db_registry.session() as session:
204
202
  with session.begin():
205
203
  name_to_id, id_to_name = self._resolve_tools(
206
204
  session,
@@ -356,7 +354,7 @@ class AgentManager:
356
354
  new_idents = set(agent_update.identity_ids or [])
357
355
  new_tags = set(agent_update.tags or [])
358
356
 
359
- with self.session_maker() as session, session.begin():
357
+ with db_registry.session() as session, session.begin():
360
358
 
361
359
  agent: AgentModel = AgentModel.read(db_session=session, identifier=agent_id, actor=actor)
362
360
  agent.updated_at = datetime.now(timezone.utc)
@@ -503,7 +501,7 @@ class AgentManager:
503
501
  Returns:
504
502
  List[PydanticAgentState]: The filtered list of matching agents.
505
503
  """
506
- with self.session_maker() as session:
504
+ with db_registry.session() as session:
507
505
  query = select(AgentModel).distinct(AgentModel.created_at, AgentModel.id)
508
506
  query = AgentModel.apply_access_predicate(query, actor, ["read"], AccessType.ORGANIZATION)
509
507
 
@@ -541,7 +539,7 @@ class AgentManager:
541
539
  Returns:
542
540
  List[PydanticAgentState: The filtered list of matching agents.
543
541
  """
544
- with self.session_maker() as session:
542
+ with db_registry.session() as session:
545
543
  query = select(AgentModel).where(AgentModel.organization_id == actor.organization_id)
546
544
 
547
545
  if match_all:
@@ -569,20 +567,20 @@ class AgentManager:
569
567
  """
570
568
  Get the total count of agents for the given user.
571
569
  """
572
- with self.session_maker() as session:
570
+ with db_registry.session() as session:
573
571
  return AgentModel.size(db_session=session, actor=actor)
574
572
 
575
573
  @enforce_types
576
574
  def get_agent_by_id(self, agent_id: str, actor: PydanticUser) -> PydanticAgentState:
577
575
  """Fetch an agent by its ID."""
578
- with self.session_maker() as session:
576
+ with db_registry.session() as session:
579
577
  agent = AgentModel.read(db_session=session, identifier=agent_id, actor=actor)
580
578
  return agent.to_pydantic()
581
579
 
582
580
  @enforce_types
583
581
  def get_agent_by_name(self, agent_name: str, actor: PydanticUser) -> PydanticAgentState:
584
582
  """Fetch an agent by its ID."""
585
- with self.session_maker() as session:
583
+ with db_registry.session() as session:
586
584
  agent = AgentModel.read(db_session=session, name=agent_name, actor=actor)
587
585
  return agent.to_pydantic()
588
586
 
@@ -599,7 +597,7 @@ class AgentManager:
599
597
  Raises:
600
598
  NoResultFound: If agent doesn't exist
601
599
  """
602
- with self.session_maker() as session:
600
+ with db_registry.session() as session:
603
601
  # Retrieve the agent
604
602
  logger.debug(f"Hard deleting Agent with ID: {agent_id} with actor={actor}")
605
603
  agent = AgentModel.read(db_session=session, identifier=agent_id, actor=actor)
@@ -635,7 +633,7 @@ class AgentManager:
635
633
 
636
634
  @enforce_types
637
635
  def serialize(self, agent_id: str, actor: PydanticUser) -> AgentSchema:
638
- with self.session_maker() as session:
636
+ with db_registry.session() as session:
639
637
  agent = AgentModel.read(db_session=session, identifier=agent_id, actor=actor)
640
638
  schema = MarshmallowAgentSchema(session=session, actor=actor)
641
639
  data = schema.dump(agent)
@@ -665,7 +663,7 @@ class AgentManager:
665
663
 
666
664
  serialized_agent_dict[MarshmallowAgentSchema.FIELD_MESSAGE_IDS] = message_ids
667
665
 
668
- with self.session_maker() as session:
666
+ with db_registry.session() as session:
669
667
  schema = MarshmallowAgentSchema(session=session, actor=actor)
670
668
  agent = schema.load(serialized_agent_dict, session=session)
671
669
 
@@ -728,7 +726,7 @@ class AgentManager:
728
726
  Returns:
729
727
  PydanticAgentState: The updated agent as a Pydantic model.
730
728
  """
731
- with self.session_maker() as session:
729
+ with db_registry.session() as session:
732
730
  # Retrieve the agent
733
731
  agent = AgentModel.read(db_session=session, identifier=agent_id, actor=actor)
734
732
 
@@ -767,7 +765,7 @@ class AgentManager:
767
765
 
768
766
  @enforce_types
769
767
  def list_groups(self, agent_id: str, actor: PydanticUser, manager_type: Optional[str] = None) -> List[PydanticGroup]:
770
- with self.session_maker() as session:
768
+ with db_registry.session() as session:
771
769
  agent = AgentModel.read(db_session=session, identifier=agent_id, actor=actor)
772
770
  if manager_type:
773
771
  return [group.to_pydantic() for group in agent.groups if group.manager_type == manager_type]
@@ -908,7 +906,7 @@ class AgentManager:
908
906
  Returns:
909
907
  PydanticAgentState: The updated agent state with no linked messages.
910
908
  """
911
- with self.session_maker() as session:
909
+ with db_registry.session() as session:
912
910
  # Retrieve the existing agent (will raise NoResultFound if invalid)
913
911
  agent = AgentModel.read(db_session=session, identifier=agent_id, actor=actor)
914
912
 
@@ -985,6 +983,17 @@ class AgentManager:
985
983
  )
986
984
  return agent_state
987
985
 
986
+ @enforce_types
987
+ async def refresh_memory_async(self, agent_state: PydanticAgentState, actor: PydanticUser) -> PydanticAgentState:
988
+ block_ids = [b.id for b in agent_state.memory.blocks]
989
+ if not block_ids:
990
+ return agent_state
991
+
992
+ agent_state.memory.blocks = await self.block_manager.get_all_blocks_by_ids_async(
993
+ block_ids=[b.id for b in agent_state.memory.blocks], actor=actor
994
+ )
995
+ return agent_state
996
+
988
997
  # ======================================================================================================================
989
998
  # Source Management
990
999
  # ======================================================================================================================
@@ -1003,7 +1012,7 @@ class AgentManager:
1003
1012
  IntegrityError: If the source is already attached to the agent
1004
1013
  """
1005
1014
 
1006
- with self.session_maker() as session:
1015
+ with db_registry.session() as session:
1007
1016
  # Verify both agent and source exist and user has permission to access them
1008
1017
  agent = AgentModel.read(db_session=session, identifier=agent_id, actor=actor)
1009
1018
 
@@ -1056,7 +1065,7 @@ class AgentManager:
1056
1065
  Returns:
1057
1066
  List[str]: List of source IDs attached to the agent
1058
1067
  """
1059
- with self.session_maker() as session:
1068
+ with db_registry.session() as session:
1060
1069
  # Verify agent exists and user has permission to access it
1061
1070
  agent = AgentModel.read(db_session=session, identifier=agent_id, actor=actor)
1062
1071
 
@@ -1073,7 +1082,7 @@ class AgentManager:
1073
1082
  source_id: ID of the source to detach
1074
1083
  actor: User performing the action
1075
1084
  """
1076
- with self.session_maker() as session:
1085
+ with db_registry.session() as session:
1077
1086
  # Verify agent exists and user has permission to access it
1078
1087
  agent = AgentModel.read(db_session=session, identifier=agent_id, actor=actor)
1079
1088
 
@@ -1101,7 +1110,7 @@ class AgentManager:
1101
1110
  actor: PydanticUser,
1102
1111
  ) -> PydanticBlock:
1103
1112
  """Gets a block attached to an agent by its label."""
1104
- with self.session_maker() as session:
1113
+ with db_registry.session() as session:
1105
1114
  agent = AgentModel.read(db_session=session, identifier=agent_id, actor=actor)
1106
1115
  for block in agent.core_memory:
1107
1116
  if block.label == block_label:
@@ -1117,7 +1126,7 @@ class AgentManager:
1117
1126
  actor: PydanticUser,
1118
1127
  ) -> PydanticAgentState:
1119
1128
  """Updates which block is assigned to a specific label for an agent."""
1120
- with self.session_maker() as session:
1129
+ with db_registry.session() as session:
1121
1130
  agent = AgentModel.read(db_session=session, identifier=agent_id, actor=actor)
1122
1131
  new_block = BlockModel.read(db_session=session, identifier=new_block_id, actor=actor)
1123
1132
 
@@ -1135,7 +1144,7 @@ class AgentManager:
1135
1144
  @enforce_types
1136
1145
  def attach_block(self, agent_id: str, block_id: str, actor: PydanticUser) -> PydanticAgentState:
1137
1146
  """Attaches a block to an agent."""
1138
- with self.session_maker() as session:
1147
+ with db_registry.session() as session:
1139
1148
  agent = AgentModel.read(db_session=session, identifier=agent_id, actor=actor)
1140
1149
  block = BlockModel.read(db_session=session, identifier=block_id, actor=actor)
1141
1150
 
@@ -1151,7 +1160,7 @@ class AgentManager:
1151
1160
  actor: PydanticUser,
1152
1161
  ) -> PydanticAgentState:
1153
1162
  """Detaches a block from an agent."""
1154
- with self.session_maker() as session:
1163
+ with db_registry.session() as session:
1155
1164
  agent = AgentModel.read(db_session=session, identifier=agent_id, actor=actor)
1156
1165
  original_length = len(agent.core_memory)
1157
1166
 
@@ -1171,7 +1180,7 @@ class AgentManager:
1171
1180
  actor: PydanticUser,
1172
1181
  ) -> PydanticAgentState:
1173
1182
  """Detaches a block with the specified label from an agent."""
1174
- with self.session_maker() as session:
1183
+ with db_registry.session() as session:
1175
1184
  agent = AgentModel.read(db_session=session, identifier=agent_id, actor=actor)
1176
1185
  original_length = len(agent.core_memory)
1177
1186
 
@@ -1215,7 +1224,7 @@ class AgentManager:
1215
1224
  embedded_text = np.array(embedded_text)
1216
1225
  embedded_text = np.pad(embedded_text, (0, MAX_EMBEDDING_DIM - embedded_text.shape[0]), mode="constant").tolist()
1217
1226
 
1218
- with self.session_maker() as session:
1227
+ with db_registry.session() as session:
1219
1228
  # Start with base query for source passages
1220
1229
  source_passages = None
1221
1230
  if not agent_only: # Include source passages
@@ -1389,7 +1398,7 @@ class AgentManager:
1389
1398
  agent_only: bool = False,
1390
1399
  ) -> List[PydanticPassage]:
1391
1400
  """Lists all passages attached to an agent."""
1392
- with self.session_maker() as session:
1401
+ with db_registry.session() as session:
1393
1402
  main_query = self._build_passage_query(
1394
1403
  actor=actor,
1395
1404
  agent_id=agent_id,
@@ -1447,7 +1456,7 @@ class AgentManager:
1447
1456
  agent_only: bool = False,
1448
1457
  ) -> int:
1449
1458
  """Returns the count of passages matching the given criteria."""
1450
- with self.session_maker() as session:
1459
+ with db_registry.session() as session:
1451
1460
  main_query = self._build_passage_query(
1452
1461
  actor=actor,
1453
1462
  agent_id=agent_id,
@@ -1487,7 +1496,7 @@ class AgentManager:
1487
1496
  Returns:
1488
1497
  PydanticAgentState: The updated agent state.
1489
1498
  """
1490
- with self.session_maker() as session:
1499
+ with db_registry.session() as session:
1491
1500
  # Verify the agent exists and user has permission to access it
1492
1501
  agent = AgentModel.read(db_session=session, identifier=agent_id, actor=actor)
1493
1502
 
@@ -1522,7 +1531,7 @@ class AgentManager:
1522
1531
  Returns:
1523
1532
  PydanticAgentState: The updated agent state.
1524
1533
  """
1525
- with self.session_maker() as session:
1534
+ with db_registry.session() as session:
1526
1535
  # Verify the agent exists and user has permission to access it
1527
1536
  agent = AgentModel.read(db_session=session, identifier=agent_id, actor=actor)
1528
1537
 
@@ -1551,7 +1560,7 @@ class AgentManager:
1551
1560
  Returns:
1552
1561
  List[PydanticTool]: List of tools attached to the agent.
1553
1562
  """
1554
- with self.session_maker() as session:
1563
+ with db_registry.session() as session:
1555
1564
  agent = AgentModel.read(db_session=session, identifier=agent_id, actor=actor)
1556
1565
  return [tool.to_pydantic() for tool in agent.tools]
1557
1566
 
@@ -1574,7 +1583,7 @@ class AgentManager:
1574
1583
  Returns:
1575
1584
  List[str]: List of all tags.
1576
1585
  """
1577
- with self.session_maker() as session:
1586
+ with db_registry.session() as session:
1578
1587
  query = (
1579
1588
  session.query(AgentsTags.tag)
1580
1589
  .join(AgentModel, AgentModel.id == AgentsTags.agent_id)