massgen 0.1.3__py3-none-any.whl → 0.1.5__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.

Potentially problematic release.


This version of massgen might be problematic. Click here for more details.

Files changed (90) hide show
  1. massgen/__init__.py +1 -1
  2. massgen/api_params_handler/_chat_completions_api_params_handler.py +4 -0
  3. massgen/api_params_handler/_claude_api_params_handler.py +4 -0
  4. massgen/api_params_handler/_gemini_api_params_handler.py +4 -0
  5. massgen/api_params_handler/_response_api_params_handler.py +4 -0
  6. massgen/backend/base_with_custom_tool_and_mcp.py +25 -5
  7. massgen/backend/docs/permissions_and_context_files.md +2 -2
  8. massgen/backend/response.py +2 -0
  9. massgen/chat_agent.py +340 -20
  10. massgen/cli.py +326 -19
  11. massgen/configs/README.md +92 -41
  12. massgen/configs/memory/gpt5mini_gemini_baseline_research_to_implementation.yaml +94 -0
  13. massgen/configs/memory/gpt5mini_gemini_context_window_management.yaml +187 -0
  14. massgen/configs/memory/gpt5mini_gemini_research_to_implementation.yaml +127 -0
  15. massgen/configs/memory/gpt5mini_high_reasoning_gemini.yaml +107 -0
  16. massgen/configs/memory/single_agent_compression_test.yaml +64 -0
  17. massgen/configs/tools/custom_tools/crawl4ai_example.yaml +55 -0
  18. massgen/configs/tools/custom_tools/multimodal_tools/text_to_file_generation_multi.yaml +61 -0
  19. massgen/configs/tools/custom_tools/multimodal_tools/text_to_file_generation_single.yaml +29 -0
  20. massgen/configs/tools/custom_tools/multimodal_tools/text_to_image_generation_multi.yaml +51 -0
  21. massgen/configs/tools/custom_tools/multimodal_tools/text_to_image_generation_single.yaml +33 -0
  22. massgen/configs/tools/custom_tools/multimodal_tools/text_to_speech_generation_multi.yaml +55 -0
  23. massgen/configs/tools/custom_tools/multimodal_tools/text_to_speech_generation_single.yaml +33 -0
  24. massgen/configs/tools/custom_tools/multimodal_tools/text_to_video_generation_multi.yaml +47 -0
  25. massgen/configs/tools/custom_tools/multimodal_tools/text_to_video_generation_single.yaml +29 -0
  26. massgen/configs/tools/custom_tools/multimodal_tools/understand_audio.yaml +1 -1
  27. massgen/configs/tools/custom_tools/multimodal_tools/understand_file.yaml +1 -1
  28. massgen/configs/tools/custom_tools/multimodal_tools/understand_image.yaml +1 -1
  29. massgen/configs/tools/custom_tools/multimodal_tools/understand_video.yaml +1 -1
  30. massgen/configs/tools/custom_tools/multimodal_tools/youtube_video_analysis.yaml +1 -1
  31. massgen/filesystem_manager/_filesystem_manager.py +1 -0
  32. massgen/filesystem_manager/_path_permission_manager.py +148 -0
  33. massgen/memory/README.md +277 -0
  34. massgen/memory/__init__.py +26 -0
  35. massgen/memory/_base.py +193 -0
  36. massgen/memory/_compression.py +237 -0
  37. massgen/memory/_context_monitor.py +211 -0
  38. massgen/memory/_conversation.py +255 -0
  39. massgen/memory/_fact_extraction_prompts.py +333 -0
  40. massgen/memory/_mem0_adapters.py +257 -0
  41. massgen/memory/_persistent.py +687 -0
  42. massgen/memory/docker-compose.qdrant.yml +36 -0
  43. massgen/memory/docs/DESIGN.md +388 -0
  44. massgen/memory/docs/QUICKSTART.md +409 -0
  45. massgen/memory/docs/SUMMARY.md +319 -0
  46. massgen/memory/docs/agent_use_memory.md +408 -0
  47. massgen/memory/docs/orchestrator_use_memory.md +586 -0
  48. massgen/memory/examples.py +237 -0
  49. massgen/message_templates.py +160 -12
  50. massgen/orchestrator.py +223 -7
  51. massgen/tests/memory/test_agent_compression.py +174 -0
  52. massgen/{configs/tools → tests}/memory/test_context_window_management.py +30 -30
  53. massgen/tests/memory/test_force_compression.py +154 -0
  54. massgen/tests/memory/test_simple_compression.py +147 -0
  55. massgen/tests/test_agent_memory.py +534 -0
  56. massgen/tests/test_binary_file_blocking.py +274 -0
  57. massgen/tests/test_case_studies.md +12 -12
  58. massgen/tests/test_conversation_memory.py +382 -0
  59. massgen/tests/test_multimodal_size_limits.py +407 -0
  60. massgen/tests/test_orchestrator_memory.py +620 -0
  61. massgen/tests/test_persistent_memory.py +435 -0
  62. massgen/token_manager/token_manager.py +6 -0
  63. massgen/tool/_manager.py +7 -2
  64. massgen/tool/_multimodal_tools/image_to_image_generation.py +293 -0
  65. massgen/tool/_multimodal_tools/text_to_file_generation.py +455 -0
  66. massgen/tool/_multimodal_tools/text_to_image_generation.py +222 -0
  67. massgen/tool/_multimodal_tools/text_to_speech_continue_generation.py +226 -0
  68. massgen/tool/_multimodal_tools/text_to_speech_transcription_generation.py +217 -0
  69. massgen/tool/_multimodal_tools/text_to_video_generation.py +223 -0
  70. massgen/tool/_multimodal_tools/understand_audio.py +19 -1
  71. massgen/tool/_multimodal_tools/understand_file.py +6 -1
  72. massgen/tool/_multimodal_tools/understand_image.py +112 -8
  73. massgen/tool/_multimodal_tools/understand_video.py +32 -5
  74. massgen/tool/_web_tools/crawl4ai_tool.py +718 -0
  75. massgen/tool/docs/multimodal_tools.md +589 -0
  76. massgen/tools/__init__.py +8 -0
  77. massgen/tools/_planning_mcp_server.py +520 -0
  78. massgen/tools/planning_dataclasses.py +434 -0
  79. {massgen-0.1.3.dist-info → massgen-0.1.5.dist-info}/METADATA +142 -82
  80. {massgen-0.1.3.dist-info → massgen-0.1.5.dist-info}/RECORD +84 -41
  81. massgen/configs/tools/custom_tools/crawl4ai_mcp_example.yaml +0 -67
  82. massgen/configs/tools/custom_tools/crawl4ai_multi_agent_example.yaml +0 -68
  83. massgen/configs/tools/memory/README.md +0 -199
  84. massgen/configs/tools/memory/gpt5mini_gemini_context_window_management.yaml +0 -131
  85. massgen/configs/tools/memory/gpt5mini_gemini_no_persistent_memory.yaml +0 -133
  86. massgen/configs/tools/multimodal/gpt5mini_gpt5nano_documentation_evolution.yaml +0 -97
  87. {massgen-0.1.3.dist-info → massgen-0.1.5.dist-info}/WHEEL +0 -0
  88. {massgen-0.1.3.dist-info → massgen-0.1.5.dist-info}/entry_points.txt +0 -0
  89. {massgen-0.1.3.dist-info → massgen-0.1.5.dist-info}/licenses/LICENSE +0 -0
  90. {massgen-0.1.3.dist-info → massgen-0.1.5.dist-info}/top_level.txt +0 -0
massgen/__init__.py CHANGED
@@ -68,7 +68,7 @@ from .chat_agent import (
68
68
  from .message_templates import MessageTemplates, get_templates
69
69
  from .orchestrator import Orchestrator, create_orchestrator
70
70
 
71
- __version__ = "0.1.3"
71
+ __version__ = "0.1.5"
72
72
  __author__ = "MassGen Contributors"
73
73
 
74
74
 
@@ -24,6 +24,10 @@ class ChatCompletionsAPIParamsHandler(APIParamsHandlerBase):
24
24
  "allowed_tools",
25
25
  "exclude_tools",
26
26
  "custom_tools", # Custom tools configuration (processed separately)
27
+ "enable_file_generation", # Internal flag for file generation (used in system messages only)
28
+ "enable_image_generation", # Internal flag for image generation (used in system messages only)
29
+ "enable_audio_generation", # Internal flag for audio generation (used in system messages only)
30
+ "enable_video_generation", # Internal flag for video generation (used in system messages only)
27
31
  },
28
32
  )
29
33
 
@@ -24,6 +24,10 @@ class ClaudeAPIParamsHandler(APIParamsHandlerBase):
24
24
  "exclude_tools",
25
25
  "custom_tools", # Custom tools configuration (processed separately)
26
26
  "_has_files_api_files",
27
+ "enable_file_generation", # Internal flag for file generation (used in system messages only)
28
+ "enable_image_generation", # Internal flag for image generation (used in system messages only)
29
+ "enable_audio_generation", # Internal flag for audio generation (used in system messages only)
30
+ "enable_video_generation", # Internal flag for video generation (used in system messages only)
27
31
  },
28
32
  )
29
33
 
@@ -19,6 +19,10 @@ class GeminiAPIParamsHandler(APIParamsHandlerBase):
19
19
  "allowed_tools",
20
20
  "exclude_tools",
21
21
  "custom_tools",
22
+ "enable_file_generation", # Internal flag for file generation (used in system messages only)
23
+ "enable_image_generation", # Internal flag for image generation (used in system messages only)
24
+ "enable_audio_generation", # Internal flag for audio generation (used in system messages only)
25
+ "enable_video_generation", # Internal flag for video generation (used in system messages only)
22
26
  }
23
27
  return set(base) | extra
24
28
 
@@ -24,6 +24,10 @@ class ResponseAPIParamsHandler(APIParamsHandlerBase):
24
24
  "exclude_tools",
25
25
  "custom_tools", # Custom tools configuration (processed separately)
26
26
  "_has_file_search_files", # Internal flag for file search tracking
27
+ "enable_file_generation", # Internal flag for file generation (used in system messages only)
28
+ "enable_image_generation", # Internal flag for image generation (used in system messages only)
29
+ "enable_audio_generation", # Internal flag for audio generation (used in system messages only)
30
+ "enable_video_generation", # Internal flag for video generation (used in system messages only)
27
31
  },
28
32
  )
29
33
 
@@ -284,9 +284,19 @@ class CustomToolAndMCPBackend(LLMBackend):
284
284
 
285
285
  # Register each function with its corresponding values
286
286
  for i, func in enumerate(functions):
287
+ # Inject agent_cwd into preset_args if filesystem_manager is available
288
+ final_preset_args = preset_args_list[i].copy() if preset_args_list[i] else {}
289
+ if self.filesystem_manager and self.filesystem_manager.cwd:
290
+ final_preset_args["agent_cwd"] = self.filesystem_manager.cwd
291
+ logger.info(f"Injecting agent_cwd for {func}: {self.filesystem_manager.cwd}")
292
+ elif self.filesystem_manager:
293
+ logger.warning(f"filesystem_manager exists but cwd is None for {func}")
294
+ else:
295
+ logger.warning(f"No filesystem_manager available for {func}")
296
+
287
297
  # Load the function first if custom name is needed
288
298
  if names[i] and names[i] != func:
289
- # Need to load function and apply custom name
299
+ # Load function to apply custom name
290
300
  if path:
291
301
  loaded_func = self.custom_tool_manager._load_function_from_path(path, func)
292
302
  else:
@@ -296,7 +306,6 @@ class CustomToolAndMCPBackend(LLMBackend):
296
306
  logger.error(f"Could not load function '{func}' from path: {path}")
297
307
  continue
298
308
 
299
- # Apply custom name by modifying __name__ attribute
300
309
  loaded_func.__name__ = names[i]
301
310
 
302
311
  # Register with loaded function (no path needed)
@@ -304,7 +313,7 @@ class CustomToolAndMCPBackend(LLMBackend):
304
313
  path=None,
305
314
  func=loaded_func,
306
315
  category=category,
307
- preset_args=preset_args_list[i],
316
+ preset_args=final_preset_args,
308
317
  description=descriptions[i],
309
318
  )
310
319
  else:
@@ -313,7 +322,7 @@ class CustomToolAndMCPBackend(LLMBackend):
313
322
  path=path,
314
323
  func=func,
315
324
  category=category,
316
- preset_args=preset_args_list[i],
325
+ preset_args=final_preset_args,
317
326
  description=descriptions[i],
318
327
  )
319
328
 
@@ -404,9 +413,19 @@ class CustomToolAndMCPBackend(LLMBackend):
404
413
  """
405
414
  import json
406
415
 
416
+ # Parse arguments
417
+ arguments = json.loads(call["arguments"]) if isinstance(call["arguments"], str) else call["arguments"]
418
+
419
+ # Ensure agent_cwd is always injected if filesystem_manager is available
420
+ # This provides a fallback in case preset_args didn't work during registration
421
+ if self.filesystem_manager and self.filesystem_manager.cwd:
422
+ if "agent_cwd" not in arguments or arguments.get("agent_cwd") is None:
423
+ arguments["agent_cwd"] = self.filesystem_manager.cwd
424
+ logger.info(f"Dynamically injected agent_cwd at execution time: {self.filesystem_manager.cwd}")
425
+
407
426
  tool_request = {
408
427
  "name": call["name"],
409
- "input": json.loads(call["arguments"]) if isinstance(call["arguments"], str) else call["arguments"],
428
+ "input": arguments,
410
429
  }
411
430
 
412
431
  result_text = ""
@@ -1120,6 +1139,7 @@ class CustomToolAndMCPBackend(LLMBackend):
1120
1139
  **kwargs,
1121
1140
  ) -> AsyncGenerator[StreamChunk, None]:
1122
1141
  """Simple passthrough streaming without MCP processing."""
1142
+
1123
1143
  agent_id = kwargs.get("agent_id", None)
1124
1144
  all_params = {**self.config, **kwargs}
1125
1145
  processed_messages = await self._process_upload_files(messages, all_params)
@@ -1067,8 +1067,8 @@ Files delivered:
1067
1067
  - **Multi-Turn Design**: `docs/dev_notes/multi_turn_filesystem_design.md` - Detailed architecture for session persistence and turn-based workflows
1068
1068
  - **MCP Integration**: `docs/dev_notes/gemini_filesystem_mcp_design.md` - How filesystem access works through Model Context Protocol
1069
1069
  - **Context Sharing**: `docs/dev_notes/v0.0.14-context.md` - Original context sharing design
1070
- - **User Context Paths**: `docs/case_studies/user-context-path-support-with-copy-mcp.md` - Case study on adding user-specified paths
1071
- - **Claude Code Workspace**: `docs/case_studies/claude-code-workspace-management.md` - Native filesystem integration patterns
1070
+ - **User Context Paths**: `docs/source/examples/case_studies/user-context-path-support-with-copy-mcp.md` - Case study on adding user-specified paths
1071
+ - **Claude Code Workspace**: `docs/source/examples/case_studies/claude-code-workspace-management.md` - Native filesystem integration patterns
1072
1072
 
1073
1073
  ## Conclusion
1074
1074
 
@@ -57,6 +57,7 @@ class ResponseBackend(CustomToolAndMCPBackend):
57
57
 
58
58
  Wraps parent implementation to ensure File Search cleanup happens after streaming completes.
59
59
  """
60
+
60
61
  try:
61
62
  async for chunk in super().stream_with_tools(messages, tools, **kwargs):
62
63
  yield chunk
@@ -145,6 +146,7 @@ class ResponseBackend(CustomToolAndMCPBackend):
145
146
  **kwargs,
146
147
  ) -> AsyncGenerator[StreamChunk, None]:
147
148
  """Recursively stream MCP responses, executing function calls as needed."""
149
+
148
150
  agent_id = kwargs.get("agent_id")
149
151
 
150
152
  # Build API params for this iteration
massgen/chat_agent.py CHANGED
@@ -14,6 +14,8 @@ from abc import ABC, abstractmethod
14
14
  from typing import Any, AsyncGenerator, Dict, List, Optional
15
15
 
16
16
  from .backend.base import LLMBackend, StreamChunk
17
+ from .logger_config import logger
18
+ from .memory import ConversationMemory, PersistentMemoryBase
17
19
  from .stream_chunk import ChunkType
18
20
  from .utils import CoordinationStage
19
21
 
@@ -26,10 +28,19 @@ class ChatAgent(ABC):
26
28
  providing a unified way to interact with any type of agent system.
27
29
  """
28
30
 
29
- def __init__(self, session_id: Optional[str] = None):
31
+ def __init__(
32
+ self,
33
+ session_id: Optional[str] = None,
34
+ conversation_memory: Optional[ConversationMemory] = None,
35
+ persistent_memory: Optional[PersistentMemoryBase] = None,
36
+ ):
30
37
  self.session_id = session_id or f"chat_session_{uuid.uuid4().hex[:8]}"
31
38
  self.conversation_history: List[Dict[str, Any]] = []
32
39
 
40
+ # Memory components
41
+ self.conversation_memory = conversation_memory
42
+ self.persistent_memory = persistent_memory
43
+
33
44
  @abstractmethod
34
45
  async def chat(
35
46
  self,
@@ -132,6 +143,9 @@ class SingleAgent(ChatAgent):
132
143
  agent_id: Optional[str] = None,
133
144
  system_message: Optional[str] = None,
134
145
  session_id: Optional[str] = None,
146
+ conversation_memory: Optional[ConversationMemory] = None,
147
+ persistent_memory: Optional[PersistentMemoryBase] = None,
148
+ context_monitor: Optional[Any] = None,
135
149
  ):
136
150
  """
137
151
  Initialize single agent.
@@ -141,11 +155,43 @@ class SingleAgent(ChatAgent):
141
155
  agent_id: Optional agent identifier
142
156
  system_message: Optional system message for the agent
143
157
  session_id: Optional session identifier
158
+ conversation_memory: Optional conversation memory instance
159
+ persistent_memory: Optional persistent memory instance
160
+ context_monitor: Optional context window monitor for tracking token usage
144
161
  """
145
- super().__init__(session_id)
162
+ super().__init__(session_id, conversation_memory, persistent_memory)
146
163
  self.backend = backend
147
164
  self.agent_id = agent_id or f"agent_{uuid.uuid4().hex[:8]}"
148
165
  self.system_message = system_message
166
+ self.context_monitor = context_monitor
167
+ self._turn_number = 0
168
+
169
+ # Track orchestrator turn number (for turn-aware memory)
170
+ self._orchestrator_turn = None
171
+
172
+ # Track if compression has occurred (for smart retrieval)
173
+ self._compression_has_occurred = False
174
+
175
+ # Retrieval configuration (defaults, can be overridden from config)
176
+ self._retrieval_limit = 5 # Number of memory facts to retrieve from mem0
177
+ self._retrieval_exclude_recent = True # Don't retrieve before compression (avoid duplicates)
178
+
179
+ # Track previous winning agents for shared memory retrieval
180
+ # Format: [{"agent_id": "agent_b", "turn": 1}, {"agent_id": "agent_a", "turn": 2}]
181
+ self._previous_winners = []
182
+
183
+ # Create context compressor if monitor and conversation_memory exist
184
+ self.context_compressor = None
185
+ if self.context_monitor and self.conversation_memory:
186
+ from .memory._compression import ContextCompressor
187
+ from .token_manager.token_manager import TokenCostCalculator
188
+
189
+ self.context_compressor = ContextCompressor(
190
+ token_calculator=TokenCostCalculator(),
191
+ conversation_memory=self.conversation_memory,
192
+ persistent_memory=self.persistent_memory,
193
+ )
194
+ logger.info(f"🗜️ Context compressor created for {self.agent_id}")
149
195
 
150
196
  # Add system message to history if provided
151
197
  if self.system_message:
@@ -174,6 +220,11 @@ class SingleAgent(ChatAgent):
174
220
  assistant_response = ""
175
221
  tool_calls = []
176
222
  complete_message = None
223
+ messages_to_record = []
224
+
225
+ # Accumulate all chunks for complete memory recording
226
+ reasoning_chunks = [] # Accumulate reasoning content
227
+ reasoning_summaries = [] # Accumulate reasoning summaries
177
228
 
178
229
  try:
179
230
  async for chunk in backend_stream:
@@ -185,6 +236,16 @@ class SingleAgent(ChatAgent):
185
236
  chunk_tool_calls = getattr(chunk, "tool_calls", []) or []
186
237
  tool_calls.extend(chunk_tool_calls)
187
238
  yield chunk
239
+ elif chunk_type == "reasoning":
240
+ # Accumulate reasoning chunks for memory
241
+ if hasattr(chunk, "content") and chunk.content:
242
+ reasoning_chunks.append(chunk.content)
243
+ yield chunk
244
+ elif chunk_type == "reasoning_summary":
245
+ # Accumulate reasoning summaries
246
+ if hasattr(chunk, "content") and chunk.content:
247
+ reasoning_summaries.append(chunk.content)
248
+ yield chunk
188
249
  elif chunk_type == "complete_message":
189
250
  # Backend provided the complete message structure
190
251
  complete_message = chunk.complete_message
@@ -207,24 +268,161 @@ class SingleAgent(ChatAgent):
207
268
  yield StreamChunk(type="tool_calls", tool_calls=response_tool_calls)
208
269
  # Complete response is for internal use - don't yield it
209
270
  elif chunk_type == "done":
210
- # Add complete response to history
271
+ # Debug: Log what we have before assembling
272
+ logger.debug(
273
+ f"🔍 [done] complete_message type: {type(complete_message)}, has_output: {isinstance(complete_message, dict) and 'output' in complete_message if complete_message else False}",
274
+ )
275
+ logger.debug(f"🔍 [done] assistant_response length: {len(assistant_response)}, reasoning: {len(reasoning_chunks)}, summaries: {len(reasoning_summaries)}")
276
+
277
+ # Assemble complete memory from all accumulated chunks
278
+ messages_to_record = []
279
+
280
+ # 1. Add reasoning if present (full context for memory)
281
+ if reasoning_chunks:
282
+ combined_reasoning = "\n".join(reasoning_chunks)
283
+ messages_to_record.append(
284
+ {
285
+ "role": "assistant",
286
+ "content": f"[Reasoning]\n{combined_reasoning}",
287
+ },
288
+ )
289
+
290
+ # 2. Add reasoning summaries if present
291
+ if reasoning_summaries:
292
+ combined_summary = "\n".join(reasoning_summaries)
293
+ messages_to_record.append(
294
+ {
295
+ "role": "assistant",
296
+ "content": f"[Reasoning Summary]\n{combined_summary}",
297
+ },
298
+ )
299
+
300
+ # 3. Add final text response (MCP tools not included - they're implementation details)
211
301
  if complete_message:
212
302
  # For Responses API: complete_message is the response object with 'output' array
213
- # Each item in output should be added to conversation history individually
214
303
  if isinstance(complete_message, dict) and "output" in complete_message:
304
+ # Store raw output for orchestrator (needs full format)
215
305
  self.conversation_history.extend(complete_message["output"])
306
+
307
+ # Debug: Log what's in the output array
308
+ logger.debug(f"🔍 [done] complete_message['output'] has {len(complete_message['output'])} items")
309
+ for i, item in enumerate(complete_message["output"][:3]): # Show first 3
310
+ item_type = item.get("type") if isinstance(item, dict) else type(item).__name__
311
+ logger.debug(f" [{i}] type={item_type}")
312
+
313
+ # Extract text from output items
314
+ for output_item in complete_message["output"]:
315
+ if not isinstance(output_item, dict):
316
+ continue
317
+
318
+ output_type = output_item.get("type")
319
+
320
+ # Skip function_call (workflow tools - not conversation content)
321
+ if output_type == "function_call":
322
+ continue
323
+
324
+ # Extract text content from various formats
325
+ if output_type == "output_text":
326
+ # Responses API format
327
+ text_content = output_item.get("text", "")
328
+ elif output_type == "message":
329
+ # Standard message format
330
+ text_content = output_item.get("content", "")
331
+ elif output_type == "reasoning":
332
+ # Reasoning chunks are already captured above, skip duplicate
333
+ continue
334
+ else:
335
+ # Unknown type - try to get content/text
336
+ text_content = output_item.get("content") or output_item.get("text", "")
337
+ logger.debug(f" ⚠️ Unknown output type '{output_type}', extracted: {bool(text_content)}")
338
+
339
+ if text_content:
340
+ logger.debug(f" ✅ Extracted text ({len(text_content)} chars) from type={output_type}")
341
+ messages_to_record.append(
342
+ {
343
+ "role": "assistant",
344
+ "content": text_content,
345
+ },
346
+ )
347
+ else:
348
+ logger.debug(f" ⚠️ No text content found in type={output_type}")
216
349
  else:
217
350
  # Fallback if it's already in message format
218
351
  self.conversation_history.append(complete_message)
219
- elif assistant_response.strip() or tool_calls:
220
- # Fallback for legacy backends
352
+ if isinstance(complete_message, dict) and complete_message.get("content"):
353
+ messages_to_record.append(complete_message)
354
+ elif assistant_response.strip():
355
+ # Fallback for legacy backends - use accumulated text
221
356
  message_data = {
222
357
  "role": "assistant",
223
358
  "content": assistant_response.strip(),
224
359
  }
225
- if tool_calls:
226
- message_data["tool_calls"] = tool_calls
227
360
  self.conversation_history.append(message_data)
361
+ messages_to_record.append(message_data)
362
+
363
+ # Record to memories
364
+ logger.debug(f"📋 [done chunk] messages_to_record has {len(messages_to_record)} message(s)")
365
+
366
+ if messages_to_record:
367
+ logger.debug(f"✅ Will record {len(messages_to_record)} message(s) to memory")
368
+ # Add to conversation memory (use formatted messages, not raw output)
369
+ if self.conversation_memory:
370
+ try:
371
+ await self.conversation_memory.add(messages_to_record)
372
+ logger.debug(f"📝 Added {len(messages_to_record)} message(s) to conversation memory")
373
+ except Exception as e:
374
+ # Log but don't fail if memory add fails
375
+ logger.warning(f"⚠️ Failed to add response to conversation memory: {e}")
376
+ # Record to persistent memory with turn metadata
377
+ if self.persistent_memory:
378
+ try:
379
+ # Include turn number in metadata for temporal filtering
380
+ logger.debug(f"📝 Recording {len(messages_to_record)} messages to persistent memory (turn {self._orchestrator_turn})")
381
+ await self.persistent_memory.record(
382
+ messages_to_record,
383
+ metadata={"turn": self._orchestrator_turn} if self._orchestrator_turn else None,
384
+ )
385
+ logger.debug("✅ Successfully recorded to persistent memory")
386
+ except NotImplementedError:
387
+ # Memory backend doesn't support record
388
+ logger.warning("⚠️ Persistent memory doesn't support record()")
389
+ except Exception as e:
390
+ # Log but don't fail if memory record fails
391
+ logger.warning(f"⚠️ Failed to record to persistent memory: {e}")
392
+ else:
393
+ logger.warning("⚠️ [done chunk] messages_to_record is EMPTY - nothing to record!")
394
+
395
+ # Log context usage after response (if monitor enabled)
396
+ if self.context_monitor:
397
+ # Use conversation history for accurate token count
398
+ current_history = self.conversation_history if not self.conversation_memory else await self.conversation_memory.get_messages()
399
+ usage_info = self.context_monitor.log_context_usage(current_history, turn_number=self._turn_number)
400
+
401
+ # Compress if needed
402
+ if self.context_compressor and usage_info.get("should_compress"):
403
+ logger.info(
404
+ f"🔄 Attempting compression for {self.agent_id} " f"({usage_info['current_tokens']:,} → {usage_info['target_tokens']:,} tokens)",
405
+ )
406
+ compression_stats = await self.context_compressor.compress_if_needed(
407
+ messages=current_history,
408
+ current_tokens=usage_info["current_tokens"],
409
+ target_tokens=usage_info["target_tokens"],
410
+ should_compress=True,
411
+ )
412
+
413
+ # Update conversation_history if compression occurred
414
+ if compression_stats and self.conversation_memory:
415
+ # Reload from conversation memory (it was updated by compressor)
416
+ self.conversation_history = await self.conversation_memory.get_messages()
417
+ # Mark that compression has occurred
418
+ self._compression_has_occurred = True
419
+ logger.info(
420
+ f"✅ Conversation history updated after compression: " f"{len(self.conversation_history)} messages",
421
+ )
422
+ elif usage_info.get("should_compress") and not self.context_compressor:
423
+ logger.warning(
424
+ f"⚠️ Should compress but compressor not available " f"(monitor={self.context_monitor is not None}, " f"conv_mem={self.conversation_memory is not None})",
425
+ )
228
426
  yield chunk
229
427
  else:
230
428
  yield chunk
@@ -242,12 +440,28 @@ class SingleAgent(ChatAgent):
242
440
  reset_chat: bool = False,
243
441
  clear_history: bool = False,
244
442
  current_stage: CoordinationStage = None,
443
+ orchestrator_turn: Optional[int] = None,
444
+ previous_winners: Optional[List[Dict[str, Any]]] = None,
245
445
  ) -> AsyncGenerator[StreamChunk, None]:
246
- # print("Agent: ", self.agent_id)
247
- # for message in messages:
248
- # print(f"Message: {message}\n")
249
- # print("Messages End. \n")
250
- """Process messages through single backend with tool support."""
446
+ """
447
+ Process messages through single backend with tool support.
448
+
449
+ Args:
450
+ orchestrator_turn: Current orchestrator turn number (for turn-aware memory)
451
+ previous_winners: List of previous winning agents with turns
452
+ Format: [{"agent_id": "agent_b", "turn": 1}, ...]
453
+ """
454
+ # Update orchestrator turn if provided
455
+ if orchestrator_turn is not None:
456
+ logger.debug(f"🔍 [chat] Setting orchestrator_turn={orchestrator_turn} for {self.agent_id}")
457
+ self._orchestrator_turn = orchestrator_turn
458
+
459
+ # Update previous winners if provided
460
+ if previous_winners is not None:
461
+ logger.debug(f"🔍 [chat] Setting previous_winners={previous_winners} for {self.agent_id}")
462
+ self._previous_winners = previous_winners
463
+ else:
464
+ logger.debug(f"🔍 [chat] No previous_winners provided to {self.agent_id} (current: {self._previous_winners})")
251
465
  if clear_history:
252
466
  # Clear history but keep system message if it exists
253
467
  system_messages = [msg for msg in self.conversation_history if msg.get("role") == "system"]
@@ -255,28 +469,121 @@ class SingleAgent(ChatAgent):
255
469
  # Clear backend history while maintaining session
256
470
  if self.backend.is_stateful():
257
471
  await self.backend.clear_history()
472
+ # Clear conversation memory if available
473
+ if self.conversation_memory:
474
+ await self.conversation_memory.clear()
258
475
 
259
476
  if reset_chat:
477
+ # Skip pre-restart recording - messages are already recorded via done chunks
478
+ # Pre-restart would duplicate content and include orchestrator system prompts (noise)
479
+ # The conversation_memory contains:
480
+ # 1. User messages - will be in new conversation after reset
481
+ # 2. Agent responses - already recorded to persistent_memory via done chunks
482
+ # 3. System messages - orchestrator prompts, don't want in long-term memory
483
+ logger.debug(f"🔄 Resetting chat for {self.agent_id} (skipping pre-restart recording - already captured via done chunks)")
484
+
260
485
  # Reset conversation history to the provided messages
261
486
  self.conversation_history = messages.copy()
262
487
  # Reset backend state completely
263
488
  if self.backend.is_stateful():
264
489
  await self.backend.reset_state()
490
+ # Reset conversation memory
491
+ if self.conversation_memory:
492
+ await self.conversation_memory.clear()
493
+ await self.conversation_memory.add(messages)
265
494
  backend_messages = self.conversation_history.copy()
266
495
  else:
267
496
  # Regular conversation - append new messages to agent's history
268
497
  self.conversation_history.extend(messages)
269
- # Handle stateful vs stateless backends differently
270
- if self.backend.is_stateful():
271
- # Stateful: only send new messages, backend maintains context
272
- backend_messages = messages.copy()
273
- else:
274
- # Stateless: send full conversation history
275
- backend_messages = self.conversation_history.copy()
498
+ # Add to conversation memory
499
+ if self.conversation_memory:
500
+ try:
501
+ await self.conversation_memory.add(messages)
502
+ except Exception as e:
503
+ # Log but don't fail if memory add fails
504
+ logger.warning(f"Failed to add messages to conversation memory: {e}")
505
+ backend_messages = self.conversation_history.copy()
506
+
507
+ # Retrieve relevant persistent memories if available
508
+ # ALWAYS retrieve on reset_chat (to restore recent context after restart)
509
+ # Otherwise, only retrieve if compression has occurred (to avoid duplicating recent context)
510
+ memory_context = ""
511
+ should_retrieve = self.persistent_memory and (reset_chat or self._compression_has_occurred or not self._retrieval_exclude_recent) # Always retrieve on reset to restore context
512
+
513
+ if should_retrieve:
514
+ try:
515
+ # Log retrieval reason and scope
516
+ if reset_chat:
517
+ logger.info(
518
+ f"🔄 Retrieving memories after reset for {self.agent_id} " f"(restoring recent context + {len(self._previous_winners) if self._previous_winners else 0} winner(s))...",
519
+ )
520
+ elif self._previous_winners:
521
+ logger.info(
522
+ f"🔍 Retrieving memories for {self.agent_id} + {len(self._previous_winners)} previous winner(s) " f"(limit={self._retrieval_limit}/agent)...",
523
+ )
524
+ logger.debug(f" Previous winners: {self._previous_winners}")
525
+ else:
526
+ logger.info(
527
+ f"🔍 Retrieving memories for {self.agent_id} " f"(limit={self._retrieval_limit}, compressed={self._compression_has_occurred})...",
528
+ )
529
+
530
+ memory_context = await self.persistent_memory.retrieve(
531
+ messages,
532
+ limit=self._retrieval_limit,
533
+ previous_winners=self._previous_winners if self._previous_winners else None,
534
+ )
535
+
536
+ if memory_context:
537
+ memory_lines = memory_context.strip().split("\n")
538
+ logger.info(
539
+ f"💭 Retrieved {len(memory_lines)} memory fact(s) from mem0",
540
+ )
541
+ # Show preview at INFO level (truncate to first 300 chars for readability)
542
+ preview = memory_context[:300] + "..." if len(memory_context) > 300 else memory_context
543
+ logger.info(f" 📝 Preview:\n{preview}")
544
+ else:
545
+ logger.info(" ℹ️ No relevant memories found")
546
+ except NotImplementedError:
547
+ logger.debug(" Persistent memory doesn't support retrieval")
548
+ except Exception as e:
549
+ logger.warning(f"⚠️ Failed to retrieve from persistent memory: {e}")
550
+ elif self.persistent_memory and self._retrieval_exclude_recent:
551
+ logger.debug(
552
+ f"⏭️ Skipping retrieval for {self.agent_id} " f"(no compression yet, all context in conversation_memory)",
553
+ )
554
+
555
+ # Handle stateful vs stateless backends differently
556
+ if self.backend.is_stateful():
557
+ # Stateful: only send new messages, backend maintains context
558
+ backend_messages = messages.copy()
559
+ # Inject memory context before user messages if available
560
+ if memory_context:
561
+ memory_msg = {
562
+ "role": "system",
563
+ "content": f"Relevant memories:\n{memory_context}",
564
+ }
565
+ backend_messages.insert(0, memory_msg)
566
+ else:
567
+ # Stateless: send full conversation history
568
+ backend_messages = self.conversation_history.copy()
569
+ # Inject memory context after system message but before conversation
570
+ if memory_context:
571
+ memory_msg = {
572
+ "role": "system",
573
+ "content": f"Relevant memories:\n{memory_context}",
574
+ }
575
+ # Insert after existing system messages
576
+ system_count = sum(1 for msg in backend_messages if msg.get("role") == "system")
577
+ backend_messages.insert(system_count, memory_msg)
276
578
 
277
579
  if current_stage:
278
580
  self.backend.set_stage(current_stage)
279
581
 
582
+ # Log context usage before processing (if monitor enabled)
583
+ self._turn_number += 1
584
+ if self.context_monitor:
585
+ self.context_monitor.log_context_usage(backend_messages, turn_number=self._turn_number)
586
+
280
587
  # Create backend stream and process it
281
588
  backend_stream = self.backend.stream_with_tools(
282
589
  messages=backend_messages,
@@ -311,6 +618,10 @@ class SingleAgent(ChatAgent):
311
618
  if self.backend.is_stateful():
312
619
  await self.backend.reset_state()
313
620
 
621
+ # Clear conversation memory (not persistent memory)
622
+ if self.conversation_memory:
623
+ await self.conversation_memory.clear()
624
+
314
625
  # Re-add system message if it exists
315
626
  if self.system_message:
316
627
  self.conversation_history.append({"role": "system", "content": self.system_message})
@@ -356,6 +667,9 @@ class ConfigurableAgent(SingleAgent):
356
667
  config, # AgentConfig - avoid circular import
357
668
  backend: LLMBackend,
358
669
  session_id: Optional[str] = None,
670
+ conversation_memory: Optional[ConversationMemory] = None,
671
+ persistent_memory: Optional[PersistentMemoryBase] = None,
672
+ context_monitor: Optional[Any] = None,
359
673
  ):
360
674
  """
361
675
  Initialize configurable agent.
@@ -364,6 +678,9 @@ class ConfigurableAgent(SingleAgent):
364
678
  config: AgentConfig with all settings
365
679
  backend: LLM backend
366
680
  session_id: Optional session identifier
681
+ conversation_memory: Optional conversation memory instance
682
+ persistent_memory: Optional persistent memory instance
683
+ context_monitor: Optional context window monitor for tracking token usage
367
684
  """
368
685
  # Extract system message without triggering deprecation warning
369
686
  system_message = None
@@ -375,6 +692,9 @@ class ConfigurableAgent(SingleAgent):
375
692
  agent_id=config.agent_id,
376
693
  system_message=system_message,
377
694
  session_id=session_id,
695
+ conversation_memory=conversation_memory,
696
+ persistent_memory=persistent_memory,
697
+ context_monitor=context_monitor,
378
698
  )
379
699
  self.config = config
380
700