Knowledge Graph Research Excavation Report¶

# 1. Use UNWIND for batch operations (dramatically faster)
query = """
UNWIND $batch as row
CREATE (n:Node {id: row.id, data: row.data})
"""
driver.execute_query(query, batch=data_batch)

# 2. Create indexes for frequently filtered properties
driver.execute_query("CREATE INDEX node_name IF NOT EXISTS FOR (n:Node) ON (n.name)")

# 3. Use parameters (never concatenate)
query = "MATCH (n:Node {id: $id}) RETURN n"
result = driver.execute_query(query, id=node_id)

# 4. Install Rust extension for 3-10x speedup
# pip install neo4j-rust-ext

# .amplihack/memory/config.json
{
  "neo4j": {
    "uri": "bolt://localhost:7687",
    "database": "project_memory_<hash>",  # Isolated per project
    "auth": ("neo4j", "generated_password")
  }
}

# Start project-specific container
docker run -d \
  --name neo4j-project-abc123 \
  -p 7687:7687 \
  -v ./amplihack/memory/data:/data \
  neo4j:latest

import json
from neo4j import GraphDatabase

def export_graph_to_json(driver, database="neo4j"):
    """Export entire graph to JSON (for version control)"""
    nodes_query = "MATCH (n) RETURN id(n) as id, labels(n) as labels, properties(n) as props"
    rels_query = "MATCH ()-[r]->() RETURN id(r) as id, type(r) as type, properties(r) as props, id(startNode(r)) as from, id(endNode(r)) as to"

    with driver.session(database=database) as session:
        nodes = list(session.run(nodes_query))
        rels = list(session.run(rels_query))

    return {
        "nodes": [dict(n) for n in nodes],
        "relationships": [dict(r) for r in rels]
    }

# Usage: commit to git
backup = export_graph_to_json(driver)
with open('.amplihack/memory/graph_backup.json', 'w') as f:
    json.dump(backup, f, indent=2)

from blarify import CodeGraph

# Initialize graph generator
graph = CodeGraph(
    codebase_path="./src",
    languages=["python", "javascript"],
    use_scip=True  # Auto-detected if available
)

# Generate graph
graph.build()

# Export to Neo4j
graph.export_to_neo4j(
    uri="bolt://localhost:7687",
    auth=("neo4j", "password")
)

# Query the graph
results = graph.query("""
    MATCH (f1:Function)-[:CALLS*1..3]->(f2:Function)
    WHERE f1.name = 'main'
    RETURN f1.name, f2.name, length(path) as depth
""")

from blarify import CodeGraph, NodeExtractor

class CustomExtractor(NodeExtractor):
    def extract(self, ast_node):
        if ast_node.type == "decorator_definition":
            return {
                "node_type": "Decorator",
                "properties": {
                    "name": ast_node.child_by_field_name("name").text,
                    "file": self.current_file
                }
            }
        return None

# Register custom extractor
graph = CodeGraph(codebase_path="./src")
graph.add_extractor(CustomExtractor())

# Export to custom graph store
def export_to_custom_store(graph_data):
    # graph_data: {"nodes": [...], "edges": [...]}
    for node in graph_data["nodes"]:
        custom_store.add_node(node)
    for edge in graph_data["edges"]:
        custom_store.add_edge(edge)

graph.export(format="dict", callback=export_to_custom_store)

class KnowledgeGraph:
    def __init__(self):
        self.core_memory = CoreMemoryStore()       # Persistent facts
        self.episodic = EpisodicMemoryStore()      # Time-stamped events
        self.semantic = SemanticMemoryStore()      # Entity relationships
        self.procedural = ProceduralMemoryStore()  # How-to knowledge
        self.resources = ResourceMemoryStore()     # Documents, code

    def update(self, event):
        # Route to appropriate memory store
        if event.type == "conversation":
            self.episodic.add(event)
            entities = self.extract_entities(event)
            self.semantic.add_entities(entities)
        elif event.type == "code_change":
            self.resources.add(event)
            procedures = self.extract_procedures(event)
            self.procedural.add(procedures)

# Vector embeddings + Graph structure
class HybridKnowledgeStore:
    def __init__(self):
        self.graph = Neo4jDriver()
        self.vectors = VectorDatabase()  # FAISS, Pinecone, etc.

    def add_entity(self, entity):
        # Store in graph
        self.graph.create_node(entity)

        # Store embedding
        embedding = self.embed(entity.description)
        self.vectors.add(entity.id, embedding)

    def search(self, query):
        # 1. Vector similarity search (semantic)
        candidates = self.vectors.search(query, top_k=50)

        # 2. Graph traversal (structural)
        related = self.graph.expand(candidates, depth=2)

        # 3. Rerank by combined score
        return self.rerank(candidates + related)

def process_event(event, kg):
    # Extract entities
    entities = extract_entities(event.content)

    # Resolve to existing entities (deduplication)
    for entity in entities:
        existing = kg.find_similar(entity, threshold=0.85)
        if existing:
            kg.merge(existing, entity)
        else:
            kg.create(entity)

    # Extract relationships
    relationships = extract_relationships(entities)
    kg.add_relationships(relationships)

    # Update community structure
    kg.recompute_communities()

# Multi-modal retrieval
def retrieve_context(query, kg):
    # Semantic search
    entities = kg.semantic.search(query, top_k=10)

    # Find related entities (graph traversal)
    expanded = kg.graph.expand(entities, depth=2)

    # Find relevant episodes (temporal)
    episodes = kg.episodic.search(
        entities=expanded,
        time_range=(now - 30_days, now)
    )

    # Find procedures (how-to)
    procedures = kg.procedural.find_by_entities(expanded)

    return {
        "entities": expanded,
        "episodes": episodes,
        "procedures": procedures
    }

def handle_contradiction(kg, new_fact, old_fact):
    # Temporal invalidation (Zep pattern)
    if new_fact.contradicts(old_fact):
        # Don't delete old fact, mark as invalidated
        old_fact.t_invalid = new_fact.t_valid
        old_fact.invalidated_by = new_fact.id
        kg.update(old_fact)
        kg.add(new_fact)

class EpisodicMemory:
    def add_episode(self, event):
        episode = {
            "id": generate_id(),
            "timestamp": event.timestamp,
            "type": event.type,
            "summary": summarize(event.content),
            "details": event.content,
            "actor": event.user_id
        }

        # Store in graph
        self.db.create_node("Episode", episode)

        # Extract and link entities
        entities = extract_entities(event.content)
        for entity in entities:
            self.db.create_relationship(
                episode["id"],
                "MENTIONS",
                entity.id
            )

    def search(self, query, time_range=None):
        # Temporal + semantic search
        results = self.db.query("""
            MATCH (e:Episode)-[:MENTIONS]->(entity:Entity)
            WHERE entity.name CONTAINS $query
              AND e.timestamp > $start
              AND e.timestamp < $end
            RETURN e
            ORDER BY e.timestamp DESC
        """, query=query, start=time_range[0], end=time_range[1])
        return results

class SemanticMemory:
    def add_entity(self, entity):
        # Check for duplicates
        existing = self.find_similar(entity)
        if existing:
            self.merge(existing, entity)
        else:
            self.db.create_node("Entity", entity)

    def find_similar(self, entity, threshold=0.85):
        # Hybrid search: embedding + graph
        embedding = self.embed(entity.name + " " + entity.summary)

        # Vector similarity
        candidates = self.vector_db.search(embedding, top_k=10)

        # Rerank by name similarity and context
        for candidate in candidates:
            score = self.similarity_score(entity, candidate)
            if score > threshold:
                return candidate
        return None

    def query(self, query_text):
        # Natural language query to graph traversal
        entities = self.extract_entities(query_text)

        # Graph traversal
        results = self.db.query("""
            MATCH (e1:Entity)-[r]-(e2:Entity)
            WHERE e1.name IN $entities
            RETURN e1, r, e2
        """, entities=[e.name for e in entities])

        return results

class ProceduralMemory:
    def add_procedure(self, procedure):
        proc = {
            "id": generate_id(),
            "name": procedure.name,
            "description": procedure.description,
            "steps": procedure.steps,
            "trigger_pattern": procedure.trigger,
            "success_rate": 0.5,  # Initial neutral
            "times_used": 0
        }
        self.db.create_node("Procedure", proc)

    def find_procedure(self, context):
        # Match context to procedure triggers
        results = self.db.query("""
            MATCH (p:Procedure)
            WHERE p.trigger_pattern =~ $pattern
            RETURN p
            ORDER BY p.success_rate DESC, p.times_used DESC
        """, pattern=self.extract_pattern(context))
        return results

    def record_outcome(self, procedure_id, success):
        # Update success rate
        proc = self.db.get_node(procedure_id)
        proc.times_used += 1

        # Exponential moving average
        alpha = 0.1
        proc.success_rate = (
            alpha * (1 if success else 0) +
            (1 - alpha) * proc.success_rate
        )
        proc.last_used = datetime.now()
        self.db.update_node(proc)

class WorkingMemory:
    def __init__(self):
        self.current_file = None
        self.current_function = None
        self.recent_context = deque(maxlen=10)  # Last 10 interactions
        self.active_entities = set()  # Entities in current focus
        self.active_goal = None

    def update(self, event):
        # Add to recent context
        self.recent_context.append(event)

        # Extract entities mentioned
        entities = extract_entities(event.content)
        self.active_entities.update(entities)

        # Detect goal/intent
        if event.type == "user_request":
            self.active_goal = extract_goal(event.content)

    def get_context(self):
        # Return current context for LLM
        return {
            "current_file": self.current_file,
            "current_function": self.current_function,
            "recent_interactions": list(self.recent_context),
            "active_entities": list(self.active_entities),
            "goal": self.active_goal
        }

class ProspectiveMemory:
    def add_intention(self, goal, trigger=None, due_date=None):
        intention = {
            "id": generate_id(),
            "goal": goal,
            "created_at": datetime.now(),
            "due_date": due_date,
            "trigger_condition": trigger,
            "status": "planned"
        }
        self.db.create_node("Intention", intention)

    def check_triggers(self, current_state):
        # Check if any intentions should activate
        results = self.db.query("""
            MATCH (i:Intention)
            WHERE i.status = 'planned'
              AND (i.due_date < datetime() OR i.trigger_condition IN $events)
            RETURN i
        """, events=current_state.recent_events)

        for intention in results:
            self.activate(intention)

    def activate(self, intention):
        # Move intention to working memory
        intention.status = "in_progress"
        self.db.update_node(intention)
        notify_user(f"Ready to work on: {intention.goal}")

class UnifiedMemoryGraph:
    def __init__(self, neo4j_driver):
        self.driver = neo4j_driver
        self.vector_db = VectorDatabase()

    def ingest_event(self, event):
        # 1. Store raw episode (non-lossy)
        episode = self.create_episode(event)

        # 2. Extract entities (semantic layer)
        entities = self.extract_entities(event.content)
        for entity in entities:
            existing = self.find_or_create_entity(entity)
            # Link episode to entity
            self.link(episode, "MENTIONS", existing)

        # 3. Extract relationships (semantic layer)
        relationships = self.extract_relationships(entities)
        for rel in relationships:
            self.create_relationship(rel)

        # 4. Update communities (incremental)
        affected_communities = self.find_communities(entities)
        for community in affected_communities:
            self.update_community_summary(community)

    def retrieve(self, query, top_k=10):
        # Multi-stage retrieval

        # Stage 1: Semantic search (vector similarity)
        embedding = self.embed(query)
        candidate_entities = self.vector_db.search(embedding, top_k=50)

        # Stage 2: Graph traversal (structural)
        expanded = self.graph_expand(candidate_entities, depth=2)

        # Stage 3: Temporal filtering
        recent = self.filter_by_recency(expanded)

        # Stage 4: Episode retrieval
        episodes = self.get_episodes(recent)

        # Stage 5: Reranking
        reranked = self.rerank(
            entities=recent,
            episodes=episodes,
            query=query
        )

        return reranked[:top_k]

class MultiModalMemory:
    """
    Six specialized memory components + meta-manager
    """
    def __init__(self):
        # Memory components
        self.core = CoreMemory()           # Persistent facts (agent + user)
        self.episodic = EpisodicMemory()   # Time-stamped events
        self.semantic = SemanticMemory()   # Entity relationships
        self.procedural = ProceduralMemory()  # Workflows and procedures
        self.resource = ResourceMemory()   # Documents, code files
        self.knowledge_vault = KnowledgeVault()  # Sensitive data

        # Meta-manager (routing)
        self.meta_manager = MetaMemoryManager()

    def process_event(self, event):
        # Meta-manager routes to appropriate memory stores
        routing = self.meta_manager.route(event)

        for component, instructions in routing.items():
            memory = getattr(self, component)
            memory.update(event, instructions)

    def retrieve(self, query):
        # Generate current topic
        topic = self.meta_manager.extract_topic(query)

        # Retrieve from all components in parallel
        results = {
            "core": self.core.retrieve(topic),
            "episodic": self.episodic.retrieve(topic),
            "semantic": self.semantic.retrieve(topic),
            "procedural": self.procedural.retrieve(topic),
            "resource": self.resource.retrieve(topic)
            # knowledge_vault requires explicit permission
        }

        # Tag by source
        tagged = self.tag_sources(results)

        # Format for LLM
        return self.format_for_llm(tagged)

class CoreMemory:
    """High-priority persistent information"""
    def __init__(self):
        self.persona = {}  # Agent identity
        self.human = {}    # User facts
        self.max_tokens = 2048

    def update(self, key, value):
        if self.compute_tokens() > self.max_tokens * 0.9:
            self.rewrite()  # Compress and summarize
        self.persona[key] = value

class EpisodicMemory:
    """Time-stamped event log"""
    def add(self, event):
        episode = {
            "id": generate_id(),
            "timestamp": event.timestamp,
            "event_type": event.type,
            "summary": summarize(event),
            "details": event.content,
            "actor": event.actor
        }
        self.db.create_node("Episode", episode)

class SemanticMemory:
    """Entity knowledge graph"""
    def add_entity(self, entity):
        # Deduplication via embedding similarity
        existing = self.find_similar(entity)
        if existing:
            self.merge(existing, entity)
        else:
            self.db.create_node("Entity", entity)

class ProceduralMemory:
    """How-to knowledge"""
    def add_procedure(self, procedure):
        proc = {
            "type": "workflow|guide|script",
            "description": procedure.description,
            "steps": procedure.steps,
            "success_rate": 0.5
        }
        self.db.create_node("Procedure", proc)

class ResourceMemory:
    """Documents and files"""
    def add_resource(self, resource):
        res = {
            "title": resource.title,
            "type": resource.type,
            "summary": summarize(resource.content),
            "content_excerpt": resource.content[:1000],
            "file_path": resource.path
        }
        self.db.create_node("Resource", res)

class KnowledgeVault:
    """Sensitive information with access control"""
    def add(self, key, value, sensitivity="high"):
        # Encrypted storage
        encrypted = self.encrypt(value)
        self.db.create_node("Secret", {
            "key": key,
            "value": encrypted,
            "sensitivity": sensitivity,
            "requires_permission": True
        })

def hybrid_search(query, kg):
    # Semantic search (vector similarity)
    query_embedding = embed(query)
    semantic_results = kg.vector_search(query_embedding, top_k=50)

    # Structural search (graph traversal)
    entities = extract_entities(query)
    structural_results = kg.graph_query("""
        MATCH (e:Entity)-[*1..2]-(related)
        WHERE e.name IN $entities
        RETURN related
    """, entities=entities)

    # Combine and rerank
    combined = semantic_results + structural_results
    reranked = reciprocal_rank_fusion(combined)
    return reranked[:10]

def temporal_query(query, kg, time_context):
    # Extract temporal aspects
    temporal_refs = extract_temporal_references(query)
    # Examples: "yesterday", "last week", "when we started the project"

    # Convert to absolute time ranges
    time_ranges = resolve_temporal_refs(temporal_refs, time_context)

    # Query with temporal constraints
    results = kg.query("""
        MATCH (e:Episode)-[:MENTIONS]->(entity:Entity)
        WHERE e.timestamp > $start AND e.timestamp < $end
        RETURN e, entity
        ORDER BY e.timestamp DESC
    """, start=time_ranges["start"], end=time_ranges["end"])

    return results

def multi_hop_reasoning(query, kg, max_hops=3):
    # Start with seed entities
    seed_entities = extract_entities(query)

    # Iteratively expand
    results = []
    current_entities = seed_entities

    for hop in range(max_hops):
        # Find related entities
        related = kg.query("""
            MATCH (e:Entity)-[r]-(related:Entity)
            WHERE e.id IN $entities
            RETURN related, r, e
        """, entities=[e.id for e in current_entities])

        # Score by relevance (decay by hop distance)
        for rel in related:
            rel.score *= (0.7 ** hop)

        results.extend(related)
        current_entities = [r.related for r in related]

    # Rerank by combined score
    return sorted(results, key=lambda x: x.score, reverse=True)

def detect_contradictions(new_fact, kg):
    # Find facts about same entities
    related_facts = kg.query("""
        MATCH (e1:Entity)<-[r1]-(f1:Fact)-[r2]->(e2:Entity)
        WHERE e1.id = $entity1 AND e2.id = $entity2
          AND f1.t_invalid IS NULL  // Only active facts
        RETURN f1
    """, entity1=new_fact.entity1, entity2=new_fact.entity2)

    # Check for contradictions
    for old_fact in related_facts:
        if contradicts(new_fact, old_fact):
            # Temporal invalidation (don't delete, mark invalid)
            old_fact.t_invalid = new_fact.t_valid
            old_fact.invalidated_by = new_fact.id
            kg.update(old_fact)

    # Add new fact
    kg.add(new_fact)

class EpisodicMemory:
    def record_conversation(self, messages, outcome=None):
        episode = {
            "id": generate_id(),
            "timestamp": datetime.now(),
            "type": "conversation",
            "summary": summarize(messages),
            "details": json.dumps(messages),
            "actor": messages[0].user_id,
            "outcome": outcome
        }

        # Create episode node
        self.db.create_node("Episode", episode)

        # Extract and link entities
        entities = extract_entities(messages)
        for entity in entities:
            self.link_entity(episode["id"], entity)

    def record_code_change(self, commit):
        episode = {
            "id": generate_id(),
            "timestamp": commit.timestamp,
            "type": "commit",
            "summary": commit.message,
            "details": commit.diff,
            "actor": commit.author,
            "files_modified": commit.files
        }
        self.db.create_node("Episode", episode)

        # Link to modified files
        for file in commit.files:
            self.db.create_relationship(
                episode["id"], "MODIFIES", file
            )

class SemanticMemory:
    def index_codebase(self, codebase_path):
        # Use blarify or tree-sitter for parsing
        from blarify import CodeGraph

        graph = CodeGraph(
            codebase_path=codebase_path,
            use_scip=True  # 330x faster
        )
        graph.build()

        # Export to Neo4j
        graph.export_to_neo4j(
            uri=self.neo4j_uri,
            auth=self.neo4j_auth
        )

    def update_on_file_change(self, file_path, content):
        # Incremental update (don't rebuild entire graph)

        # 1. Parse changed file
        ast = parse_file(file_path, content)

        # 2. Extract entities
        functions = extract_functions(ast)
        classes = extract_classes(ast)

        # 3. Update graph
        for func in functions:
            existing = self.db.find_node("Function", name=func.name, file=file_path)
            if existing:
                self.db.update_node(existing.id, func)
            else:
                self.db.create_node("Function", func)

        # 4. Update relationships (calls, imports)
        calls = extract_calls(ast)
        for call in calls:
            self.db.create_relationship(
                call.caller, "CALLS", call.callee, {"line": call.line}
            )

class ProceduralMemory:
    def learn_procedure(self, problem, solution_steps, outcome):
        # Extract procedure from successful resolution
        procedure = {
            "id": generate_id(),
            "name": extract_procedure_name(problem),
            "description": summarize(problem),
            "trigger_pattern": extract_trigger(problem),
            "steps": solution_steps,
            "success_rate": 1.0 if outcome == "success" else 0.0,
            "times_used": 1
        }

        self.db.create_node("Procedure", procedure)

        # Link to error types
        error_type = extract_error_type(problem)
        self.db.create_relationship(
            procedure["id"], "FIXES", error_type
        )

    def find_procedure(self, problem):
        # Match problem to procedure triggers
        error_type = extract_error_type(problem)

        results = self.db.query("""
            MATCH (p:Procedure)-[:FIXES]->(e:Error)
            WHERE e.type = $error_type
            RETURN p
            ORDER BY p.success_rate DESC, p.times_used DESC
        """, error_type=error_type)

        return results[0] if results else None

    def update_success_rate(self, procedure_id, success):
        # Exponential moving average
        proc = self.db.get_node(procedure_id)
        alpha = 0.1
        proc.success_rate = (
            alpha * (1 if success else 0) +
            (1 - alpha) * proc.success_rate
        )
        proc.times_used += 1
        self.db.update_node(proc)

class CodeMemoryIntegration:
    def __init__(self, codebase_path):
        self.episodic = EpisodicMemory()
        self.semantic = SemanticMemory()
        self.procedural = ProceduralMemory()
        self.code_graph = CodeGraph(codebase_path)

    def on_code_change(self, file_path, content):
        # 1. Update code graph
        self.code_graph.update_file(file_path, content)

        # 2. Create episode
        episode = self.episodic.record_code_change({
            "file": file_path,
            "content": content,
            "timestamp": datetime.now()
        })

        # 3. Update semantic entities
        entities = self.code_graph.extract_entities(file_path)
        for entity in entities:
            self.semantic.update_entity(entity)

    def on_error(self, error):
        # 1. Record episode
        episode = self.episodic.record_error(error)

        # 2. Check for known procedure
        procedure = self.procedural.find_procedure(error)

        if procedure:
            return {
                "procedure": procedure,
                "confidence": procedure.success_rate
            }
        else:
            # Search for similar past errors
            similar = self.episodic.find_similar_episodes(
                error, type="error"
            )
            return {
                "similar_cases": similar,
                "confidence": 0.5
            }

from neo4j import GraphDatabase
from blarify import CodeGraph
import datetime

class CodingAssistantMemory:
    def __init__(self, codebase_path, neo4j_uri, neo4j_auth):
        self.driver = GraphDatabase.driver(neo4j_uri, auth=neo4j_auth)
        self.codebase_path = codebase_path

        # Initialize memory components
        self.episodic = EpisodicMemory(self.driver)
        self.semantic = SemanticMemory(self.driver)
        self.procedural = ProceduralMemory(self.driver)

        # Initialize code graph
        self.code_graph = CodeGraph(codebase_path, use_scip=True)

    def setup(self):
        """Initial setup: index codebase"""
        print("Indexing codebase...")
        self.code_graph.build()
        self.code_graph.export_to_neo4j(
            uri=self.driver._pool.address[0],
            auth=self.driver._pool.auth
        )
        print("Codebase indexed successfully")

    def record_conversation(self, messages, files_mentioned=None):
        """Record a conversation with user"""
        episode_id = self.episodic.record({
            "type": "conversation",
            "messages": messages,
            "timestamp": datetime.datetime.now(),
            "files": files_mentioned or []
        })

        # Extract entities and link
        entities = self.extract_entities(messages)
        for entity in entities:
            self.semantic.link_episode_to_entity(episode_id, entity)

        return episode_id

    def retrieve_context(self, query, max_results=10):
        """Retrieve relevant context for query"""

        # Multi-modal retrieval

        # 1. Semantic search (code entities)
        code_entities = self.semantic.search(query, type="code")

        # 2. Find related code via graph traversal
        related_code = self.code_graph.expand(code_entities, depth=2)

        # 3. Search episodes (conversations, errors)
        episodes = self.episodic.search(query, time_range="30d")

        # 4. Find applicable procedures
        procedures = self.procedural.search(query)

        # 5. Rerank and combine
        context = self.rerank({
            "code": related_code,
            "episodes": episodes,
            "procedures": procedures
        }, query)

        return context[:max_results]

    def on_error(self, error_info):
        """Handle error occurrence"""

        # 1. Record episode
        episode_id = self.episodic.record({
            "type": "error",
            "error": error_info,
            "timestamp": datetime.datetime.now()
        })

        # 2. Find similar past errors
        similar = self.episodic.find_similar(error_info, type="error")

        # 3. Find applicable procedure
        procedure = self.procedural.find_by_error_type(error_info.type)

        # 4. Assemble suggestions
        suggestions = {
            "procedure": procedure,
            "similar_cases": similar,
            "confidence": procedure.success_rate if procedure else 0.3
        }

        return suggestions

    def learn_from_resolution(self, error_id, resolution_steps, success):
        """Learn from error resolution"""

        # 1. Update episode with resolution
        self.episodic.update(error_id, {
            "resolution": resolution_steps,
            "success": success
        })

        # 2. Update or create procedure
        procedure = self.procedural.learn(
            error_id=error_id,
            steps=resolution_steps,
            success=success
        )

        # 3. Update success rates
        if procedure:
            self.procedural.update_success_rate(procedure.id, success)

    def close(self):
        self.driver.close()


# Usage
memory = CodingAssistantMemory(
    codebase_path="./src",
    neo4j_uri="bolt://localhost:7687",
    neo4j_auth=("neo4j", "password")
)

# Initial setup
memory.setup()

# Record conversation
messages = [
    {"role": "user", "content": "How do I fix this ImportError?"},
    {"role": "assistant", "content": "Let me check..."}
]
episode_id = memory.record_conversation(messages, files_mentioned=["auth.py"])

# Retrieve context for new query
context = memory.retrieve_context("authentication bug")

# Handle error
error = {
    "type": "ImportError",
    "message": "Module 'requests' not found",
    "file": "auth.py",
    "line": 10
}
suggestions = memory.on_error(error)

# Learn from resolution
memory.learn_from_resolution(
    error_id=episode_id,
    resolution_steps=["pip install requests", "verify PYTHONPATH"],
    success=True
)

memory.close()