Shared Memory in Machine Learning

from typing import TypedDict, Annotated
from langgraph.graph import StateGraph, END
from langgraph.checkpoint.redis import RedisSaver
import operator

# Define shared state schema with reducer functions
class SharedMemory(TypedDict):
    # append reducer: each agent's findings are accumulated
    research_findings: Annotated[list[str], operator.add]
    # append reducer: conversation history grows
    messages: Annotated[list[dict], operator.add]
    # last-writer-wins: only the latest summary matters
    current_summary: str
    # append reducer: action items from all agents
    action_items: Annotated[list[str], operator.add]
    # metadata for coordination
    iteration_count: int

# Agent functions that read/write shared state
def researcher_agent(state: SharedMemory) -> dict:
    """Reads existing findings, adds new ones."""
    existing = state.get("research_findings", [])
    # Agent uses existing findings to avoid duplicates
    new_finding = f"Finding #{len(existing) + 1}: Market size is $4.2B"
    return {
        "research_findings": [new_finding],
        "messages": [{"role": "researcher", "content": new_finding}],
    }

def analyst_agent(state: SharedMemory) -> dict:
    """Reads research findings, produces analysis."""
    findings = state.get("research_findings", [])
    analysis = f"Analysis of {len(findings)} findings: growth trend detected"
    return {
        "current_summary": analysis,
        "messages": [{"role": "analyst", "content": analysis}],
        "action_items": ["Validate growth trend with secondary sources"],
    }

def writer_agent(state: SharedMemory) -> dict:
    """Reads summary and findings, writes final output."""
    summary = state.get("current_summary", "")
    findings = state.get("research_findings", [])
    report = f"Report: Based on {len(findings)} findings. {summary}"
    return {
        "messages": [{"role": "writer", "content": report}],
        "iteration_count": state.get("iteration_count", 0) + 1,
    }

# Build the graph with shared memory
graph = StateGraph(SharedMemory)
graph.add_node("researcher", researcher_agent)
graph.add_node("analyst", analyst_agent)
graph.add_node("writer", writer_agent)
graph.add_edge("researcher", "analyst")
graph.add_edge("analyst", "writer")
graph.add_edge("writer", END)
graph.set_entry_point("researcher")

# Compile with Redis checkpointer for persistence
checkpointer = RedisSaver(redis_url="redis://localhost:6379")
app = graph.compile(checkpointer=checkpointer)

# Run with thread-level persistence
config = {"configurable": {"thread_id": "session-001"}}
result = app.invoke(
    {"research_findings": [], "messages": [], "action_items": [], "iteration_count": 0},
    config=config,
)
print(result["current_summary"])
print(f"Total action items: {len(result['action_items'])}")

import redis
import json
import time
import hashlib
from typing import Any, Optional
from dataclasses import dataclass, field

@dataclass
class StateEntry:
    value: Any
    version: int
    author: str
    timestamp: float
    checksum: str = ""

class SharedMemoryStore:
    """Production shared memory with optimistic concurrency control."""

    def __init__(self, redis_url: str = "redis://localhost:6379", namespace: str = "agent"):
        self.redis = redis.from_url(redis_url, decode_responses=True)
        self.ns = namespace

    def _key(self, session_id: str, field: str) -> str:
        return f"{self.ns}:{session_id}:{field}"

    def read(self, session_id: str, field: str) -> Optional[StateEntry]:
        """Read a field from shared memory."""
        raw = self.redis.get(self._key(session_id, field))
        if raw is None:
            return None
        data = json.loads(raw)
        return StateEntry(**data)

    def write(
        self,
        session_id: str,
        field: str,
        value: Any,
        author: str,
        expected_version: Optional[int] = None,
    ) -> bool:
        """Write with optimistic concurrency control.
        
        Returns True if write succeeded, False if version conflict.
        Uses Redis WATCH/MULTI for atomic check-and-set.
        """
        key = self._key(session_id, field)
        with self.redis.pipeline() as pipe:
            try:
                pipe.watch(key)
                current_raw = pipe.get(key)
                current_version = 0
                if current_raw:
                    current = json.loads(current_raw)
                    current_version = current["version"]
                
                # Optimistic concurrency check
                if expected_version is not None and current_version != expected_version:
                    return False  # Conflict detected
                
                new_entry = StateEntry(
                    value=value,
                    version=current_version + 1,
                    author=author,
                    timestamp=time.time(),
                    checksum=hashlib.sha256(
                        json.dumps(value, sort_keys=True).encode()
                    ).hexdigest()[:16],
                )
                pipe.multi()
                pipe.set(key, json.dumps(new_entry.__dict__))
                pipe.execute()
                # Publish state change event for reactive agents
                self.redis.publish(
                    f"{self.ns}:events:{session_id}",
                    json.dumps({"field": field, "author": author, "version": new_entry.version}),
                )
                return True
            except redis.WatchError:
                return False  # Another agent modified the key

    def append_to_list(
        self, session_id: str, field: str, items: list, author: str
    ) -> int:
        """Atomic append to a list field. Returns new list length."""
        key = self._key(session_id, field)
        # Use Lua script for atomic read-modify-write
        lua_script = """
        local current = redis.call('GET', KEYS[1])
        local data
        if current then
            data = cjson.decode(current)
        else
            data = {value={}, version=0, author='', timestamp=0, checksum=''}
        end
        local new_items = cjson.decode(ARGV[1])
        for _, item in ipairs(new_items) do
            table.insert(data.value, item)
        end
        data.version = data.version + 1
        data.author = ARGV[2]
        data.timestamp = tonumber(ARGV[3])
        redis.call('SET', KEYS[1], cjson.encode(data))
        return #data.value
        """
        return self.redis.eval(
            lua_script, 1, key, json.dumps(items), author, str(time.time())
        )

    def checkpoint(self, session_id: str) -> str:
        """Snapshot all fields for a session. Returns checkpoint ID."""
        pattern = f"{self.ns}:{session_id}:*"
        keys = list(self.redis.scan_iter(match=pattern))
        snapshot = {}
        for key in keys:
            field_name = key.split(":")[-1]
            snapshot[field_name] = self.redis.get(key)
        checkpoint_id = f"cp:{session_id}:{int(time.time())}"
        self.redis.set(checkpoint_id, json.dumps(snapshot), ex=86400)  # 24h TTL
        return checkpoint_id

    def restore(self, checkpoint_id: str) -> bool:
        """Restore shared memory from a checkpoint."""
        raw = self.redis.get(checkpoint_id)
        if not raw:
            return False
        snapshot = json.loads(raw)
        # Extract session_id from checkpoint_id
        session_id = checkpoint_id.split(":")[1]
        for field_name, value in snapshot.items():
            key = self._key(session_id, field_name)
            self.redis.set(key, value)
        return True


# Usage example
store = SharedMemoryStore(redis_url="redis://localhost:6379")

# Agent 1 writes research findings
store.append_to_list("session-42", "findings", ["GDP growth at 7.2%"], "researcher")

# Agent 2 reads and writes analysis (with version check)
entry = store.read("session-42", "findings")
if entry:
    store.write("session-42", "analysis", 
                f"Analyzed {len(entry.value)} findings", 
                author="analyst", expected_version=None)

# Checkpoint before risky operation
cp_id = store.checkpoint("session-42")
print(f"Checkpointed: {cp_id}")

from typing import Any, Callable
from dataclasses import dataclass, field
import threading
import time

@dataclass
class BlackboardEntry:
    key: str
    value: Any
    source_agent: str
    confidence: float
    timestamp: float = field(default_factory=time.time)

class Blackboard:
    """Classical blackboard architecture for multi-agent coordination."""

    def __init__(self):
        self._state: dict[str, BlackboardEntry] = {}
        self._subscribers: dict[str, list[Callable]] = {}
        self._lock = threading.RLock()
        self._history: list[tuple[str, BlackboardEntry]] = []

    def write(self, key: str, value: Any, source: str, confidence: float = 1.0):
        """Post information to the blackboard."""
        with self._lock:
            entry = BlackboardEntry(
                key=key, value=value, source_agent=source,
                confidence=confidence, timestamp=time.time(),
            )
            # Conflict resolution: higher confidence wins
            if key in self._state:
                existing = self._state[key]
                if existing.confidence > confidence:
                    return  # Existing entry has higher confidence
            self._state[key] = entry
            self._history.append(("write", entry))
        # Notify subscribers outside the lock
        for callback in self._subscribers.get(key, []):
            callback(entry)

    def read(self, key: str) -> BlackboardEntry | None:
        """Read from the blackboard."""
        with self._lock:
            return self._state.get(key)

    def read_all(self, prefix: str = "") -> dict[str, BlackboardEntry]:
        """Read all entries matching a prefix."""
        with self._lock:
            return {
                k: v for k, v in self._state.items()
                if k.startswith(prefix)
            }

    def subscribe(self, key: str, callback: Callable):
        """Subscribe to changes on a specific key."""
        self._subscribers.setdefault(key, []).append(callback)

    def get_history(self, last_n: int = 10) -> list[tuple[str, BlackboardEntry]]:
        """Get recent history for debugging / time-travel."""
        return self._history[-last_n:]


# Define knowledge sources (agents)
class KnowledgeSource:
    def __init__(self, name: str, blackboard: Blackboard):
        self.name = name
        self.bb = blackboard

class ResearchAgent(KnowledgeSource):
    def act(self, query: str):
        # Simulate research
        result = f"Research on '{query}': Found 3 relevant papers"
        self.bb.write(f"research:{query}", result, self.name, confidence=0.85)
        self.bb.write("status", f"{self.name} completed research", self.name)

class AnalysisAgent(KnowledgeSource):
    def act(self):
        # Read all research entries from the blackboard
        research = self.bb.read_all(prefix="research:")
        analysis = f"Synthesized {len(research)} research entries"
        self.bb.write("analysis:summary", analysis, self.name, confidence=0.90)

class QualityAgent(KnowledgeSource):
    def act(self):
        # Review analysis quality
        analysis = self.bb.read("analysis:summary")
        if analysis and analysis.confidence < 0.95:
            self.bb.write(
                "review:needed", True, self.name, confidence=1.0
            )


# Orchestrate via the blackboard
bb = Blackboard()
researcher = ResearchAgent("researcher-1", bb)
analyst = AnalysisAgent("analyst-1", bb)
reviewer = QualityAgent("reviewer-1", bb)

# Agents interact through the shared blackboard
researcher.act("transformer architectures")
researcher.act("multi-agent memory")
analyst.act()
reviewer.act()

# Inspect the blackboard state
for key, entry in bb._state.items():
    print(f"[{entry.source_agent}] {key}: {entry.value} (conf: {entry.confidence})")

from mem0 import Memory

# Initialize Mem0 with shared memory configuration
config = {
    "llm": {
        "provider": "openai",
        "config": {"model": "gpt-4o-mini", "temperature": 0.1},
    },
    "embedder": {
        "provider": "openai",
        "config": {"model": "text-embedding-3-small"},
    },
    "vector_store": {
        "provider": "qdrant",
        "config": {"collection_name": "agent_memory", "host": "localhost", "port": 6333},
    },
}

memory = Memory.from_config(config)

# Agent 1: Customer research agent stores findings
memory.add(
    "Customer prefers budget smartphones under INR 15,000 with good camera",
    user_id="customer-123",
    agent_id="research-agent",
    metadata={"category": "preference", "confidence": 0.92},
)

memory.add(
    "Customer previously purchased Redmi Note 12 in March 2025",
    user_id="customer-123",
    agent_id="research-agent",
    metadata={"category": "history", "confidence": 0.99},
)

# Agent 2: Recommendation agent reads shared memory
relevant_memories = memory.search(
    "What phones should I recommend to this customer?",
    user_id="customer-123",
    limit=5,
)

for mem in relevant_memories["results"]:
    print(f"Memory: {mem['memory']} (score: {mem['score']:.2f})")

# Agent 3: Follow-up agent adds new information
memory.add(
    "Customer mentioned interest in 5G connectivity for Jio network",
    user_id="customer-123",
    agent_id="followup-agent",
    metadata={"category": "preference", "confidence": 0.88},
)

# All agents now have access to the combined memory
all_memories = memory.get_all(user_id="customer-123")
print(f"Total shared memories for customer: {len(all_memories['results'])}")

# LangGraph + Redis shared memory configuration
langgraph:
  state_type: SharedMemory
  checkpointer:
    backend: redis
    redis_url: redis://localhost:6379
    key_prefix: "agent:state:"
    ttl_seconds: 86400  # 24 hours
  reducers:
    messages: append      # list concatenation
    findings: append      # list concatenation  
    summary: replace      # last-writer-wins
    metadata: merge       # dict merge
  view_projection:
    max_tokens_per_agent: 4000
    strategy: recency_and_relevance
    summarize_after: 20   # summarize state after 20 entries
  conflict_resolution:
    strategy: optimistic_concurrency
    max_retries: 3
    retry_backoff_ms: 50