Delt minne og kommunikasjon mellom agenter

# Simple in-memory blackboard using a dictionary
from threading import Lock

class Blackboard:
    def __init__(self):
        self._data = {}
        self._lock = Lock()  # thread-safe for parallel agents

    def write(self, key: str, value, agent_id: str):
        with self._lock:
            self._data[key] = {'value': value, 'written_by': agent_id}
            print(f'[{agent_id}] wrote: {key}')

    def read(self, key: str):
        with self._lock:
            return self._data.get(key, {}).get('value')

    def keys(self):
        with self._lock:
            return list(self._data.keys())

blackboard = Blackboard()

import redis
import json

class RedisSharedMemory:
    def __init__(self, prefix='agent:'):
        self.redis = redis.Redis(host='localhost', port=6379, decode_responses=True)
        self.prefix = prefix

    def set(self, key: str, value, ttl_seconds=3600):
        full_key = self.prefix + key
        self.redis.setex(full_key, ttl_seconds, json.dumps(value))

    def get(self, key: str):
        full_key = self.prefix + key
        raw = self.redis.get(full_key)
        return json.loads(raw) if raw else None

    def append_to_list(self, key: str, item):
        full_key = self.prefix + key
        self.redis.rpush(full_key, json.dumps(item))

    def get_list(self, key: str):
        full_key = self.prefix + key
        return [json.loads(x) for x in self.redis.lrange(full_key, 0, -1)]

memory = RedisSharedMemory(prefix='research_project:')

class NamespacedMemory:
    def __init__(self, project_id: str, agent_id: str, redis_client):
        self.base = f'{project_id}:{agent_id}'
        self.redis = redis_client

    def write_finding(self, topic: str, content: str):
        key = f'{self.base}:findings:{topic}'
        self.redis.set(key, content)

    def read_all_findings(self, project_id: str):
        # Read findings from ALL agents in this project
        pattern = f'{project_id}:*:findings:*'
        keys = self.redis.keys(pattern)
        return {k: self.redis.get(k) for k in keys}

# Usage
researcher_memory = NamespacedMemory('proj_123', 'researcher', redis_client)
researcher_memory.write_finding('competitors', 'OpenAI, Anthropic, Google...')

writer_memory = NamespacedMemory('proj_123', 'writer', redis_client)
all_findings = writer_memory.read_all_findings('proj_123')

from pydantic import BaseModel
from typing import Optional, list
from datetime import datetime

class ResearchFinding(BaseModel):
    topic: str
    summary: str
    sources: list[str]
    confidence: float  # 0.0 to 1.0
    written_by: str
    timestamp: datetime

# Validated write - bad data is caught before it enters shared memory
def write_finding(memory, finding_dict: dict):
    finding = ResearchFinding(**finding_dict)  # validates on creation
    memory.set(f'findings:{finding.topic}', finding.model_dump())
    print(f'Validated finding written for topic: {finding.topic}')

import redis

# Publisher (researcher agent)
def researcher_agent(topic, redis_client):
    findings = do_research(topic)
    redis_client.set(f'findings:{topic}', findings)
    
    # Notify all subscribers that research is done
    redis_client.publish('agent_events', f'research_complete:{topic}')
    print(f'Research complete, published event for topic: {topic}')

# Subscriber (writer agent) - runs in separate process
def writer_agent_listener(redis_client):
    pubsub = redis_client.pubsub()
    pubsub.subscribe('agent_events')
    
    for message in pubsub.listen():
        if message['type'] == 'message':
            event = message['data']
            if event.startswith('research_complete:'):
                topic = event.split(':')[1]
                findings = redis_client.get(f'findings:{topic}')
                write_draft(findings)  # start writing immediately

from langgraph.graph import StateGraph
from typing import TypedDict

class SharedState(TypedDict):
    # All shared data lives here - every node can read any field
    query: str
    research_findings: str   # written by researcher, read by writer
    written_draft: str       # written by writer, read by reviewer
    review_notes: str        # written by reviewer, read by writer (loop)
    final_output: str        # written by synthesizer

# Researcher writes to 'research_findings'
def researcher(state: SharedState) -> dict:
    findings = search_and_summarize(state['query'])
    return {'research_findings': findings}  # partial state update

# Writer reads 'research_findings', writes 'written_draft'
def writer(state: SharedState) -> dict:
    draft = write_from_findings(state['research_findings'])  # reads researcher output
    return {'written_draft': draft}

# Optimistic locking with Redis
def safe_write(redis_client, key, new_value, expected_version):
    with redis_client.pipeline() as pipe:
        try:
            pipe.watch(key + ':version')  # watch for concurrent modification
            current_version = int(pipe.get(key + ':version') or 0)
            
            if current_version != expected_version:
                raise ValueError(f'Version conflict: expected {expected_version}, got {current_version}')
            
            pipe.multi()  # start transaction
            pipe.set(key, new_value)
            pipe.set(key + ':version', current_version + 1)
            pipe.execute()  # atomic commit
            print('Write successful')
        except redis.WatchError:
            print('Conflict detected, retry write')

def cleanup_project_memory(redis_client, project_id: str):
    pattern = f'{project_id}:*'
    keys = redis_client.keys(pattern)
    if keys:
        redis_client.delete(*keys)
        print(f'Cleaned up {len(keys)} memory entries for project {project_id}')

# Set TTL when writing
def write_with_ttl(redis_client, key, value, ttl_hours=2):
    redis_client.setex(
        key,
        ttl_hours * 3600,  # convert to seconds
        json.dumps(value)
    )

# Register cleanup callback when workflow completes
def on_workflow_complete(project_id):
    cleanup_project_memory(redis_client, project_id)
    print(f'Workflow {project_id} complete, memory cleaned up')

import time
from dataclasses import dataclass

@dataclass
class AgentAction:
    agent_id: str
    action_type: str     # 'research', 'write', 'review', 'tool_call'
    input_summary: str
    output_summary: str
    timestamp: float
    success: bool

def log_action(memory, action: AgentAction):
    key = f'action_log:{action.agent_id}:{action.timestamp}'
    memory.set(key, vars(action))

# Usage in an agent
def researcher_with_logging(state, memory):
    start = time.time()
    findings = do_research(state['query'])
    log_action(memory, AgentAction(
        agent_id='researcher',
        action_type='research',
        input_summary=state['query'][:100],
        output_summary=findings[:100],
        timestamp=start,
        success=True
    ))
    return findings