# kg-memory-mcp

An in-memory knowledge graph MCP server that gives coding agents structural and semantic recall over codebases. Indexes Python source, ADR documents, and project configuration via tree-sitter parsing and optional LLM extraction, then exposes 7 tools over the Model Context Protocol for search, traversal, context retrieval, and natural-language Q&A.

## Key Features

- **Structural indexing** via tree-sitter -- extracts modules, classes, functions, imports, call graphs, and inheritance hierarchies from Python source
- **ADR indexing** -- parses architectural decision records into first-class graph nodes with GOVERNS edges linking decisions to the code they affect
- **LLM semantic extraction** -- enriches the graph with concepts, trade-offs, design principles, and cross-references that only exist in prose
- **3-tier entity reconciliation** -- deduplicates nodes via exact match, fuzzy match, and LLM-based semantic similarity
- **7 MCP tools** -- search, relate, context, path, ask, reindex, status
- **Interactive graph visualization** -- React + Sigma.js explorer with filtering, search, and node navigation
- **Fast restarts** -- persists the graph to `.kg-index.json`; cold start loads from cache in milliseconds
- **Per-project configuration** -- drop a `.mcp.json` in any project root to enable

## Architecture

```
Codebase Files
     |
     v
+------------------+     +------------------+
| tree-sitter      |     | LLM Extraction   |
| (Python, MD)     |     | (litellm)        |
| Structural parse |     | Semantic parse   |
+--------+---------+     +--------+---------+
         |                        |
         v                        v
    +----+------------------------+----+
    |      In-Memory Knowledge Graph   |
    |   9 node types, 14 edge types    |
    |   5 O(1) indexes                 |
    +----+-------------------+---------+
         |                   |
         v                   v
  +------+------+    +------+------+
  | 7 MCP Tools |    | Viz Server  |
  | (stdio)     |    | (HTTP)      |
  +-------------+    +-------------+
```

## Quick Start

### 1. Install

```bash
cd kg-memory-mcp
python -m venv .venv && source .venv/bin/activate
pip install -e .

# Optional: file watcher for auto-reindex on save
pip install -e ".[watcher]"
```

### 2. Configure for a project

Create `.mcp.json` in your project root:

```json
{
  "mcpServers": {
    "kg-memory": {
      "command": "/path/to/kg-memory-mcp/.venv/bin/kg-memory-mcp",
      "env": {
        "KG_PROJECT_ROOT": "/path/to/your/project"
      }
    }
  }
}
```

### 3. Start Claude Code

```bash
cd /path/to/your/project
claude
```

The MCP server launches automatically, indexes the codebase, and the 7 tools become available.

### 4. (Optional) Enable LLM extraction

Add to the `env` section of `.mcp.json`:

```json
{
  "KG_LLM_ENABLED": "true",
  "KG_LLM_MODEL": "openai/gpt-4o-mini",
  "OPENAI_API_KEY": "sk-..."
}
```

This adds concepts, decisions, and trade-offs extracted from ADRs and documentation.

### 5. (Optional) Graph visualization

```bash
cd viz && npm install && npm run build && cd ..
kg-memory-viz --project-root /path/to/your/project
# Open http://localhost:8765
```

---

## Indexing Pipeline

The pipeline has four phases, orchestrated by `server.py:reindex()`.

### Phase 1: Structural Indexing (tree-sitter)

Three indexers run without any LLM calls.

#### Python Indexer

Parses all `.py` files using `tree-sitter-python`. Extracts:

| Node Type | Source | Example ID |
|-----------|--------|------------|
| `file` | Each `.py` file | `file:src/myapp/scoring.py` |
| `module` | Directories with `__init__.py` | `module:src.myapp.domain` |
| `class` | `class` definitions | `class:src.myapp.scoring.Scorer` |
| `function` | `def` definitions (top-level and methods) | `func:src.myapp.scoring.Scorer.compute` |

Edges created from code structure:

| Edge Type | Meaning | Detection |
|-----------|---------|-----------|
| `contains` | file/module/class contains child | CST parent-child nesting |
| `imports` | file imports another file | `import` / `from...import` statements |
| `calls` | function calls another function | CST call expressions with name resolution |
| `inherits` | class extends base class | `class Foo(Bar)` superclass list |
| `method_of` | method belongs to class | Method nested inside class body |

Properties captured per node: `path`, `line`, `docstring`, `params`, `decorators`, `bases`, `return type`.

#### ADR Indexer

Parses `docs/adr/*.md` files using `tree-sitter-markdown`. Uses a two-pass approach:

**First pass** -- for each ADR file:
- Creates an `adr` node with title, status, decision summary, and principles
- Extracts `concept` nodes from Decision subsection headings
- Creates `GOVERNS` edges via two strategies:
  - **Path-based**: backtick code spans matching module paths are resolved against the graph's path index
  - **Identifier-based**: PascalCase/snake_case identifiers in backticks are resolved against the graph's name index (prefers class nodes)
- Creates `DEFINES` edges linking the ADR to its concepts

**Second pass** -- resolves `REFERENCES` edges between ADRs (cross-references like `ADR-0003`).

#### Config Indexer

Parses `CLAUDE.md` (or configured project instructions file). Creates `config` nodes from h2 sections and extracts `concept` nodes from markdown tables.

### Phase 2: LLM Semantic Extraction (optional)

When `KG_LLM_ENABLED=true`, the `LLMExtractor` sends each ADR's full text to the LLM and extracts:

| Extracted Type | Node Type | Edge to ADR | What it captures |
|----------------|-----------|-------------|------------------|
| Decisions | `decision` | `decided_in` | Statement, rationale, confidence score |
| Concepts | `concept` | `defines` | Name, definition, category (algorithm/pattern/principle/term) |
| Trade-offs | `trade_off` | `considered` | Option chosen, option rejected, reasoning |
| Principles | (property) | -- | Stored on the ADR node |
| Module governance | -- | `governs` | Additional file paths parsed from prose |

All ADR extractions run concurrently via `asyncio.gather()`. Per-item Pydantic validation ensures one bad extraction doesn't drop the entire batch. Node IDs use SHA-256 hashes for deterministic, idempotent re-extraction.

**What LLM extraction adds that structural indexing misses:**

Structural indexing sees code and heading text. LLM extraction reads prose -- the "why" behind architecture. It captures rationale ("We chose Redis over Kafka because..."), design constraints ("Must support at-least-once delivery"), and trade-off reasoning that only exists in natural language. For a project with 15 ADRs, this typically adds 50-100+ semantic nodes and 100-200+ edges.

### Phase 3: Entity Reconciliation

Deduplicates nodes created by different indexers without destructive merging.

**Tier 1 -- Exact match**: Normalizes names (lowercase, collapse whitespace, strip separators). Applies an alias dictionary (`"kg"` -> `"knowledge graph"`, `"mcp"` -> `"model context protocol"`, etc.). Creates `SAME_AS` edges with confidence 1.0.

**Tier 2 -- Fuzzy match**: Uses `difflib.SequenceMatcher` across 4 name/canonical combinations. Thresholds: >= 0.85 -> `SAME_AS`, >= 0.70 -> `RELATED_TO`.

**Tier 3 -- LLM semantic match** (only with LLM enabled): Pre-filters to concept/decision pairs sharing significant words. Batches 10 pairs per LLM call, up to 5 concurrent calls, hard-capped at 200 pairs. Thresholds: >= 0.85 -> `SAME_AS`, >= 0.60 -> `RELATED_TO`.

Transitive closure via Union-Find groups all `SAME_AS` pairs into clusters. Both nodes are preserved -- linked, not merged.

### Phase 4: Persistence

The graph serializes to `.kg-index.json` via `orjson` with atomic write (temp file + rename). On next startup, the cache loads in milliseconds, skipping the full indexing pipeline.

---

## Graph Schema

### 9 Node Types

| Type | Source | Description |
|------|--------|-------------|
| `module` | Python indexer | Python package (directory with `__init__.py`) |
| `class` | Python indexer | Class definition |
| `function` | Python indexer | Function or method definition |
| `file` | Python indexer | Source file |
| `adr` | ADR indexer | Architectural Decision Record |
| `decision` | LLM extraction | Architectural decision with rationale |
| `concept` | ADR indexer + LLM | Technical concept or term |
| `trade_off` | LLM extraction | Evaluated trade-off with chosen/rejected options |
| `config` | Config indexer | Project configuration section |

### 14 Edge Types

| Type | Direction | Source |
|------|-----------|--------|
| `contains` | parent -> child | Structural |
| `imports` | file -> file | Structural |
| `calls` | function -> function | Structural |
| `inherits` | class -> base class | Structural |
| `method_of` | method -> class | Structural |
| `depends_on` | module -> module | Structural |
| `references` | ADR -> ADR | Structural |
| `governs` | ADR/config -> file/class/function | Structural + LLM |
| `decided_in` | decision -> ADR | LLM extraction |
| `defines` | ADR/config -> concept | Structural + LLM |
| `implements` | code -> decision | LLM extraction |
| `considered` | trade_off -> ADR | LLM extraction |
| `same_as` | node -> node | Reconciliation |
| `related_to` | node -> node | Reconciliation |

---

## MCP Tools Deep-Dive

All tools return markdown-formatted text designed for LLM consumption.

### kg_search

Search the graph by keyword or concept name.

| Parameter | Type | Required | Default | Description |
|-----------|------|----------|---------|-------------|
| `query` | string | yes | -- | Search term |
| `node_types` | string[] | no | all | Filter by node type |
| `limit` | integer | no | 20 | Max results (1-100) |

**Scoring algorithm** -- each node is scored against the query on a 4-tier scale:

| Score | Condition |
|-------|-----------|
| 1.0 | Exact name match (case-insensitive) |
| 0.8 | Name contains query string |
| 0.6 | Query words appear in name (proportional to overlap) |
| 0.4 | Query found in any property value |

Results are sorted by score descending. Type-specific details are included: file path for files, title for ADRs, definition for concepts, path and line number for classes/functions.

**Example**: `kg_search("scoring", node_types=["class", "function"])` returns the `Scorer` class, `compute_decay_score` function, and other scoring-related symbols -- ranked by relevance.

### kg_relate

Expand the neighborhood of a node via BFS traversal.

| Parameter | Type | Required | Default | Description |
|-----------|------|----------|---------|-------------|
| `node_id` | string | yes | -- | Starting node (e.g., `file:src/scoring.py`) |
| `edge_types` | string[] | no | all | Filter by edge type |
| `depth` | integer | no | 1 | Traversal depth (1-5) |
| `limit` | integer | no | 30 | Max nodes (1-100) |

Traverses both outgoing and incoming edges. Results are grouped by node type with edges shown as `source --[TYPE]--> target`.

**Example**: `kg_relate("adr:0008", depth=2)` returns the ADR's decisions, trade-offs, concepts, governed files, and the classes/functions in those files -- a 2-hop view of everything ADR-0008 touches.

### kg_context

Get full structured context for a specific file.

| Parameter | Type | Required | Default | Description |
|-----------|------|----------|---------|-------------|
| `file_path` | string | yes | -- | Relative file path |

Unlike generic BFS, this follows specific semantic paths:
1. **Structure**: classes and functions contained in the file (`CONTAINS`)
2. **Imports**: what this file imports (`IMPORTS` outgoing)
3. **Dependents**: who imports this file (`IMPORTS` incoming)
4. **Governing ADRs**: which ADRs govern this file (`GOVERNS` incoming)
5. **Decisions**: architectural decisions from those ADRs (`DECIDED_IN`)
6. **Related concepts**: concepts defined by those ADRs (`DEFINES`)
7. **Entity links**: reconciled duplicates (`SAME_AS`)

**Example**: `kg_context("src/context_graph/domain/scoring.py")` returns the file's classes, functions, what ADRs govern it (ADR-0008), key decisions like "Use Ebbinghaus decay scoring", related concepts, and who depends on this file.

### kg_path

Find the shortest path between two nodes.

| Parameter | Type | Required | Default | Description |
|-----------|------|----------|---------|-------------|
| `source` | string | yes | -- | Source node ID |
| `target` | string | yes | -- | Target node ID |
| `max_depth` | integer | no | 5 | Max hops (1-10) |

Uses BFS over both edge directions. Returns the path as an indented chain showing each node's name, type, and ID.

**Example**: `kg_path("concept:ebbinghaus_decay", "file:src/scoring.py")` might return: `concept:ebbinghaus_decay` -> `adr:0008` -> `file:src/scoring.py` -- showing how the concept connects to the implementation through the ADR.

### kg_ask

Natural-language Q&A grounded in the knowledge graph. Requires `KG_LLM_ENABLED=true`.

| Parameter | Type | Required | Default | Description |
|-----------|------|----------|---------|-------------|
| `question` | string | yes | -- | Question about the codebase |

**How it works**:
1. Searches the graph for the 10 most relevant nodes
2. Expands the top 5 results by 1 hop to gather surrounding context
3. Formats node properties (docstrings, decisions, definitions) as context
4. Sends to the LLM with a structured prompt: WHY (motivation) -> WHAT (how it works) -> WHERE (code references)
5. Appends a Sources section with provenance

**Example**: `kg_ask("How does memory consolidation work?")` searches for consolidation-related functions, ADRs, and concepts, then synthesizes a narrative explanation citing specific files and decisions.

### kg_reindex

Trigger a full rebuild of the knowledge graph.

| Parameter | Type | Required | Default | Description |
|-----------|------|----------|---------|-------------|
| `llm` | boolean | no | true | Run LLM extraction |

Executes all 4 pipeline phases: structural indexing -> LLM extraction -> reconciliation -> persistence. Returns node/edge counts and breakdown by type.

### kg_status

Show current graph statistics and server state. Takes no parameters.

Returns: total nodes/edges, counts by type, file watcher status, LLM model/enabled state, last indexed timestamp, reconciliation cluster count.

---

## Graph Visualization

The `viz/` directory contains a React + Sigma.js interactive graph explorer.

**Tech stack**: Vite, React 18, TypeScript, Tailwind CSS, Sigma.js v3 (WebGL), graphology, ForceAtlas2 layout.

**Features**:
- ForceAtlas2 force-directed layout (handles 1000+ nodes via WebGL)
- Filter by node type and edge type with instant toggle (Sigma reducers, no re-render)
- 4 filter presets: Architecture, Call Graph, ADRs, Show All
- Full-text search with type-priority matching and camera auto-center
- Click any node to see properties, connections grouped by edge type, and navigate to neighbors
- Node drag-and-drop, zoom controls, keyboard shortcuts (Escape, `/` to search)
- Search-selected nodes reveal their neighbors even when filtered out

**Running**:
```bash
cd viz && npm install && npm run build && cd ..
kg-memory-viz --project-root /path/to/project
# Open http://localhost:8765
```

The viz server is a zero-dependency Python HTTP server that reads `.kg-index.json` and serves the React build.

---

## Configuration Reference

| Environment Variable | Default | Description |
|---------------------|---------|-------------|
| `KG_PROJECT_ROOT` | cwd | Project root to index |
| `KG_SOURCE_DIRS` | `src` | Comma-separated source directories |
| `KG_ADR_DIR` | `docs/adr` | ADR directory (`none` to disable) |
| `KG_CONFIG_PATH` | `CLAUDE.md` | Project config file (`none` to disable) |
| `KG_INDEX_FILE` | `.kg-index.json` | Graph persistence file |
| `KG_ADR_PATTERN` | `^\d{4}-.+\.md$` | Regex for ADR filename matching |
| `KG_WATCH_ENABLED` | `true` | Enable file watcher for live reindex |
| `KG_LLM_ENABLED` | `false` | Enable LLM extraction |
| `KG_LLM_MODEL` | `claude-sonnet-4-6` | LLM model (any litellm-supported model) |

Configuration can also be set in `kg-memory.toml` at the project root, with optional `[aliases]` for entity reconciliation.

---

## Project Structure

```
kg-memory-mcp/
  src/kg_memory/
    server.py              MCP server entry point, startup, reindex orchestration
    graph.py               KnowledgeGraph, Node, Edge, 9 NodeTypes, 14 EdgeTypes
    retrieval.py           Search, BFS traversal, shortest path, file context
    extraction.py          LLMExtractor: ADR/config extraction, Q&A, reconciliation
    reconciliation.py      3-tier entity reconciliation (exact, fuzzy, LLM)
    persistence.py         orjson serialize/deserialize to .kg-index.json
    config.py              ProjectConfig from env vars / kg-memory.toml
    watcher.py             Watchdog file watcher for incremental reindex
    tools.py               7 MCP tool definitions and handlers
    viz_server.py          HTTP server for graph visualization
    indexers/
      __init__.py          Orchestrates all indexers, incremental reindex routing
      python_indexer.py    tree-sitter Python: modules, classes, functions, edges
      adr_indexer.py       tree-sitter Markdown: ADRs, concepts, GOVERNS edges
      config_indexer.py    tree-sitter Markdown: config sections, table concepts
  viz/
    src/
      App.tsx              Main layout, state management, keyboard shortcuts
      types.ts             TypeScript types, node/edge color maps
      hooks/useGraphData.ts  Fetch API, build graphology instance
      components/
        GraphCanvas.tsx    Sigma.js container, ForceAtlas2, reducers, drag
        FilterPanel.tsx    Node/edge type toggles, presets
        SearchBar.tsx      Debounced search, type-priority results
        NodeDetail.tsx     Selected node properties and connection browser
        StatsBar.tsx       Graph statistics bar
  tests/
    test_graph.py          Graph data model tests
    test_extraction.py     LLM extraction tests
    test_reconciliation.py Entity reconciliation tests
    test_retrieval.py      Search and traversal tests
    test_persistence.py    Serialization tests
    test_config.py         Configuration tests
  pyproject.toml           Package config, dependencies, CLI entry points
```

## License

MIT