flock-core 0.4.0b49__py3-none-any.whl → 0.4.1__py3-none-any.whl

flock/modules/mem0/mem0_module.py

@@ -1,5 +1,8 @@
 from typing import Any
 
+# from mem0.client.main import MemoryClient
+# from mem0.memory.main import Memory
+from mem0 import Memory, MemoryClient
 from pydantic import Field
 
 from flock.core.context.context import FlockContext
@@ -11,35 +14,46 @@ from flock.core.logging.logging import get_logger
 logger = get_logger("module.mem0")
 
 
-class Mem0ModuleConfig(FlockModuleConfig):
-    qdrant_host: str = Field(default="localhost", description="Qdrant hostname")
-    qdrant_port: int = Field(default=6333, description="Qdrant port")
-    collection: str = Field(default="flock_memories", description="Vector collection")
-    embedder_provider: str = Field(default="openai", description="'openai' or 'local'")
-    embedder_model: str = Field(default="text-embedding-ada-002",
-                                description="Model name for embeddings")
-    # Optional: allow separate LLM for reflection/summarisation
-    llm_provider: str | None = Field(default=None)
-    llm_model: str | None = Field(default=None)
-    top_k: int = Field(default=5, description="Number of memories to retrieve")
+config = {
+    "vector_store": {
+        "provider": "chroma",
+        "config": {
+            "collection_name": "flock_memory",
+            "path": ".flock/memory",
+        }
+    }
+}
+
 
+class Mem0ModuleConfig(FlockModuleConfig):
+    top_k: int = Field(default=10, description="Number of memories to retrieve")
+    user_id: str = Field(default="flock", description="User ID the memories will be associated with")
+    agent_id: str = Field(default="flock", description="Agent ID the memories will be associated with")
+    memory_input_key: str | None = Field(default=None, description="Input key to use for memory, if none the description of the agent will be used")
+    api_key: str | None = Field(default=None, description="API key for mem0 Platform")
+    config: dict[str, Any] = Field(default=config, description="Configuration for mem0")
 
 
 @flock_component(config_class=Mem0ModuleConfig)
 class Mem0Module(FlockModule):
-    """Module that adds Zep capabilities to a Flock agent."""
 
     name: str = "mem0"
     config: Mem0ModuleConfig = Mem0ModuleConfig()
-    session_id: str | None = None
-    user_id: str | None = None
+
 
     def __init__(self, name, config: Mem0ModuleConfig) -> None:
+        global memory
         """Initialize Mem0 module."""
         super().__init__(name=name, config=config)
         logger.debug("Initializing Mem0 module")
 
 
+
+
+    def dict_to_str_repr(self,d: dict) -> str:
+        return repr(d)
+
+
     async def on_post_evaluate(
         self,
         agent: FlockAgent,
@@ -47,6 +61,26 @@ class Mem0Module(FlockModule):
         context: FlockContext | None = None,
         result: dict[str, Any] | None = None,
     ) -> dict[str, Any]:
+        if self.config.api_key:
+            memory = MemoryClient(api_key=self.config.api_key)
+        else:
+            memory = Memory.from_config(config_dict=self.config.config)
+
+        agent_id = self.config.agent_id if self.config.agent_id else agent.name
+
+        # get the result without the inputs
+        filtered_result = {k: v for k, v in result.items() if k not in inputs}
+        # get the inputs without memory
+        filtered_inputs = {k: v for k, v in inputs.items() if k not in [self.config.memory_input_key]}
+
+        # add memories about the user inputs
+        added_user_memory =  memory.add(self.dict_to_str_repr(filtered_inputs), user_id=self.config.user_id)
+        logger.info(f"Added caller memory: {added_user_memory}")
+
+        # add memories about the agent result
+        added_agent_memory =  memory.add(self.dict_to_str_repr(filtered_result), agent_id=agent_id)
+        logger.info(f"Added agent memory: {added_agent_memory}")
+
 
         return result
 
@@ -56,8 +90,37 @@ class Mem0Module(FlockModule):
         inputs: dict[str, Any],
         context: FlockContext | None = None,
     ) -> dict[str, Any]:
-
-
+        if self.config.api_key:
+            memory = MemoryClient(api_key=self.config.api_key)
+        else:
+            memory = Memory.from_config(config_dict=self.config.config)
+
+        message = self.dict_to_str_repr(inputs)
+        agent_id = self.config.agent_id if self.config.agent_id else agent.name
+
+        relevant_agent_memories = memory.search(query=message, agent_id=agent_id, limit=self.config.top_k)
+        logger.info(f"Relevant agent memories: {relevant_agent_memories}")
+
+        relevant_user_memories = memory.search(query=message, user_id=self.config.user_id, limit=self.config.top_k)
+        logger.info(f"Relevant user memories: {relevant_user_memories}")
+
+        if relevant_agent_memories or relevant_user_memories:
+            memories_str = ''
+            if "results" in relevant_agent_memories:
+                memories_str = "\n".join(f"- {entry['memory']}" for entry in relevant_agent_memories["results"])
+            else:
+                memories_str = "\n".join(f"- {entry}" for entry in relevant_agent_memories)
+
+            if "results" in relevant_user_memories:
+                memories_str = memories_str + "\n" + "\n".join(f"- {entry['memory']}" for entry in relevant_user_memories["results"])
+            else:
+                memories_str = memories_str + "\n" + "\n".join(f"- {entry}" for entry in relevant_user_memories)
+
+            if memories_str:
+                if self.config.memory_input_key:
+                    inputs[self.config.memory_input_key] = memories_str
+                else:
+                    agent.description = agent.description + "\n\n Memories:" + memories_str
 
 
         return inputs

flock/modules/mem0_async/async_mem0_module.py

@@ -0,0 +1,126 @@
+from typing import Any
+
+# from mem0.client.main import AsyncMemoryClient, MemoryClient
+# from mem0.memory.main import AsyncMemory
+from mem0 import AsyncMemory, AsyncMemoryClient
+from pydantic import Field
+
+from flock.core.context.context import FlockContext
+from flock.core.flock_agent import FlockAgent
+from flock.core.flock_module import FlockModule, FlockModuleConfig
+from flock.core.flock_registry import flock_component
+from flock.core.logging.logging import get_logger
+
+logger = get_logger("module.mem0")
+
+
+config = {
+    "vector_store": {
+        "provider": "chroma",
+        "config": {
+            "collection_name": "flock_memory",
+            "path": ".flock/memory",
+        }
+    }
+}
+
+
+class AsyncMem0ModuleConfig(FlockModuleConfig):
+    top_k: int = Field(default=10, description="Number of memories to retrieve")
+    user_id: str = Field(default="flock", description="User ID the memories will be associated with")
+    agent_id: str = Field(default="flock", description="Agent ID the memories will be associated with")
+    memory_input_key: str | None = Field(default=None, description="Input key to use for memory, if none the description of the agent will be used")
+    api_key: str | None = Field(default=None, description="API key for mem0 Platform")
+    config: dict[str, Any] = Field(default=config, description="Configuration for mem0")
+
+
+@flock_component(config_class=AsyncMem0ModuleConfig)
+class AsyncMem0Module(FlockModule):
+
+    name: str = "mem0"
+    config: AsyncMem0ModuleConfig = AsyncMem0ModuleConfig()
+
+
+    def __init__(self, name, config: AsyncMem0ModuleConfig) -> None:
+        global memory
+        """Initialize Mem0 module."""
+        super().__init__(name=name, config=config)
+        logger.debug("Initializing Mem0 module")
+
+
+
+
+    def dict_to_str_repr(self,d: dict) -> str:
+        return repr(d)
+
+
+    async def on_post_evaluate(
+        self,
+        agent: FlockAgent,
+        inputs: dict[str, Any],
+        context: FlockContext | None = None,
+        result: dict[str, Any] | None = None,
+    ) -> dict[str, Any]:
+        if self.config.api_key:
+            memory = AsyncMemoryClient(api_key=self.config.api_key)
+        else:
+            memory = await AsyncMemory.from_config(config_dict=self.config.config)
+
+        agent_id = self.config.agent_id if self.config.agent_id else agent.name
+
+        # get the result without the inputs
+        filtered_result = {k: v for k, v in result.items() if k not in inputs}
+        # get the inputs without memory
+        filtered_inputs = {k: v for k, v in inputs.items() if k not in [self.config.memory_input_key]}
+
+        # add memories about the user inputs
+        added_user_memory = await memory.add(self.dict_to_str_repr(filtered_inputs), user_id=self.config.user_id)
+        logger.info(f"Added caller memory: {added_user_memory}")
+
+        # add memories about the agent result
+        added_agent_memory = await memory.add(self.dict_to_str_repr(filtered_result), agent_id=agent_id)
+        logger.info(f"Added agent memory: {added_agent_memory}")
+
+
+        return result
+
+    async def on_pre_evaluate(
+        self,
+        agent: FlockAgent,
+        inputs: dict[str, Any],
+        context: FlockContext | None = None,
+    ) -> dict[str, Any]:
+        if self.config.api_key:
+            memory = AsyncMemoryClient(api_key=self.config.api_key)
+        else:
+            memory = await AsyncMemory.from_config(config_dict=self.config.config)
+
+        message = self.dict_to_str_repr(inputs)
+        agent_id = self.config.agent_id if self.config.agent_id else agent.name
+
+        relevant_agent_memories = await memory.search(query=message, agent_id=agent_id, limit=self.config.top_k)
+        logger.info(f"Relevant agent memories: {relevant_agent_memories}")
+
+        relevant_user_memories = await memory.search(query=message, user_id=self.config.user_id, limit=self.config.top_k)
+        logger.info(f"Relevant user memories: {relevant_user_memories}")
+
+        if relevant_agent_memories or relevant_user_memories:
+            memories_str = ''
+            if "results" in relevant_agent_memories:
+                memories_str = "\n".join(f"- {entry['memory']}" for entry in relevant_agent_memories["results"])
+            else:
+                memories_str = "\n".join(f"- {entry}" for entry in relevant_agent_memories)
+
+            if "results" in relevant_user_memories:
+                memories_str = memories_str + "\n" + "\n".join(f"- {entry['memory']}" for entry in relevant_user_memories["results"])
+            else:
+                memories_str = memories_str + "\n" + "\n".join(f"- {entry}" for entry in relevant_user_memories)
+
+            if memories_str:
+                if self.config.memory_input_key:
+                    inputs[self.config.memory_input_key] = memories_str
+                else:
+                    agent.description = agent.description + "\n\n Memories:" + memories_str
+
+
+        return inputs

flock/modules/memory/memory_module.py

@@ -4,7 +4,6 @@ from datetime import datetime
 from typing import Any, Literal
 
 from pydantic import Field
-from tqdm import tqdm
 
 from flock.core.context.context import FlockContext
 
@@ -27,7 +26,7 @@ class MemoryModuleConfig(FlockModuleConfig):
     """
 
     folder_path: str = Field(
-        default="concept_memory/",
+        default=".flock/memory/",
         description="Directory where memory file and concept graph will be saved",
     )
     concept_graph_file: str = Field(
@@ -50,14 +49,19 @@ class MemoryModuleConfig(FlockModuleConfig):
         default=True, description="Whether to save memory after each update"
     )
     splitting_mode: Literal["summary", "semantic", "characters", "none"] = (
-        Field(default="none", description="Mode to split memory content")
+        Field(default="characters", description="Mode to split memory content")
     )
     enable_read_only_mode: bool = Field(
         default=False, description="Whether to enable read only mode"
     )
+    enable_write_only_mode: bool = Field(
+        default=False, description="Whether to enable write only mode"
+    )
     number_of_concepts_to_extract: int = Field(
         default=3, description="Number of concepts to extract from the memory"
     )
+    memory_input_key: str | None = Field(default=None, description="Input key to use for memory, if none the description of the agent will be used")
+
 
 
 @flock_component(config_class=MemoryModuleConfig)
@@ -111,6 +115,9 @@ class MemoryModule(FlockModule):
         if not self.memory_store:
             return inputs
 
+        if self.config.enable_write_only_mode:
+            return inputs
+
         inputs = await self.search_memory(agent, inputs)
 
         if "context" in inputs:
@@ -169,10 +176,11 @@ class MemoryModule(FlockModule):
 
     async def search_memory(
         self, agent: FlockAgent, query: dict[str, Any]
-    ) -> list[str]:
+    ) -> dict[str, Any]:
         """Search memory for the query."""
         if not self.memory_store:
-            return []
+            # No memory store loaded – just return the untouched input
+            return query
 
         try:
             input_text = json.dumps(query)
@@ -237,6 +245,9 @@ class MemoryModule(FlockModule):
         if not self.memory_store:
             return result
 
+        if self.config.enable_read_only_mode:
+            return result
+
         try:
             chunks = await self._get_chunks(agent, inputs, result)
             await self._store_chunks(agent, chunks)
@@ -313,7 +324,7 @@ class MemoryModule(FlockModule):
         agent: FlockAgent,
         inputs: dict[str, Any],
         result: dict[str, Any],
-    ) -> str | list[dict[str, str]]:
+    ) -> list[str]:
         """Extract information chunks using semantic mode."""
         split_signature = agent.create_dspy_signature_class(
             f"{self.name}_splitter",
@@ -327,7 +338,15 @@ class MemoryModule(FlockModule):
         splitter = agent._select_task(split_signature, "Completion")
         full_text = json.dumps(inputs) + (json.dumps(result) if result else "")
         split_result = splitter(content=full_text)
-        return split_result.chunks
+        # Flatten list[dict] into list[str] of "title: content" strings to
+        # keep downstream storage logic simple and type-safe.
+        flattened: list[str] = []
+        for chunk in split_result.chunks:
+            if isinstance(chunk, dict):
+                flattened.extend([f"{k}: {v}" for k, v in chunk.items()])
+            else:
+                flattened.append(str(chunk))
+        return flattened
 
     async def _character_splitter_mode(
         self,
@@ -394,7 +413,8 @@ class MemoryModule(FlockModule):
         if isinstance(chunks, str):
             await self._store_chunk(agent, chunks)
         elif isinstance(chunks, list):
-            for chunk in tqdm(chunks, desc="Storing chunks in memory"):
+            # Avoid tqdm in async context – simple for-loop is safer.
+            for chunk in chunks:
                 await self._store_chunk(agent, chunk)
 
     def save_memory(self) -> None:

flock/modules/performance/metrics_module.py

@@ -48,7 +48,7 @@ class MetricsModuleConfig(FlockModuleConfig):
         default="json", description="Where to store metrics"
     )
     metrics_dir: str = Field(
-        default="metrics/", description="Directory for metrics storage"
+        default=".flock/metrics/", description="Directory for metrics storage"
     )
 
     # Aggregation settings
@@ -77,6 +77,9 @@ class MetricsModuleConfig(FlockModuleConfig):
 class MetricsModule(FlockModule):
     """Module for collecting and analyzing agent performance metrics."""
 
+    # --- Singleton holder for convenient static access ---
+    _INSTANCE: "MetricsModule | None" = None
+
     name: str = "performance_metrics"
     config: MetricsModuleConfig = Field(
         default_factory=MetricsModuleConfig,
@@ -85,6 +88,8 @@ class MetricsModule(FlockModule):
 
     def __init__(self, name, config):
         super().__init__(name=name, config=config)
+        # Register singleton for static helpers
+        MetricsModule._INSTANCE = self
         self._metrics = defaultdict(list)
         self._start_time: float | None = None
         self._start_memory: int | None = None
@@ -495,3 +500,21 @@ class MetricsModule(FlockModule):
             1,
             {"agent": agent.name, "error_type": type(error).__name__},
         )
+
+    # --------------------------------------------------
+    # Public helper for external modules
+    # --------------------------------------------------
+    @classmethod
+    def record(cls, name: str, value: int | float | str, tags: dict[str, str] | None = None):
+        """Record a metric from anywhere in the codebase.
+
+        Example:
+            MetricsModule.record("custom_latency", 123, {"stage": "inference"})
+        The call will forward to the *first* instantiated MetricsModule.  If no
+        instance exists in the current run the call is a no-op so that importing
+        this helper never crashes test-code.
+        """
+        instance = cls._INSTANCE
+        if instance is None:
+            return  # silently ignore if module isn't active
+        instance._record_metric(name, value, tags or {})

flock/modules/zep/__init__.py

@@ -0,0 +1 @@
+# Package for modules

flock/modules/zep/zep_module.py

@@ -0,0 +1,192 @@
+import uuid
+from typing import Any
+
+from pydantic import Field
+from zep_python.client import Zep
+from zep_python.types import Message as ZepMessage, SessionSearchResult
+
+from flock.core.context.context import FlockContext
+from flock.core.flock_agent import FlockAgent
+from flock.core.flock_module import FlockModule, FlockModuleConfig
+from flock.core.flock_registry import flock_component
+from flock.core.logging.logging import get_logger
+
+logger = get_logger("module.zep")
+
+
+class ZepModuleConfig(FlockModuleConfig):
+    """Configuration for the Zep module."""
+
+    zep_url: str = "http://localhost:8000"
+    zep_api_key: str = "apikey"
+    min_fact_rating: float = Field(
+        default=0.7, description="Minimum rating for facts to be considered"
+    )
+    enable_read: bool = True
+    enable_write: bool = False
+    top_k: int = Field(default=10, description="Number of memories to retrieve")
+    user_id: str = Field(default="flock", description="User ID the memories will be associated with")
+    agent_id: str = Field(default="flock", description="Agent ID the memories will be associated with")
+    memory_input_key: str | None = Field(default=None, description="Input key to use for memory, if none the description of the agent will be used")
+
+
+
+@flock_component(config_class=ZepModuleConfig)
+class ZepModule(FlockModule):
+    """Module that adds Zep capabilities to a Flock agent."""
+
+    name: str = "zep"
+    config: ZepModuleConfig = ZepModuleConfig()
+    session_id: str | None = None
+    user_id: str | None = None
+
+    def __init__(self, name, config: ZepModuleConfig) -> None:
+        """Initialize Zep module."""
+        super().__init__(name=name, config=config)
+        logger.debug("Initializing Zep module")
+        zep_client = Zep(
+            base_url=self.config.zep_url, api_key=self.config.zep_api_key
+        )
+        self.user_id = self.name
+        self._setup_user(zep_client)
+        self.session_id = str(uuid.uuid4())
+        self._setup_session(zep_client)
+
+    def _setup_user(self, zep_client: Zep) -> None:
+        """Set up user in Zep."""
+        if not zep_client or not self.user_id:
+            raise ValueError("Zep service or user_id not initialized")
+
+        try:
+            user = zep_client.user.get(user_id=self.user_id)
+            if not user:
+                zep_client.user.add(user_id=self.user_id)
+        except Exception:
+            zep_client.user.add(user_id=self.user_id)
+
+    def _setup_session(self, zep_client: Zep) -> None:
+        """Set up new session."""
+        if not zep_client or not self.user_id or not self.session_id:
+            raise ValueError(
+                "Zep service, user_id, or session_id not initialized"
+            )
+
+        zep_client.memory.add_session(
+            user_id=self.user_id,
+            session_id=self.session_id,
+        )
+
+    def get_client(self) -> Zep:
+        """Get Zep client."""
+        return Zep(
+            base_url=self.config.zep_url, api_key=self.config.zep_api_key
+        )
+
+    def get_memory(self, zep_client: Zep) -> str | None:
+        """Get memory for the current session."""
+        if not zep_client or not self.session_id:
+            logger.error("Zep service or session_id not initialized")
+            return None
+
+        try:
+            memory = zep_client.memory.get(
+                self.session_id, min_rating=self.config.min_fact_rating
+            )
+            if memory:
+                return f"{memory.relevant_facts}"
+        except Exception as e:
+            logger.error(f"Error fetching memory: {e}")
+            return None
+
+        return None
+
+    def split_text(
+        self, text: str | None, max_length: int = 1000
+    ) -> list[ZepMessage]:
+        """Split text into smaller chunks."""
+        result: list[ZepMessage] = []
+        if not text:
+            return result
+        if len(text) <= max_length:
+            return [ZepMessage(role="user", content=text, role_type="user")]
+        for i in range(0, len(text), max_length):
+            result.append(
+                ZepMessage(
+                    role="user",
+                    content=text[i : i + max_length],
+                    role_type="user",
+                )
+            )
+        return result
+
+    def add_to_memory(self, text: str, zep_client: Zep) -> None:
+        """Add text to memory."""
+        if not zep_client or not self.session_id:
+            logger.error("Zep service or session_id not initialized")
+            return
+
+        messages = self.split_text(text)
+        zep_client.memory.add(session_id=self.session_id, messages=messages)
+
+    def search_memory(
+        self, query: str, zep_client: Zep
+    ) -> list[SessionSearchResult]:
+        """Search memory for a query."""
+        if not zep_client or not self.user_id:
+            logger.error("Zep service or user_id not initialized")
+            return []
+
+        response = zep_client.memory.search_sessions(
+            text=query,
+            user_id=self.user_id,
+            search_scope="facts",
+            min_fact_rating=self.config.min_fact_rating,
+        )
+        if not response.results:
+            return []
+        return response.results
+
+    async def on_post_evaluate(
+        self,
+        agent: FlockAgent,
+        inputs: dict[str, Any],
+        context: FlockContext | None = None,
+        result: dict[str, Any] | None = None,
+    ) -> dict[str, Any]:
+        """Format and display the output."""
+        if not self.config.enable_write:
+            return result
+        logger.debug("Saving data to memory")
+        zep_client = Zep(
+            base_url=self.config.zep_url, api_key=self.config.zep_api_key
+        )
+        self.add_to_memory(str(result), zep_client)
+        return result
+
+    async def on_pre_evaluate(
+        self,
+        agent: FlockAgent,
+        inputs: dict[str, Any],
+        context: FlockContext | None = None,
+    ) -> dict[str, Any]:
+        """Format and display the output."""
+        if not self.config.enable_read:
+            return inputs
+
+        zep_client = Zep(
+            base_url=self.config.zep_url, api_key=self.config.zep_api_key
+        )
+
+        logger.debug("Searching memory")
+        facts = self.search_memory(str(inputs), zep_client)
+
+        # Add memory to inputs
+        facts_str = ""
+        if facts:
+            for fact in facts:
+                facts_str += fact.fact.fact + "\n"
+            logger.debug("Found facts in memory: {}", facts_str)
+            agent.input = agent.input + ", memory"
+            inputs["memory"] = facts_str
+
+        return inputs