Code Context Injection at Memory Retrieval¶

Status: Implemented (Week 5-6) Feature: Enrich memory retrieval with related code files and functions

Overview¶

Code context injection enriches memory retrieval by automatically including information about related code files, functions, and classes when memories are retrieved. This enables agents to have immediate access to code structure information alongside memories.

Implementation Summary¶

1. Memory Retrieval API - MemoryCoordinator.retrieve()¶

File: src/amplihack/memory/coordinator.py

Added include_code_context parameter to RetrievalQuery:

@dataclass
class RetrievalQuery:
    """Query to retrieve memories.

    Args:
        query_text: Search text
        token_budget: Max tokens to return (default 8000)
        memory_types: Filter by specific types
        time_range: Filter by time range (start, end)
        include_code_context: Include related code files/functions (default False)
    """

    query_text: str
    token_budget: int = 8000
    memory_types: list[MemoryType] | None = None
    time_range: tuple[datetime, datetime] | None = None
    include_code_context: bool = False

2. Code Context Enrichment¶

Implementation: MemoryCoordinator._enrich_with_code_context()

The enrichment process:

Check backend capabilities: Verify backend supports graph queries
Get code graph instance: Access Kuzu code graph integration
Query relationships: Find RELATES_TO_FILE_* and RELATES_TO_FUNCTION_* links
Format context: Convert code information to LLM-readable format
Add to metadata: Inject code_context key into memory metadata

async def _enrich_with_code_context(
    self, memories: list[MemoryEntry]
) -> list[MemoryEntry]:
    """Enrich memories with related code context.

    Queries Kuzu graph fer code files and functions linked to each memory
    via RELATES_TO_FILE_* and RELATES_TO_FUNCTION_* relationships.

    Args:
        memories: List of memory entries to enrich

    Returns:
        Same memories with code_context added to metadata

    Performance: Must complete under 100ms total (not per memory)
    """

3. Code Context Format¶

Code context is formatted as markdown-style text for LLM consumption:

**Related Files:**

- src/amplihack/memory/coordinator.py (python)
- src/amplihack/memory/backends/kuzu_backend.py (python)

**Related Functions:**

- `async def retrieve(self, query: RetrievalQuery) -> list[MemoryEntry]`
  Retrieve memories matching query.
  (complexity: 12.5)

**Related Classes:**

- amplihack.memory.coordinator.MemoryCoordinator
  Coordinates memory storage and retrieval with quality control.

4. Backend Integration¶

File: src/amplihack/memory/backends/kuzu_backend.py

Added public accessor for code graph:

def get_code_graph(self) -> KuzuCodeGraph | None:
    """Get code graph instance for querying code-memory relationships.

    Returns:
        KuzuCodeGraph instance if available, None if code graph not initialized

    Example:
        >>> backend = KuzuBackend()
        >>> backend.initialize()
        >>> code_graph = backend.get_code_graph()
        >>> if code_graph:
        ...     context = code_graph.query_code_context(memory_id)
    """

Usage Examples¶

Basic Usage¶

from amplihack.memory.coordinator import MemoryCoordinator, RetrievalQuery
from amplihack.memory.types import MemoryType

coordinator = MemoryCoordinator()

# Retrieve WITHOUT code context (default)
query = RetrievalQuery(
    query_text="Kuzu backend implementation",
    include_code_context=False,
)
memories = await coordinator.retrieve(query)

# Retrieve WITH code context
query = RetrievalQuery(
    query_text="Kuzu backend implementation",
    include_code_context=True,  # Enable code context injection
)
memories = await coordinator.retrieve(query)

# Check code context
for memory in memories:
    if "code_context" in memory.metadata:
        print(f"Memory: {memory.content[:50]}...")
        print(f"Related code:\n{memory.metadata['code_context']}")

Integration with Agent Context Building¶

# In agent initialization or context building
query = RetrievalQuery(
    query_text="current task context",
    include_code_context=True,  # Get code structure with memories
    memory_types=[MemoryType.SEMANTIC, MemoryType.PROCEDURAL],
)

relevant_memories = await coordinator.retrieve(query)

# Build agent context with both memories and code structure
context = []
for memory in relevant_memories:
    context.append(f"Memory: {memory.content}")

    # Add code context if available
    if "code_context" in memory.metadata:
        context.append(f"\nRelated Code:\n{memory.metadata['code_context']}")

Backend Support¶

Kuzu Backend¶

Full support for code context injection:

Graph queries: ✓ Native Cypher support
Code graph: ✓ Via KuzuCodeGraph
Memory-code links: ✓ RELATES_TO_FILE_* and RELATES_TO_FUNCTION_* relationships

SQLite Backend¶

Graceful fallback:

Graph queries: ✗ Not supported
Code context injection: Skipped silently
No errors, memories returned without code context

Performance¶

Requirements¶

Enrichment overhead: <100ms total (not per memory)
Retrieval total: <150ms including enrichment
Memory overhead: ~200 bytes per memory for code context metadata

Benchmarks¶

import time

# Store some memories
for i in range(10):
    await coordinator.store(StorageRequest(...))

# Benchmark retrieval with code context
start = time.time()
query = RetrievalQuery(
    query_text="test",
    include_code_context=True,
)
memories = await coordinator.retrieve(query)
elapsed_ms = (time.time() - start) * 1000

print(f"Retrieval with code context: {elapsed_ms:.1f}ms")
# Expected: <150ms

Injection Points (Priority Implementation)¶

✅ 1. Memory Retrieval API¶

File: src/amplihack/memory/coordinator.py Status: IMPLEMENTED

Added include_code_context parameter to RetrievalQuery
Implemented _enrich_with_code_context() helper method
Queries Kuzu for RELATES_TO_FILE_* and RELATES_TO_FUNCTION_* links
Formats code context into readable format
Injects into memory metadata

🚧 2. Agent Context Building (Future)¶

File: To be implemented in agent memory hooks Status: Planned

Extend memory hooks to automatically inject code context:

# In PreRequest hook or agent initialization
def inject_memory_context(agent_context: dict) -> dict:
    """Inject memories with code context into agent context."""
    query = RetrievalQuery(
        query_text=agent_context.get("task", ""),
        include_code_context=True,  # Auto-enable
    )
    memories = await coordinator.retrieve(query)
    agent_context["relevant_memories"] = memories
    return agent_context

🚧 3. Session Start (Future)¶

Status: Planned

Add code overview to session initialization:

# In SessionStart hook
async def load_code_overview(session: Session):
    """Load code structure overview into session context."""
    code_graph = backend.get_code_graph()
    if code_graph:
        stats = code_graph.get_code_stats()
        session.context["code_stats"] = stats

🚧 4. File Edit Operations (Future)¶

Status: Planned

Hook into edit operations to inject related memories:

# Before file edit
def get_file_context(file_path: str) -> dict:
    """Get memories and code context for file."""
    # Query memories linked to this file
    # Include related functions and classes
    # Return context for agent

Testing¶

Test Suite¶

File: tests/memory/test_code_context_injection.py

Comprehensive test coverage:

✓ test_retrieve_with_code_context_flag - Parameter acceptance
✓ test_code_context_enrichment - Enrichment functionality
✓ test_code_context_fallback_sqlite - SQLite fallback
✓ test_code_context_performance - Performance requirements
✓ test_code_context_format - Output format
✓ test_code_context_with_no_links - No links handling
✓ test_code_context_default_false - Default behavior

Running Tests¶

# Run all code context injection tests
python -m pytest tests/memory/test_code_context_injection.py -v

# Run specific test
python -m pytest tests/memory/test_code_context_injection.py::test_code_context_enrichment -v

# Run with coverage
python -m pytest tests/memory/test_code_context_injection.py --cov=amplihack.memory.coordinator

Demo Scripts¶

File: examples/code_context_injection_demo.py

# Run comprehensive demo
python examples/code_context_injection_demo.py

# Run simple test
python examples/simple_code_context_test.py

Success Criteria¶

✅ Completed¶

✓ Memory retrieval can include related code files/functions
✓ Agent context includes code structure information
✓ Code context formatted for LLM consumption
✓ Tests validate code injection functionality
✓ Graceful fallback for non-Kuzu backends
✓ Performance requirements met (<100ms enrichment)

🚧 Future Enhancements¶

⏳ Integrate with agent memory hooks
⏳ Add to session start/end workflows
⏳ Hook into file edit operations
⏳ Add caching for frequently queried code context
⏳ Support configurable context depth (1-hop, 2-hop traversal)

Configuration¶

Environment Variables¶

None required - feature is opt-in via include_code_context parameter.

Backend Requirements¶

Kuzu backend: Full support
SQLite backend: Graceful fallback (no errors)
Neo4j backend: Not yet implemented

Troubleshooting¶

Code context is empty¶

Possible causes:

No code graph data imported (run blarify first)
Memory not linked to code (auto-linking didn't find matches)
Code graph not initialized

Solution:

# Import code graph using blarify
amplihack memory blarify /path/to/codebase

# Check code graph stats
python -c "
from amplihack.memory.backends import create_backend
backend = create_backend('kuzu')
code_graph = backend.get_code_graph()
if code_graph:
    print(code_graph.get_code_stats())
"

Performance degradation¶