flock-core 0.4.527__py3-none-any.whl → 0.5.0b0__py3-none-any.whl

flock/{modules/enterprise_memory/enterprise_memory_module.py → components/utility/memory_utility_component.py}

@@ -1,8 +1,7 @@
-from __future__ import annotations
-
-"""Enterprise-grade memory module for Flock.
+# src/flock/components/utility/memory_utility_component.py
+"""Enterprise-grade memory utility component for Flock using unified component architecture.
 
-This module persists:
+This component persists:
 • vector embeddings in a Chroma collection (or any collection that
   implements the same API)
 • a concept graph in Neo4j/Memgraph (Cypher-compatible)
@@ -11,12 +10,14 @@ It follows the same life-cycle callbacks as the standard MemoryModule but
 is designed for large-scale, concurrent deployments.
 """
 
+from __future__ import annotations
+
 import asyncio
 import json
 import time
 import uuid
 from pathlib import Path
-from typing import Any, Literal
+from typing import TYPE_CHECKING, Any, Literal
 
 from neo4j import AsyncGraphDatabase
 from opentelemetry import trace
@@ -27,31 +28,32 @@ from flock.adapter.azure_adapter import AzureSearchAdapter
 from flock.adapter.chroma_adapter import ChromaAdapter
 from flock.adapter.faiss_adapter import FAISSAdapter
 from flock.adapter.pinecone_adapter import PineconeAdapter
-
-# Adapter imports
 from flock.adapter.vector_base import VectorAdapter
+from flock.core.component.agent_component_base import AgentComponentConfig
+from flock.core.component.utility_component import UtilityComponent
 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
-from flock.modules.performance.metrics_module import MetricsModule
+from flock.core.registry import flock_component
+
+# Conditional import for MetricsUtilityComponent to avoid circular imports
+if TYPE_CHECKING:
+    from flock.components.utility.metrics_utility_component import (
+        MetricsUtilityComponent,
+    )
+    from flock.core.flock_agent import FlockAgent
 
-logger = get_logger("enterprise_memory")
+logger = get_logger("components.utility.memory")
 tracer = trace.get_tracer(__name__)
 
 
-# ---------------------------------------------------------------------------
-# Configuration
-# ---------------------------------------------------------------------------
-class EnterpriseMemoryModuleConfig(FlockModuleConfig):
-    """Configuration for EnterpriseMemoryModule."""
+class MemoryUtilityConfig(AgentComponentConfig):
+    """Configuration for MemoryUtilityComponent."""
 
     # ---------------------
     # Vector store settings
     # ---------------------
 
-    vector_backend: Literal["chroma", "pinecone", "azure"] = Field(
+    vector_backend: Literal["chroma", "pinecone", "azure", "faiss"] = Field(
         default="chroma",
         description="Which vector backend to use (chroma | pinecone | azure)",
     )
@@ -68,29 +70,48 @@ class EnterpriseMemoryModuleConfig(FlockModuleConfig):
         default=None,
         description="If provided, connect to a remote Chroma HTTP server at this host",
     )
-    chroma_port: int = Field(default=8000, description="Remote Chroma HTTP port")
+    chroma_port: int = Field(
+        default=8000, description="Remote Chroma HTTP port"
+    )
 
     # --- Pinecone ---
-    pinecone_api_key: str | None = Field(default=None, description="Pinecone API key")
-    pinecone_env: str | None = Field(default=None, description="Pinecone environment")
-    pinecone_index: str | None = Field(default=None, description="Pinecone index name")
+    pinecone_api_key: str | None = Field(
+        default=None, description="Pinecone API key"
+    )
+    pinecone_env: str | None = Field(
+        default=None, description="Pinecone environment"
+    )
+    pinecone_index: str | None = Field(
+        default=None, description="Pinecone index name"
+    )
 
     # --- Azure Cognitive Search ---
-    azure_search_endpoint: str | None = Field(default=None, description="Azure search endpoint (https://<service>.search.windows.net)")
-    azure_search_key: str | None = Field(default=None, description="Azure search admin/key")
-    azure_search_index_name: str | None = Field(default=None, description="Azure search index name")
+    azure_search_endpoint: str | None = Field(
+        default=None,
+        description="Azure search endpoint (https://<service>.search.windows.net)",
+    )
+    azure_search_key: str | None = Field(
+        default=None, description="Azure search admin/key"
+    )
+    azure_search_index_name: str | None = Field(
+        default=None, description="Azure search index name"
+    )
 
     # Graph DB (Neo4j / Memgraph) settings
     cypher_uri: str = Field(
         default="bolt://localhost:7687", description="Bolt URI for the graph DB"
     )
     cypher_username: str = Field(default="neo4j", description="Username for DB")
-    cypher_password: str = Field(default="password", description="Password for DB")
+    cypher_password: str = Field(
+        default="password", description="Password for DB"
+    )
 
     similarity_threshold: float = Field(
         default=0.5, description="Cosine-similarity threshold for retrieval"
     )
-    max_results: int = Field(default=10, description="Maximum retrieved memories")
+    max_results: int = Field(
+        default=10, description="Maximum retrieved memories"
+    )
     number_of_concepts_to_extract: int = Field(
         default=3, description="Number of concepts extracted per chunk"
     )
@@ -109,16 +130,17 @@ class EnterpriseMemoryModuleConfig(FlockModuleConfig):
     )
 
 
-# ---------------------------------------------------------------------------
-# Storage Abstraction
-# ---------------------------------------------------------------------------
-class EnterpriseMemoryStore:
-    """Persistence layer that wraps Chroma + Cypher graph."""
+class MemoryStore:
+    """Persistence layer that wraps vector store + Cypher graph."""
 
-    def __init__(self, cfg: EnterpriseMemoryModuleConfig, metrics_module: MetricsModule | None = None):
+    def __init__(
+        self,
+        cfg: MemoryUtilityConfig,
+        metrics_component: MetricsUtilityComponent | None = None,
+    ):
         self.cfg = cfg
-        # Metrics module (DI-resolved or fallback)
-        self._metrics = metrics_module or MetricsModule  # can be either instance or class exposing .record
+        # Metrics component (DI-resolved or fallback)
+        self._metrics = metrics_component
         # Lazy initialise expensive resources
         self._embedding_model: SentenceTransformer | None = None
         self._adapter: VectorAdapter | None = None
@@ -127,15 +149,16 @@ class EnterpriseMemoryStore:
         self._write_lock = asyncio.Lock()
         self._concept_cache: set[str] | None = None  # names of known concepts
 
-    # ---------------------------------------------------------------------
-    # Connections
-    # ---------------------------------------------------------------------
     def _ensure_embedding_model(self) -> SentenceTransformer:
         if self._embedding_model is None:
             logger.debug("Loading embedding model 'all-MiniLM-L6-v2'")
-            with tracer.start_as_current_span("memory.load_embedding_model") as span:
+            with tracer.start_as_current_span(
+                "memory.load_embedding_model"
+            ) as span:
                 try:
-                    self._embedding_model = SentenceTransformer("all-MiniLM-L6-v2")
+                    self._embedding_model = SentenceTransformer(
+                        "all-MiniLM-L6-v2"
+                    )
                     span.set_attribute("model", "all-MiniLM-L6-v2")
                 except Exception as e:
                     span.record_exception(e)
@@ -183,9 +206,6 @@ class EnterpriseMemoryStore:
             )
         return self._driver
 
-    # ---------------------------------------------------------------------
-    # Public API
-    # ---------------------------------------------------------------------
     async def add_entry(
         self,
         content: str,
@@ -194,7 +214,8 @@ class EnterpriseMemoryStore:
     ) -> str:
         """Store a chunk in both vector store and graph DB and return its id."""
         with tracer.start_as_current_span("memory.add_entry") as span:
-            span.set_attribute("entry_id", str(uuid.uuid4()))
+            entry_id = str(uuid.uuid4())
+            span.set_attribute("entry_id", entry_id)
 
             # Embed
             embedding = self._ensure_embedding_model().encode(content).tolist()
@@ -207,7 +228,7 @@ class EnterpriseMemoryStore:
             start_t = time.perf_counter()
             try:
                 adapter.add(
-                    id=span.get_attribute("entry_id"),
+                    id=entry_id,
                     content=content,
                     embedding=embedding,
                     metadata=metadata,
@@ -217,18 +238,22 @@ class EnterpriseMemoryStore:
                 raise
             finally:
                 elapsed = (time.perf_counter() - start_t) * 1000  # ms
-                self._metrics.record(
-                    "memory_add_latency_ms",
-                    elapsed,
-                    {"backend": self.cfg.vector_backend},
-                )
+                if self._metrics:
+                    self._metrics.record(
+                        "memory_add_latency_ms",
+                        elapsed,
+                        {"backend": self.cfg.vector_backend},
+                    )
 
-        # Schedule graph writes (batched)
-        async with self._write_lock:
-            self._pending_writes.append((span.get_attribute("entry_id"), {"concepts": concepts}))
-            if self.cfg.save_interval and len(self._pending_writes) >= self.cfg.save_interval:
-                await self._flush_pending_graph_writes()
-        return span.get_attribute("entry_id")
+            # Schedule graph writes (batched)
+            async with self._write_lock:
+                self._pending_writes.append((entry_id, {"concepts": concepts}))
+                if (
+                    self.cfg.save_interval
+                    and len(self._pending_writes) >= self.cfg.save_interval
+                ):
+                    await self._flush_pending_graph_writes()
+            return entry_id
 
     async def search(
         self, query_text: str, threshold: float, k: int
@@ -247,9 +272,10 @@ class EnterpriseMemoryStore:
             search_start = time.perf_counter()
             vector_hits = adapter.query(embedding=embedding, k=k)
             search_elapsed = (time.perf_counter() - search_start) * 1000
-            self._metrics.record(
-                "memory_search_hits", len(vector_hits), {"backend": backend}
-            )
+            if self._metrics:
+                self._metrics.record(
+                    "memory_search_hits", len(vector_hits), {"backend": backend}
+                )
             for hit in vector_hits:
                 if hit.score < threshold:
                     continue
@@ -263,14 +289,14 @@ class EnterpriseMemoryStore:
                 )
 
             span.set_attribute("results_count", len(results))
-            self._metrics.record(
-                "memory_search_latency_ms", search_elapsed, {"backend": backend}
-            )
+            if self._metrics:
+                self._metrics.record(
+                    "memory_search_latency_ms",
+                    search_elapsed,
+                    {"backend": backend},
+                )
             return results
 
-    # ------------------------------------------------------------------
-    # Graph persistence helpers
-    # ------------------------------------------------------------------
     async def _flush_pending_graph_writes(self):
         """Commit queued node/edge creations to the Cypher store."""
         if not self._pending_writes:
@@ -304,6 +330,7 @@ class EnterpriseMemoryStore:
         """Generate and save a PNG of the concept graph."""
         try:
             import matplotlib
+
             matplotlib.use("Agg")
             import matplotlib.pyplot as plt
             import networkx as nx
@@ -312,7 +339,10 @@ class EnterpriseMemoryStore:
                 "MATCH (c1:Concept)<-[:MENTIONS]-(:Memory)-[:MENTIONS]->(c2:Concept) "
                 "RETURN DISTINCT c1.name AS source, c2.name AS target"
             )
-            edges = [(r["source"], r["target"]) for r in await records.values("source", "target")]
+            edges = [
+                (r["source"], r["target"])
+                for r in await records.values("source", "target")
+            ]
             if not edges:
                 return
 
@@ -321,7 +351,9 @@ class EnterpriseMemoryStore:
 
             pos = nx.spring_layout(G, k=0.4)
             plt.figure(figsize=(12, 9), dpi=100)
-            nx.draw_networkx_nodes(G, pos, node_color="#8fa8d6", node_size=500, edgecolors="white")
+            nx.draw_networkx_nodes(
+                G, pos, node_color="#8fa8d6", node_size=500, edgecolors="white"
+            )
             nx.draw_networkx_edges(G, pos, alpha=0.5, width=1.5)
             nx.draw_networkx_labels(G, pos, font_size=8)
             plt.axis("off")
@@ -335,6 +367,52 @@ class EnterpriseMemoryStore:
         except Exception as e:
             logger.warning("Failed to export concept graph image: %s", e)
 
+    async def _deduplicate_concepts(self, new_concepts: set[str]) -> set[str]:
+        """Return a set of concept names that merges with existing ones to avoid duplicates.
+
+        Strategy: case-insensitive equality first, then fuzzy match via difflib with cutoff 0.85.
+        """
+        await self._ensure_concept_cache()
+        assert self._concept_cache is not None
+
+        from difflib import get_close_matches
+
+        unified: set[str] = set()
+        for concept in new_concepts:
+            # Exact (case-insensitive) match
+            lower = concept.lower()
+            exact = next(
+                (c for c in self._concept_cache if c.lower() == lower), None
+            )
+            if exact:
+                unified.add(exact)
+                continue
+
+            # Fuzzy match (>=0.85 similarity)
+            close = get_close_matches(
+                concept, list(self._concept_cache), n=1, cutoff=0.85
+            )
+            if close:
+                unified.add(close[0])
+                continue
+
+            # No match – treat as new
+            unified.add(concept)
+            self._concept_cache.add(concept)
+        return unified
+
+    async def _ensure_concept_cache(self):
+        if self._concept_cache is not None:
+            return
+        driver = self._ensure_graph_driver()
+        async with driver.session() as session:
+            records = await session.run(
+                "MATCH (c:Concept) RETURN c.name AS name"
+            )
+            self._concept_cache = {
+                r["name"] for r in await records.values("name")
+            }
+
     async def close(self):
         if self._pending_writes:
             await self._flush_pending_graph_writes()
@@ -344,77 +422,51 @@ class EnterpriseMemoryStore:
             self._adapter.close()
 
 
-# ---------------------------------------------------------------------------
-# Module
-# ---------------------------------------------------------------------------
-@flock_component(config_class=EnterpriseMemoryModuleConfig)
-class EnterpriseMemoryModule(FlockModule):
-    """Enterprise-ready memory module using real datastores."""
+@flock_component(config_class=MemoryUtilityConfig)
+class MemoryUtilityComponent(UtilityComponent):
+    """Enterprise-ready memory utility component using real datastores."""
 
-    name: str = "enterprise_memory"
-    config: EnterpriseMemoryModuleConfig = Field(default_factory=EnterpriseMemoryModuleConfig)
-
-    _store: EnterpriseMemoryStore | None = None
-    _container: Any | None = None  # DI container if supplied
-    _metrics_module: MetricsModule | None = None
+    config: MemoryUtilityConfig = Field(
+        default_factory=MemoryUtilityConfig,
+        description="Memory configuration",
+    )
 
-    # ----------------------------------------------------------
-    # DI-enabled constructor
-    # ----------------------------------------------------------
     def __init__(
         self,
-        name: str = "enterprise_memory",
-        config: EnterpriseMemoryModuleConfig | None = None,
-        *,
-        container: object | None = None,
-        **kwargs,
+        name: str = "memory",
+        config: MemoryUtilityConfig | None = None,
+        **data,
     ):
-        """Create a new EnterpriseMemoryModule instance.
-
-        Parameters
-        ----------
-        container : ServiceProvider | None
-            Optional DI container used to resolve shared services.  When
-            provided, the module will attempt to resolve
-            :class:`flock.modules.performance.metrics_module.MetricsModule` from
-            it.  Falling back to the global singleton when not available keeps
-            backward-compatibility.
-        """
-        from wd.di.container import (
-            ServiceProvider,  # Local import to avoid hard dependency if wd.di is absent
-        )
-
         if config is None:
-            config = EnterpriseMemoryModuleConfig()
+            config = MemoryUtilityConfig()
+        super().__init__(name=name, config=config, **data)
 
-        super().__init__(name=name, config=config, **kwargs)
+        self._store: MemoryStore | None = None
+        self._metrics_component: MetricsUtilityComponent | None = None
 
-        self._container = container if isinstance(container, ServiceProvider) else None
-
-        # Attempt to resolve MetricsModule via DI, then via FlockModule registry
-        resolved_metrics: MetricsModule | None = None
-        if self._container is not None:
-            try:
-                resolved_metrics = self._container.get_service(MetricsModule)
-            except Exception:
-                resolved_metrics = None
-
-        if resolved_metrics is None:
-            resolved_metrics = MetricsModule._INSTANCE
+    def _get_metrics_component(self) -> MetricsUtilityComponent | None:
+        """Try to get the metrics component from the singleton if available."""
+        try:
+            from flock.components.utility.metrics_utility_component import (
+                MetricsUtilityComponent,
+            )
 
-        self._metrics_module = resolved_metrics
+            return MetricsUtilityComponent._INSTANCE
+        except ImportError:
+            return None
 
-    # ----------------------------------------------------------
-    # Life-cycle hooks
-    # ----------------------------------------------------------
     async def on_initialize(
         self,
         agent: FlockAgent,
         inputs: dict[str, Any],
         context: FlockContext | None = None,
     ) -> None:
-        self._store = EnterpriseMemoryStore(self.config, self._metrics_module)
-        logger.info("EnterpriseMemoryModule initialised", agent=agent.name)
+        """Initialize the memory store."""
+        self._metrics_component = self._get_metrics_component()
+        self._store = MemoryStore(self.config, self._metrics_component)
+        logger.info(
+            "MemoryUtilityComponent initialised for agent: %s", agent.name
+        )
 
     async def on_pre_evaluate(
         self,
@@ -422,6 +474,7 @@ class EnterpriseMemoryModule(FlockModule):
         inputs: dict[str, Any],
         context: FlockContext | None = None,
     ) -> dict[str, Any]:
+        """Retrieve relevant memories and inject into inputs."""
         if not self._store:
             return inputs
         try:
@@ -434,10 +487,15 @@ class EnterpriseMemoryModule(FlockModule):
             if matches:
                 inputs = {**inputs, "context": matches}
                 # Advertise new input key to DSPy signature if needed
-                if isinstance(agent.input, str) and "context:" not in agent.input:
+                if (
+                    isinstance(agent.input, str)
+                    and "context:" not in agent.input
+                ):
                     agent.input += ", context: list | retrieved memories"
         except Exception as e:
-            logger.warning("Enterprise memory retrieval failed: %s", e, agent=agent.name)
+            logger.warning(
+                "Memory retrieval failed for agent %s: %s", agent.name, e
+            )
         return inputs
 
     async def on_post_evaluate(
@@ -447,35 +505,40 @@ class EnterpriseMemoryModule(FlockModule):
         context: FlockContext | None = None,
         result: dict[str, Any] | None = None,
     ) -> dict[str, Any] | None:
+        """Store new memories from inputs and results."""
         if not self._store:
             return result
         try:
-            full_text = json.dumps(inputs) + (json.dumps(result) if result else "")
+            full_text = json.dumps(inputs) + (
+                json.dumps(result) if result else ""
+            )
             concepts = await self._extract_concepts(agent, full_text)
             if self._store:
                 concepts = await self._store._deduplicate_concepts(concepts)
             await self._store.add_entry(full_text, concepts)
         except Exception as e:
-            logger.warning("Enterprise memory store failed: %s", e, agent=agent.name)
+            logger.warning(
+                "Memory store failed for agent %s: %s", agent.name, e
+            )
         return result
 
     async def on_terminate(
         self,
         agent: FlockAgent,
         inputs: dict[str, Any],
-        result: dict[str, Any],
         context: FlockContext | None = None,
+        result: dict[str, Any] | None = None,
     ) -> None:
+        """Clean up memory store resources."""
         if self._store:
             await self._store.close()
 
-    # ----------------------------------------------------------
-    # Helpers (mostly copied from original module but simplified)
-    # ----------------------------------------------------------
-    async def _extract_concepts(self, agent: FlockAgent, text: str) -> set[str]:
+    async def _extract_concepts(
+        self, agent: FlockAgent, text: str
+    ) -> set[str]:
         """Use the LLM to extract concept tokens."""
         concept_signature = agent.create_dspy_signature_class(
-            f"{agent.name}_concept_extractor_enterprise",
+            f"{agent.name}_concept_extractor_memory",
             "Extract key concepts from text",
             "text: str | Input text -> concepts: list[str] | key concepts lower case",
         )
@@ -483,44 +546,3 @@ class EnterpriseMemoryModule(FlockModule):
         predictor = agent._select_task(concept_signature, "Completion")
         res = predictor(text=text)
         return set(getattr(res, "concepts", []))
-
-    # --------------------------------------------------------------
-    # Concept helpers
-    # --------------------------------------------------------------
-    async def _ensure_concept_cache(self):
-        if self._concept_cache is not None:
-            return
-        driver = self._ensure_graph_driver()
-        async with driver.session() as session:
-            records = await session.run("MATCH (c:Concept) RETURN c.name AS name")
-            self._concept_cache = {r["name"] for r in await records.values("name")}
-
-    async def _deduplicate_concepts(self, new_concepts: set[str]) -> set[str]:
-        """Return a set of concept names that merges with existing ones to avoid duplicates.
-
-        Strategy: case-insensitive equality first, then fuzzy match via difflib with cutoff 0.85.
-        """
-        await self._ensure_concept_cache()
-        assert self._concept_cache is not None
-
-        from difflib import get_close_matches
-
-        unified: set[str] = set()
-        for concept in new_concepts:
-            # Exact (case-insensitive) match
-            lower = concept.lower()
-            exact = next((c for c in self._concept_cache if c.lower() == lower), None)
-            if exact:
-                unified.add(exact)
-                continue
-
-            # Fuzzy match (>=0.85 similarity)
-            close = get_close_matches(concept, list(self._concept_cache), n=1, cutoff=0.85)
-            if close:
-                unified.add(close[0])
-                continue
-
-            # No match – treat as new
-            unified.add(concept)
-            self._concept_cache.add(concept)
-        return unified