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

massgen/formatter/_gemini_formatter.py

@@ -52,11 +52,57 @@ class GeminiFormatter(FormatterBase):
         """
         return tools or []
 
-    def format_mcp_tools(self, mcp_functions: Dict[str, Any]) -> List[Dict[str, Any]]:
+    def format_mcp_tools(
+        self,
+        mcp_functions: Dict[str, Any],
+        return_sdk_objects: bool = True,
+    ) -> List[Any]:
         """
-        MCP tools are passed via SDK sessions in stream_with_tools; not function declarations.
+        Convert MCP Function objects to Gemini FunctionDeclaration format.
+
         """
-        return []
+        if not mcp_functions:
+            return []
+
+        # Step 1: Convert MCP Function objects to Gemini dictionary format
+        gemini_dicts = []
+
+        for mcp_function in mcp_functions.values():
+            try:
+                # Extract attributes from Function object
+                name = getattr(mcp_function, "name", "")
+                description = getattr(mcp_function, "description", "")
+                parameters = getattr(mcp_function, "parameters", {})
+
+                # Build Gemini-compatible dictionary
+                gemini_dict = {
+                    "name": name,
+                    "description": description,
+                    "parameters": parameters,
+                }
+                gemini_dicts.append(gemini_dict)
+
+                logger.debug(f"[GeminiFormatter] Converted MCP tool '{name}' to dictionary format")
+
+            except Exception as e:
+                logger.error(f"[GeminiFormatter] Failed to convert MCP tool: {e}")
+                # Continue processing remaining tools instead of failing completely
+                continue
+
+        if not return_sdk_objects:
+            return gemini_dicts
+
+        # Step 2: Convert dictionaries to SDK FunctionDeclaration objects
+        function_declarations = self._convert_to_function_declarations(gemini_dicts)
+
+        # Log successful conversion
+        for func_decl in function_declarations:
+            if hasattr(func_decl, "name"):
+                logger.debug(
+                    f"[GeminiFormatter] Converted MCP tool '{func_decl.name}' to FunctionDeclaration",
+                )
+
+        return function_declarations
 
     # Coordination helpers
 
@@ -239,15 +285,14 @@ Make your decision and include the JSON at the very end of your response."""
             vote_data = structured_response.get("vote_data", {})
             return [
                 {
-                    "id": f"vote_{abs(hash(str(vote_data))) % 10000 + 1}",
-                    "type": "function",
-                    "function": {
-                        "name": "vote",
-                        "arguments": {
+                    "call_id": f"vote_{abs(hash(str(vote_data))) % 10000 + 1}",
+                    "name": "vote",
+                    "arguments": json.dumps(
+                        {
                             "agent_id": vote_data.get("agent_id", ""),
                             "reason": vote_data.get("reason", ""),
                         },
-                    },
+                    ),
                 },
             ]
 
@@ -255,12 +300,13 @@ Make your decision and include the JSON at the very end of your response."""
             answer_data = structured_response.get("answer_data", {})
             return [
                 {
-                    "id": f"new_answer_{abs(hash(str(answer_data))) % 10000 + 1}",
-                    "type": "function",
-                    "function": {
-                        "name": "new_answer",
-                        "arguments": {"content": answer_data.get("content", "")},
-                    },
+                    "call_id": f"new_answer_{abs(hash(str(answer_data))) % 10000 + 1}",
+                    "name": "new_answer",
+                    "arguments": json.dumps(
+                        {
+                            "content": answer_data.get("content", ""),
+                        },
+                    ),
                 },
             ]
 

massgen/memory/README.md

@@ -0,0 +1,277 @@
+# MassGen Memory System
+
+The MassGen memory system provides agents with the ability to store and recall information across conversations, enabling more contextual and intelligent interactions.
+
+## Architecture
+
+The memory system consists of three main components:
+
+### 1. **ConversationMemory** - Short-term Memory
+Fast in-memory storage for active conversations.
+
+```python
+from massgen.memory import ConversationMemory
+
+# Initialize
+memory = ConversationMemory()
+
+# Add messages
+await memory.add({"role": "user", "content": "Hello!"})
+await memory.add({"role": "assistant", "content": "Hi there!"})
+
+# Retrieve messages
+messages = await memory.get_messages()
+print(f"Stored {len(messages)} messages")
+
+# Get last 10 messages only
+recent = await memory.get_messages(limit=10)
+
+# Filter by role
+user_messages = await memory.get_messages_by_role("user")
+
+# Manage memory size
+await memory.truncate_to_size(100)  # Keep last 100 messages
+
+# Clear all
+await memory.clear()
+
+# Save/load state
+state = memory.state_dict()
+# ... later ...
+new_memory = ConversationMemory()
+new_memory.load_state_dict(state)
+```
+
+### 2. **PersistentMemory** - Long-term Memory
+Semantic memory storage using mem0 with vector search capabilities.
+
+```python
+from massgen.memory import PersistentMemory
+from massgen.backend import OpenAIBackend
+
+# Initialize with backends
+llm_backend = OpenAIBackend(model="gpt-4")
+embedding_backend = OpenAIBackend(model="text-embedding-3-small")
+
+memory = PersistentMemory(
+    agent_name="my_agent",
+    llm_backend=llm_backend,
+    embedding_backend=embedding_backend,
+    on_disk=True  # Persist to disk
+)
+
+# Developer Interface: Record conversation
+await memory.record([
+    {"role": "user", "content": "My name is Alice"},
+    {"role": "assistant", "content": "Nice to meet you, Alice!"}
+])
+
+# Developer Interface: Retrieve relevant memories
+relevant = await memory.retrieve("What's the user's name?")
+print(relevant)  # "Alice"
+
+# Agent Tool Interface: Agent saves important info
+result = await memory.save_to_memory(
+    thinking="User shared personal information",
+    content=["User's favorite color is blue", "User works as a teacher"]
+)
+
+# Agent Tool Interface: Agent recalls from memory
+result = await memory.recall_from_memory(
+    keywords=["favorite color", "job"],
+    limit=5
+)
+for mem in result['memories']:
+    print(mem)
+```
+
+### 3. **Memory Integration Modes**
+
+#### Mode 1: Developer-Controlled (Automatic)
+```python
+# Framework automatically manages memory
+
+class MyAgent:
+    def __init__(self):
+        self.conversation_memory = ConversationMemory()
+        self.persistent_memory = PersistentMemory(
+            agent_name="agent_1",
+            llm_backend=llm,
+            embedding_backend=embedding
+        )
+
+    async def chat(self, user_message):
+        # Auto-add to conversation memory
+        await self.conversation_memory.add({
+            "role": "user",
+            "content": user_message
+        })
+
+        # Auto-retrieve relevant long-term memories
+        context = await self.persistent_memory.retrieve(user_message)
+
+        # Build full context
+        messages = await self.conversation_memory.get_messages()
+        if context:
+            messages.insert(0, {
+                "role": "system",
+                "content": f"Relevant memories: {context}"
+            })
+
+        # Generate response...
+        response = await self.backend.chat(messages)
+
+        # Auto-record response
+        await self.conversation_memory.add({
+            "role": "assistant",
+            "content": response
+        })
+
+        # Optionally save important conversations to long-term memory
+        await self.persistent_memory.record(messages)
+
+        return response
+```
+
+#### Mode 2: Agent-Controlled (Tools)
+```python
+# Agent actively manages its own memory
+
+# Register memory tools
+tools = [
+    {
+        "name": "save_to_memory",
+        "description": "Save important information for future reference",
+        "parameters": {
+            "thinking": "Why this information is important",
+            "content": ["List of things to remember"]
+        },
+        "function": memory.save_to_memory
+    },
+    {
+        "name": "recall_from_memory",
+        "description": "Search long-term memory for information",
+        "parameters": {
+            "keywords": ["Keywords to search for"]
+        },
+        "function": memory.recall_from_memory
+    }
+]
+
+# Agent can now call these tools during conversation
+# Example agent reasoning:
+# "The user mentioned their birthday. I should save this."
+# -> Calls save_to_memory(thinking="...", content=["Birthday: March 15"])
+```
+
+## Installation
+
+The memory system requires the mem0 library for persistent memory:
+
+```bash
+pip install mem0ai
+```
+
+## Vector Store Options
+
+By default, PersistentMemory uses Qdrant for vector storage. You can configure other backends:
+
+```python
+from mem0.vector_stores.configs import VectorStoreConfig
+
+# Custom vector store
+vector_config = VectorStoreConfig(
+    provider="qdrant",  # or "chroma", "pinecone", etc.
+    config={
+        "on_disk": True,
+        "path": "./my_memories"
+    }
+)
+
+memory = PersistentMemory(
+    agent_name="agent_1",
+    llm_backend=llm,
+    embedding_backend=embedding,
+    vector_store_config=vector_config
+)
+```
+
+## Memory Organization
+
+Memories are organized by metadata:
+
+- **agent_name**: Isolate memories per agent
+- **user_name**: Track user-specific information
+- **session_name**: Separate different conversation sessions
+
+```python
+# Agent-specific memory
+agent_memory = PersistentMemory(agent_name="research_agent")
+
+# User-specific memory
+user_memory = PersistentMemory(user_name="alice")
+
+# Session-specific memory
+session_memory = PersistentMemory(
+    agent_name="agent_1",
+    session_name="session_123"
+)
+
+# Combined filtering
+memory = PersistentMemory(
+    agent_name="agent_1",
+    user_name="alice",
+    session_name="2024-01-15"
+)
+```
+
+## Best Practices
+
+1. **Short-term for speed**: Use ConversationMemory for current dialogue
+2. **Long-term for knowledge**: Use PersistentMemory for cross-session information
+3. **Truncate regularly**: Keep conversation memory manageable
+4. **Meaningful metadata**: Use descriptive agent/user/session names
+5. **Selective persistence**: Only save important information to long-term memory
+
+## Advanced: Custom mem0 Configuration
+
+For full control over mem0 behavior:
+
+```python
+from mem0.configs.base import MemoryConfig
+from mem0.configs.llms.configs import LlmConfig
+from mem0.configs.embeddings.configs import EmbedderConfig
+
+custom_mem0_config = MemoryConfig(
+    llm=LlmConfig(provider="openai", config={"model": "gpt-4"}),
+    embedder=EmbedderConfig(provider="openai", config={"model": "text-embedding-3-large"}),
+    # ... other mem0 settings
+)
+
+memory = PersistentMemory(
+    agent_name="agent_1",
+    mem0_config=custom_mem0_config
+)
+```
+
+## Troubleshooting
+
+**Issue**: `ImportError: mem0 library is required`
+- **Solution**: `pip install mem0ai`
+
+**Issue**: Memories not persisting across restarts
+- **Solution**: Ensure `on_disk=True` in PersistentMemory initialization
+
+**Issue**: Out of memory errors
+- **Solution**: Use `truncate_to_size()` on ConversationMemory regularly
+
+**Issue**: Slow retrieval
+- **Solution**: Reduce `limit` parameter in retrieve calls, or upgrade vector store
+
+## Contributing
+
+The memory system is designed to be extensible. To add new memory types:
+
+1. Inherit from `MemoryBase` or `PersistentMemoryBase`
+2. Implement required abstract methods
+3. Add to `__init__.py` exports

massgen/memory/__init__.py

@@ -0,0 +1,26 @@
+# -*- coding: utf-8 -*-
+"""
+MassGen Memory System
+
+This module provides memory capabilities for MassGen agents, supporting both
+short-term conversation memory and long-term persistent memory storage.
+
+The memory system is designed to be:
+- Asynchronous: All operations are async for optimal performance
+- Pluggable: Easy to swap different storage backends
+- Serializable: Support for saving and loading agent memory state
+"""
+
+from ._base import MemoryBase, PersistentMemoryBase
+from ._compression import CompressionStats, ContextCompressor
+from ._conversation import ConversationMemory
+from ._persistent import PersistentMemory
+
+__all__ = [
+    "MemoryBase",
+    "PersistentMemoryBase",
+    "ConversationMemory",
+    "PersistentMemory",
+    "ContextCompressor",
+    "CompressionStats",
+]

massgen/memory/_base.py

@@ -0,0 +1,193 @@
+# -*- coding: utf-8 -*-
+"""
+Base classes for MassGen memory system.
+
+This module defines the abstract interfaces that all memory implementations
+must follow, ensuring consistency across different storage backends.
+"""
+
+from abc import ABC, abstractmethod
+from typing import Any, Dict, List, Union
+
+
+class MemoryBase(ABC):
+    """
+    Abstract base class for memory storage in MassGen.
+
+    All memory implementations (conversation, persistent, etc.) should inherit
+    from this class and implement the required methods.
+    """
+
+    @abstractmethod
+    async def add(self, *args: Any, **kwargs: Any) -> None:
+        """
+        Add new items to memory.
+
+        Args:
+            *args: Variable positional arguments
+            **kwargs: Variable keyword arguments
+        """
+
+    @abstractmethod
+    async def delete(self, *args: Any, **kwargs: Any) -> None:
+        """
+        Remove items from memory.
+
+        Args:
+            *args: Variable positional arguments
+            **kwargs: Variable keyword arguments
+        """
+
+    @abstractmethod
+    async def retrieve(self, *args: Any, **kwargs: Any) -> Any:
+        """
+        Retrieve items from memory based on query.
+
+        Args:
+            *args: Variable positional arguments
+            **kwargs: Variable keyword arguments
+
+        Returns:
+            Retrieved memory content
+        """
+
+    @abstractmethod
+    async def size(self) -> int:
+        """
+        Get the current size of the memory.
+
+        Returns:
+            Number of items in memory
+        """
+
+    @abstractmethod
+    async def clear(self) -> None:
+        """Clear all content from memory."""
+
+    @abstractmethod
+    async def get_messages(self, *args: Any, **kwargs: Any) -> List[Dict[str, Any]]:
+        """
+        Get the stored messages in a format suitable for LLM consumption.
+
+        Returns:
+            List of message dictionaries
+        """
+
+    @abstractmethod
+    def state_dict(self) -> Dict[str, Any]:
+        """
+        Export memory state for serialization.
+
+        Returns:
+            Dictionary containing the memory state
+        """
+
+    @abstractmethod
+    def load_state_dict(self, state_dict: Dict[str, Any], strict: bool = True) -> None:
+        """
+        Load memory state from serialized data.
+
+        Args:
+            state_dict: Dictionary containing the memory state
+            strict: If True, raise error on missing/extra keys
+        """
+
+
+class PersistentMemoryBase(ABC):
+    """
+    Abstract base class for long-term persistent memory.
+
+    This type of memory is designed to store information across sessions,
+    with support for semantic retrieval and knowledge management.
+
+    Two modes of operation:
+    1. Developer-controlled: Use `record()` and `retrieve()` methods programmatically
+    2. Agent-controlled: Agent calls `save_to_memory()` and `recall_from_memory()` tools
+    """
+
+    async def record(
+        self,
+        messages: List[Dict[str, Any]],
+        **kwargs: Any,
+    ) -> None:
+        """
+        Developer interface: Record information to persistent memory.
+
+        This method is called by the framework to automatically save important
+        information from conversations.
+
+        Args:
+            messages: List of message dictionaries to record
+            **kwargs: Additional recording options
+        """
+        raise NotImplementedError(
+            "The `record` method is not implemented in this memory backend.",
+        )
+
+    async def retrieve(
+        self,
+        query: Union[str, Dict[str, Any], List[Dict[str, Any]]],
+        **kwargs: Any,
+    ) -> str:
+        """
+        Developer interface: Retrieve relevant information from persistent memory.
+
+        This method is called by the framework to automatically inject relevant
+        historical knowledge into the current conversation.
+
+        Args:
+            query: Query message(s) or string to search for
+            **kwargs: Additional retrieval options
+
+        Returns:
+            Retrieved information as formatted string
+        """
+        raise NotImplementedError(
+            "The `retrieve` method is not implemented in this memory backend.",
+        )
+
+    async def save_to_memory(
+        self,
+        thinking: str,
+        content: List[str],
+        **kwargs: Any,
+    ) -> Dict[str, Any]:
+        """
+        Agent tool interface: Save important information to memory.
+
+        This is a tool function that agents can call to explicitly save
+        information they deem important for future reference.
+
+        Args:
+            thinking: Agent's reasoning about why this information is important
+            content: List of information items to save
+            **kwargs: Additional save options
+
+        Returns:
+            Tool response with status and saved memory IDs
+        """
+        raise NotImplementedError(
+            "The `save_to_memory` tool is not implemented in this memory backend. " "Implement this method to allow agents to actively manage their memory.",
+        )
+
+    async def recall_from_memory(
+        self,
+        keywords: List[str],
+        **kwargs: Any,
+    ) -> Dict[str, Any]:
+        """
+        Agent tool interface: Recall information based on keywords.
+
+        This is a tool function that agents can call to retrieve relevant
+        information from their long-term memory.
+
+        Args:
+            keywords: Keywords to search for (person names, dates, topics, etc.)
+            **kwargs: Additional recall options
+
+        Returns:
+            Tool response with retrieved memories
+        """
+        raise NotImplementedError(
+            "The `recall_from_memory` tool is not implemented in this memory backend. " "Implement this method to allow agents to actively query their memory.",
+        )