massgen 0.1.4__py3-none-any.whl → 0.1.5__py3-none-any.whl

massgen/memory/examples.py

@@ -0,0 +1,237 @@
+# -*- coding: utf-8 -*-
+"""
+Usage examples for MassGen memory system.
+
+These examples demonstrate how to use conversation and persistent memory
+in your MassGen agents.
+"""
+
+import asyncio
+
+
+async def example_conversation_memory():
+    """Example: Using ConversationMemory for short-term dialogue."""
+    from massgen.memory import ConversationMemory
+
+    print("=" * 60)
+    print("Example 1: Conversation Memory (Short-term)")
+    print("=" * 60)
+
+    memory = ConversationMemory()
+
+    # Simulate a conversation
+    conversation = [
+        {"role": "user", "content": "Hello! My name is Alice."},
+        {"role": "assistant", "content": "Hi Alice! How can I help you today?"},
+        {"role": "user", "content": "I'm interested in learning about Python."},
+        {
+            "role": "assistant",
+            "content": "Great! Python is a versatile programming language...",
+        },
+    ]
+
+    # Add messages one by one
+    for msg in conversation:
+        await memory.add(msg)
+        print(f"Added: {msg['role']} - {msg['content'][:50]}...")
+
+    # Get all messages
+    await memory.get_messages()
+    print(f"\nTotal messages: {await memory.size()}")
+
+    # Get last message
+    last = await memory.get_last_message()
+    print(f"Last message role: {last['role']}")
+
+    # Filter by role
+    user_messages = await memory.get_messages_by_role("user")
+    print(f"User messages: {len(user_messages)}")
+
+    # Truncate to keep only recent messages
+    await memory.truncate_to_size(2)
+    print(f"After truncation: {await memory.size()} messages")
+
+    # Save and restore state
+    state = memory.state_dict()
+    print(f"\nState saved: {len(state['messages'])} messages")
+
+    # Create new memory from state
+    restored_memory = ConversationMemory()
+    restored_memory.load_state_dict(state)
+    print(f"State restored: {await restored_memory.size()} messages")
+
+    print("\n✅ Conversation memory example completed!\n")
+
+
+async def example_persistent_memory():
+    """
+    Example: Using PersistentMemory for long-term storage.
+
+    Note: This requires mem0 to be installed and proper backends configured.
+    This is a conceptual example - adjust backends as needed.
+    """
+    print("=" * 60)
+    print("Example 2: Persistent Memory (Long-term)")
+    print("=" * 60)
+
+    # NOTE: This is a conceptual example
+    # In practice, you need to provide actual MassGen backends
+    print("\n⚠️  This example requires actual LLM and embedding backends.")
+    print("     Uncomment and configure backends to run this example.\n")
+
+    # Conceptual usage:
+    """
+    from massgen.memory import PersistentMemory
+    from massgen.backend import OpenAIBackend  # Or your backend
+
+    # Initialize backends
+    llm_backend = OpenAIBackend(model="gpt-4")
+    embedding_backend = OpenAIBackend(model="text-embedding-3-small")
+
+    # Create persistent memory
+    memory = PersistentMemory(
+        agent_name="learning_assistant",
+        user_name="alice",
+        llm_backend=llm_backend,
+        embedding_backend=embedding_backend,
+        on_disk=True
+    )
+
+    # Record a conversation
+    await memory.record([
+        {"role": "user", "content": "I love Python programming"},
+        {"role": "assistant", "content": "That's great! Python is very versatile."}
+    ])
+    print("✓ Recorded conversation to long-term memory")
+
+    # Retrieve relevant memories
+    query = "What programming languages does the user like?"
+    relevant = await memory.retrieve(query)
+    print(f"Retrieved: {relevant}")
+
+    # Agent-controlled memory saving
+    result = await memory.save_to_memory(
+        thinking="User expressed interest in a topic",
+        content=["User likes Python", "User is a beginner"]
+    )
+    print(f"Save result: {result}")
+
+    # Agent-controlled memory recall
+    result = await memory.recall_from_memory(
+        keywords=["programming", "Python"],
+        limit=3
+    )
+    print(f"Recalled {result['count']} memories")
+    for mem in result['memories']:
+        print(f"  - {mem}")
+    """
+
+    print("✅ Persistent memory example completed!\n")
+
+
+async def example_combined_usage():
+    """Example: Using both memory types together."""
+    print("=" * 60)
+    print("Example 3: Combined Memory Usage")
+    print("=" * 60)
+
+    from massgen.memory import ConversationMemory
+
+    # Short-term memory for active conversation
+    short_term = ConversationMemory()
+
+    # Simulate ongoing conversation
+    messages = [
+        {"role": "system", "content": "You are a helpful assistant."},
+        {"role": "user", "content": "What's the weather like?"},
+        {"role": "assistant", "content": "I can help you check the weather!"},
+    ]
+
+    for msg in messages:
+        await short_term.add(msg)
+
+    print(f"Short-term memory: {await short_term.size()} messages")
+
+    # In a real agent, you would:
+    # 1. Retrieve relevant long-term memories based on current message
+    # 2. Inject them into the conversation context
+    # 3. Generate response
+    # 4. Add response to short-term memory
+    # 5. Optionally save important parts to long-term memory
+
+    print("\n💡 In production, this would be integrated with:")
+    print("   - LLM backend for generating responses")
+    print("   - Persistent memory for cross-session knowledge")
+    print("   - Tool system for agent-controlled memory")
+
+    print("\n✅ Combined usage example completed!\n")
+
+
+async def example_memory_management():
+    """Example: Memory management best practices."""
+    print("=" * 60)
+    print("Example 4: Memory Management")
+    print("=" * 60)
+
+    from massgen.memory import ConversationMemory
+
+    memory = ConversationMemory()
+
+    # Add many messages to simulate long conversation
+    for i in range(100):
+        await memory.add(
+            {
+                "role": "user" if i % 2 == 0 else "assistant",
+                "content": f"Message {i}",
+            },
+        )
+
+    print(f"Added {await memory.size()} messages")
+
+    # Best practice 1: Regular truncation
+    await memory.truncate_to_size(50)
+    print(f"After truncation: {await memory.size()} messages")
+
+    # Best practice 2: Get only recent messages
+    recent = await memory.get_messages(limit=10)
+    print(f"Retrieved last {len(recent)} messages")
+
+    # Best practice 3: Periodic cleanup
+    user_msgs = await memory.get_messages_by_role("user")
+    print(f"User sent {len(user_msgs)} messages")
+
+    # Best practice 4: Clear when starting new topic
+    await memory.clear()
+    print(f"After clearing: {await memory.size()} messages")
+
+    # Best practice 5: State persistence for crash recovery
+    await memory.add({"role": "user", "content": "Important message"})
+    state = memory.state_dict()
+    print(f"State saved with {len(state['messages'])} messages")
+
+    print("\n💾 Save this state to disk for persistence across restarts!")
+    print("\n✅ Memory management example completed!\n")
+
+
+async def main():
+    """Run all examples."""
+    print("\n🚀 MassGen Memory System Examples\n")
+
+    await example_conversation_memory()
+    await example_persistent_memory()
+    await example_combined_usage()
+    await example_memory_management()
+
+    print("=" * 60)
+    print("All examples completed! 🎉")
+    print("=" * 60)
+    print("\nNext steps:")
+    print("1. Install mem0: pip install mem0ai")
+    print("2. Configure your LLM and embedding backends")
+    print("3. Try persistent memory with real backends")
+    print("4. Integrate into your MassGen agents")
+    print("\n")
+
+
+if __name__ == "__main__":
+    asyncio.run(main())

massgen/orchestrator.py

@@ -42,6 +42,7 @@ from .logger_config import (
     log_stream_chunk,
     log_tool_call,
 )
+from .memory import ConversationMemory, PersistentMemoryBase
 from .message_templates import MessageTemplates
 from .stream_chunk import ChunkType
 from .tool import get_post_evaluation_tools, get_workflow_tools
@@ -118,6 +119,9 @@ class Orchestrator(ChatAgent):
         snapshot_storage: Optional[str] = None,
         agent_temporary_workspace: Optional[str] = None,
         previous_turns: Optional[List[Dict[str, Any]]] = None,
+        winning_agents_history: Optional[List[Dict[str, Any]]] = None,
+        shared_conversation_memory: Optional[ConversationMemory] = None,
+        shared_persistent_memory: Optional[PersistentMemoryBase] = None,
     ):
         """
         Initialize MassGen orchestrator.
@@ -130,13 +134,22 @@ class Orchestrator(ChatAgent):
             snapshot_storage: Optional path to store agent workspace snapshots
             agent_temporary_workspace: Optional path for agent temporary workspaces
             previous_turns: List of previous turn metadata for multi-turn conversations (loaded by CLI)
+            winning_agents_history: List of previous winning agents for memory sharing
+                                   Format: [{"agent_id": "agent_b", "turn": 1}, ...]
+                                   Loaded from session storage to persist across orchestrator recreations
+            shared_conversation_memory: Optional shared conversation memory for all agents
+            shared_persistent_memory: Optional shared persistent memory for all agents
         """
-        super().__init__(session_id)
+        super().__init__(session_id, shared_conversation_memory, shared_persistent_memory)
         self.orchestrator_id = orchestrator_id
         self.agents = agents
         self.agent_states = {aid: AgentState() for aid in agents.keys()}
         self.config = config or AgentConfig.create_openai_config()
 
+        # Shared memory for all agents
+        self.shared_conversation_memory = shared_conversation_memory
+        self.shared_persistent_memory = shared_persistent_memory
+
         # Get message templates from config
         self.message_templates = self.config.message_templates or MessageTemplates(
             voting_sensitivity=self.config.voting_sensitivity,
@@ -158,6 +171,14 @@ class Orchestrator(ChatAgent):
         self._selected_agent: Optional[str] = None
         self._final_presentation_content: Optional[str] = None
 
+        # Track winning agents by turn for memory sharing
+        # Format: [{"agent_id": "agent_b", "turn": 1}, {"agent_id": "agent_a", "turn": 2}]
+        # Restore from session storage if provided (for multi-turn persistence)
+        self._winning_agents_history: List[Dict[str, Any]] = winning_agents_history or []
+        if self._winning_agents_history:
+            logger.info(f"📚 Restored {len(self._winning_agents_history)} winning agent(s) from session: {self._winning_agents_history}")
+        self._current_turn: int = 0
+
         # Timeout and resource tracking
         self.total_tokens: int = 0
         self.coordination_start_time: float = 0
@@ -365,6 +386,113 @@ class Orchestrator(ChatAgent):
             "full_messages": messages,
         }
 
+    async def _inject_shared_memory_context(
+        self,
+        messages: List[Dict[str, Any]],
+        agent_id: str,
+    ) -> List[Dict[str, Any]]:
+        """
+        Inject shared memory context into agent messages.
+
+        This allows all agents to see shared memories including what other agents
+        have stored in the shared memory.
+
+        Args:
+            messages: Original messages to send to agent
+            agent_id: ID of the agent receiving the messages
+
+        Returns:
+            Messages with shared memory context injected
+        """
+        if not self.shared_conversation_memory and not self.shared_persistent_memory:
+            # No shared memory configured, return original messages
+            return messages
+
+        memory_context_parts = []
+
+        # Get conversation memory content
+        if self.shared_conversation_memory:
+            try:
+                conv_messages = await self.shared_conversation_memory.get_messages()
+                if conv_messages:
+                    memory_context_parts.append("=== SHARED CONVERSATION MEMORY ===")
+                    for msg in conv_messages[-10:]:  # Last 10 messages
+                        role = msg.get("role", "unknown")
+                        content = msg.get("content", "")
+                        agent_source = msg.get("agent_id", "unknown")
+                        memory_context_parts.append(f"[{agent_source}] {role}: {content}")
+            except Exception as e:
+                logger.warning(f"Failed to retrieve shared conversation memory: {e}")
+
+        # Get persistent memory content
+        if self.shared_persistent_memory:
+            try:
+                # Extract user message for retrieval
+                user_messages = [msg for msg in messages if msg.get("role") == "user"]
+                if user_messages:
+                    retrieved = await self.shared_persistent_memory.retrieve(user_messages)
+                    if retrieved:
+                        memory_context_parts.append("\n=== SHARED PERSISTENT MEMORY ===")
+                        memory_context_parts.append(retrieved)
+            except NotImplementedError:
+                # Memory backend doesn't support retrieve
+                pass
+            except Exception as e:
+                logger.warning(f"Failed to retrieve shared persistent memory: {e}")
+
+        # Inject memory context if we have any
+        if memory_context_parts:
+            memory_message = {
+                "role": "system",
+                "content": ("You have access to shared memory that all agents can see and contribute to.\n" + "\n".join(memory_context_parts)),
+            }
+
+            # Insert after existing system messages but before user messages
+            system_count = sum(1 for msg in messages if msg.get("role") == "system")
+            modified_messages = messages.copy()
+            modified_messages.insert(system_count, memory_message)
+            return modified_messages
+
+        return messages
+
+    async def _record_to_shared_memory(
+        self,
+        agent_id: str,
+        content: str,
+        role: str = "assistant",
+    ) -> None:
+        """
+        Record agent's contribution to shared memory.
+
+        Args:
+            agent_id: ID of the agent contributing
+            content: Content to record
+            role: Role of the message (default: "assistant")
+        """
+        message = {
+            "role": role,
+            "content": content,
+            "agent_id": agent_id,
+            "timestamp": time.time(),
+        }
+
+        # Add to conversation memory
+        if self.shared_conversation_memory:
+            try:
+                await self.shared_conversation_memory.add(message)
+            except Exception as e:
+                logger.warning(f"Failed to add to shared conversation memory: {e}")
+
+        # Record to persistent memory
+        if self.shared_persistent_memory:
+            try:
+                await self.shared_persistent_memory.record([message])
+            except NotImplementedError:
+                # Memory backend doesn't support record
+                pass
+            except Exception as e:
+                logger.warning(f"Failed to record to shared persistent memory: {e}")
+
     def save_coordination_logs(self):
         """Public method to save coordination logs after final presentation is complete."""
         # End the coordination session
@@ -798,6 +926,18 @@ Your answer:"""
         current_answers = {aid: state.answer for aid, state in self.agent_states.items() if state.answer}
         self._selected_agent = self._determine_final_agent_from_votes(votes, current_answers)
 
+        # Track winning agent for memory sharing in future turns
+        self._current_turn += 1
+        if self._selected_agent:
+            winner_entry = {
+                "agent_id": self._selected_agent,
+                "turn": self._current_turn,
+            }
+            self._winning_agents_history.append(winner_entry)
+            logger.info(
+                f"🏆 Turn {self._current_turn} winner: {self._selected_agent} " f"(tracked for memory sharing)",
+            )
+
         log_coordination_step(
             "Final agent selected",
             {"selected_agent": self._selected_agent, "votes": votes},
@@ -1806,6 +1946,13 @@ Your answer:"""
                 {"role": "system", "content": conversation["system_message"]},
                 {"role": "user", "content": conversation["user_message"]},
             ]
+
+            # Inject shared memory context
+            conversation_messages = await self._inject_shared_memory_context(
+                conversation_messages,
+                agent_id,
+            )
+
             enforcement_msg = self.message_templates.enforcement_message()
 
             # Update agent status to STREAMING
@@ -1832,20 +1979,42 @@ Your answer:"""
                     # First attempt: orchestrator provides initial conversation
                     # But we need the agent to have this in its history for subsequent calls
                     # First attempt: provide complete conversation and reset agent's history
-                    chat_stream = agent.chat(conversation_messages, self.workflow_tools, reset_chat=True, current_stage=CoordinationStage.INITIAL_ANSWER)
+                    # Pass current turn and previous winners for memory sharing
+                    chat_stream = agent.chat(
+                        conversation_messages,
+                        self.workflow_tools,
+                        reset_chat=True,
+                        current_stage=CoordinationStage.INITIAL_ANSWER,
+                        orchestrator_turn=self._current_turn + 1,  # Next turn number
+                        previous_winners=self._winning_agents_history.copy(),
+                    )
                 else:
                     # Subsequent attempts: send enforcement message (set by error handling)
 
                     if isinstance(enforcement_msg, list):
                         # Tool message array
-                        chat_stream = agent.chat(enforcement_msg, self.workflow_tools, reset_chat=False, current_stage=CoordinationStage.ENFORCEMENT)
+                        chat_stream = agent.chat(
+                            enforcement_msg,
+                            self.workflow_tools,
+                            reset_chat=False,
+                            current_stage=CoordinationStage.ENFORCEMENT,
+                            orchestrator_turn=self._current_turn + 1,
+                            previous_winners=self._winning_agents_history.copy(),
+                        )
                     else:
                         # Single user message
                         enforcement_message = {
                             "role": "user",
                             "content": enforcement_msg,
                         }
-                        chat_stream = agent.chat([enforcement_message], self.workflow_tools, reset_chat=False, current_stage=CoordinationStage.ENFORCEMENT)
+                        chat_stream = agent.chat(
+                            [enforcement_message],
+                            self.workflow_tools,
+                            reset_chat=False,
+                            current_stage=CoordinationStage.ENFORCEMENT,
+                            orchestrator_turn=self._current_turn + 1,
+                            previous_winners=self._winning_agents_history.copy(),
+                        )
                 response_text = ""
                 tool_calls = []
                 workflow_tool_found = False
@@ -2101,6 +2270,14 @@ Your answer:"""
                                 "reason": reason,
                             }
 
+                            # Record vote to shared memory
+                            vote_message = f"Voted for {voted_agent}. Reason: {reason}"
+                            await self._record_to_shared_memory(
+                                agent_id=agent_id,
+                                content=vote_message,
+                                role="assistant",
+                            )
+
                             # Send tool result - orchestrator will decide if vote is accepted
                             # Vote submitted (result will be shown by orchestrator)
                             yield (
@@ -2193,6 +2370,14 @@ Your answer:"""
                                         return
                             # Send successful tool result back to agent
                             # Answer recorded (result will be shown by orchestrator)
+
+                            # Record to shared memory
+                            await self._record_to_shared_memory(
+                                agent_id=agent_id,
+                                content=content,
+                                role="assistant",
+                            )
+
                             yield ("result", ("answer", content))
                             yield ("done", None)
                             return
@@ -2623,7 +2808,13 @@ INSTRUCTIONS FOR NEXT ATTEMPT:
 
         try:
             # Track final round iterations (each chunk is like an iteration)
-            async for chunk in agent.chat(presentation_messages, reset_chat=True, current_stage=CoordinationStage.PRESENTATION):
+            async for chunk in agent.chat(
+                presentation_messages,
+                reset_chat=True,
+                current_stage=CoordinationStage.PRESENTATION,
+                orchestrator_turn=self._current_turn,
+                previous_winners=self._winning_agents_history.copy(),
+            ):
                 chunk_type = self._get_chunk_type_value(chunk)
                 # Start new iteration for this chunk
                 self.coordination_tracker.start_new_iteration()
@@ -2872,7 +3063,14 @@ Then call either submit(confirmed=True) if the answer is satisfactory, or restar
         try:
             timeout_seconds = self.config.timeout_config.orchestrator_timeout_seconds
             async with asyncio.timeout(timeout_seconds):
-                async for chunk in agent.chat(messages=evaluation_messages, tools=post_eval_tools, reset_chat=True, current_stage=CoordinationStage.POST_EVALUATION):
+                async for chunk in agent.chat(
+                    messages=evaluation_messages,
+                    tools=post_eval_tools,
+                    reset_chat=True,
+                    current_stage=CoordinationStage.POST_EVALUATION,
+                    orchestrator_turn=self._current_turn,
+                    previous_winners=self._winning_agents_history.copy(),
+                ):
                     chunk_type = self._get_chunk_type_value(chunk)
 
                     if chunk_type == "content" and chunk.content:
@@ -3103,7 +3301,8 @@ Then call either submit(confirmed=True) if the answer is satisfactory, or restar
         Get final result for session persistence.
 
         Returns:
-            Dict with final_answer, winning_agent_id, and workspace_path, or None if not available
+            Dict with final_answer, winning_agent_id, workspace_path, and winning_agents_history,
+            or None if not available
         """
         if not self._selected_agent or not self._final_presentation_content:
             return None
@@ -3117,6 +3316,7 @@ Then call either submit(confirmed=True) if the answer is satisfactory, or restar
             "final_answer": self._final_presentation_content,
             "winning_agent_id": self._selected_agent,
             "workspace_path": workspace_path,
+            "winning_agents_history": self._winning_agents_history.copy(),  # For cross-turn memory sharing
         }
 
     def get_status(self) -> Dict[str, Any]: