From 2a429b3160ae2030635c88dfba60f3f9d54ab1ae Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 7 Sep 2026 22:59:14 +0000 Subject: [PATCH 1/2] Add knowledge management: LightRAG-backed GraphRAG knowledge base Persist what a run learns (files, AST entities, dependency graph, user stories) into a queryable knowledge base instead of only writing a JSON report. GraphRAG was chosen over an LLM-wiki or vanilla RAG because a codebase is a graph: multi-hop and global questions need entity/relation retrieval, and the exact dependency graph from static analysis can seed the knowledge graph without spending LLM tokens. LightRAG (MIT) was selected for its incremental inserts, custom-KG import, pluggable Anthropic/OpenAI/ Ollama adapters and storage back-ends (NetworkX/Neo4j/Postgres/Memgraph, nano-vectordb/pgvector/Qdrant/Milvus/FAISS, JSON/Postgres/Redis/Mongo). - New workflow node `knowledge_base` (opt-in: `archaeo --knowledge-base`, `knowledge_base.enabled` in config) with deterministic indexing and an optional LLM entity-extraction pass (`--kb-extract`) - New MCP server `knowledge-base` (kb_query, kb_retrieve, kb_context, kb_graph, kb_stats, kb_add_documents, kb_export_wiki), registered in the config and the in-process MCP client - New CLI `archaeo-kb` (query, retrieve, context, graph, stats, add, wiki) - Obsidian-compatible Markdown wiki export as the human-facing projection - Offline mode (byte tokenizer, hashing embeddings) used by the new tests - AST entities now carry `file_path` so members link to their file - Docs: docs/KNOWLEDGE_MANAGEMENT.md decision record, README, quickstart, MCP usage, changelog; example under examples/knowledge_base; v0.4.0 Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_011hwnef5pyjkdpQ6deZJYDD --- .env.example | 17 + .gitignore | 2 + QUICKSTART.md | 10 + README.md | 45 +- config/langgraph_config.yaml | 54 + config/mcp_servers_config.yaml | 25 + data/knowledge_base/.gitkeep | 0 docs/CHANGELOG.md | 16 + docs/KNOWLEDGE_MANAGEMENT.md | 182 +++ docs/MCP_ARCHITECTURE.md | 2 +- docs/MCP_SERVERS_USAGE.md | 66 +- examples/knowledge_base/__init__.py | 0 .../knowledge_base/query_knowledge_base.py | 76 + pyproject.toml | 5 +- requirements.txt | 7 + src/mcp_servers/knowledge_base/__init__.py | 18 + .../knowledge_base/knowledge_base.py | 1240 +++++++++++++++++ .../knowledge_base/knowledge_mcp_server.py | 242 ++++ src/orchestration/cli.py | 61 + src/orchestration/graph.py | 16 +- src/orchestration/graph_mcp.py | 10 +- src/orchestration/kb_cli.py | 179 +++ src/orchestration/nodes/ast_node.py | 2 + src/orchestration/nodes/knowledge_node.py | 89 ++ src/orchestration/state/graph_state.py | 4 + src/orchestration/utils/mcp_client.py | 16 + tests/test_knowledge_base.py | 338 +++++ 27 files changed, 2703 insertions(+), 19 deletions(-) create mode 100644 data/knowledge_base/.gitkeep create mode 100644 docs/KNOWLEDGE_MANAGEMENT.md create mode 100644 examples/knowledge_base/__init__.py create mode 100644 examples/knowledge_base/query_knowledge_base.py create mode 100644 src/mcp_servers/knowledge_base/__init__.py create mode 100644 src/mcp_servers/knowledge_base/knowledge_base.py create mode 100644 src/mcp_servers/knowledge_base/knowledge_mcp_server.py create mode 100644 src/orchestration/kb_cli.py create mode 100644 src/orchestration/nodes/knowledge_node.py create mode 100644 tests/test_knowledge_base.py diff --git a/.env.example b/.env.example index 0a4a4d6..760c09d 100644 --- a/.env.example +++ b/.env.example @@ -35,6 +35,23 @@ USER_STORY_TEMPERATURE=0.3 # NEO4J_USER=neo4j # NEO4J_PASSWORD=password +# Knowledge Base (Optional) — LightRAG-backed GraphRAG, enabled per run with +# `archaeo --knowledge-base`. Defaults are file-based (no infrastructure). +# KB_WORKING_DIR=data/knowledge_base +# KB_WORKSPACE= # partition inside shared back-ends (one per analyzed system) +# KB_GRAPH_STORAGE=NetworkXStorage # Neo4JStorage | PGGraphStorage | MemgraphStorage | MongoGraphStorage +# KB_VECTOR_STORAGE=NanoVectorDBStorage # PGVectorStorage | QdrantVectorDBStorage | MilvusVectorDBStorage | FaissVectorDBStorage +# KB_KV_STORAGE=JsonKVStorage # PGKVStorage | RedisKVStorage | MongoKVStorage +# KB_EMBEDDING_PROVIDER=sentence-transformers # openai | ollama | hashing (tests only) +# KB_EMBEDDING_MODEL=all-MiniLM-L6-v2 +# KB_LLM_PROVIDER=auto # anthropic | openai | ollama | none +# KB_EXTRACT_ENTITIES=false # LLM entity extraction over source files (costs tokens) +# KB_TOKENIZER=tiktoken # bytes = fully offline +# Scale-out back-end credentials (LightRAG names): +# NEO4J_USERNAME=neo4j # falls back to NEO4J_USER +# POSTGRES_HOST=localhost POSTGRES_PORT=5432 POSTGRES_USER= POSTGRES_PASSWORD= POSTGRES_DATABASE= +# QDRANT_URL=http://localhost:6333 + # Workflow Settings ENABLE_CHECKPOINTS=true ENABLE_HITL=false diff --git a/.gitignore b/.gitignore index 575181f..4bf2e31 100644 --- a/.gitignore +++ b/.gitignore @@ -51,6 +51,8 @@ data/graph_db/* !data/graph_db/.gitkeep data/metadata/* !data/metadata/.gitkeep +data/knowledge_base/* +!data/knowledge_base/.gitkeep # Logs *.log diff --git a/QUICKSTART.md b/QUICKSTART.md index f70b1e1..b8fffa9 100644 --- a/QUICKSTART.md +++ b/QUICKSTART.md @@ -60,6 +60,15 @@ echo "LANGSMITH_PROJECT=archaeocode" >> .env ``` Then view traces at https://smith.langchain.com. Details: [docs/LANGSMITH_SETUP.md](docs/LANGSMITH_SETUP.md) +**Knowledge base** — keep what a run learned and ask questions about it: + +```bash +python archaeo --source sample_data/cobol --source-lang cobol --knowledge-base --kb-wiki ./wiki +archaeo-kb retrieve "PAYMENT" --mode local # works without an API key +archaeo-kb query "Which programs touch the customer file?" # needs an LLM key +``` +Local files by default; Neo4j/Postgres/Qdrant for large systems. Details: [docs/KNOWLEDGE_MANAGEMENT.md](docs/KNOWLEDGE_MANAGEMENT.md) + **Smalltalk grammars** — required before analyzing Smalltalk code: ```bash @@ -76,4 +85,5 @@ python scripts/build_smalltalk_grammar.py - [README](README.md) — feature overview and Python API - [examples/user_story_extraction/basic_usage.py](examples/user_story_extraction/basic_usage.py) — programmatic usage +- [examples/knowledge_base/query_knowledge_base.py](examples/knowledge_base/query_knowledge_base.py) — build and query the knowledge base - [docs/ROADMAP.md](docs/ROADMAP.md) — where the project is heading diff --git a/README.md b/README.md index b319890..380556e 100644 --- a/README.md +++ b/README.md @@ -12,7 +12,8 @@ Most reverse-engineering tools stop at syntax: they parse code and draw diagrams - **User stories from code** — the unique feature: an LLM analyzes each source file and produces user stories with roles, capabilities, benefits, acceptance criteria, priority, and confidence scores. Legacy knowledge becomes a product backlog. - **Legacy-first language support** — COBOL, Fortran, Pascal, and Smalltalk alongside Java, Python, JavaScript, and TypeScript. See the [language matrix](#language-support) for per-language capabilities. -- **MCP architecture** — analysis tools (AST parsing, dependency graphs, RAG) are [Model Context Protocol](https://modelcontextprotocol.io/) servers, so any MCP-capable agent can reuse them independently of this workflow. +- **Knowledge management built in** — every run can be persisted into a [GraphRAG knowledge base](docs/KNOWLEDGE_MANAGEMENT.md) (powered by [LightRAG](https://github.com/HKUDS/LightRAG)): files, entities, dependency edges and user stories become a queryable knowledge graph, exportable as an Obsidian-style wiki. Ask it questions with `archaeo-kb` or from any agent via MCP. Runs on local files or on Neo4j / Postgres / Qdrant for large estates. +- **MCP architecture** — analysis tools (AST parsing, dependency graphs, RAG, knowledge base) are [Model Context Protocol](https://modelcontextprotocol.io/) servers, so any MCP-capable agent can reuse them independently of this workflow. - **Observable by design** — every workflow run can be traced in [LangSmith](https://docs.smith.langchain.com) (token costs, latency, state snapshots). ## How it works @@ -20,13 +21,13 @@ Most reverse-engineering tools stop at syntax: they parse code and draw diagrams ``` LangGraph Orchestration │ -[Discovery] ──► [AST Analysis] ──► [Dependency Mapping] ──► [User Stories] - │ │ │ │ -file catalog tree-sitter graph + cycle Claude / GPT - parsing detection +[Discovery] ──► [AST Analysis] ──► [Dependency Mapping] ──► [User Stories] ──► [Knowledge Base] + │ │ │ │ │ +file catalog tree-sitter graph + cycle Claude / GPT LightRAG GraphRAG + parsing detection (optional, --knowledge-base) ``` -Each analysis step is an MCP server under `src/mcp_servers/` (static analysis, graph DB, RAG pipeline); the LangGraph workflow under `src/orchestration/` wires them together with checkpointing and state management. +Each analysis step is an MCP server under `src/mcp_servers/` (static analysis, graph DB, RAG pipeline, knowledge base); the LangGraph workflow under `src/orchestration/` wires them together with checkpointing and state management. ## Quick start @@ -56,6 +57,20 @@ python archaeo --source sample_data/java --source-lang java --target-lang python You get a console summary plus a JSON report: file catalog, language breakdown, dependency edges/cycles/layers, and (with an API key) generated user stories. +Keep the knowledge instead of just reporting it: + +```bash +# persist the run into the knowledge base (+ an Obsidian-compatible wiki) +python archaeo --source sample_data/cobol --source-lang cobol --knowledge-base --kb-wiki ./wiki + +# then ask it questions (query needs an LLM key; retrieve/graph/stats do not) +archaeo-kb query "Which programs touch the customer master file?" +archaeo-kb retrieve "PAYMENT" --mode local +archaeo-kb graph "PAYMENT.cob" +``` + +See [docs/KNOWLEDGE_MANAGEMENT.md](docs/KNOWLEDGE_MANAGEMENT.md) for the design decision (GraphRAG via LightRAG), query modes, and how to scale it to Neo4j / Postgres / Qdrant. + Or use it from Python: ```python @@ -86,8 +101,11 @@ A complete runnable example is in [`examples/user_story_extraction/basic_usage.p | LangSmith tracing | ✅ working | | Checkpointing / resumable workflows | ✅ working | | Smalltalk grammars (standard + Cincom) | ✅ working (grammar build required, see below) | -| RAG semantic code search (MCP server) | 🚧 functional, not yet wired into the workflow | -| Neo4j-backed dependency graphs | 🚧 in-memory fallback works; live Neo4j optional | +| Knowledge base: GraphRAG over files, entities, dependencies, stories (LightRAG, MCP server, `archaeo-kb`) | ✅ working (opt-in per run) | +| Knowledge base: LLM extraction of business entities | ✅ working (`--kb-extract`, needs API key) | +| Knowledge base: Obsidian-compatible wiki export | ✅ working | +| RAG semantic code search (MCP server) | 🚧 mock; superseded by the knowledge base | +| Neo4j-backed dependency graphs | 🚧 in-memory fallback works; live Neo4j optional (the knowledge base can use Neo4j today) | | Code generation to target language | 🎯 planned | | HP NonStop COBOL extensions (TMF, Pathway) | 🎯 planned | @@ -95,7 +113,7 @@ A complete runnable example is in [`examples/user_story_extraction/basic_usage.p - Python 3.10+ - An Anthropic or OpenAI API key for user-story extraction (everything else runs without one) -- Optional: Neo4j 5.x if you want persistent dependency graphs, LangSmith account for tracing +- Optional: Neo4j 5.x / PostgreSQL+pgvector / Qdrant to scale the knowledge base beyond local files, LangSmith account for tracing ## Language support @@ -116,18 +134,20 @@ The AST MCP server additionally parses C, C++, C#, Go, Rust, Ruby, PHP, and Bash ## Project structure ``` -├── archaeo # CLI entry point +├── archaeo # CLI entry point (archaeo-kb: knowledge base CLI) ├── src/ │ ├── orchestration/ # LangGraph workflow │ │ ├── graph.py # Direct workflow (in-process nodes) │ │ ├── graph_mcp.py # MCP-backed workflow -│ │ ├── nodes/ # Discovery, AST, dependency, user-story nodes +│ │ ├── kb_cli.py # archaeo-kb: query the knowledge base +│ │ ├── nodes/ # Discovery, AST, dependency, user-story, knowledge-base nodes │ │ ├── state/ # Workflow state schema │ │ └── utils/ # MCP client, LangSmith tracing │ ├── mcp_servers/ │ │ ├── static_analysis/ # tree-sitter AST analysis + custom grammars │ │ ├── graph_db/ # dependency graph (Neo4j / in-memory) -│ │ └── rag_pipeline/ # chunking, embeddings, semantic search +│ │ ├── rag_pipeline/ # chunking, embeddings, semantic search +│ │ └── knowledge_base/ # GraphRAG knowledge base (LightRAG adapter + MCP server) │ └── parsers/ # language-specific parser extensions ├── config/ # workflow + MCP server configuration ├── sample_data/ # COBOL, Java, Smalltalk, Fortran, Pascal, Python samples @@ -152,6 +172,7 @@ The suite runs in [GitHub Actions](.github/workflows/ci.yml) on Python 3.10–3. - [MCP server architecture](docs/MCP_ARCHITECTURE.md) and [usage](docs/MCP_SERVERS_USAGE.md) - [LangSmith setup](docs/LANGSMITH_SETUP.md) - [User story extraction](docs/USER_STORIES.md) +- [Knowledge management (GraphRAG knowledge base)](docs/KNOWLEDGE_MANAGEMENT.md) - [Roadmap](docs/ROADMAP.md) · [Changelog](docs/CHANGELOG.md) ## Contributing diff --git a/config/langgraph_config.yaml b/config/langgraph_config.yaml index c207ff2..eff7109 100644 --- a/config/langgraph_config.yaml +++ b/config/langgraph_config.yaml @@ -78,6 +78,45 @@ tools: enabled: true query_timeout: 60 +# Knowledge Base (knowledge management) +# -------------------------------------- +# Persists every run's findings (files, AST entities, dependency graph, user +# stories) into a GraphRAG store powered by LightRAG (https://github.com/HKUDS/LightRAG). +# Opt-in per run: `archaeo --knowledge-base` or `enabled: true` here. +# Environment variables KB_* override these values; see .env.example. +knowledge_base: + enabled: false + working_dir: "data/knowledge_base" + workspace: "" # one workspace per analyzed system on shared back-ends + # Storage — file-based defaults need no infrastructure; swap for scale: + # graph_storage: NetworkXStorage | Neo4JStorage | PGGraphStorage | MemgraphStorage | MongoGraphStorage + # vector_storage: NanoVectorDBStorage | PGVectorStorage | QdrantVectorDBStorage | MilvusVectorDBStorage | FaissVectorDBStorage + # kv_storage: JsonKVStorage | PGKVStorage | RedisKVStorage | MongoKVStorage + graph_storage: "NetworkXStorage" + vector_storage: "NanoVectorDBStorage" + kv_storage: "JsonKVStorage" + doc_status_storage: "JsonDocStatusStorage" + # Embeddings: sentence-transformers (local, default) | openai | ollama | hashing (tests) + embedding_provider: "sentence-transformers" + embedding_model: "all-MiniLM-L6-v2" + # LLM for answers / keyword extraction / entity extraction: auto | anthropic | openai | ollama | none + llm_provider: "auto" + llm_model: null # defaults to MODEL_NAME + # Chunking (tokens of the active tokenizer; "bytes" tokenizer works offline) + chunk_token_size: 1200 + chunk_overlap_token_size: 100 + tokenizer: "tiktoken" + # LLM-based entity/relation extraction of business concepts (costs tokens) + extract_entities: false + max_extract_files: 0 # 0 = all files + # Retrieval defaults + default_query_mode: "mix" # local | global | hybrid | naive | mix + top_k: 40 + chunk_top_k: 20 + # Wiki export (Obsidian-compatible Markdown), null = off + wiki_dir: null + max_graph_nodes: 5000 + # MCP Server Configuration # MCP (Model Context Protocol) servers expose tools for AI agents mcp_servers: @@ -123,6 +162,21 @@ mcp_servers: - get_graph_metrics - clear_graph + knowledge-base: + enabled: true + command: "python" + args: + - "src/mcp_servers/knowledge_base/knowledge_mcp_server.py" + description: "GraphRAG knowledge base of the analyzed system (LightRAG)" + tools: + - kb_query + - kb_retrieve + - kb_context + - kb_graph + - kb_stats + - kb_add_documents + - kb_export_wiki + # Performance Optimization performance: parallel_execution: true diff --git a/config/mcp_servers_config.yaml b/config/mcp_servers_config.yaml index 850ef6e..b814862 100644 --- a/config/mcp_servers_config.yaml +++ b/config/mcp_servers_config.yaml @@ -65,6 +65,31 @@ rag_pipeline: separators: ["\n\n", "\n", " ", ""] function_boundary_aware: true +# Knowledge Base Configuration (LightRAG GraphRAG) +# The MCP server reads KB_* environment variables (see .env.example); these +# profiles document the recommended back-end combinations per scale. +knowledge_base: + profiles: + laptop: # one developer, up to ~100k LOC, zero infrastructure + graph_storage: "NetworkXStorage" + vector_storage: "NanoVectorDBStorage" + kv_storage: "JsonKVStorage" + embedding_provider: "sentence-transformers" + team: # shared Neo4j + Postgres, millions of LOC, several analysts + graph_storage: "Neo4JStorage" + vector_storage: "PGVectorStorage" + kv_storage: "PGKVStorage" + doc_status_storage: "PGDocStatusStorage" + embedding_provider: "openai" + embedding_model: "text-embedding-3-large" + enterprise: # dedicated vector cluster, Redis KV, many systems (one workspace each) + graph_storage: "Neo4JStorage" + vector_storage: "QdrantVectorDBStorage" # or MilvusVectorDBStorage + kv_storage: "RedisKVStorage" + doc_status_storage: "RedisDocStatusStorage" + embedding_provider: "openai" + embedding_model: "text-embedding-3-large" + # Graph Database Configuration graph_db: provider: "neo4j" # Options: "neo4j", "memgraph" diff --git a/data/knowledge_base/.gitkeep b/data/knowledge_base/.gitkeep new file mode 100644 index 0000000..e69de29 diff --git a/docs/CHANGELOG.md b/docs/CHANGELOG.md index 9c32085..df3bb71 100644 --- a/docs/CHANGELOG.md +++ b/docs/CHANGELOG.md @@ -1,5 +1,21 @@ # Changelog +## [0.4.0] - 2026-09-07 + +### Added +- **Knowledge management**: a persistent, queryable knowledge base built on [LightRAG](https://github.com/HKUDS/LightRAG) (GraphRAG: knowledge graph + vector retrieval). See [docs/KNOWLEDGE_MANAGEMENT.md](KNOWLEDGE_MANAGEMENT.md) for the decision record. + - New workflow node `knowledge_base` (opt-in via `--knowledge-base` / `knowledge_base.enabled`) that writes files, AST entities, dependency edges (incl. circular-dependency markers) and user stories into the graph **without any LLM call**; `--kb-extract` adds LLM extraction of business entities. + - New MCP server `src/mcp_servers/knowledge_base/knowledge_mcp_server.py` with `kb_query`, `kb_retrieve`, `kb_context`, `kb_graph`, `kb_stats`, `kb_add_documents`, `kb_export_wiki`; registered in `config/langgraph_config.yaml` and the in-process MCP client. + - New CLI `archaeo-kb` (`query`, `retrieve`, `context`, `graph`, `stats`, `add`, `wiki`). + - Obsidian-compatible Markdown wiki export (`--kb-wiki DIR`, `archaeo-kb wiki`). + - Storage back-ends selectable by configuration: NetworkX / Neo4j / Postgres+AGE / Memgraph (graph), nano-vectordb / FAISS / pgvector / Qdrant / Milvus (vectors), JSON / Postgres / Redis / Mongo (KV); workspaces isolate several systems in one back-end. Embeddings via sentence-transformers (default, local), OpenAI, Ollama; LLM via Anthropic, OpenAI, Ollama. + - Offline mode (`KB_TOKENIZER=bytes`, `KB_EMBEDDING_PROVIDER=hashing`) used by the new test module `tests/test_knowledge_base.py`. + - Example `examples/knowledge_base/query_knowledge_base.py`; new `knowledge_base` state field and report section. + +### Changed +- AST entities now carry `file_path`, so classes/functions/methods can be linked to their file. +- `requirements.txt` adds `lightrag-hku`. + ## [0.3.0] - 2026-08-27 First public release, under the new name **archaeocode** (previously `agentic-reverse-engineering`). diff --git a/docs/KNOWLEDGE_MANAGEMENT.md b/docs/KNOWLEDGE_MANAGEMENT.md new file mode 100644 index 0000000..e20670f --- /dev/null +++ b/docs/KNOWLEDGE_MANAGEMENT.md @@ -0,0 +1,182 @@ +# Knowledge Management + +archaeocode excavates a lot of knowledge from a legacy system — file catalogs, AST +entities, dependency graphs, user stories — but until v0.4 every run threw it +away after writing a JSON report. The **knowledge base** keeps it: a persistent, +queryable store that people query with `archaeo-kb` and that agents query +through the `knowledge-base` MCP server. + +This document records what was chosen, why, and how to use and scale it. + +## Decision: GraphRAG (LightRAG) with a wiki projection + +Three families were on the table: + +| Approach | What it is | Fit for legacy-system knowledge | +|---|---|---| +| **LLM wiki** (Karpathy-style: an LLM maintains a Markdown wiki) | Human-readable pages, incrementally curated by the model | Great *output* format, but no retrieval layer: does not scale to millions of LOC, every update is an LLM call, and questions like "what depends on copybook X" cannot be answered without re-reading pages | +| **Vanilla RAG** (chunks + vector search) | Embed text chunks, retrieve top-k | Good for "find code that looks like this", poor for multi-hop and global questions ("which business processes touch the tariff tables?"), and it ignores the exact dependency graph static analysis already produced | +| **GraphRAG** (knowledge graph + vectors) | Entities and relations in a graph, chunks and entities embedded, retrieval walks both | Matches the data: a codebase *is* a graph. Multi-hop questions, global summaries, and exact structural facts coexist; the graph can be seeded deterministically from the AST/dependency analysis so the LLM is only spent on business concepts | + +**Choice: GraphRAG, with the wiki as a projection of the graph** rather than as +the system of record. You get the human-facing "LLM wiki" experience (one page +per entity, `[[wikilinks]]`, works in Obsidian) without giving up a queryable, +scalable backbone. + +### Why LightRAG and not the alternatives + +The requirement was "reuse a fitting solution that integrates best and scales by +all dimensions". Candidates were assessed on integration fit (Python, LangGraph, +MCP, Neo4j already in the stack, Anthropic first), incremental updates, storage +back-ends, cost control, query capabilities, licence and maturity. + +| Candidate | Verdict | +|---|---| +| **[LightRAG](https://github.com/HKUDS/LightRAG)** (HKUDS, MIT) | **Chosen.** Pure-Python library, async, pluggable LLM/embedding functions (Anthropic, OpenAI, Ollama, HuggingFace adapters built in), storage back-ends for NetworkX/Neo4j/Memgraph/Postgres+AGE (graph), nano-vectordb/FAISS/pgvector/Qdrant/Milvus (vectors), JSON/Postgres/Redis/Mongo (KV). Incremental inserts with document tracking and deletion. **`insert_custom_kg`** lets us write the exact dependency graph without any LLM. Query modes local/global/hybrid/naive/mix, structured retrieval without LLM (`query_data`), graph export. Workspaces isolate multiple systems in one back-end. | +| Microsoft GraphRAG | Batch indexer built around community summaries: expensive to (re)index, weak incremental story, parquet/LanceDB-centric, hard to reuse the deterministic graph. Best for static document corpora, not for a codebase that keeps changing. | +| Cognee | Solid "AI memory" engine with graph + vector and an MCP server, but more opinionated pipelines ("cognify"), telemetry on by default, and less direct control over inserting a prebuilt graph. Close second. | +| Graphiti (Zep) | Temporal knowledge graph aimed at agent memory; Neo4j/FalkorDB mandatory, extraction is LLM-first. Excellent for evolving facts, over-engineered for code corpora. | +| neo4j-graphraph (official Neo4j) | Clean and Neo4j-native, but Neo4j is mandatory (no zero-infrastructure mode for laptops/CI) and there is no global/community retrieval. | +| LlamaIndex PropertyGraphIndex | Capable, but pulls a second orchestration framework next to LangGraph. | + +## Architecture + +``` +archaeo run + Discovery ─► AST ─► Dependencies ─► User stories ─► Knowledge Base node + │ + deterministic (no LLM) ────────────────┤ + files, classes, functions, methods │ + import/call edges, circular deps │ LightRAG + user stories ⟷ files │ ┌────────────┐ + ├─►│ graph store │ NetworkX | Neo4j | Postgres/AGE | Memgraph + optional LLM extraction ───────────────┤ │ vector db │ nano-vectordb | pgvector | Qdrant | Milvus | FAISS + business entities & relations │ │ kv / status │ JSON | Postgres | Redis | Mongo + (customers, invoices, tariffs, ...) │ └────────────┘ + │ + archaeo-kb CLI ◄──── query / retrieve / graph ◄─────┤ + knowledge-base MCP server ◄─────────────────────────┤ + Obsidian wiki (index.md + entities/*.md) ◄──────────┘ +``` + +Code lives in: + +- `src/mcp_servers/knowledge_base/knowledge_base.py` — the adapter (`KnowledgeBase`, `KnowledgeBaseConfig`, `build_custom_kg`) +- `src/mcp_servers/knowledge_base/knowledge_mcp_server.py` — MCP tools +- `src/orchestration/nodes/knowledge_node.py` — workflow node +- `src/orchestration/kb_cli.py` — `archaeo-kb` + +### What gets indexed + +| Source in workflow state | Graph entities | Relations | Embedded chunks | +|---|---|---|---| +| `code_artifacts` | one `file` entity per file (language, size, complexity, defined members) | | source windows of ~`chunk_token_size` tokens with a `File: path (lines a-b)` header | +| `parsed_entities` (tree-sitter) | `class` / `function` / `method` entities named `path::Name` | `file —contains→ member` | (linked to the chunk containing the definition) | +| `dependency_graph` | `external_dependency` entities for targets outside the sources | `file —imports/calls→ file/external`, `circular-dependency` keyword on cycle edges | | +| `user_stories` | `user_story` entities with role, capability, benefit, criteria | `story —derived_from→ file` | the story as Markdown | +| documents added later (`archaeo-kb add`, `kb_add_documents`) | `document` entities | via LLM extraction (optional) | document windows | + +Everything in the table is written with `insert_custom_kg`: exact, cheap, +idempotent (re-indexing upserts). LLM extraction (`--kb-extract`) additionally +runs LightRAG's entity/relation extraction over the sources and merges the +result into the same graph — that is where business vocabulary appears. + +## Usage + +### Index during a run + +```bash +# zero infrastructure: NetworkX + nano-vectordb under data/knowledge_base, +# local sentence-transformers embeddings, no LLM needed +python archaeo --source ./legacy --source-lang cobol --knowledge-base --kb-wiki ./wiki + +# with an API key: also let the LLM extract business entities +python archaeo --source ./legacy --source-lang cobol --knowledge-base --kb-extract +``` + +Flags: `--knowledge-base/--kb`, `--kb-dir`, `--kb-workspace`, `--kb-extract`, +`--kb-wiki DIR`, `--kb-embeddings {sentence-transformers,openai,ollama,hashing}`. +The same settings live under `knowledge_base:` in `config/langgraph_config.yaml` +and as `KB_*` environment variables (`.env.example`). + +### Ask questions + +```bash +archaeo-kb query "Which programs update the customer master file, and what business rules apply?" +archaeo-kb query "Summarize the payment domain" --mode global +archaeo-kb retrieve "PROCESS-PAYMENT" --mode local # no LLM: entities, relations, chunks +archaeo-kb context "SQLCA" # prompt-ready evidence block +archaeo-kb graph "PAYMENT.cob" --depth 2 +archaeo-kb stats +archaeo-kb add docs/runbook.md docs/interview-notes.md # enrich with non-code knowledge +archaeo-kb wiki ./wiki # Obsidian-compatible export +``` + +Query modes (LightRAG): `local` (entity neighbourhood), `global` (relations / +themes), `hybrid`, `naive` (chunks only), `mix` (default: graph + chunks). +`query` needs an LLM (Anthropic → OpenAI → Ollama, picked from the keys present); +`retrieve`, `context`, `graph`, `stats`, `wiki` never do. + +### From Python + +```python +from src.mcp_servers.knowledge_base import KnowledgeBase, KnowledgeBaseConfig + +config = KnowledgeBaseConfig(working_dir="data/knowledge_base") + +async def ask(kb: KnowledgeBase): + return await kb.query("Which files implement invoice calculation?", mode="mix") + +print(KnowledgeBase.run_with(ask, config)) +``` + +### From an agent (MCP) + +`python src/mcp_servers/knowledge_base/knowledge_mcp_server.py` exposes +`kb_query`, `kb_retrieve`, `kb_context`, `kb_graph`, `kb_stats`, +`kb_add_documents`, `kb_export_wiki`. Add it to Claude Desktop / Claude Code like +the other servers (see `docs/MCP_SERVERS_USAGE.md`); set `KB_WORKING_DIR` or pass +`working_dir` per call. + +## Scaling + +The adapter never touches back-end specifics; scale is a configuration change. +Recommended profiles (also listed in `config/mcp_servers_config.yaml`): + +| Profile | Graph | Vectors | KV / status | Embeddings | Typical size | +|---|---|---|---|---|---| +| laptop / CI | `NetworkXStorage` | `NanoVectorDBStorage` | JSON files | sentence-transformers (local) | up to ~100k LOC, one user | +| team | `Neo4JStorage` | `PGVectorStorage` | `PGKVStorage` | OpenAI `text-embedding-3-large` | millions of LOC, several analysts, shared | +| enterprise | `Neo4JStorage` | `QdrantVectorDBStorage` / `MilvusVectorDBStorage` | `RedisKVStorage` | OpenAI / self-hosted | many systems, one `workspace` each | + +Set them via environment (`KB_GRAPH_STORAGE=Neo4JStorage`, `NEO4J_URI`, +`NEO4J_USERNAME` (falls back to `NEO4J_USER`), `NEO4J_PASSWORD`, +`POSTGRES_*`, `QDRANT_URL`, ...) or in the YAML section. Other scale levers: + +- **Workspaces** — `--kb-workspace billing-core` keeps several analyzed systems + apart inside one Neo4j/Postgres/Qdrant. +- **Batching** — custom-KG inserts are streamed in batches of files, so a + million-line codebase does not need to fit in one write. +- **Cost** — deterministic indexing is free of LLM tokens. `--kb-extract` is the + only expensive step; cap it with `max_extract_files`, and LightRAG caches LLM + responses so re-runs do not pay twice. Local embeddings via sentence-transformers + or Ollama keep data on-premise. +- **Offline** — `KB_TOKENIZER=bytes` and `KB_EMBEDDING_PROVIDER=hashing` run the + whole pipeline without network access (used by the test suite). If the + sentence-transformers model cannot be downloaded the adapter falls back to + hashing embeddings with a warning. +- **Observability** — LightRAG calls go through the same Anthropic/OpenAI SDKs; + the workflow node is a normal LangGraph node, so it shows up in LangSmith. + +## Limitations and next steps + +- AST entities are only available for tree-sitter-supported languages; COBOL, + Fortran and Pascal contribute files, dependencies and stories until their + grammars land (see the language matrix in the README). +- The dependency mapper still returns illustrative edges for some languages; + the knowledge base faithfully indexes whatever it gets — improving the mapper + improves the graph automatically. +- Planned: incremental re-indexing keyed by file hash (LightRAG already tracks + document ids), community summaries for very large systems, and feeding + retrieved context back into the user-story prompt so stories are written with + cross-file knowledge. diff --git a/docs/MCP_ARCHITECTURE.md b/docs/MCP_ARCHITECTURE.md index d797b04..3768c81 100644 --- a/docs/MCP_ARCHITECTURE.md +++ b/docs/MCP_ARCHITECTURE.md @@ -58,7 +58,7 @@ result = mcp_client.call_tool("parse_file", { ## 🏗️ Our MCP Architecture -We're implementing **3 MCP servers** for our reverse engineering system: +We're implementing **3 analysis MCP servers** for our reverse engineering system, plus a fourth, the **knowledge base server** (GraphRAG via LightRAG) that persists what the analysis found — see [KNOWLEDGE_MANAGEMENT.md](KNOWLEDGE_MANAGEMENT.md) and [MCP_SERVERS_USAGE.md](MCP_SERVERS_USAGE.md#4-knowledge-base-mcp-server): ``` ┌──────────────────────────────────────────────────────────────┐ diff --git a/docs/MCP_SERVERS_USAGE.md b/docs/MCP_SERVERS_USAGE.md index 896b8b4..4376075 100644 --- a/docs/MCP_SERVERS_USAGE.md +++ b/docs/MCP_SERVERS_USAGE.md @@ -4,11 +4,12 @@ This document explains how to use the MCP (Model Context Protocol) servers imple ## Overview -The project implements **three MCP servers** using the official [Anthropic MCP SDK](https://github.com/anthropics/mcp): +The project implements **four MCP servers** using the official [Anthropic MCP SDK](https://github.com/anthropics/mcp): 1. **AST Analysis Server** - Code parsing and static analysis using tree-sitter 2. **RAG Pipeline Server** - Code chunking, embeddings, and semantic search 3. **Neo4j Graph Database Server** - Dependency graph operations +4. **Knowledge Base Server** - GraphRAG knowledge base of an analyzed system (LightRAG) All servers follow the official MCP protocol and can be used standalone or integrated with Claude Desktop, IDEs, or custom applications. @@ -552,6 +553,69 @@ python src/mcp_servers/graph_db/neo4j_mcp_server.py --- +## 4. Knowledge Base MCP Server + +**Location:** `src/mcp_servers/knowledge_base/knowledge_mcp_server.py` + +Exposes the knowledge base that `archaeo --knowledge-base` builds (see +[KNOWLEDGE_MANAGEMENT.md](KNOWLEDGE_MANAGEMENT.md)). Backed by +[LightRAG](https://github.com/HKUDS/LightRAG): knowledge graph + vector indexes, +file-based by default, Neo4j/Postgres/Qdrant/... via `KB_*` environment variables. + +All tools accept optional `working_dir` (default `$KB_WORKING_DIR` or +`data/knowledge_base`) and `workspace` (default `$KB_WORKSPACE`) arguments and +return `{"success": true, ...}` or `{"success": false, "error": "..."}`. + +### Available Tools + +#### `kb_query` +Natural-language question → LLM-synthesized answer with citations. Needs an LLM +(Anthropic / OpenAI / Ollama, picked from the keys present). + +**Parameters:** `question` (required), `mode` (`local` | `global` | `hybrid` | `naive` | `mix`), `response_type` + +#### `kb_retrieve` +Structured retrieval without an LLM: matching entities, relationships, chunks and references. + +**Parameters:** `question` (required), `mode`, `top_k`, `chunk_top_k` + +**Returns:** +```json +{ + "success": true, + "question": "PROCESS-PAYMENT", + "mode": "local", + "entities": [{"entity_name": "PAYMENT.cob::PROCESS-PAYMENT", "entity_type": "function", "description": "...", "file_path": "PAYMENT.cob"}], + "relationships": [{"src_id": "PAYMENT.cob", "tgt_id": "PAYMENT.cob::PROCESS-PAYMENT", "keywords": "contains", "description": "..."}], + "chunks": [{"file_path": "PAYMENT.cob", "content": "File: PAYMENT.cob (lines 1-36, cobol)\n..."}], + "references": [], + "metadata": {"query_mode": "local", "keywords": {"high_level": [], "low_level": ["PROCESS-PAYMENT"]}} +} +``` + +#### `kb_context` +The prompt-ready context block LightRAG would hand to an LLM — for agents that reason over the evidence themselves. + +#### `kb_graph` +Neighbourhood of an entity (`label`, `max_depth`, `max_nodes`); `label="*"` returns the whole (capped) graph. + +#### `kb_stats` +Entity count, entity-type breakdown, document processing status, and the active storage / embedding / LLM configuration. + +#### `kb_add_documents` +Ingest additional documents (`documents: [{"path", "content"}]`, optional `extract` for LLM entity extraction). + +#### `kb_export_wiki` +Write an Obsidian-compatible Markdown wiki (`output_dir`, optional `max_nodes`). + +### Running the Knowledge Base Server + +```bash +KB_WORKING_DIR=data/knowledge_base python src/mcp_servers/knowledge_base/knowledge_mcp_server.py +``` + +--- + ## Integration with Claude Desktop To use these MCP servers with Claude Desktop, add them to your Claude Desktop configuration file: diff --git a/examples/knowledge_base/__init__.py b/examples/knowledge_base/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/examples/knowledge_base/query_knowledge_base.py b/examples/knowledge_base/query_knowledge_base.py new file mode 100644 index 0000000..e042685 --- /dev/null +++ b/examples/knowledge_base/query_knowledge_base.py @@ -0,0 +1,76 @@ +""" +Knowledge base example: analyze the bundled COBOL sample, persist the findings +into the LightRAG-backed knowledge base, then query it. + +Runs fully offline (no API key needed) with local embeddings. With +ANTHROPIC_API_KEY (or OPENAI_API_KEY) set, the last step also asks the LLM a +question and you can pass --extract to let it mine business entities. + +Run from the project root: + python examples/knowledge_base/query_knowledge_base.py [--extract] +""" + +import os +import sys + +project_root = os.path.join(os.path.dirname(__file__), "..", "..") +sys.path.insert(0, project_root) + +from src.orchestration.graph import create_graph # noqa: E402 +from src.mcp_servers.knowledge_base import KnowledgeBase, KnowledgeBaseConfig, KnowledgeBaseError # noqa: E402 + + +def main() -> None: + extract = "--extract" in sys.argv + kb_dir = os.path.join(project_root, "data", "knowledge_base", "example") + + graph = create_graph( + { + "checkpoints": {"enabled": False}, + "hitl": {"enabled": False}, + "knowledge_base": { + "enabled": True, + "working_dir": kb_dir, + "extract_entities": extract, + "wiki_dir": os.path.join(kb_dir, "wiki"), + }, + } + ) + result = graph.run( + source_language="cobol", + target_language="java", + source_path=os.path.join(project_root, "sample_data", "cobol"), + ) + stats = result["knowledge_base"] + print(f"\nIndexed {stats['entities']} entities / {stats['relationships']} relations into {stats['working_dir']}") + + config = KnowledgeBaseConfig.from_env(working_dir=kb_dir) + + async def explore(kb: KnowledgeBase): + hits = await kb.retrieve("customer payment", mode="hybrid") + print("\nEntities matching 'customer payment':") + for entity in hits["entities"][:5]: + print(f" - {entity['entity_name']} [{entity['entity_type']}]") + + neighbourhood = await kb.graph("PAYMENT.cob", max_depth=1) + print("\nNeighbourhood of PAYMENT.cob:") + for edge in neighbourhood["edges"]: + print(f" {edge['source']} —{edge['keywords']}→ {edge['target']}") + + if kb.llm_available: + print("\nAsking the LLM...") + print(await kb.query("What does the payment program do, and which other programs does it rely on?")) + else: + print("\n(no LLM key configured — skipping the natural-language answer)") + + try: + KnowledgeBase.run_with(explore, config) + except KnowledgeBaseError as exc: + print(f"Knowledge base error: {exc}") + + print(f"\nWiki: {os.path.join(kb_dir, 'wiki', 'index.md')}") + print(f"CLI: archaeo-kb --dir {kb_dir} retrieve \"payment\"") + + +if __name__ == "__main__": + main() diff --git a/pyproject.toml b/pyproject.toml index 1d0c6f9..fccdb89 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,13 +4,13 @@ build-backend = "setuptools.build_meta" [project] name = "archaeocode" -version = "0.3.0" +version = "0.4.0" description = "Automated software archaeology: excavate user stories, dependency maps, and structure from legacy code with LangGraph + MCP agents" readme = "README.md" requires-python = ">=3.10" license = { text = "MIT" } authors = [{ name = "Oliver Sick" }] -keywords = ["reverse-engineering", "legacy-code", "langgraph", "mcp", "tree-sitter", "cobol", "smalltalk", "migration"] +keywords = ["reverse-engineering", "legacy-code", "langgraph", "mcp", "tree-sitter", "cobol", "smalltalk", "migration", "graphrag", "knowledge-graph"] classifiers = [ "Development Status :: 3 - Alpha", "Intended Audience :: Developers", @@ -26,6 +26,7 @@ dynamic = ["dependencies"] [project.scripts] archaeo = "src.orchestration.cli:main" +archaeo-kb = "src.orchestration.kb_cli:main" [project.urls] Homepage = "https://github.com/osick/archaeocode" diff --git a/requirements.txt b/requirements.txt index e5e13de..56906e8 100644 --- a/requirements.txt +++ b/requirements.txt @@ -18,6 +18,13 @@ mcp>=1.0.0 # Config parsing PyYAML>=6.0 +# Knowledge management (GraphRAG): knowledge graph + vector retrieval. +# File-based back-ends (NetworkX, nano-vectordb) work out of the box; scale-out +# back-ends only need their client library: +# Neo4j -> neo4j (already below) Postgres/pgvector -> asyncpg +# Qdrant -> qdrant-client Milvus -> pymilvus Redis -> redis +lightrag-hku>=1.5.0 + # Vector store (choose one or more) chromadb>=0.5.0 # Embedded vector DB # qdrant-client>=1.11.0 # Alternative: Qdrant diff --git a/src/mcp_servers/knowledge_base/__init__.py b/src/mcp_servers/knowledge_base/__init__.py new file mode 100644 index 0000000..11e884e --- /dev/null +++ b/src/mcp_servers/knowledge_base/__init__.py @@ -0,0 +1,18 @@ +""" +Knowledge Base MCP Server +========================= + +Persistent, queryable knowledge about an analyzed legacy system, backed by +`LightRAG `_ (graph + vector retrieval). + +See ``knowledge_base.py`` for the adapter and ``knowledge_mcp_server.py`` for +the MCP tool surface. +""" + +from .knowledge_base import ( # noqa: F401 + KnowledgeBase, + KnowledgeBaseConfig, + KnowledgeBaseError, + build_custom_kg, + is_available, +) diff --git a/src/mcp_servers/knowledge_base/knowledge_base.py b/src/mcp_servers/knowledge_base/knowledge_base.py new file mode 100644 index 0000000..d2569d9 --- /dev/null +++ b/src/mcp_servers/knowledge_base/knowledge_base.py @@ -0,0 +1,1240 @@ +""" +Knowledge Base — LightRAG adapter +================================= + +archaeocode's knowledge-management layer. It turns the artifacts a workflow run +produces (file catalog, AST entities, dependency graph, user stories) into a +persistent knowledge base that can be queried by humans and agents. + +The heavy lifting is delegated to `LightRAG `_, +a graph-augmented retrieval engine (GraphRAG) that stores a knowledge graph plus +vector indexes for entities, relations and text chunks. archaeocode contributes: + +* **Deterministic knowledge** — files, classes, functions, dependency edges and + user stories are written straight into the graph with ``insert_custom_kg``. + No LLM call is needed; static analysis already produced exact facts. +* **Optional LLM enrichment** — with an API key, raw documents can additionally + go through LightRAG's entity/relation extraction to surface *business* + concepts (customers, invoices, tariffs, ...) that live inside the code. +* **Pluggable storage** — the same adapter runs on local files (NetworkX + + nano-vectordb, zero infrastructure) or on Neo4j / PostgreSQL+pgvector / + Qdrant / Milvus / Redis / MongoDB for large systems, selected purely through + configuration. +* **Wiki projection** — the graph can be exported as an Obsidian-compatible + Markdown wiki (one page per entity, ``[[wikilinks]]``) for humans. + +LightRAG is an optional dependency: everything in here degrades to a clear +``KnowledgeBaseError`` when it is not installed. +""" + +from __future__ import annotations + +import asyncio +import hashlib +import logging +import os +import re +import threading +from dataclasses import dataclass, field, fields +from datetime import datetime +from pathlib import Path +from typing import Any, Awaitable, Callable, Dict, List, Optional, Tuple, TypeVar + +logger = logging.getLogger(__name__) + +T = TypeVar("T") + +DEFAULT_WORKING_DIR = "data/knowledge_base" + +# Storage back-ends understood by LightRAG (see lightrag.kg.STORAGE_IMPLEMENTATIONS). +GRAPH_STORAGES = ["NetworkXStorage", "Neo4JStorage", "PGGraphStorage", "MemgraphStorage", "MongoGraphStorage"] +VECTOR_STORAGES = [ + "NanoVectorDBStorage", + "FaissVectorDBStorage", + "PGVectorStorage", + "QdrantVectorDBStorage", + "MilvusVectorDBStorage", + "MongoVectorDBStorage", +] +KV_STORAGES = ["JsonKVStorage", "PGKVStorage", "RedisKVStorage", "MongoKVStorage"] +DOC_STATUS_STORAGES = ["JsonDocStatusStorage", "PGDocStatusStorage", "RedisDocStatusStorage", "MongoDocStatusStorage"] + +EMBEDDING_PROVIDERS = ["sentence-transformers", "openai", "ollama", "hashing"] +LLM_PROVIDERS = ["auto", "anthropic", "openai", "ollama", "none"] +QUERY_MODES = ["local", "global", "hybrid", "naive", "mix"] + +HASHING_EMBEDDING_DIM = 256 + + +class KnowledgeBaseError(RuntimeError): + """Raised for configuration or availability problems of the knowledge base.""" + + +def is_available() -> bool: + """Return True when the LightRAG library can be imported.""" + try: + import lightrag # noqa: F401 + + return True + except Exception: # pragma: no cover - depends on the environment + return False + + +# --------------------------------------------------------------------------- +# Configuration +# --------------------------------------------------------------------------- + + +def _env_bool(name: str, default: bool) -> bool: + value = os.getenv(name) + if value is None: + return default + return value.strip().lower() in {"1", "true", "yes", "on"} + + +@dataclass +class KnowledgeBaseConfig: + """ + Configuration of the knowledge base. + + Every field can be set from ``config/langgraph_config.yaml`` (section + ``knowledge_base``), from the CLI, or from ``KB_*`` environment variables. + """ + + # Where LightRAG keeps its local files (also used by the file back-ends). + working_dir: str = DEFAULT_WORKING_DIR + # Logical partition inside shared back-ends (Neo4j, Postgres, ...): one + # workspace per analyzed system keeps knowledge bases apart. + workspace: str = "" + + # Storage back-ends — file-based defaults need no infrastructure. + graph_storage: str = "NetworkXStorage" + vector_storage: str = "NanoVectorDBStorage" + kv_storage: str = "JsonKVStorage" + doc_status_storage: str = "JsonDocStatusStorage" + + # Embeddings: local sentence-transformers by default (no API key needed). + embedding_provider: str = "sentence-transformers" + embedding_model: str = "all-MiniLM-L6-v2" + embedding_dim: Optional[int] = None + # Minimum cosine similarity for vector hits (None = LightRAG default). + similarity_threshold: Optional[float] = None + + # LLM used for answer synthesis, keyword extraction and (optional) entity + # extraction. "auto" picks Anthropic, then OpenAI, then nothing, based on + # the API keys that are present. + llm_provider: str = "auto" + llm_model: Optional[str] = None + llm_max_async: int = 4 + enable_llm_cache: bool = True + + # Chunking of source files (in tokens of the active tokenizer). + chunk_token_size: int = 1200 + chunk_overlap_token_size: int = 100 + # "tiktoken" needs a one-time download of the encoding; "bytes" works offline. + tokenizer: str = "tiktoken" + + # LLM-based entity/relation extraction over raw documents (costs tokens). + extract_entities: bool = False + # Cap the number of files sent through LLM extraction (0 = all). + max_extract_files: int = 0 + + # Retrieval defaults. + default_query_mode: str = "mix" + top_k: int = 40 + chunk_top_k: int = 20 + + # Graph export limits. + max_graph_nodes: int = 5000 + + verbose: bool = False + extra: Dict[str, Any] = field(default_factory=dict) + + @classmethod + def from_dict(cls, data: Optional[Dict[str, Any]]) -> "KnowledgeBaseConfig": + """Build a config from a plain dict (unknown keys go to ``extra``).""" + data = dict(data or {}) + known = {f.name for f in fields(cls)} + kwargs = {k: v for k, v in data.items() if k in known and k != "extra"} + extra = {k: v for k, v in data.items() if k not in known} + extra.update(data.get("extra") or {}) + return cls(**kwargs, extra=extra) + + @classmethod + def from_env(cls, **overrides: Any) -> "KnowledgeBaseConfig": + """Build a config from ``KB_*`` environment variables plus overrides.""" + env = { + "working_dir": os.getenv("KB_WORKING_DIR"), + "workspace": os.getenv("KB_WORKSPACE"), + "graph_storage": os.getenv("KB_GRAPH_STORAGE"), + "vector_storage": os.getenv("KB_VECTOR_STORAGE"), + "kv_storage": os.getenv("KB_KV_STORAGE"), + "doc_status_storage": os.getenv("KB_DOC_STATUS_STORAGE"), + "embedding_provider": os.getenv("KB_EMBEDDING_PROVIDER"), + "embedding_model": os.getenv("KB_EMBEDDING_MODEL"), + "embedding_dim": os.getenv("KB_EMBEDDING_DIM"), + "llm_provider": os.getenv("KB_LLM_PROVIDER"), + "llm_model": os.getenv("KB_LLM_MODEL"), + "tokenizer": os.getenv("KB_TOKENIZER"), + "default_query_mode": os.getenv("KB_QUERY_MODE"), + "chunk_token_size": os.getenv("KB_CHUNK_TOKEN_SIZE"), + "chunk_overlap_token_size": os.getenv("KB_CHUNK_OVERLAP_TOKEN_SIZE"), + } + data: Dict[str, Any] = {k: v for k, v in env.items() if v not in (None, "")} + for int_key in ("embedding_dim", "chunk_token_size", "chunk_overlap_token_size"): + if int_key in data: + data[int_key] = int(data[int_key]) + if os.getenv("KB_EXTRACT_ENTITIES") is not None: + data["extract_entities"] = _env_bool("KB_EXTRACT_ENTITIES", False) + data.update({k: v for k, v in overrides.items() if v is not None}) + return cls.from_dict(data) + + def validate(self) -> None: + if self.graph_storage not in GRAPH_STORAGES: + raise KnowledgeBaseError(f"Unknown graph_storage '{self.graph_storage}'. Options: {GRAPH_STORAGES}") + if self.vector_storage not in VECTOR_STORAGES: + raise KnowledgeBaseError(f"Unknown vector_storage '{self.vector_storage}'. Options: {VECTOR_STORAGES}") + if self.kv_storage not in KV_STORAGES: + raise KnowledgeBaseError(f"Unknown kv_storage '{self.kv_storage}'. Options: {KV_STORAGES}") + if self.embedding_provider not in EMBEDDING_PROVIDERS: + raise KnowledgeBaseError( + f"Unknown embedding_provider '{self.embedding_provider}'. Options: {EMBEDDING_PROVIDERS}" + ) + if self.llm_provider not in LLM_PROVIDERS: + raise KnowledgeBaseError(f"Unknown llm_provider '{self.llm_provider}'. Options: {LLM_PROVIDERS}") + if self.default_query_mode not in QUERY_MODES: + raise KnowledgeBaseError(f"Unknown query mode '{self.default_query_mode}'. Options: {QUERY_MODES}") + + def to_dict(self) -> Dict[str, Any]: + return {f.name: getattr(self, f.name) for f in fields(self)} + + +# --------------------------------------------------------------------------- +# Small helpers +# --------------------------------------------------------------------------- + + +def run_sync(coro_factory: Callable[[], Awaitable[T]]) -> T: + """ + Run an async knowledge-base operation from synchronous code. + + LightRAG binds its storages to the event loop that initialized them, so a + whole unit of work (initialize → operate → close) must run on one loop. When + no loop is running we use ``asyncio.run``; inside a running loop (async + tests, MCP servers) we hop to a worker thread with its own loop. + """ + try: + asyncio.get_running_loop() + except RuntimeError: + return asyncio.run(coro_factory()) + + result: Dict[str, Any] = {} + + def _target() -> None: + try: + result["value"] = asyncio.run(coro_factory()) + except BaseException as exc: # pragma: no cover - re-raised below + result["error"] = exc + + thread = threading.Thread(target=_target, name="archaeocode-kb", daemon=True) + thread.start() + thread.join() + if "error" in result: + raise result["error"] + return result["value"] + + +class ByteTokenizer: + """ + Offline tokenizer: one token per UTF-8 byte. + + Used when tiktoken's encoding files cannot be downloaded. Token counts are + roughly 4x tiktoken's, so chunk sizes are scaled accordingly by the adapter. + Stateless, therefore thread-safe and deep-copyable (LightRAG requires both). + """ + + SCALE = 4 + + def encode(self, content: str) -> List[int]: + return list(content.encode("utf-8")) + + def decode(self, tokens: List[int]) -> str: + return bytes(tokens).decode("utf-8", errors="ignore") + + +def _hash_embedding(text: str, dim: int = HASHING_EMBEDDING_DIM): + """Deterministic bag-of-identifiers embedding (tests / offline smoke runs).""" + import numpy as np + + vec = np.zeros(dim, dtype=np.float32) + for token in re.findall(r"[A-Za-z_][A-Za-z0-9_]*", text.lower()): + bucket = int(hashlib.md5(token.encode("utf-8")).hexdigest(), 16) % dim + vec[bucket] += 1.0 + norm = float(np.linalg.norm(vec)) + return vec / norm if norm else vec + + +def _slugify(name: str) -> str: + slug = re.sub(r"[^A-Za-z0-9._-]+", "-", name).strip("-.") + slug = re.sub(r"-{2,}", "-", slug) + if not slug: + slug = "entity" + if len(slug) > 80: + slug = slug[:72] + "-" + hashlib.md5(name.encode("utf-8")).hexdigest()[:7] + return slug + + +def _relative_path(path: str, root: Optional[str]) -> str: + if not root: + return path + try: + rel = os.path.relpath(path, root) + except ValueError: # different drives on Windows + return path + return path if rel.startswith("..") else rel.replace(os.sep, "/") + + +def keywords_from_question(question: str, limit: int = 12) -> List[str]: + """Cheap keyword extraction used when no LLM is available for it.""" + seen: List[str] = [] + for token in re.findall(r"[A-Za-z_][A-Za-z0-9_.-]{2,}", question): + if token.lower() in _STOPWORDS: + continue + if token not in seen: + seen.append(token) + if len(seen) >= limit: + break + return seen + + +_STOPWORDS = { + "the", "and", "for", "with", "that", "this", "what", "which", "where", "when", "how", "does", "did", + "are", "was", "were", "from", "into", "about", "there", "their", "they", "them", "than", "then", + "have", "has", "had", "not", "but", "can", "could", "should", "would", "will", "who", "whom", "why", + "all", "any", "some", "each", "every", "one", "two", "use", "used", "uses", "using", "show", "list", +} + + +# --------------------------------------------------------------------------- +# Building the deterministic knowledge graph from workflow state +# --------------------------------------------------------------------------- + + +def _chunk_lines( + lines: List[str], max_chars: int, overlap_chars: int +) -> List[Tuple[int, int]]: + """Split ``lines`` into (start, end) line windows bounded by ``max_chars``.""" + windows: List[Tuple[int, int]] = [] + n = len(lines) + start = 0 + while start < n: + size = 0 + end = start + while end < n and (size + len(lines[end]) + 1 <= max_chars or end == start): + size += len(lines[end]) + 1 + end += 1 + windows.append((start, end)) + if end >= n: + break + # Walk back for the overlap, but always make progress. + back = end + overlap = 0 + while back > start + 1 and overlap + len(lines[back - 1]) + 1 <= overlap_chars: + back -= 1 + overlap += len(lines[back]) + 1 + start = max(back, start + 1) + return windows + + +def _story_markdown(story: Dict[str, Any]) -> str: + md = [f"# User Story: {story.get('title', 'Untitled')}", ""] + md.append( + f"Priority: {story.get('priority', 'Medium')} | Complexity: {story.get('complexity', '?')} points" + f" | Confidence: {story.get('confidence', 0.0)}" + ) + md.append("") + md.append(f"As a **{story.get('user_role', 'user')}**, I want to **{story.get('capability', '')}**, " + f"so that **{story.get('benefit', '')}**.") + criteria = story.get("acceptance_criteria") or [] + if criteria: + md.append("") + md.append("Acceptance criteria:") + for i, criterion in enumerate(criteria, 1): + md.append(f"{i}. {criterion}") + files = story.get("code_files") or [] + if files: + md.append("") + md.append("Code: " + ", ".join(files)) + return "\n".join(md) + + +def build_custom_kg( + state: Dict[str, Any], + config: Optional[KnowledgeBaseConfig] = None, + chars_per_token: int = 4, +) -> Tuple[Dict[str, Any], Dict[str, Any]]: + """ + Translate a workflow state into LightRAG's custom-KG format. + + The result contains ``chunks`` (embedded source windows and user stories), + ``entities`` (files, classes, functions, methods, external dependencies, + user stories) and ``relationships`` (contains / imports / derived-from). + Everything is derived from static analysis — no LLM involved — so the + graph is exact and cheap to rebuild. + + Returns: + (custom_kg, stats) + """ + config = config or KnowledgeBaseConfig() + root = state.get("source_path") + max_chars = max(400, config.chunk_token_size * chars_per_token) + overlap_chars = max(0, config.chunk_overlap_token_size * chars_per_token) + + chunks: List[Dict[str, Any]] = [] + entities: Dict[str, Dict[str, Any]] = {} + relationships: Dict[Tuple[str, str], Dict[str, Any]] = {} + + def add_entity(name: str, **data: Any) -> None: + entities[name] = {"entity_name": name, **data} + + def add_relationship(src: str, tgt: str, **data: Any) -> None: + if src == tgt: + return + key = tuple(sorted((src, tgt))) + if key in relationships: + # Merge keyword sets of parallel edges. + existing = relationships[key] + kws = set(existing.get("keywords", "").split(", ")) | set(data.get("keywords", "").split(", ")) + existing["keywords"] = ", ".join(sorted(k for k in kws if k)) + return + relationships[key] = {"src_id": src, "tgt_id": tgt, **data} + + # --- files and their text chunks ------------------------------------ + artifacts = state.get("code_artifacts") or [] + rel_by_path: Dict[str, str] = {} + first_chunk_of_file: Dict[str, str] = {} + chunk_ranges_of_file: Dict[str, List[Tuple[int, int, str]]] = {} + + for artifact in artifacts: + path = artifact.get("path") or artifact.get("id") or "unknown" + rel = _relative_path(path, root) + rel_by_path[path] = rel + rel_by_path[rel] = rel + language = artifact.get("language", "unknown") + content = artifact.get("content") or "" + lines = content.splitlines() + metadata = artifact.get("metadata") or {} + + ranges: List[Tuple[int, int, str]] = [] + for index, (start, end) in enumerate(_chunk_lines(lines, max_chars, overlap_chars) or [(0, 0)]): + alias = f"{rel}#chunk{index}" + body = "\n".join(lines[start:end]) + header = f"File: {rel} (lines {start + 1}-{end}, {language})\n" + chunks.append( + { + "content": header + body, + "source_id": alias, + "file_path": rel, + "chunk_order_index": index, + } + ) + ranges.append((start + 1, end, alias)) + first_chunk_of_file[rel] = ranges[0][2] + chunk_ranges_of_file[rel] = ranges + + description = f"{language} source file {rel}" + details = [] + if metadata.get("line_count") is not None: + details.append(f"{metadata['line_count']} lines") + if artifact.get("complexity_score") is not None: + details.append(f"complexity {artifact['complexity_score']}") + if details: + description += " (" + ", ".join(details) + ")" + add_entity( + rel, + entity_type="file", + description=description, + source_id=first_chunk_of_file[rel], + file_path=rel, + ) + + def chunk_alias_for(rel: str, line: Optional[int]) -> str: + ranges = chunk_ranges_of_file.get(rel) or [] + if line is not None: + for start, end, alias in ranges: + if start <= line <= end: + return alias + return first_chunk_of_file.get(rel, ranges[0][2] if ranges else "") + + # --- AST entities ----------------------------------------------------- + members_of_file: Dict[str, List[str]] = {} + parsed = state.get("parsed_entities") or {} + for kind in ("classes", "functions", "methods"): + for item in parsed.get(kind, []) or []: + name = item.get("name") + file_path = item.get("file_path") + if not name or not file_path: + continue + rel = rel_by_path.get(file_path) or _relative_path(file_path, root) + if rel not in first_chunk_of_file: + continue + entity_type = item.get("type") or kind.rstrip("s") + qualified = f"{rel}::{name}" + line_start = item.get("line_start") + span = f"lines {line_start}-{item.get('line_end')}" if line_start else "unknown position" + add_entity( + qualified, + entity_type=entity_type, + description=f"{entity_type} {name} defined in {rel} ({span})", + source_id=chunk_alias_for(rel, line_start), + file_path=rel, + ) + add_relationship( + rel, + qualified, + description=f"{rel} contains {entity_type} {name}", + keywords="contains", + weight=1.0, + source_id=chunk_alias_for(rel, line_start), + file_path=rel, + ) + members_of_file.setdefault(rel, []).append(f"{entity_type} {name}") + + for rel, members in members_of_file.items(): + entity = entities.get(rel) + if entity: + shown = ", ".join(members[:25]) + more = f" and {len(members) - 25} more" if len(members) > 25 else "" + entity["description"] += f". Defines: {shown}{more}" + + # --- dependency graph ------------------------------------------------- + cyclic_edges = set() + for cycle in state.get("circular_dependencies") or []: + for a, b in zip(cycle, cycle[1:]): + cyclic_edges.add((a, b)) + + known_basenames: Dict[str, str] = {} + for rel in first_chunk_of_file: + base = os.path.basename(rel) + known_basenames.setdefault(base.lower(), rel) + known_basenames.setdefault(os.path.splitext(base)[0].lower(), rel) + + def resolve_target(target: str) -> Tuple[str, bool]: + """Map a dependency target to a file entity if possible.""" + if target in rel_by_path: + return rel_by_path[target], True + rel = _relative_path(target, root) + if rel in first_chunk_of_file: + return rel, True + hit = known_basenames.get(target.lower()) + if hit: + return hit, True + return target, False + + external_count = 0 + for edge in state.get("dependency_graph") or []: + source = edge.get("source") + target = edge.get("target") + if not source or not target: + continue + src_rel = rel_by_path.get(source) or _relative_path(source, root) + if src_rel not in first_chunk_of_file: + continue + tgt_name, is_file = resolve_target(target) + rel_type = edge.get("relationship_type") or "depends_on" + if not is_file and tgt_name not in entities: + language = (edge.get("metadata") or {}).get("language", "") + add_entity( + tgt_name, + entity_type="external_dependency", + description=f"External {language} dependency {tgt_name} (referenced by {src_rel}, not part of the analyzed sources)", + source_id=first_chunk_of_file[src_rel], + file_path=src_rel, + ) + external_count += 1 + keywords = rel_type + if (source, target) in cyclic_edges or (src_rel, tgt_name) in cyclic_edges: + keywords += ", circular-dependency" + add_relationship( + src_rel, + tgt_name, + description=f"{src_rel} {rel_type} {tgt_name}", + keywords=keywords, + weight=1.0, + source_id=first_chunk_of_file[src_rel], + file_path=src_rel, + ) + + # --- user stories ----------------------------------------------------- + story_count = 0 + for index, story in enumerate(state.get("user_stories") or []): + title = story.get("title") or f"Story {index + 1}" + alias = f"story#{index}" + markdown = _story_markdown(story) + story_files = [rel_by_path.get(f) or _relative_path(f, root) for f in story.get("code_files") or []] + file_path = story_files[0] if story_files else "user_stories" + chunks.append({"content": markdown, "source_id": alias, "file_path": file_path, "chunk_order_index": 0}) + entity_name = f"Story: {title}" + criteria = "; ".join(story.get("acceptance_criteria") or []) + add_entity( + entity_name, + entity_type="user_story", + description=( + f"User story '{title}': as a {story.get('user_role', 'user')}, I want to " + f"{story.get('capability', '')}, so that {story.get('benefit', '')}. " + f"Priority {story.get('priority', 'Medium')}, complexity {story.get('complexity', '?')}, " + f"confidence {story.get('confidence', 0.0)}." + + (f" Acceptance criteria: {criteria}" if criteria else "") + ), + source_id=alias, + file_path=file_path, + ) + for rel in story_files: + if rel in first_chunk_of_file: + add_relationship( + entity_name, + rel, + description=f"User story '{title}' was derived from {rel}", + keywords="derived_from, user_story", + weight=1.0, + source_id=alias, + file_path=rel, + ) + story_count += 1 + + custom_kg = { + "chunks": chunks, + "entities": list(entities.values()), + "relationships": list(relationships.values()), + } + stats = { + "files": len(first_chunk_of_file), + "chunks": len(chunks), + "entities": len(entities), + "relationships": len(relationships), + "external_dependencies": external_count, + "user_stories": story_count, + } + return custom_kg, stats + + +# --------------------------------------------------------------------------- +# The knowledge base +# --------------------------------------------------------------------------- + + +class KnowledgeBase: + """ + Thin, opinionated wrapper around a LightRAG instance. + + Usage (async):: + + kb = KnowledgeBase(KnowledgeBaseConfig(working_dir="data/knowledge_base")) + await kb.initialize() + await kb.index_state(workflow_state) + answer = await kb.query("Which programs update the customer master file?") + await kb.close() + + Synchronous callers can use :func:`run_sync` around a coroutine that does + all three steps, or the convenience class method :meth:`run_with`. + """ + + def __init__(self, config: Optional[KnowledgeBaseConfig] = None): + self.config = config or KnowledgeBaseConfig.from_env() + self.config.validate() + self.rag = None + self.llm_available = False + self.llm_description = "none" + self.embedding_description = "" + self.tokenizer_description = "" + self._token_scale = 1 + + # -- lifecycle -------------------------------------------------------- + + async def initialize(self) -> "KnowledgeBase": + """Create the LightRAG instance and open its storages.""" + if self.rag is not None: + return self + if not is_available(): + raise KnowledgeBaseError( + "LightRAG is not installed. Run `pip install lightrag-hku` (it is listed in requirements.txt)." + ) + + from lightrag import LightRAG + from lightrag.kg.shared_storage import initialize_pipeline_status + + if not self.config.verbose: + logging.getLogger("lightrag").setLevel(logging.WARNING) + + self._apply_backend_env() + Path(self.config.working_dir).mkdir(parents=True, exist_ok=True) + + embedding_func = self._make_embedding_func() + llm_func, llm_name, llm_kwargs = self._make_llm_func() + tokenizer = self._make_tokenizer() + + kwargs: Dict[str, Any] = dict( + working_dir=self.config.working_dir, + workspace=self.config.workspace, + kv_storage=self.config.kv_storage, + vector_storage=self.config.vector_storage, + graph_storage=self.config.graph_storage, + doc_status_storage=self.config.doc_status_storage, + embedding_func=embedding_func, + llm_model_func=llm_func, + llm_model_name=llm_name, + llm_model_kwargs=llm_kwargs, + llm_model_max_async=self.config.llm_max_async, + enable_llm_cache=self.config.enable_llm_cache, + chunk_token_size=self.config.chunk_token_size * self._token_scale, + chunk_overlap_token_size=self.config.chunk_overlap_token_size * self._token_scale, + max_graph_nodes=self.config.max_graph_nodes, + ) + if tokenizer is not None: + kwargs["tokenizer"] = tokenizer + threshold = self.config.similarity_threshold + if threshold is None and self.config.embedding_provider == "hashing": + threshold = 0.05 + if threshold is not None: + kwargs["cosine_better_than_threshold"] = threshold + kwargs["vector_db_storage_cls_kwargs"] = {"cosine_better_than_threshold": threshold} + + self.rag = LightRAG(**kwargs) + await self.rag.initialize_storages() + await initialize_pipeline_status(self.config.workspace or None) + logger.info( + "Knowledge base ready (graph=%s, vectors=%s, embeddings=%s, llm=%s, tokenizer=%s)", + self.config.graph_storage, + self.config.vector_storage, + self.embedding_description, + self.llm_description, + self.tokenizer_description, + ) + return self + + async def close(self) -> None: + """Flush and close all storages.""" + if self.rag is not None: + await self.rag.finalize_storages() + self.rag = None + + async def __aenter__(self) -> "KnowledgeBase": + return await self.initialize() + + async def __aexit__(self, *exc: Any) -> None: + await self.close() + + @classmethod + def run_with( + cls, + operation: Callable[["KnowledgeBase"], Awaitable[T]], + config: Optional[KnowledgeBaseConfig] = None, + ) -> T: + """Open a knowledge base, run ``operation`` on it, close it — from sync code.""" + + async def _job() -> T: + async with cls(config) as kb: + return await operation(kb) + + return run_sync(_job) + + # -- factories -------------------------------------------------------- + + def _apply_backend_env(self) -> None: + """Bridge archaeocode's environment names to LightRAG's.""" + if self.config.graph_storage == "Neo4JStorage": + if os.getenv("NEO4J_USER") and not os.getenv("NEO4J_USERNAME"): + os.environ["NEO4J_USERNAME"] = os.environ["NEO4J_USER"] + if self.config.workspace and not os.getenv("NEO4J_WORKSPACE"): + os.environ["NEO4J_WORKSPACE"] = self.config.workspace + for key, value in (self.config.extra.get("env") or {}).items(): + os.environ.setdefault(str(key), str(value)) + + def _make_tokenizer(self): + from lightrag.utils import Tokenizer + + choice = (self.config.tokenizer or "tiktoken").lower() + if choice == "tiktoken": + try: + import tiktoken + + encoding = tiktoken.encoding_for_model("gpt-4o-mini") + self.tokenizer_description = "tiktoken/o200k_base" + return Tokenizer(model_name="gpt-4o-mini", tokenizer=encoding) + except Exception as exc: # network blocked, encoding not cached, ... + logger.warning("tiktoken unavailable (%s); falling back to the offline byte tokenizer", exc) + self._token_scale = ByteTokenizer.SCALE + self.tokenizer_description = "bytes (offline)" + return Tokenizer(model_name="bytes", tokenizer=ByteTokenizer()) + + def _make_embedding_func(self): + from lightrag.utils import EmbeddingFunc + + provider = self.config.embedding_provider + model = self.config.embedding_model + + if provider == "sentence-transformers": + try: + from sentence_transformers import SentenceTransformer + + st_model = SentenceTransformer(model) + dim = st_model.get_sentence_embedding_dimension() + except Exception as exc: + logger.warning( + "sentence-transformers model '%s' unavailable (%s); using hashing embeddings", model, exc + ) + provider = "sentence-transformers-fallback" + else: + lock = threading.Lock() + + async def st_embed(texts: List[str], **_: Any): + def _encode(): + with lock: + return st_model.encode(list(texts), convert_to_numpy=True, normalize_embeddings=True) + + return await asyncio.to_thread(_encode) + + self.embedding_description = f"sentence-transformers/{model} ({dim}d)" + return EmbeddingFunc(embedding_dim=dim, max_token_size=8192, func=st_embed, model_name=model) + + if provider == "openai": + from functools import partial + + from lightrag.llm.openai import openai_embed + + dim = self.config.embedding_dim or (3072 if "large" in model else 1536) + self.embedding_description = f"openai/{model} ({dim}d)" + return EmbeddingFunc( + embedding_dim=dim, + max_token_size=8191, + func=partial(openai_embed.func, model=model), + model_name=model, + ) + + if provider == "ollama": + from functools import partial + + from lightrag.llm.ollama import ollama_embed + + dim = self.config.embedding_dim or 1024 + host = self.config.extra.get("ollama_host") or os.getenv("OLLAMA_HOST", "http://localhost:11434") + self.embedding_description = f"ollama/{model} ({dim}d)" + return EmbeddingFunc( + embedding_dim=dim, + max_token_size=8192, + func=partial(ollama_embed.func, embed_model=model, host=host), + model_name=model, + ) + + # "hashing" — deterministic, dependency-free (tests, offline smoke runs). + import numpy as np + + dim = self.config.embedding_dim or HASHING_EMBEDDING_DIM + + async def hash_embed(texts: List[str], **_: Any): + return np.stack([_hash_embedding(t, dim) for t in texts]) + + self.embedding_description = f"hashing ({dim}d)" + if provider == "sentence-transformers-fallback": + self.embedding_description += f", fallback because sentence-transformers/{model} is unavailable" + return EmbeddingFunc(embedding_dim=dim, max_token_size=1_000_000, func=hash_embed, model_name="hashing") + + def _make_llm_func(self): + """Pick the LLM adapter; returns (func, model_name, kwargs).""" + provider = self.config.llm_provider + if provider == "auto": + if os.getenv("ANTHROPIC_API_KEY", "").startswith("sk-ant-") and "your-api-key" not in os.getenv( + "ANTHROPIC_API_KEY", "" + ): + provider = "anthropic" + elif os.getenv("OPENAI_API_KEY", "").startswith("sk-"): + provider = "openai" + elif os.getenv("OLLAMA_HOST"): + provider = "ollama" + else: + provider = "none" + + if provider == "anthropic": + from lightrag.llm.anthropic import anthropic_complete + + model = self.config.llm_model or os.getenv("MODEL_NAME", "claude-sonnet-5") + self.llm_available = True + self.llm_description = f"anthropic/{model}" + return anthropic_complete, model, {} + + if provider == "openai": + from lightrag.llm.openai import openai_complete + + model = self.config.llm_model or os.getenv("MODEL_NAME", "gpt-4o-mini") + if not model.startswith(("gpt", "o1", "o3", "o4")): + model = "gpt-4o-mini" + self.llm_available = True + self.llm_description = f"openai/{model}" + return openai_complete, model, {} + + if provider == "ollama": + from lightrag.llm.ollama import ollama_model_complete + + model = self.config.llm_model or os.getenv("OLLAMA_MODEL", "qwen2.5-coder:7b") + host = self.config.extra.get("ollama_host") or os.getenv("OLLAMA_HOST", "http://localhost:11434") + self.llm_available = True + self.llm_description = f"ollama/{model}" + return ollama_model_complete, model, {"host": host, "options": {"num_ctx": 32768}} + + async def no_llm(prompt: str, *args: Any, **kwargs: Any) -> str: + # Returning an empty string makes LightRAG fall back to the raw + # query for keyword extraction; answer synthesis is guarded in query(). + return "" + + self.llm_available = False + self.llm_description = "none" + return no_llm, "none", {} + + # -- indexing --------------------------------------------------------- + + def _require(self): + if self.rag is None: + raise KnowledgeBaseError("KnowledgeBase.initialize() must be awaited first") + return self.rag + + async def insert_custom_kg(self, custom_kg: Dict[str, Any], batch_files: int = 200) -> None: + """Insert a custom KG, batching the chunks so huge codebases stream in.""" + rag = self._require() + chunks = custom_kg.get("chunks") or [] + entities = custom_kg.get("entities") or [] + relationships = custom_kg.get("relationships") or [] + + if len(chunks) <= batch_files * 4: + await rag.ainsert_custom_kg(custom_kg) + return + + # Group chunks by file so an entity always lands in the same batch as + # the chunk its source_id refers to. + by_file: Dict[str, List[Dict[str, Any]]] = {} + for chunk in chunks: + by_file.setdefault(chunk.get("file_path", ""), []).append(chunk) + alias_to_file = {c["source_id"]: c.get("file_path", "") for c in chunks} + files = list(by_file) + for i in range(0, len(files), batch_files): + batch_paths = set(files[i : i + batch_files]) + batch_chunks = [c for f in batch_paths for c in by_file[f]] + batch_entities = [e for e in entities if alias_to_file.get(e.get("source_id", "")) in batch_paths] + batch_rels = [r for r in relationships if alias_to_file.get(r.get("source_id", "")) in batch_paths] + await rag.ainsert_custom_kg( + {"chunks": batch_chunks, "entities": batch_entities, "relationships": batch_rels} + ) + + async def index_state(self, state: Dict[str, Any]) -> Dict[str, Any]: + """ + Index a workflow state: deterministic graph first, LLM extraction second. + + Returns statistics about what was written. + """ + rag = self._require() + custom_kg, stats = build_custom_kg(state, self.config, chars_per_token=4) + started = datetime.now() + await self.insert_custom_kg(custom_kg) + stats["extraction"] = "skipped" + + if self.config.extract_entities: + if not self.llm_available: + stats["extraction"] = "skipped (no LLM configured)" + else: + artifacts = state.get("code_artifacts") or [] + limit = self.config.max_extract_files or len(artifacts) + docs = [] + ids = [] + paths = [] + for artifact in artifacts[:limit]: + content = artifact.get("content") or "" + if not content.strip(): + continue + rel = _relative_path(artifact.get("path", ""), state.get("source_path")) + docs.append(content) + ids.append("doc-" + hashlib.md5(rel.encode("utf-8")).hexdigest()) + paths.append(rel) + if docs: + await rag.ainsert(docs, ids=ids, file_paths=paths) + stats["extraction"] = f"{len(docs)} documents" + else: + stats["extraction"] = "no documents" + + stats["seconds"] = round((datetime.now() - started).total_seconds(), 2) + stats["working_dir"] = self.config.working_dir + stats["workspace"] = self.config.workspace + stats["graph_storage"] = self.config.graph_storage + stats["vector_storage"] = self.config.vector_storage + stats["embeddings"] = self.embedding_description + stats["llm"] = self.llm_description + return stats + + async def add_documents( + self, documents: List[Dict[str, Any]], extract: Optional[bool] = None + ) -> Dict[str, Any]: + """ + Add free-form documents (design docs, tickets, runbooks, interview notes). + + With ``extract`` (default: config.extract_entities) and an LLM, the + documents go through LightRAG's entity extraction; otherwise they are + only chunked and embedded, which still makes them searchable. + """ + rag = self._require() + extract = self.config.extract_entities if extract is None else extract + texts, ids, paths = [], [], [] + for index, doc in enumerate(documents): + content = (doc.get("content") or "").strip() + if not content: + continue + path = doc.get("path") or doc.get("id") or f"document-{index}" + texts.append(content) + ids.append(doc.get("id") or "doc-" + hashlib.md5(path.encode("utf-8")).hexdigest()) + paths.append(path) + if not texts: + return {"documents": 0, "mode": "none"} + + if extract and self.llm_available: + track_id = await rag.ainsert(texts, ids=ids, file_paths=paths) + return {"documents": len(texts), "mode": "llm-extraction", "track_id": track_id} + + custom_kg = {"chunks": [], "entities": [], "relationships": []} + for text, doc_id, path in zip(texts, ids, paths): + lines = text.splitlines() + windows = _chunk_lines(lines, self.config.chunk_token_size * 4, self.config.chunk_overlap_token_size * 4) + for index, (start, end) in enumerate(windows): + custom_kg["chunks"].append( + { + "content": "\n".join(lines[start:end]), + "source_id": f"{doc_id}#chunk{index}", + "file_path": path, + "chunk_order_index": index, + } + ) + custom_kg["entities"].append( + { + "entity_name": path, + "entity_type": "document", + "description": f"Document {path}: " + " ".join(lines[:3])[:300], + "source_id": f"{doc_id}#chunk0", + "file_path": path, + } + ) + await rag.ainsert_custom_kg(custom_kg) + return {"documents": len(texts), "mode": "chunks-only", "chunks": len(custom_kg["chunks"])} + + # -- retrieval -------------------------------------------------------- + + def _query_param(self, mode: Optional[str], **overrides: Any): + from lightrag import QueryParam + + mode = mode or self.config.default_query_mode + if mode not in QUERY_MODES: + raise KnowledgeBaseError(f"Unknown query mode '{mode}'. Options: {QUERY_MODES}") + params: Dict[str, Any] = dict( + mode=mode, + top_k=self.config.top_k, + chunk_top_k=self.config.chunk_top_k, + enable_rerank=False, + ) + params.update({k: v for k, v in overrides.items() if v is not None}) + return QueryParam(**params) + + async def retrieve( + self, + question: str, + mode: Optional[str] = None, + top_k: Optional[int] = None, + chunk_top_k: Optional[int] = None, + ) -> Dict[str, Any]: + """ + Structured retrieval without answer synthesis. + + Returns LightRAG's ``query_data`` payload: matching entities, relations, + chunks and references. Works without any LLM (keywords are then derived + from the question directly). + """ + rag = self._require() + overrides: Dict[str, Any] = {"top_k": top_k, "chunk_top_k": chunk_top_k} + if not self.llm_available: + keywords = keywords_from_question(question) or [question] + overrides["ll_keywords"] = keywords + overrides["hl_keywords"] = keywords + param = self._query_param(mode, **overrides) + result = await rag.aquery_data(question, param) + data = result.get("data") or {} + return { + "question": question, + "mode": param.mode, + "entities": data.get("entities", []), + "relationships": data.get("relationships", []), + "chunks": data.get("chunks", []), + "references": data.get("references", []), + "metadata": result.get("metadata", {}), + } + + async def context(self, question: str, mode: Optional[str] = None) -> str: + """The exact context block LightRAG would hand to the LLM (prompt-ready).""" + rag = self._require() + overrides: Dict[str, Any] = {"only_need_context": True} + if not self.llm_available: + keywords = keywords_from_question(question) or [question] + overrides["ll_keywords"] = keywords + overrides["hl_keywords"] = keywords + result = await rag.aquery(question, self._query_param(mode, **overrides)) + return result if isinstance(result, str) else str(result) + + async def query( + self, + question: str, + mode: Optional[str] = None, + response_type: str = "Multiple Paragraphs", + system_prompt: Optional[str] = None, + ) -> str: + """Ask a question and get an LLM-synthesized answer with citations.""" + rag = self._require() + if not self.llm_available: + raise KnowledgeBaseError( + "No LLM configured for answer synthesis. Set ANTHROPIC_API_KEY / OPENAI_API_KEY " + "(or KB_LLM_PROVIDER=ollama), or use retrieve()/context() which need no LLM." + ) + param = self._query_param(mode, response_type=response_type) + result = await rag.aquery(question, param, system_prompt=system_prompt) + return result if isinstance(result, str) else "".join([chunk async for chunk in result]) + + # -- graph access ----------------------------------------------------- + + async def graph(self, label: str = "*", max_depth: int = 2, max_nodes: Optional[int] = None) -> Dict[str, Any]: + """Neighbourhood of an entity (or the whole graph with ``*``) as plain dicts.""" + rag = self._require() + kg = await rag.get_knowledge_graph(label, max_depth=max_depth, max_nodes=max_nodes) + return { + "label": label, + "is_truncated": bool(getattr(kg, "is_truncated", False)), + "nodes": [ + {"id": n.id, "labels": list(n.labels), **(n.properties or {})} for n in kg.nodes + ], + "edges": [ + {"source": e.source, "target": e.target, **(e.properties or {})} for e in kg.edges + ], + } + + async def labels(self) -> List[str]: + rag = self._require() + return list(await rag.get_graph_labels()) + + async def stats(self) -> Dict[str, Any]: + """Size of the knowledge base and how it is configured.""" + rag = self._require() + labels = await rag.get_graph_labels() + kg = await rag.get_knowledge_graph("*", max_depth=1, max_nodes=self.config.max_graph_nodes) + types: Dict[str, int] = {} + for node in kg.nodes: + entity_type = (node.properties or {}).get("entity_type", "UNKNOWN") + types[entity_type] = types.get(entity_type, 0) + 1 + try: + processing = await rag.get_processing_status() + except Exception: + processing = {} + return { + "entities": len(labels), + "edges_sampled": len(kg.edges), + "sample_truncated": bool(getattr(kg, "is_truncated", False)), + "entity_types": dict(sorted(types.items(), key=lambda kv: -kv[1])), + "documents": processing, + "working_dir": self.config.working_dir, + "workspace": self.config.workspace, + "graph_storage": self.config.graph_storage, + "vector_storage": self.config.vector_storage, + "kv_storage": self.config.kv_storage, + "embeddings": self.embedding_description, + "llm": self.llm_description, + "tokenizer": self.tokenizer_description, + } + + # -- wiki export ------------------------------------------------------ + + async def export_wiki(self, output_dir: str, max_nodes: Optional[int] = None) -> Dict[str, Any]: + """ + Project the knowledge graph onto an Obsidian-compatible Markdown wiki. + + One page per entity (``[[wikilinks]]`` for relations), grouped index + page, and LightRAG's own tabular export next to it. This is the + human-facing "LLM wiki" view; the graph stays the source of truth. + """ + rag = self._require() + graph = await self.graph("*", max_depth=1, max_nodes=max_nodes or self.config.max_graph_nodes) + out = Path(output_dir) + pages_dir = out / "entities" + pages_dir.mkdir(parents=True, exist_ok=True) + + nodes = {n["id"]: n for n in graph["nodes"]} + slugs: Dict[str, str] = {} + used: set = set() + for node_id in sorted(nodes): + slug = _slugify(node_id) + if slug in used: + slug = f"{slug}-{hashlib.md5(node_id.encode('utf-8')).hexdigest()[:6]}" + used.add(slug) + slugs[node_id] = slug + + neighbours: Dict[str, List[Dict[str, Any]]] = {} + for edge in graph["edges"]: + neighbours.setdefault(edge["source"], []).append({"other": edge["target"], **edge}) + neighbours.setdefault(edge["target"], []).append({"other": edge["source"], **edge}) + + by_type: Dict[str, List[str]] = {} + for node_id, node in nodes.items(): + entity_type = node.get("entity_type", "UNKNOWN") + by_type.setdefault(entity_type, []).append(node_id) + lines = [ + f"# {node_id}", + "", + f"- **Type**: {entity_type}", + f"- **Source**: `{node.get('file_path', 'n/a')}`", + "", + "## Description", + "", + str(node.get("description", "")).replace("", "\n\n"), + "", + ] + links = neighbours.get(node_id, []) + if links: + lines += ["## Relations", ""] + for link in sorted(links, key=lambda l: l["other"]): + other = link["other"] + target = slugs.get(other, _slugify(other)) + keywords = link.get("keywords", "") + description = link.get("description", "") + lines.append(f"- [[{target}|{other}]] — {keywords}: {description}") + lines.append("") + (pages_dir / f"{slugs[node_id]}.md").write_text("\n".join(lines), encoding="utf-8") + + index = [ + "# Knowledge Base", + "", + f"Generated {datetime.now().isoformat(timespec='seconds')} from LightRAG " + f"({self.config.graph_storage}, workspace '{self.config.workspace or 'default'}').", + "", + f"{len(nodes)} entities, {len(graph['edges'])} relations" + + (" (truncated view)" if graph["is_truncated"] else "") + + ".", + "", + ] + for entity_type in sorted(by_type): + index += [f"## {entity_type} ({len(by_type[entity_type])})", ""] + for node_id in sorted(by_type[entity_type]): + index.append(f"- [[entities/{slugs[node_id]}|{node_id}]]") + index.append("") + (out / "index.md").write_text("\n".join(index), encoding="utf-8") + + try: + await rag.aexport_data(str(out / "graph_export.md"), file_format="md") + except Exception as exc: # pragma: no cover - export formats vary by version + logger.warning("LightRAG tabular export failed: %s", exc) + + return { + "output_dir": str(out), + "pages": len(nodes), + "relations": len(graph["edges"]), + "truncated": graph["is_truncated"], + "index": str(out / "index.md"), + } diff --git a/src/mcp_servers/knowledge_base/knowledge_mcp_server.py b/src/mcp_servers/knowledge_base/knowledge_mcp_server.py new file mode 100644 index 0000000..329e52b --- /dev/null +++ b/src/mcp_servers/knowledge_base/knowledge_mcp_server.py @@ -0,0 +1,242 @@ +""" +Knowledge Base MCP Server +========================= + +Exposes the archaeocode knowledge base (LightRAG-backed GraphRAG) as MCP tools, +so any MCP-capable agent — Claude Desktop, Claude Code, an IDE, or the +archaeocode workflow itself — can ask questions about an analyzed legacy +system, walk its knowledge graph, add documents, and export a Markdown wiki. + +Tools: +- kb_query natural-language question → LLM-synthesized answer with citations +- kb_retrieve structured retrieval (entities, relations, chunks) without an LLM +- kb_context the prompt-ready context block for a question +- kb_graph neighbourhood of an entity (or the whole graph) +- kb_stats size and configuration of the knowledge base +- kb_add_documents ingest additional documents (design docs, tickets, runbooks) +- kb_export_wiki write an Obsidian-compatible Markdown wiki + +Run standalone (stdio transport): + python src/mcp_servers/knowledge_base/knowledge_mcp_server.py +""" + +from __future__ import annotations + +import asyncio +import json +import os +import sys +from typing import Any, Optional + +try: # mcp >= 2.0 + from mcp.server import MCPServer +except ImportError: # mcp 1.x + from mcp.server.fastmcp import FastMCP as MCPServer +from mcp import types + +# Allow running the server file directly from a source checkout +_PROJECT_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..", "..")) +if _PROJECT_ROOT not in sys.path: + sys.path.insert(0, _PROJECT_ROOT) + +from src.mcp_servers.knowledge_base.knowledge_base import ( # noqa: E402 + KnowledgeBase, + KnowledgeBaseConfig, + KnowledgeBaseError, +) + +server = MCPServer("knowledge-base") + +# One knowledge base per (working_dir, workspace, event loop). LightRAG binds +# its storages to the loop that opened them, so a different loop gets a fresh +# instance. +_instances: dict[tuple, KnowledgeBase] = {} + + +async def get_knowledge_base(working_dir: Optional[str] = None, workspace: Optional[str] = None) -> KnowledgeBase: + """Return an initialized knowledge base for the given location (cached).""" + config = KnowledgeBaseConfig.from_env(working_dir=working_dir, workspace=workspace) + key = (config.working_dir, config.workspace, id(asyncio.get_running_loop())) + kb = _instances.get(key) + if kb is None: + kb = KnowledgeBase(config) + await kb.initialize() + _instances[key] = kb + return kb + + +def _ok(payload: dict[str, Any]) -> list[types.TextContent]: + payload.setdefault("success", True) + return [types.TextContent(type="text", text=json.dumps(payload, indent=2, default=str))] + + +def _fail(exc: Exception, **context: Any) -> list[types.TextContent]: + payload = {"success": False, "error": f"{type(exc).__name__}: {exc}", **context} + return [types.TextContent(type="text", text=json.dumps(payload, indent=2, default=str))] + + +@server.tool() +async def kb_query( + question: str, + mode: Optional[str] = None, + response_type: Optional[str] = None, + working_dir: Optional[str] = None, + workspace: Optional[str] = None, +) -> list[types.TextContent]: + """ + Ask the knowledge base a question and get an LLM-synthesized answer. + + Args: + question: Natural-language question about the analyzed system + mode: Retrieval mode — local (entity-centric), global (relation/theme-centric), + hybrid, naive (chunks only) or mix (graph + chunks; default) + response_type: Free-text answer format hint, e.g. "Bullet Points", "Single Paragraph" + working_dir: Knowledge base directory (default: KB_WORKING_DIR or data/knowledge_base) + workspace: Logical partition inside shared back-ends (default: KB_WORKSPACE) + + Returns: + The answer text plus the retrieval mode used + """ + try: + kb = await get_knowledge_base(working_dir, workspace) + answer = await kb.query(question, mode=mode, response_type=response_type or "Multiple Paragraphs") + return _ok({"question": question, "mode": mode or kb.config.default_query_mode, "answer": answer}) + except Exception as exc: + return _fail(exc, question=question) + + +@server.tool() +async def kb_retrieve( + question: str, + mode: Optional[str] = None, + top_k: Optional[int] = None, + chunk_top_k: Optional[int] = None, + working_dir: Optional[str] = None, + workspace: Optional[str] = None, +) -> list[types.TextContent]: + """ + Structured retrieval without answer synthesis (no LLM required). + + Args: + question: Search query (natural language, identifier, file name, ...) + mode: local / global / hybrid / naive / mix (default from config) + top_k: Maximum entities/relations to retrieve + chunk_top_k: Maximum text chunks to retrieve + working_dir: Knowledge base directory + workspace: Logical partition inside shared back-ends + + Returns: + Matching entities, relationships, text chunks and references + """ + try: + kb = await get_knowledge_base(working_dir, workspace) + result = await kb.retrieve(question, mode=mode, top_k=top_k, chunk_top_k=chunk_top_k) + return _ok(result) + except Exception as exc: + return _fail(exc, question=question) + + +@server.tool() +async def kb_context( + question: str, + mode: Optional[str] = None, + working_dir: Optional[str] = None, + workspace: Optional[str] = None, +) -> list[types.TextContent]: + """ + Return the prompt-ready context block the knowledge base would give an LLM. + + Useful for agents that want to reason over the evidence themselves. + """ + try: + kb = await get_knowledge_base(working_dir, workspace) + context = await kb.context(question, mode=mode) + return _ok({"question": question, "mode": mode or kb.config.default_query_mode, "context": context}) + except Exception as exc: + return _fail(exc, question=question) + + +@server.tool() +async def kb_graph( + label: str = "*", + max_depth: int = 2, + max_nodes: Optional[int] = None, + working_dir: Optional[str] = None, + workspace: Optional[str] = None, +) -> list[types.TextContent]: + """ + Get the neighbourhood of an entity in the knowledge graph. + + Args: + label: Entity name (file path, "path::Function", "Story: ...") or "*" for everything + max_depth: Traversal depth from the entity + max_nodes: Cap on returned nodes + """ + try: + kb = await get_knowledge_base(working_dir, workspace) + return _ok(await kb.graph(label, max_depth=max_depth, max_nodes=max_nodes)) + except Exception as exc: + return _fail(exc, label=label) + + +@server.tool() +async def kb_stats( + working_dir: Optional[str] = None, + workspace: Optional[str] = None, +) -> list[types.TextContent]: + """Size, entity-type breakdown and configuration of the knowledge base.""" + try: + kb = await get_knowledge_base(working_dir, workspace) + return _ok(await kb.stats()) + except Exception as exc: + return _fail(exc) + + +@server.tool() +async def kb_add_documents( + documents: list[dict[str, Any]], + extract: Optional[bool] = None, + working_dir: Optional[str] = None, + workspace: Optional[str] = None, +) -> list[types.TextContent]: + """ + Add documents (design docs, tickets, runbooks, interview notes) to the knowledge base. + + Args: + documents: List of {"path": str, "content": str} objects + extract: Run LLM entity/relation extraction (needs an LLM); default from config + """ + try: + kb = await get_knowledge_base(working_dir, workspace) + return _ok(await kb.add_documents(documents, extract=extract)) + except Exception as exc: + return _fail(exc) + + +@server.tool() +async def kb_export_wiki( + output_dir: str, + max_nodes: Optional[int] = None, + working_dir: Optional[str] = None, + workspace: Optional[str] = None, +) -> list[types.TextContent]: + """ + Export the knowledge graph as an Obsidian-compatible Markdown wiki. + + Args: + output_dir: Directory to write index.md and entities/*.md into + max_nodes: Cap on exported entities + """ + try: + kb = await get_knowledge_base(working_dir, workspace) + return _ok(await kb.export_wiki(output_dir, max_nodes=max_nodes)) + except Exception as exc: + return _fail(exc, output_dir=output_dir) + + +if __name__ == "__main__": + try: + server.run() + except KnowledgeBaseError as exc: # pragma: no cover + print(f"Knowledge base error: {exc}", file=sys.stderr) + sys.exit(1) diff --git a/src/orchestration/cli.py b/src/orchestration/cli.py index 1288210..77ee547 100644 --- a/src/orchestration/cli.py +++ b/src/orchestration/cli.py @@ -7,6 +7,7 @@ Usage: archaeo --source ./sample_data --source-lang cobol --target-lang java archaeo --source ./legacy_code --source-lang smalltalk --target-lang kotlin + archaeo --source ./legacy_code --source-lang cobol --knowledge-base --kb-wiki ./wiki """ import argparse @@ -79,6 +80,41 @@ def parse_args(): help="Generate workflow graph visualization" ) + kb = parser.add_argument_group( + "knowledge base", + "Persist the run into a queryable knowledge base (LightRAG GraphRAG). " + "Query it afterwards with `archaeo-kb` or the knowledge-base MCP server." + ) + kb.add_argument( + "--knowledge-base", + "--kb", + action="store_true", + help="Index files, entities, dependencies and user stories into the knowledge base" + ) + kb.add_argument( + "--kb-dir", + help="Knowledge base directory (default: $KB_WORKING_DIR or data/knowledge_base)" + ) + kb.add_argument( + "--kb-workspace", + help="Workspace name inside shared back-ends such as Neo4j (default: $KB_WORKSPACE)" + ) + kb.add_argument( + "--kb-extract", + action="store_true", + help="Also run LLM entity/relation extraction over the sources (needs an API key, costs tokens)" + ) + kb.add_argument( + "--kb-wiki", + metavar="DIR", + help="Export the knowledge graph as an Obsidian-compatible Markdown wiki into DIR" + ) + kb.add_argument( + "--kb-embeddings", + choices=["sentence-transformers", "openai", "ollama", "hashing"], + help="Embedding provider (default: $KB_EMBEDDING_PROVIDER or sentence-transformers)" + ) + return parser.parse_args() @@ -114,6 +150,10 @@ def save_report(state: dict, report_path: str): report["user_stories"] = state.get("user_stories", []) report["statistics"]["user_stories_generated"] = len(state.get("user_stories", [])) + # Knowledge base (only present when --knowledge-base was used) + if state.get("knowledge_base"): + report["knowledge_base"] = state["knowledge_base"] + # Save to file with open(report_path, 'w') as f: json.dump(report, f, indent=2) @@ -166,6 +206,15 @@ def print_summary(state: dict): if len(user_stories) > 3: print(f" ... and {len(user_stories) - 3} more") + kb = state.get("knowledge_base") + if kb: + print(f"\n🧠 Knowledge Base: {kb.get('working_dir')}") + print(f" Entities: {kb.get('entities')} | Relations: {kb.get('relationships')} | Chunks: {kb.get('chunks')}") + print(f" Storage: {kb.get('graph_storage')} / {kb.get('vector_storage')} | LLM extraction: {kb.get('extraction')}") + if kb.get("wiki"): + print(f" Wiki: {kb['wiki'].get('output_dir')} ({kb['wiki'].get('pages')} pages)") + print(f" Query it: archaeo-kb --dir {kb.get('working_dir')} query \"\"") + print("\n" + "="*70 + "\n") @@ -186,6 +235,18 @@ def main(): "verbose": args.verbose } + if args.knowledge_base: + kb_config = {"enabled": True, "extract_entities": bool(args.kb_extract)} + if args.kb_dir: + kb_config["working_dir"] = args.kb_dir + if args.kb_workspace: + kb_config["workspace"] = args.kb_workspace + if args.kb_wiki: + kb_config["wiki_dir"] = args.kb_wiki + if args.kb_embeddings: + kb_config["embedding_provider"] = args.kb_embeddings + config["knowledge_base"] = kb_config + try: # Create graph print("🔧 Initializing LangGraph workflow...\n") diff --git a/src/orchestration/graph.py b/src/orchestration/graph.py index 22ac358..4f574d6 100644 --- a/src/orchestration/graph.py +++ b/src/orchestration/graph.py @@ -15,6 +15,7 @@ from src.orchestration.nodes.ast_node import ASTAnalysisNode from src.orchestration.nodes.dependency_node import DependencyMappingNode from src.orchestration.nodes.user_story_node import UserStoryExtractionNode +from src.orchestration.nodes.knowledge_node import KnowledgeBaseNode def _setup_langsmith_tracing(): @@ -63,7 +64,11 @@ class ReverseMigrationGraph: Main orchestrator for the reverse engineering workflow. Graph structure: - START -> Discovery -> AST Analysis -> Dependency Mapping -> User Story Extraction -> END + START -> Discovery -> AST Analysis -> Dependency Mapping -> User Story Extraction + -> [Knowledge Base] -> END + + The knowledge-base step is opt-in (``config["knowledge_base"]["enabled"]``) + and persists the run's findings into a LightRAG-backed GraphRAG store. """ def __init__(self, config: Dict[str, Any]): @@ -81,6 +86,7 @@ def _build_graph(self): ast_node = ASTAnalysisNode(self.config) dependency_node = DependencyMappingNode(self.config) user_story_node = UserStoryExtractionNode(self.config) + knowledge_node = KnowledgeBaseNode(self.config) # Create state graph workflow = StateGraph(MigrationState) @@ -96,7 +102,13 @@ def _build_graph(self): workflow.add_edge("discovery", "ast_analysis") workflow.add_edge("ast_analysis", "dependency_mapping") workflow.add_edge("dependency_mapping", "user_story_extraction") - workflow.add_edge("user_story_extraction", END) + + if knowledge_node.enabled: + workflow.add_node("knowledge_base", knowledge_node) + workflow.add_edge("user_story_extraction", "knowledge_base") + workflow.add_edge("knowledge_base", END) + else: + workflow.add_edge("user_story_extraction", END) # Add conditional edges based on state # workflow.add_conditional_edges( diff --git a/src/orchestration/graph_mcp.py b/src/orchestration/graph_mcp.py index 9debb13..16bb880 100644 --- a/src/orchestration/graph_mcp.py +++ b/src/orchestration/graph_mcp.py @@ -34,6 +34,7 @@ from .nodes.ast_node_mcp import ast_analysis_node_mcp from .nodes.dependency_node_mcp import dependency_mapping_node_mcp from .nodes.user_story_node import user_story_extraction_node +from .nodes.knowledge_node import KnowledgeBaseNode from .utils.mcp_client import configure_mcp_manager, initialize_mcp_servers, shutdown_mcp_servers from .utils.tracing import is_tracing_enabled, create_run_metadata @@ -147,7 +148,14 @@ def _build_graph(self): workflow.add_edge("discovery", "ast_analysis") workflow.add_edge("ast_analysis", "dependency_mapping") workflow.add_edge("dependency_mapping", "user_story_extraction") - workflow.add_edge("user_story_extraction", END) + + knowledge_node = KnowledgeBaseNode(self.config) + if knowledge_node.enabled: + workflow.add_node("knowledge_base", knowledge_node) + workflow.add_edge("user_story_extraction", "knowledge_base") + workflow.add_edge("knowledge_base", END) + else: + workflow.add_edge("user_story_extraction", END) # Compile graph self.graph = workflow.compile(checkpointer=self.checkpointer) diff --git a/src/orchestration/kb_cli.py b/src/orchestration/kb_cli.py new file mode 100644 index 0000000..86db435 --- /dev/null +++ b/src/orchestration/kb_cli.py @@ -0,0 +1,179 @@ +""" +archaeo-kb — query the knowledge base +===================================== + +Companion CLI to ``archaeo``. After a run with ``--knowledge-base`` the +knowledge base lives on disk (or in Neo4j/Postgres/...); this tool asks it +questions, inspects the graph, adds documents, and exports a wiki. + +Usage: + archaeo-kb query "Which programs update the customer master file?" + archaeo-kb retrieve "PAYMENT" --mode local --json + archaeo-kb graph "CUSTMGMT.cob" --depth 2 + archaeo-kb stats + archaeo-kb add docs/runbook.md docs/interview-notes.md + archaeo-kb wiki ./wiki +""" + +from __future__ import annotations + +import argparse +import json +import sys +from pathlib import Path +from typing import Any, Dict + +from src.mcp_servers.knowledge_base.knowledge_base import ( + QUERY_MODES, + KnowledgeBase, + KnowledgeBaseConfig, + KnowledgeBaseError, +) + + +def build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(prog="archaeo-kb", description="Query the archaeocode knowledge base") + parser.add_argument("--dir", "-d", help="Knowledge base directory (default: $KB_WORKING_DIR or data/knowledge_base)") + parser.add_argument("--workspace", "-w", help="Workspace inside shared back-ends (default: $KB_WORKSPACE)") + parser.add_argument("--embeddings", choices=["sentence-transformers", "openai", "ollama", "hashing"], + help="Embedding provider (must match the one used for indexing)") + parser.add_argument("--llm", choices=["auto", "anthropic", "openai", "ollama", "none"], help="LLM provider") + parser.add_argument("--json", action="store_true", help="Print raw JSON") + parser.add_argument("--verbose", "-v", action="store_true") + + sub = parser.add_subparsers(dest="command", required=True) + + q = sub.add_parser("query", help="Ask a question (needs an LLM)") + q.add_argument("question") + q.add_argument("--mode", choices=QUERY_MODES) + q.add_argument("--response-type", default="Multiple Paragraphs") + + r = sub.add_parser("retrieve", help="Structured retrieval without an LLM") + r.add_argument("question") + r.add_argument("--mode", choices=QUERY_MODES) + r.add_argument("--top-k", type=int) + r.add_argument("--chunk-top-k", type=int) + + c = sub.add_parser("context", help="Print the prompt-ready context for a question") + c.add_argument("question") + c.add_argument("--mode", choices=QUERY_MODES) + + g = sub.add_parser("graph", help="Show the neighbourhood of an entity") + g.add_argument("label", nargs="?", default="*") + g.add_argument("--depth", type=int, default=2) + g.add_argument("--max-nodes", type=int) + + sub.add_parser("stats", help="Knowledge base statistics") + + a = sub.add_parser("add", help="Add documents (markdown, text, ...)") + a.add_argument("files", nargs="+") + a.add_argument("--extract", action="store_true", help="Run LLM entity extraction on the documents") + + w = sub.add_parser("wiki", help="Export an Obsidian-compatible Markdown wiki") + w.add_argument("output_dir") + w.add_argument("--max-nodes", type=int) + + return parser + + +def _config_from_args(args: argparse.Namespace) -> KnowledgeBaseConfig: + return KnowledgeBaseConfig.from_env( + working_dir=args.dir, + workspace=args.workspace, + embedding_provider=args.embeddings, + llm_provider=args.llm, + verbose=True if args.verbose else None, + ) + + +def _print_retrieval(result: Dict[str, Any]) -> None: + print(f"\n🔎 {result['question']} (mode: {result['mode']})") + entities = result.get("entities", []) + if entities: + print(f"\nEntities ({len(entities)}):") + for entity in entities[:20]: + print(f" - {entity.get('entity_name')} [{entity.get('entity_type')}]: {str(entity.get('description', ''))[:120]}") + relations = result.get("relationships", []) + if relations: + print(f"\nRelations ({len(relations)}):") + for rel in relations[:20]: + print(f" - {rel.get('src_id')} —{rel.get('keywords', '')}→ {rel.get('tgt_id')}") + chunks = result.get("chunks", []) + if chunks: + print(f"\nChunks ({len(chunks)}):") + for chunk in chunks[:10]: + first_line = str(chunk.get("content", "")).strip().splitlines()[:1] + print(f" - {chunk.get('file_path')}: {first_line[0] if first_line else ''}") + if not (entities or relations or chunks): + print(" (nothing found)") + + +def main(argv=None) -> int: + args = build_parser().parse_args(argv) + try: + config = _config_from_args(args) + except KnowledgeBaseError as exc: + print(f"❌ {exc}", file=sys.stderr) + return 2 + + async def run(kb: KnowledgeBase) -> Any: + if args.command == "query": + return {"question": args.question, "answer": await kb.query(args.question, mode=args.mode, response_type=args.response_type)} + if args.command == "retrieve": + return await kb.retrieve(args.question, mode=args.mode, top_k=args.top_k, chunk_top_k=args.chunk_top_k) + if args.command == "context": + return {"question": args.question, "context": await kb.context(args.question, mode=args.mode)} + if args.command == "graph": + return await kb.graph(args.label, max_depth=args.depth, max_nodes=args.max_nodes) + if args.command == "stats": + return await kb.stats() + if args.command == "add": + documents = [] + for file_name in args.files: + path = Path(file_name) + documents.append({"path": str(path), "content": path.read_text(encoding="utf-8", errors="ignore")}) + return await kb.add_documents(documents, extract=args.extract or None) + if args.command == "wiki": + return await kb.export_wiki(args.output_dir, max_nodes=args.max_nodes) + raise KnowledgeBaseError(f"Unknown command {args.command}") + + try: + result = KnowledgeBase.run_with(run, config) + except KnowledgeBaseError as exc: + print(f"❌ {exc}", file=sys.stderr) + return 1 + except KeyboardInterrupt: + return 130 + + if args.json: + print(json.dumps(result, indent=2, default=str)) + return 0 + + if args.command == "query": + print(f"\n❓ {result['question']}\n") + print(result["answer"]) + elif args.command == "retrieve": + _print_retrieval(result) + elif args.command == "context": + print(result["context"]) + elif args.command == "graph": + print(f"\n🕸️ {result['label']}: {len(result['nodes'])} nodes, {len(result['edges'])} edges" + + (" (truncated)" if result.get("is_truncated") else "")) + for node in result["nodes"][:50]: + print(f" • {node['id']} [{node.get('entity_type', '?')}]") + for edge in result["edges"][:50]: + print(f" {edge['source']} —{edge.get('keywords', '')}→ {edge['target']}") + elif args.command == "stats": + print("\n📚 Knowledge base") + for key, value in result.items(): + print(f" {key:16} {value}") + elif args.command == "add": + print(f"\n✓ Added {result.get('documents', 0)} document(s) ({result.get('mode')})") + elif args.command == "wiki": + print(f"\n✓ Wiki written to {result['output_dir']}: {result['pages']} pages, {result['relations']} relations") + print(f" Open {result['index']} (Obsidian-compatible)") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/src/orchestration/nodes/ast_node.py b/src/orchestration/nodes/ast_node.py index 26e1600..652ae3b 100644 --- a/src/orchestration/nodes/ast_node.py +++ b/src/orchestration/nodes/ast_node.py @@ -463,6 +463,8 @@ def __call__(self, state: MigrationState) -> MigrationState: # Extract entities entities = self.extract_entities(ast, artifact["language"]) for entity_type, entity_list in entities.items(): + for entity in entity_list: + entity.setdefault("file_path", artifact["path"]) parsed_entities[entity_type].extend(entity_list) # Calculate complexity diff --git a/src/orchestration/nodes/knowledge_node.py b/src/orchestration/nodes/knowledge_node.py new file mode 100644 index 0000000..e85b769 --- /dev/null +++ b/src/orchestration/nodes/knowledge_node.py @@ -0,0 +1,89 @@ +""" +Knowledge Base Node +=================== + +Final workflow step: persist everything the run has learned (files, AST +entities, dependency graph, user stories) into the knowledge base so it can be +queried later — by people, by the ``archaeo-kb`` CLI, or by agents through the +knowledge-base MCP server. + +The node is opt-in (``config["knowledge_base"]["enabled"]``). Deterministic +indexing needs no LLM; LLM-based entity extraction is a separate switch +(``extract_entities``) because it costs tokens proportional to the codebase. +""" + +from __future__ import annotations + +from typing import Any, Dict, Optional + +from src.orchestration.state.graph_state import MigrationState, AnalysisPhase +from src.mcp_servers.knowledge_base.knowledge_base import ( + KnowledgeBase, + KnowledgeBaseConfig, + KnowledgeBaseError, + is_available, +) + + +class KnowledgeBaseNode: + """Index the workflow state into the LightRAG-backed knowledge base.""" + + def __init__(self, config: Optional[Dict[str, Any]] = None): + self.config = config or {} + self.kb_settings: Dict[str, Any] = dict(self.config.get("knowledge_base") or {}) + self.enabled = bool(self.kb_settings.pop("enabled", False)) + self.wiki_dir: Optional[str] = self.kb_settings.pop("wiki_dir", None) + self.kb_settings.setdefault("verbose", bool(self.config.get("verbose", False))) + self.kb_config: Optional[KnowledgeBaseConfig] = None + if self.enabled: + self.kb_config = KnowledgeBaseConfig.from_env(**self.kb_settings) + + def __call__(self, state: MigrationState) -> MigrationState: + if not self.enabled: + return state + + print("\n🧠 Building knowledge base...") + + if not is_available(): + message = "Knowledge base requested but LightRAG is not installed (pip install lightrag-hku)" + print(f" ❌ {message}") + state["errors"].append(message) + state["phase"] = AnalysisPhase.COMPLETE + return state + + async def _job(kb: KnowledgeBase) -> Dict[str, Any]: + print(f" Storage: {kb.config.graph_storage} / {kb.config.vector_storage}") + print(f" Embeddings: {kb.embedding_description}") + print(f" LLM: {kb.llm_description}") + if kb.config.extract_entities and not kb.llm_available: + print(" ⚠️ extract_entities requested but no LLM configured — deterministic indexing only") + stats = await kb.index_state(state) + if self.wiki_dir: + stats["wiki"] = await kb.export_wiki(self.wiki_dir) + return stats + + try: + stats = KnowledgeBase.run_with(_job, self.kb_config) + state["knowledge_base"] = stats + print( + f" ✓ Indexed {stats['files']} files, {stats['chunks']} chunks, " + f"{stats['entities']} entities, {stats['relationships']} relations " + f"(LLM extraction: {stats['extraction']})" + ) + if "wiki" in stats: + print(f" ✓ Wiki exported to {stats['wiki']['output_dir']} ({stats['wiki']['pages']} pages)") + print(f" 📚 Knowledge base: {stats['working_dir']}") + except KnowledgeBaseError as exc: + state["errors"].append(f"Knowledge base failed: {exc}") + print(f" ❌ {exc}") + except Exception as exc: # keep the analysis results even if indexing dies + state["errors"].append(f"Knowledge base failed: {type(exc).__name__}: {exc}") + print(f" ❌ {type(exc).__name__}: {exc}") + + state["phase"] = AnalysisPhase.COMPLETE + return state + + +def knowledge_base_node(state: MigrationState, config: Optional[Dict[str, Any]] = None) -> MigrationState: + """Function wrapper around :class:`KnowledgeBaseNode`.""" + return KnowledgeBaseNode(config)(state) diff --git a/src/orchestration/state/graph_state.py b/src/orchestration/state/graph_state.py index f773e71..5c7bbd1 100644 --- a/src/orchestration/state/graph_state.py +++ b/src/orchestration/state/graph_state.py @@ -102,6 +102,9 @@ class MigrationState(TypedDict): # User stories user_stories: List[Dict[str, Any]] + # Knowledge base (LightRAG-backed GraphRAG): indexing statistics + location + knowledge_base: Dict[str, Any] + # Human-in-the-loop pending_approvals: List[Dict[str, Any]] human_feedback: List[Dict[str, Any]] @@ -175,6 +178,7 @@ def create_initial_state( generated_artifacts=[], migration_plan=None, user_stories=[], + knowledge_base={}, pending_approvals=[], human_feedback=[], errors=[], diff --git a/src/orchestration/utils/mcp_client.py b/src/orchestration/utils/mcp_client.py index 1068d28..6395745 100644 --- a/src/orchestration/utils/mcp_client.py +++ b/src/orchestration/utils/mcp_client.py @@ -143,6 +143,22 @@ async def _invoke_tool_direct(self, tool_name: str, arguments: Dict[str, Any]) - "get_graph_metrics": get_graph_metrics, "clear_graph": clear_graph } + elif self.config.name == "knowledge-base": + from src.mcp_servers.knowledge_base.knowledge_mcp_server import ( + kb_query, kb_retrieve, kb_context, kb_graph, kb_stats, + kb_add_documents, kb_export_wiki + ) + + tool_map = { + "kb_query": kb_query, + "kb_retrieve": kb_retrieve, + "kb_context": kb_context, + "kb_graph": kb_graph, + "kb_stats": kb_stats, + "kb_add_documents": kb_add_documents, + "kb_export_wiki": kb_export_wiki + } + else: raise ValueError(f"Unknown MCP server: {self.config.name}") diff --git a/tests/test_knowledge_base.py b/tests/test_knowledge_base.py new file mode 100644 index 0000000..17e75eb --- /dev/null +++ b/tests/test_knowledge_base.py @@ -0,0 +1,338 @@ +""" +Knowledge Base Tests +==================== + +Exercises the LightRAG-backed knowledge base end to end without any network +access or API key: hashing embeddings, the offline byte tokenizer, file-based +storages, and no LLM (retrieval-only paths). +""" + +import json +import os +import sys + +import pytest + +project_root = os.path.join(os.path.dirname(__file__), "..") +if project_root not in sys.path: + sys.path.insert(0, project_root) + +from src.orchestration.state.graph_state import create_initial_state # noqa: E402 +from src.orchestration.nodes.discovery_node import CodeDiscoveryNode # noqa: E402 +from src.mcp_servers.knowledge_base.knowledge_base import ( # noqa: E402 + KnowledgeBase, + KnowledgeBaseConfig, + KnowledgeBaseError, + build_custom_kg, + is_available, + keywords_from_question, +) + +requires_lightrag = pytest.mark.skipif(not is_available(), reason="lightrag-hku not installed") + +COBOL_SAMPLE = os.path.join(project_root, "sample_data", "cobol") + + +@pytest.fixture(autouse=True) +def _no_llm_keys(monkeypatch): + """Make sure tests never pick up a real API key from the developer's .env.""" + for key in ("ANTHROPIC_API_KEY", "OPENAI_API_KEY", "OLLAMA_HOST"): + monkeypatch.delenv(key, raising=False) + for key in list(os.environ): + if key.startswith("KB_"): + monkeypatch.delenv(key, raising=False) + + +def offline_config(tmp_path, **overrides) -> KnowledgeBaseConfig: + settings = dict( + working_dir=str(tmp_path / "kb"), + embedding_provider="hashing", + tokenizer="bytes", + llm_provider="none", + chunk_token_size=300, + chunk_overlap_token_size=30, + ) + settings.update(overrides) + return KnowledgeBaseConfig(**settings) + + +def sample_state(): + """Discovery over the COBOL sample plus hand-made downstream results.""" + state = create_initial_state("cobol", "java", COBOL_SAMPLE) + state = CodeDiscoveryNode({})(state) + assert len(state["code_artifacts"]) >= 2 + paths = sorted(a["path"] for a in state["code_artifacts"]) + custmgmt = next(p for p in paths if p.endswith("CUSTMGMT.cob")) + payment = next(p for p in paths if p.endswith("PAYMENT.cob")) + + state["parsed_entities"] = { + "classes": [], + "functions": [ + {"name": "PROCESS-PAYMENT", "type": "function", "line_start": 5, "line_end": 40, "file_path": payment}, + ], + "methods": [], + "imports": [], + "variables": [], + } + state["dependency_graph"] = [ + {"source": payment, "target": custmgmt, "relationship_type": "calls", "metadata": {"language": "cobol"}}, + {"source": payment, "target": "SQLCA", "relationship_type": "imports", "metadata": {"language": "cobol"}}, + {"source": custmgmt, "target": payment, "relationship_type": "calls", "metadata": {"language": "cobol"}}, + ] + state["circular_dependencies"] = [[payment, custmgmt, payment]] + state["user_stories"] = [ + { + "title": "Record a customer payment", + "user_role": "accounts clerk", + "capability": "post an incoming payment against a customer account", + "benefit": "outstanding balances stay accurate", + "acceptance_criteria": ["Given a valid customer, when a payment is posted, then the balance is reduced"], + "code_files": [payment], + "complexity": 5, + "priority": "High", + "confidence": 0.8, + } + ] + return state + + +# --------------------------------------------------------------------------- +# Pure functions (no LightRAG needed) +# --------------------------------------------------------------------------- + + +def test_build_custom_kg_from_state(): + state = sample_state() + kg, stats = build_custom_kg(state, KnowledgeBaseConfig(chunk_token_size=300, chunk_overlap_token_size=30)) + + names = {e["entity_name"] for e in kg["entities"]} + assert {"PAYMENT.cob", "CUSTMGMT.cob", "PAYMENT.cob::PROCESS-PAYMENT", "SQLCA", "Story: Record a customer payment"} <= names + + types = {e["entity_name"]: e["entity_type"] for e in kg["entities"]} + assert types["SQLCA"] == "external_dependency" + assert types["Story: Record a customer payment"] == "user_story" + assert types["PAYMENT.cob::PROCESS-PAYMENT"] == "function" + + rels = {(r["src_id"], r["tgt_id"]): r for r in kg["relationships"]} + assert ("PAYMENT.cob", "PAYMENT.cob::PROCESS-PAYMENT") in rels + assert ("PAYMENT.cob", "SQLCA") in rels + # The two-way call between the files collapses into one undirected edge flagged as circular. + circular = [r for r in kg["relationships"] if "circular-dependency" in r["keywords"]] + assert len(circular) == 1 and "calls" in circular[0]["keywords"] + # Story is linked back to its source file. + assert ("Story: Record a customer payment", "PAYMENT.cob") in rels + + # Every entity/relationship source_id refers to a chunk alias in the same payload. + aliases = {c["source_id"] for c in kg["chunks"]} + assert all(e["source_id"] in aliases for e in kg["entities"]) + assert all(r["source_id"] in aliases for r in kg["relationships"]) + # Files are chunked into several windows at this small chunk size. + assert stats["chunks"] > stats["files"] + assert stats == {**stats, "files": 2, "user_stories": 1, "external_dependencies": 1} + + +def test_config_from_dict_and_env(monkeypatch): + monkeypatch.setenv("KB_GRAPH_STORAGE", "Neo4JStorage") + monkeypatch.setenv("KB_EXTRACT_ENTITIES", "true") + monkeypatch.setenv("KB_CHUNK_TOKEN_SIZE", "800") + config = KnowledgeBaseConfig.from_env(working_dir="/tmp/x", embedding_provider="hashing") + assert config.graph_storage == "Neo4JStorage" + assert config.extract_entities is True + assert config.chunk_token_size == 800 + assert config.working_dir == "/tmp/x" + config.validate() + + with pytest.raises(KnowledgeBaseError): + KnowledgeBaseConfig(vector_storage="PineconeStorage").validate() + + loaded = KnowledgeBaseConfig.from_dict({"working_dir": "kb", "unknown_key": 1, "extra": {"env": {"A": "1"}}}) + assert loaded.extra == {"unknown_key": 1, "env": {"A": "1"}} + + +def test_keywords_from_question(): + assert keywords_from_question("Which programs update the CUSTOMER-MASTER file?") == [ + "programs", "update", "CUSTOMER-MASTER", "file" + ] + + +# --------------------------------------------------------------------------- +# LightRAG-backed behaviour (offline: hashing embeddings, byte tokenizer, no LLM) +# --------------------------------------------------------------------------- + + +@requires_lightrag +def test_index_retrieve_graph_and_wiki(tmp_path): + state = sample_state() + config = offline_config(tmp_path) + + async def job(kb: KnowledgeBase): + assert kb.llm_available is False + stats = await kb.index_state(state) + naive = await kb.retrieve("customer payment balance", mode="naive") + local = await kb.retrieve("PROCESS-PAYMENT", mode="local") + hybrid = await kb.retrieve("payment", mode="hybrid") + context = await kb.context("SQLCA", mode="mix") + graph = await kb.graph("SQLCA", max_depth=1) + summary = await kb.stats() + wiki = await kb.export_wiki(str(tmp_path / "wiki")) + with pytest.raises(KnowledgeBaseError): + await kb.query("needs an llm") + return stats, naive, local, hybrid, context, graph, summary, wiki + + stats, naive, local, hybrid, context, graph, summary, wiki = KnowledgeBase.run_with(job, config) + + assert stats["files"] == 2 and stats["entities"] >= 5 and stats["extraction"] == "skipped" + assert stats["embeddings"].startswith("hashing") + + assert naive["chunks"], "vector search over chunks returned nothing" + assert {c["file_path"] for c in naive["chunks"]} <= {"PAYMENT.cob", "CUSTMGMT.cob"} + + local_names = [e["entity_name"] for e in local["entities"]] + assert "PAYMENT.cob::PROCESS-PAYMENT" in local_names + assert local["relationships"], "graph neighbourhood of the function is empty" + + assert any(e["entity_name"] == "PAYMENT.cob" for e in hybrid["entities"]) + assert "SQLCA" in context + + assert {n["id"] for n in graph["nodes"]} == {"SQLCA", "PAYMENT.cob"} + assert graph["edges"][0]["keywords"] == "imports" + + assert summary["entities"] >= 5 + assert summary["entity_types"]["file"] == 2 + assert summary["entity_types"]["user_story"] == 1 + + assert wiki["pages"] == summary["entities"] + index = (tmp_path / "wiki" / "index.md").read_text(encoding="utf-8") + assert "## user_story (1)" in index and "[[entities/" in index + story_page = next((tmp_path / "wiki" / "entities").glob("Story-*.md")).read_text(encoding="utf-8") + assert "[[PAYMENT.cob|PAYMENT.cob]]" in story_page + + # Reopening the same directory sees the persisted graph (second event loop). + async def reopen(kb: KnowledgeBase): + return await kb.labels() + + assert "SQLCA" in KnowledgeBase.run_with(reopen, config) + + +@requires_lightrag +def test_add_documents_without_llm(tmp_path): + config = offline_config(tmp_path) + + async def job(kb: KnowledgeBase): + added = await kb.add_documents( + [{"path": "docs/runbook.md", "content": "Nightly batch runs PAYMENT after CUSTMGMT.\nRestart the job via OPC."}] + ) + hits = await kb.retrieve("nightly batch restart", mode="naive") + return added, hits + + added, hits = KnowledgeBase.run_with(job, config) + assert added["mode"] == "chunks-only" and added["documents"] == 1 + assert hits["chunks"] and hits["chunks"][0]["file_path"] == "runbook.md" + + +@requires_lightrag +def test_mcp_server_tools(tmp_path): + from src.mcp_servers.knowledge_base import knowledge_mcp_server as mcp + + state = sample_state() + config = offline_config(tmp_path) + KnowledgeBase.run_with(lambda kb: kb.index_state(state), config) + + async def call_tools(): + # The MCP server builds its config from KB_* variables + explicit args. + os.environ["KB_EMBEDDING_PROVIDER"] = "hashing" + os.environ["KB_TOKENIZER"] = "bytes" + os.environ["KB_LLM_PROVIDER"] = "none" + try: + stats = json.loads((await mcp.kb_stats(working_dir=config.working_dir))[0].text) + retrieved = json.loads( + (await mcp.kb_retrieve("PROCESS-PAYMENT", mode="local", working_dir=config.working_dir))[0].text + ) + graph = json.loads((await mcp.kb_graph("PAYMENT.cob", max_depth=1, working_dir=config.working_dir))[0].text) + query = json.loads((await mcp.kb_query("anything", working_dir=config.working_dir))[0].text) + wiki = json.loads( + (await mcp.kb_export_wiki(str(tmp_path / "mcp-wiki"), working_dir=config.working_dir))[0].text + ) + for kb in list(mcp._instances.values()): + await kb.close() + mcp._instances.clear() + finally: + for key in ("KB_EMBEDDING_PROVIDER", "KB_TOKENIZER", "KB_LLM_PROVIDER"): + os.environ.pop(key, None) + return stats, retrieved, graph, query, wiki + + import asyncio + + stats, retrieved, graph, query, wiki = asyncio.run(call_tools()) + assert stats["success"] and stats["entities"] >= 5 + assert retrieved["success"] and any(e["entity_name"] == "PAYMENT.cob::PROCESS-PAYMENT" for e in retrieved["entities"]) + assert graph["success"] and len(graph["nodes"]) >= 3 + assert query["success"] is False and "No LLM configured" in query["error"] + assert wiki["success"] and wiki["pages"] == stats["entities"] + + +@requires_lightrag +def test_workflow_with_knowledge_base_node(tmp_path): + """Full LangGraph run on the Python sample with the knowledge-base step enabled.""" + from src.orchestration.graph import create_graph + + config = { + "checkpoints": {"enabled": False}, + "hitl": {"enabled": False}, + "knowledge_base": { + "enabled": True, + "working_dir": str(tmp_path / "kb"), + "wiki_dir": str(tmp_path / "wiki"), + "embedding_provider": "hashing", + "tokenizer": "bytes", + "llm_provider": "none", + }, + } + graph = create_graph(config) + assert "knowledge_base" in graph.graph.get_graph().nodes + + result = graph.run( + source_language="python", + target_language="java", + source_path=os.path.join(project_root, "sample_data", "python"), + ) + assert not result["errors"], result["errors"] + kb = result["knowledge_base"] + assert kb["files"] == result["total_files"] + assert kb["entities"] > kb["files"], "AST entities should have been indexed alongside files" + assert kb["wiki"]["pages"] == kb["entities"] + assert (tmp_path / "wiki" / "index.md").exists() + + # Every AST entity now carries the file it came from (needed for the graph). + for kind in ("classes", "functions"): + for entity in result["parsed_entities"].get(kind, []): + assert entity["file_path"].endswith("sample.py") + + # Without the flag the node is not part of the graph. + plain = create_graph({"checkpoints": {"enabled": False}, "hitl": {"enabled": False}}) + assert "knowledge_base" not in plain.graph.get_graph().nodes + + +@requires_lightrag +def test_kb_cli_retrieve_and_stats(tmp_path, capsys): + from src.orchestration.kb_cli import main as kb_main + + state = sample_state() + config = offline_config(tmp_path) + KnowledgeBase.run_with(lambda kb: kb.index_state(state), config) + + common = ["--dir", config.working_dir, "--embeddings", "hashing", "--llm", "none"] + os.environ["KB_TOKENIZER"] = "bytes" + try: + assert kb_main(common + ["stats"]) == 0 + out = capsys.readouterr().out + assert "Knowledge base" in out and "hashing" in out + + assert kb_main(common + ["--json", "retrieve", "PROCESS-PAYMENT", "--mode", "local"]) == 0 + payload = json.loads(capsys.readouterr().out) + assert any(e["entity_name"] == "PAYMENT.cob::PROCESS-PAYMENT" for e in payload["entities"]) + + assert kb_main(common + ["query", "needs an llm"]) == 1 + assert "No LLM configured" in capsys.readouterr().err + finally: + os.environ.pop("KB_TOKENIZER", None) From b476f1baeb32156fe4a0aaaebdb4cf41441478b8 Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 17 Sep 2026 10:31:40 +0000 Subject: [PATCH 2/2] Rewrite README, quickstart and knowledge-base docs in the archaeology voice Docs-only: same commands, tables and links, told as a dig. The README is a field guide, the quickstart a first dig, the knowledge-base doc the site archive; the changelog and MCP usage guide get matching intros. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_011hwnef5pyjkdpQ6deZJYDD --- QUICKSTART.md | 61 +++++++++++------- README.md | 70 +++++++++++--------- docs/CHANGELOG.md | 4 +- docs/KNOWLEDGE_MANAGEMENT.md | 122 +++++++++++++++++++++-------------- docs/MCP_SERVERS_USAGE.md | 6 +- 5 files changed, 156 insertions(+), 107 deletions(-) diff --git a/QUICKSTART.md b/QUICKSTART.md index b8fffa9..f2bcd5d 100644 --- a/QUICKSTART.md +++ b/QUICKSTART.md @@ -1,46 +1,47 @@ -# Quick Start Guide +# Quick Start: Your First Dig -Get from clone to your first analysis in about five minutes. +From clone to your first excavated user story in about five minutes. No prior fieldwork required. -## 1. Install +## 1. Pack the tools ```bash git clone https://github.com/osick/archaeocode.git cd archaeocode pip install -r requirements.txt +pip install -e . # optional: puts `archaeo` and `archaeo-kb` on your PATH ``` -## 2. Configure (optional, for user stories) +## 2. Hire an interpreter (optional, for user stories and questions) ```bash cp .env.example .env # edit .env and set ANTHROPIC_API_KEY (or OPENAI_API_KEY) ``` -Without an API key the workflow still runs discovery, AST analysis, and dependency mapping — only user-story extraction is skipped. +Without an API key the expedition still runs discovery, AST analysis, dependency mapping, and knowledge-base indexing. Only the two steps that need someone to *read* the code for meaning are skipped: user-story extraction and natural-language answers. -## 3. Run +## 3. Break ground ```bash -# Analyze the bundled COBOL samples +# The bundled COBOL site python archaeo --source sample_data/cobol --source-lang cobol --target-lang java # Java, with a custom report path and verbose output python archaeo --source sample_data/java --source-lang java --target-lang python \ --report report.json --verbose -# Any codebase of yours +# Your own site python archaeo --source /path/to/legacy/code --source-lang cobol --target-lang java ``` Supported source languages: `cobol`, `smalltalk`, `java`, `python`, `javascript`, `fortran`, `pascal`. -## 4. What happens +## 4. What happens down there -1. **Discovery** — finds and catalogs all source files -2. **AST analysis** — parses files with tree-sitter (entities, complexity) -3. **Dependency mapping** — builds the dependency graph, detects cycles, computes layers -4. **User-story extraction** — an LLM turns each file's business logic into user stories (needs API key) +1. **Discovery** — surveys the site and catalogs every source file +2. **AST analysis** — brushes off the dirt with tree-sitter (entities, complexity) +3. **Dependency mapping** — reconstructs who called whom: graph, cycles, layers +4. **User-story extraction** — an LLM asks "what was this *for*?" and writes user stories (needs API key) 5. **Report** — a JSON report is written (default: `workflow_report.json`) View the report: @@ -49,7 +50,26 @@ View the report: python -m json.tool workflow_report.json ``` -## 5. Optional extras +## 5. Keep the finds (the knowledge base) + +Add one flag and the run is filed into a knowledge graph instead of evaporating into a JSON file: + +```bash +python archaeo --source sample_data/cobol --source-lang cobol --knowledge-base --kb-wiki ./wiki +``` + +Then interrogate the archive: + +```bash +archaeo-kb retrieve "PAYMENT" --mode local # entities, relations, code — no API key needed +archaeo-kb graph "PAYMENT.cob" # what does this program touch? +archaeo-kb query "Which programs touch the customer file?" # a real answer — needs an LLM key +archaeo-kb stats # how big is the museum? +``` + +Open `./wiki/index.md` in Obsidian (or any Markdown viewer) for one page per file, function, dependency, and user story, cross-linked. Add `--kb-extract` to also let the LLM dig out business vocabulary. By default everything lives in `data/knowledge_base/` as plain files; for a site the size of a core-banking system, point it at Neo4j / Postgres / Qdrant via `KB_*` variables. The full story: [docs/KNOWLEDGE_MANAGEMENT.md](docs/KNOWLEDGE_MANAGEMENT.md). + +## 6. Optional extras **LangSmith tracing** — full observability of every run: @@ -60,15 +80,6 @@ echo "LANGSMITH_PROJECT=archaeocode" >> .env ``` Then view traces at https://smith.langchain.com. Details: [docs/LANGSMITH_SETUP.md](docs/LANGSMITH_SETUP.md) -**Knowledge base** — keep what a run learned and ask questions about it: - -```bash -python archaeo --source sample_data/cobol --source-lang cobol --knowledge-base --kb-wiki ./wiki -archaeo-kb retrieve "PAYMENT" --mode local # works without an API key -archaeo-kb query "Which programs touch the customer file?" # needs an LLM key -``` -Local files by default; Neo4j/Postgres/Qdrant for large systems. Details: [docs/KNOWLEDGE_MANAGEMENT.md](docs/KNOWLEDGE_MANAGEMENT.md) - **Smalltalk grammars** — required before analyzing Smalltalk code: ```bash @@ -80,10 +91,12 @@ python scripts/build_smalltalk_grammar.py - **Import errors**: run from the project root; the CLI adds the right paths itself. - **Missing dependencies**: `pip install -r requirements.txt` - **No user stories generated**: check that your API key is set in `.env` and valid. +- **`archaeo-kb query` says no LLM configured**: same fix, or use `retrieve` / `graph` / `stats`, which never need one. +- **Knowledge base can't download models** (air-gapped site): set `KB_TOKENIZER=bytes` and `KB_EMBEDDING_PROVIDER=hashing`, or point `KB_EMBEDDING_PROVIDER=ollama` at a local Ollama. ## Next steps - [README](README.md) — feature overview and Python API - [examples/user_story_extraction/basic_usage.py](examples/user_story_extraction/basic_usage.py) — programmatic usage - [examples/knowledge_base/query_knowledge_base.py](examples/knowledge_base/query_knowledge_base.py) — build and query the knowledge base -- [docs/ROADMAP.md](docs/ROADMAP.md) — where the project is heading +- [docs/ROADMAP.md](docs/ROADMAP.md) — where the expedition is heading diff --git a/README.md b/README.md index 380556e..3f93185 100644 --- a/README.md +++ b/README.md @@ -4,30 +4,37 @@ [![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](LICENSE) [![Python 3.10+](https://img.shields.io/badge/python-3.10+-blue.svg)](https://www.python.org/downloads/) -**Automated software archaeology.** Point archaeocode at a legacy codebase — get back business-readable user stories, a dependency map, and structural analysis: the raw material for a migration backlog, excavated from code whose authors are long gone. +**Automated software archaeology.** Somewhere in your company there is a COBOL program older than most of your colleagues. Its author retired, its documentation is a rumour, and it still moves money every night. archaeocode is the expedition you send in: point it at a legacy codebase and get back business-readable user stories, a dependency map, structural analysis — and, since v0.4, a **knowledge base that remembers everything the dig uncovered**, so you never have to excavate the same site twice. -Most reverse-engineering tools stop at syntax: they parse code and draw diagrams. This project goes one step further and recovers the *business intent* hidden in legacy code. An AI agent workflow (LangGraph) orchestrates static-analysis tools (exposed as MCP servers) to turn COBOL, Smalltalk, Fortran, Pascal, or Java code into artifacts that stakeholders can actually read — the raw material for a migration backlog. +Most reverse-engineering tools stop at syntax: they parse code and draw diagrams. That is like cataloguing potsherds without asking what the pot was for. archaeocode goes one step further and recovers the *business intent* buried in legacy code. An AI agent workflow (LangGraph) orchestrates static-analysis tools (exposed as MCP servers) to turn COBOL, Smalltalk, Fortran, Pascal, or Java into artifacts that stakeholders can actually read: the raw material for a migration backlog. ## What makes it different -- **User stories from code** — the unique feature: an LLM analyzes each source file and produces user stories with roles, capabilities, benefits, acceptance criteria, priority, and confidence scores. Legacy knowledge becomes a product backlog. -- **Legacy-first language support** — COBOL, Fortran, Pascal, and Smalltalk alongside Java, Python, JavaScript, and TypeScript. See the [language matrix](#language-support) for per-language capabilities. -- **Knowledge management built in** — every run can be persisted into a [GraphRAG knowledge base](docs/KNOWLEDGE_MANAGEMENT.md) (powered by [LightRAG](https://github.com/HKUDS/LightRAG)): files, entities, dependency edges and user stories become a queryable knowledge graph, exportable as an Obsidian-style wiki. Ask it questions with `archaeo-kb` or from any agent via MCP. Runs on local files or on Neo4j / Postgres / Qdrant for large estates. -- **MCP architecture** — analysis tools (AST parsing, dependency graphs, RAG, knowledge base) are [Model Context Protocol](https://modelcontextprotocol.io/) servers, so any MCP-capable agent can reuse them independently of this workflow. -- **Observable by design** — every workflow run can be traced in [LangSmith](https://docs.smith.langchain.com) (token costs, latency, state snapshots). +- **User stories from code** — the signature find. An LLM reads each source file and writes user stories with roles, capabilities, benefits, acceptance criteria, priority, and confidence scores. Forty-year-old business rules become a product backlog your product owner can argue about. +- **A knowledge base that doesn't forget** — every run can be persisted into a [GraphRAG knowledge base](docs/KNOWLEDGE_MANAGEMENT.md) (powered by [LightRAG](https://github.com/HKUDS/LightRAG)). Files, entities, dependency edges, and user stories become a knowledge graph you can *interrogate*: "which programs touch the customer master file?", "what breaks if we retire this copybook?". Ask from the terminal with `archaeo-kb`, from any agent via MCP, or browse the exported Obsidian-style wiki like a museum catalogue. Runs on plain files on your laptop, or on Neo4j / Postgres / Qdrant when the site is the size of a bank. +- **Legacy-first language support** — COBOL, Fortran, Pascal, and Smalltalk get first-class treatment, next to Java, Python, JavaScript, and TypeScript. See the [language matrix](#language-support). +- **MCP architecture** — every tool (AST parsing, dependency graphs, RAG, knowledge base) is a [Model Context Protocol](https://modelcontextprotocol.io/) server. Claude Desktop, Claude Code, or your own agent can borrow the trowel without adopting the whole expedition. +- **Observable by design** — every run can be traced in [LangSmith](https://docs.smith.langchain.com): token costs, latency, state snapshots. No mystery boxes on a mystery-box project. -## How it works +## How the dig works ``` LangGraph Orchestration │ [Discovery] ──► [AST Analysis] ──► [Dependency Mapping] ──► [User Stories] ──► [Knowledge Base] │ │ │ │ │ -file catalog tree-sitter graph + cycle Claude / GPT LightRAG GraphRAG - parsing detection (optional, --knowledge-base) + survey the tree-sitter graph + cycle Claude / GPT LightRAG GraphRAG + site parsing detection (what was it for?) (the site archive, + --knowledge-base) ``` -Each analysis step is an MCP server under `src/mcp_servers/` (static analysis, graph DB, RAG pipeline, knowledge base); the LangGraph workflow under `src/orchestration/` wires them together with checkpointing and state management. +1. **Discovery** surveys the site: every source file is catalogued with language and size. +2. **AST analysis** brushes off the dirt: tree-sitter turns files into classes, functions, methods, imports, and complexity scores. +3. **Dependency mapping** reconstructs who talked to whom: import/call edges, circular dependencies, and layering. +4. **User story extraction** asks the important question, *what was this for?*, and writes the answer as user stories. +5. **Knowledge base** (opt-in) files every find in a knowledge graph so the next question doesn't need a new dig. + +Each step is an MCP server under `src/mcp_servers/`; the LangGraph workflow under `src/orchestration/` wires them together with checkpointing and state management. ## Quick start @@ -36,14 +43,14 @@ git clone https://github.com/osick/archaeocode.git cd archaeocode pip install -r requirements.txt -# Optional: install the `archaeo` command onto your PATH +# Optional: install the `archaeo` and `archaeo-kb` commands onto your PATH pip install -e . # Optional but recommended: enable user-story extraction cp .env.example .env # add your ANTHROPIC_API_KEY (or OPENAI_API_KEY) ``` -Run the workflow against the bundled samples: +Break ground on the bundled samples: ```bash # COBOL analysis @@ -57,21 +64,24 @@ python archaeo --source sample_data/java --source-lang java --target-lang python You get a console summary plus a JSON report: file catalog, language breakdown, dependency edges/cycles/layers, and (with an API key) generated user stories. -Keep the knowledge instead of just reporting it: +### Keep the finds: the knowledge base + +A JSON report is a photo of the excavation. The knowledge base is the museum. ```bash # persist the run into the knowledge base (+ an Obsidian-compatible wiki) python archaeo --source sample_data/cobol --source-lang cobol --knowledge-base --kb-wiki ./wiki -# then ask it questions (query needs an LLM key; retrieve/graph/stats do not) -archaeo-kb query "Which programs touch the customer master file?" -archaeo-kb retrieve "PAYMENT" --mode local -archaeo-kb graph "PAYMENT.cob" +# then interrogate it +archaeo-kb query "Which programs touch the customer master file?" # needs an LLM key +archaeo-kb retrieve "PAYMENT" --mode local # no LLM needed +archaeo-kb graph "PAYMENT.cob" # walk the neighbourhood +archaeo-kb add docs/runbook.md # add non-code lore ``` -See [docs/KNOWLEDGE_MANAGEMENT.md](docs/KNOWLEDGE_MANAGEMENT.md) for the design decision (GraphRAG via LightRAG), query modes, and how to scale it to Neo4j / Postgres / Qdrant. +Indexing needs no LLM at all: files, entities, and dependency edges come straight from static analysis, so the graph is exact and free. Add `--kb-extract` and the LLM also mines the code for *business* concepts (customers, tariffs, invoice runs) and files them next to the technical ones. Details, query modes, and the scaling ladder: [docs/KNOWLEDGE_MANAGEMENT.md](docs/KNOWLEDGE_MANAGEMENT.md). -Or use it from Python: +### From Python ```python from src.orchestration.graph import create_graph @@ -88,7 +98,7 @@ for story in result["user_stories"]: print("-", story["title"]) ``` -A complete runnable example is in [`examples/user_story_extraction/basic_usage.py`](examples/user_story_extraction/basic_usage.py). +Runnable examples: [`examples/user_story_extraction/basic_usage.py`](examples/user_story_extraction/basic_usage.py) and [`examples/knowledge_base/query_knowledge_base.py`](examples/knowledge_base/query_knowledge_base.py). ## Feature status @@ -105,15 +115,15 @@ A complete runnable example is in [`examples/user_story_extraction/basic_usage.p | Knowledge base: LLM extraction of business entities | ✅ working (`--kb-extract`, needs API key) | | Knowledge base: Obsidian-compatible wiki export | ✅ working | | RAG semantic code search (MCP server) | 🚧 mock; superseded by the knowledge base | -| Neo4j-backed dependency graphs | 🚧 in-memory fallback works; live Neo4j optional (the knowledge base can use Neo4j today) | +| Neo4j-backed dependency graphs | 🚧 in-memory fallback works; the knowledge base can already use Neo4j | | Code generation to target language | 🎯 planned | | HP NonStop COBOL extensions (TMF, Pathway) | 🎯 planned | ## Requirements - Python 3.10+ -- An Anthropic or OpenAI API key for user-story extraction (everything else runs without one) -- Optional: Neo4j 5.x / PostgreSQL+pgvector / Qdrant to scale the knowledge base beyond local files, LangSmith account for tracing +- An Anthropic or OpenAI API key for user-story extraction and knowledge-base questions (everything else runs without one) +- Optional: Neo4j 5.x / PostgreSQL+pgvector / Qdrant to scale the knowledge base beyond local files; a LangSmith account for tracing ## Language support @@ -134,12 +144,12 @@ The AST MCP server additionally parses C, C++, C#, Go, Rust, Ruby, PHP, and Bash ## Project structure ``` -├── archaeo # CLI entry point (archaeo-kb: knowledge base CLI) +├── archaeo # CLI entry point (`archaeo-kb` after pip install -e .) ├── src/ │ ├── orchestration/ # LangGraph workflow │ │ ├── graph.py # Direct workflow (in-process nodes) │ │ ├── graph_mcp.py # MCP-backed workflow -│ │ ├── kb_cli.py # archaeo-kb: query the knowledge base +│ │ ├── kb_cli.py # archaeo-kb: interrogate the knowledge base │ │ ├── nodes/ # Discovery, AST, dependency, user-story, knowledge-base nodes │ │ ├── state/ # Workflow state schema │ │ └── utils/ # MCP client, LangSmith tracing @@ -163,21 +173,21 @@ pip install pytest pytest-asyncio pytest ``` -The suite runs in [GitHub Actions](.github/workflows/ci.yml) on Python 3.10–3.12. Smalltalk tests are skipped automatically unless the grammar has been built. +The suite runs in [GitHub Actions](.github/workflows/ci.yml) on Python 3.10–3.12. Smalltalk tests are skipped automatically unless the grammar has been built. The knowledge-base tests run fully offline (no API key, no downloads), so CI never has to phone home. ## Documentation -- [Quick Start](QUICKSTART.md) +- [Quick Start](QUICKSTART.md) — your first dig in five minutes +- [Knowledge management](docs/KNOWLEDGE_MANAGEMENT.md) — the site archive: why GraphRAG, why LightRAG, how to query and scale it - [Architecture wireframe](docs/LANGGRAPH_WIREFRAME.md) - [MCP server architecture](docs/MCP_ARCHITECTURE.md) and [usage](docs/MCP_SERVERS_USAGE.md) - [LangSmith setup](docs/LANGSMITH_SETUP.md) - [User story extraction](docs/USER_STORIES.md) -- [Knowledge management (GraphRAG knowledge base)](docs/KNOWLEDGE_MANAGEMENT.md) - [Roadmap](docs/ROADMAP.md) · [Changelog](docs/CHANGELOG.md) ## Contributing -Issues and pull requests are welcome. Please run `pytest` before submitting, and open an issue first for larger changes. +Issues and pull requests are welcome. Please run `pytest` before submitting, and open an issue first for larger changes. Bring your own legacy horror stories; the sample data can always use another stratum. ## License diff --git a/docs/CHANGELOG.md b/docs/CHANGELOG.md index df3bb71..52b5d14 100644 --- a/docs/CHANGELOG.md +++ b/docs/CHANGELOG.md @@ -1,6 +1,8 @@ # Changelog -## [0.4.0] - 2026-09-07 +## [0.4.0] - 2026-09-07 — "The Site Archive" + +Until now every dig ended with a JSON report and amnesia. This release gives archaeocode a memory. ### Added - **Knowledge management**: a persistent, queryable knowledge base built on [LightRAG](https://github.com/HKUDS/LightRAG) (GraphRAG: knowledge graph + vector retrieval). See [docs/KNOWLEDGE_MANAGEMENT.md](KNOWLEDGE_MANAGEMENT.md) for the decision record. diff --git a/docs/KNOWLEDGE_MANAGEMENT.md b/docs/KNOWLEDGE_MANAGEMENT.md index e20670f..20fd639 100644 --- a/docs/KNOWLEDGE_MANAGEMENT.md +++ b/docs/KNOWLEDGE_MANAGEMENT.md @@ -1,43 +1,51 @@ -# Knowledge Management +# Knowledge Management: The Site Archive -archaeocode excavates a lot of knowledge from a legacy system — file catalogs, AST -entities, dependency graphs, user stories — but until v0.4 every run threw it -away after writing a JSON report. The **knowledge base** keeps it: a persistent, -queryable store that people query with `archaeo-kb` and that agents query -through the `knowledge-base` MCP server. +Every archaeological dig produces two things: a pile of finds, and a *record* of +where each find was, what it sat next to, and what the team thinks it was for. +Lose the record and you have a box of interesting rocks. -This document records what was chosen, why, and how to use and scale it. +Until v0.4, archaeocode was excellent at the first part and terrible at the +second. Each run surveyed the site, catalogued every file, reconstructed the +dependency graph, and wrote user stories — then packed it all into a JSON +report and forgot about it. Ask a follow-up question next week and you were +back in the trench with a shovel. -## Decision: GraphRAG (LightRAG) with a wiki projection +The **knowledge base** is the site archive. Every run can now be filed into a +persistent knowledge graph that people interrogate with `archaeo-kb` and that +agents interrogate through the `knowledge-base` MCP server. This document +records what was chosen, why, and how to use and scale it. -Three families were on the table: +## The decision: GraphRAG, with a wiki as the museum gift shop + +Three ways of remembering were on the table. | Approach | What it is | Fit for legacy-system knowledge | |---|---|---| -| **LLM wiki** (Karpathy-style: an LLM maintains a Markdown wiki) | Human-readable pages, incrementally curated by the model | Great *output* format, but no retrieval layer: does not scale to millions of LOC, every update is an LLM call, and questions like "what depends on copybook X" cannot be answered without re-reading pages | -| **Vanilla RAG** (chunks + vector search) | Embed text chunks, retrieve top-k | Good for "find code that looks like this", poor for multi-hop and global questions ("which business processes touch the tariff tables?"), and it ignores the exact dependency graph static analysis already produced | -| **GraphRAG** (knowledge graph + vectors) | Entities and relations in a graph, chunks and entities embedded, retrieval walks both | Matches the data: a codebase *is* a graph. Multi-hop questions, global summaries, and exact structural facts coexist; the graph can be seeded deterministically from the AST/dependency analysis so the LLM is only spent on business concepts | +| **LLM wiki** (Karpathy-style: an LLM maintains a Markdown wiki) | Human-readable pages, incrementally curated by the model | Lovely to read, terrible to *ask*. No retrieval layer, so it does not scale to millions of LOC; every update is an LLM call; and "what depends on copybook X?" means re-reading pages until you find it. A wiki is a gift shop, not a catalogue. | +| **Vanilla RAG** (chunks + vector search) | Embed text chunks, retrieve top-k | Fine for "find code that looks like this". Poor for multi-hop and global questions ("which business processes touch the tariff tables?"), and it throws away the one thing static analysis handed us for free: the exact dependency graph. Like searching a dig by the colour of the dirt. | +| **GraphRAG** (knowledge graph + vectors) | Entities and relations in a graph, chunks and entities embedded, retrieval walks both | Matches the material. A codebase *is* a graph. Multi-hop questions, global summaries, and exact structural facts coexist, and the graph can be seeded deterministically from the AST/dependency analysis so the LLM is only spent on the interesting part: business meaning. | **Choice: GraphRAG, with the wiki as a projection of the graph** rather than as -the system of record. You get the human-facing "LLM wiki" experience (one page -per entity, `[[wikilinks]]`, works in Obsidian) without giving up a queryable, -scalable backbone. +the system of record. You still get the human-facing "LLM wiki" experience (one +page per entity, `[[wikilinks]]`, opens in Obsidian) without giving up a +queryable, scalable backbone. The graph is the archive; the wiki is printed +from it whenever you like. -### Why LightRAG and not the alternatives +### Why LightRAG and not one of the other shovels -The requirement was "reuse a fitting solution that integrates best and scales by -all dimensions". Candidates were assessed on integration fit (Python, LangGraph, +The brief was: *reuse a fitting solution that integrates best and scales by all +dimensions.* Candidates were assessed on integration fit (Python, LangGraph, MCP, Neo4j already in the stack, Anthropic first), incremental updates, storage back-ends, cost control, query capabilities, licence and maturity. | Candidate | Verdict | |---|---| -| **[LightRAG](https://github.com/HKUDS/LightRAG)** (HKUDS, MIT) | **Chosen.** Pure-Python library, async, pluggable LLM/embedding functions (Anthropic, OpenAI, Ollama, HuggingFace adapters built in), storage back-ends for NetworkX/Neo4j/Memgraph/Postgres+AGE (graph), nano-vectordb/FAISS/pgvector/Qdrant/Milvus (vectors), JSON/Postgres/Redis/Mongo (KV). Incremental inserts with document tracking and deletion. **`insert_custom_kg`** lets us write the exact dependency graph without any LLM. Query modes local/global/hybrid/naive/mix, structured retrieval without LLM (`query_data`), graph export. Workspaces isolate multiple systems in one back-end. | -| Microsoft GraphRAG | Batch indexer built around community summaries: expensive to (re)index, weak incremental story, parquet/LanceDB-centric, hard to reuse the deterministic graph. Best for static document corpora, not for a codebase that keeps changing. | +| **[LightRAG](https://github.com/HKUDS/LightRAG)** (HKUDS, MIT) | **Chosen.** Pure-Python library, async, pluggable LLM/embedding functions (Anthropic, OpenAI, Ollama, HuggingFace adapters built in), storage back-ends for NetworkX/Neo4j/Memgraph/Postgres+AGE (graph), nano-vectordb/FAISS/pgvector/Qdrant/Milvus (vectors), JSON/Postgres/Redis/Mongo (KV). Incremental inserts with document tracking and deletion. **`insert_custom_kg`** lets us write the exact dependency graph without any LLM — the feature that decided it. Query modes local/global/hybrid/naive/mix, structured retrieval without LLM (`query_data`), graph export. Workspaces isolate multiple systems in one back-end. | +| Microsoft GraphRAG | A batch indexer built around community summaries: expensive to (re)index, weak incremental story, parquet/LanceDB-centric, and no clean way to hand it a graph we already computed. Superb for a static document corpus, wrong shape for a codebase that keeps changing. | | Cognee | Solid "AI memory" engine with graph + vector and an MCP server, but more opinionated pipelines ("cognify"), telemetry on by default, and less direct control over inserting a prebuilt graph. Close second. | -| Graphiti (Zep) | Temporal knowledge graph aimed at agent memory; Neo4j/FalkorDB mandatory, extraction is LLM-first. Excellent for evolving facts, over-engineered for code corpora. | -| neo4j-graphraph (official Neo4j) | Clean and Neo4j-native, but Neo4j is mandatory (no zero-infrastructure mode for laptops/CI) and there is no global/community retrieval. | -| LlamaIndex PropertyGraphIndex | Capable, but pulls a second orchestration framework next to LangGraph. | +| Graphiti (Zep) | Temporal knowledge graph aimed at agent memory; Neo4j/FalkorDB mandatory, extraction is LLM-first. Excellent for facts that change over time, over-engineered for code corpora. | +| neo4j-graphrag (official Neo4j) | Clean and Neo4j-native, but Neo4j is mandatory (no zero-infrastructure mode for laptops or CI) and there is no global/community retrieval. | +| LlamaIndex PropertyGraphIndex | Capable, but drags a second orchestration framework in next to LangGraph. One expedition leader is enough. | ## Architecture @@ -59,6 +67,17 @@ archaeo run Obsidian wiki (index.md + entities/*.md) ◄──────────┘ ``` +Two streams feed the archive: + +- **The deterministic stream** is the field record. Files, classes, functions, + methods, import/call edges, circular dependencies, and user stories are + written straight into the graph. No LLM is involved, so it is exact, free, + and idempotent: re-run the dig and the record is upserted, not duplicated. +- **The optional LLM stream** (`--kb-extract`) is the interpretation. LightRAG's + entity/relation extraction reads the sources and files *business* concepts + (customers, invoice runs, tariff tables) next to the technical ones, in the + same graph. This is the only step that costs tokens, and it is off by default. + Code lives in: - `src/mcp_servers/knowledge_base/knowledge_base.py` — the adapter (`KnowledgeBase`, `KnowledgeBaseConfig`, `build_custom_kg`) @@ -66,7 +85,7 @@ Code lives in: - `src/orchestration/nodes/knowledge_node.py` — workflow node - `src/orchestration/kb_cli.py` — `archaeo-kb` -### What gets indexed +### What gets filed | Source in workflow state | Graph entities | Relations | Embedded chunks | |---|---|---|---| @@ -76,14 +95,14 @@ Code lives in: | `user_stories` | `user_story` entities with role, capability, benefit, criteria | `story —derived_from→ file` | the story as Markdown | | documents added later (`archaeo-kb add`, `kb_add_documents`) | `document` entities | via LLM extraction (optional) | document windows | -Everything in the table is written with `insert_custom_kg`: exact, cheap, -idempotent (re-indexing upserts). LLM extraction (`--kb-extract`) additionally -runs LightRAG's entity/relation extraction over the sources and merges the -result into the same graph — that is where business vocabulary appears. +That last row matters more than it looks. Runbooks, ticket exports, and the +notes from your interview with the one person who still remembers why the +batch job runs at 02:17 all go into the same graph as the code. The archive +holds the folklore as well as the finds. ## Usage -### Index during a run +### File a dig ```bash # zero infrastructure: NetworkX + nano-vectordb under data/knowledge_base, @@ -99,7 +118,7 @@ Flags: `--knowledge-base/--kb`, `--kb-dir`, `--kb-workspace`, `--kb-extract`, The same settings live under `knowledge_base:` in `config/langgraph_config.yaml` and as `KB_*` environment variables (`.env.example`). -### Ask questions +### Interrogate the archive ```bash archaeo-kb query "Which programs update the customer master file, and what business rules apply?" @@ -108,14 +127,16 @@ archaeo-kb retrieve "PROCESS-PAYMENT" --mode local # no LLM: entities, re archaeo-kb context "SQLCA" # prompt-ready evidence block archaeo-kb graph "PAYMENT.cob" --depth 2 archaeo-kb stats -archaeo-kb add docs/runbook.md docs/interview-notes.md # enrich with non-code knowledge -archaeo-kb wiki ./wiki # Obsidian-compatible export +archaeo-kb add docs/runbook.md docs/interview-notes.md # file the folklore +archaeo-kb wiki ./wiki # print the museum catalogue ``` -Query modes (LightRAG): `local` (entity neighbourhood), `global` (relations / -themes), `hybrid`, `naive` (chunks only), `mix` (default: graph + chunks). -`query` needs an LLM (Anthropic → OpenAI → Ollama, picked from the keys present); -`retrieve`, `context`, `graph`, `stats`, `wiki` never do. +Query modes (LightRAG): `local` (entity neighbourhood: "tell me about this +thing"), `global` (relations and themes: "tell me about this area"), `hybrid`, +`naive` (chunks only, classic RAG), `mix` (default: graph + chunks). +`query` needs an LLM (Anthropic → OpenAI → Ollama, picked from the keys +present); `retrieve`, `context`, `graph`, `stats`, `wiki` never do, so the +archive is useful even on a machine with no keys and no network. ### From Python @@ -136,47 +157,48 @@ print(KnowledgeBase.run_with(ask, config)) `kb_query`, `kb_retrieve`, `kb_context`, `kb_graph`, `kb_stats`, `kb_add_documents`, `kb_export_wiki`. Add it to Claude Desktop / Claude Code like the other servers (see `docs/MCP_SERVERS_USAGE.md`); set `KB_WORKING_DIR` or pass -`working_dir` per call. +`working_dir` per call. An agent planning a migration can now ask the archive +instead of re-reading the code, which is the whole point. -## Scaling +## Scaling: from shoebox to national museum The adapter never touches back-end specifics; scale is a configuration change. Recommended profiles (also listed in `config/mcp_servers_config.yaml`): | Profile | Graph | Vectors | KV / status | Embeddings | Typical size | |---|---|---|---|---|---| -| laptop / CI | `NetworkXStorage` | `NanoVectorDBStorage` | JSON files | sentence-transformers (local) | up to ~100k LOC, one user | +| shoebox (laptop / CI) | `NetworkXStorage` | `NanoVectorDBStorage` | JSON files | sentence-transformers (local) | up to ~100k LOC, one user | | team | `Neo4JStorage` | `PGVectorStorage` | `PGKVStorage` | OpenAI `text-embedding-3-large` | millions of LOC, several analysts, shared | -| enterprise | `Neo4JStorage` | `QdrantVectorDBStorage` / `MilvusVectorDBStorage` | `RedisKVStorage` | OpenAI / self-hosted | many systems, one `workspace` each | +| national museum | `Neo4JStorage` | `QdrantVectorDBStorage` / `MilvusVectorDBStorage` | `RedisKVStorage` | OpenAI / self-hosted | many systems, one `workspace` each | Set them via environment (`KB_GRAPH_STORAGE=Neo4JStorage`, `NEO4J_URI`, `NEO4J_USERNAME` (falls back to `NEO4J_USER`), `NEO4J_PASSWORD`, `POSTGRES_*`, `QDRANT_URL`, ...) or in the YAML section. Other scale levers: - **Workspaces** — `--kb-workspace billing-core` keeps several analyzed systems - apart inside one Neo4j/Postgres/Qdrant. + apart inside one Neo4j/Postgres/Qdrant. One archive per site, one building. - **Batching** — custom-KG inserts are streamed in batches of files, so a million-line codebase does not need to fit in one write. - **Cost** — deterministic indexing is free of LLM tokens. `--kb-extract` is the only expensive step; cap it with `max_extract_files`, and LightRAG caches LLM responses so re-runs do not pay twice. Local embeddings via sentence-transformers - or Ollama keep data on-premise. + or Ollama keep data on-premise, which matters when the code *is* the bank. - **Offline** — `KB_TOKENIZER=bytes` and `KB_EMBEDDING_PROVIDER=hashing` run the - whole pipeline without network access (used by the test suite). If the - sentence-transformers model cannot be downloaded the adapter falls back to - hashing embeddings with a warning. + whole pipeline without network access (this is what the test suite does). If + the sentence-transformers model cannot be downloaded, the adapter falls back + to hashing embeddings with a warning rather than dying in the trench. - **Observability** — LightRAG calls go through the same Anthropic/OpenAI SDKs; the workflow node is a normal LangGraph node, so it shows up in LangSmith. -## Limitations and next steps +## Known gaps and the next season's dig plan - AST entities are only available for tree-sitter-supported languages; COBOL, Fortran and Pascal contribute files, dependencies and stories until their grammars land (see the language matrix in the README). - The dependency mapper still returns illustrative edges for some languages; - the knowledge base faithfully indexes whatever it gets — improving the mapper - improves the graph automatically. + the knowledge base faithfully indexes whatever it gets, so improving the + mapper improves the graph automatically. - Planned: incremental re-indexing keyed by file hash (LightRAG already tracks document ids), community summaries for very large systems, and feeding - retrieved context back into the user-story prompt so stories are written with - cross-file knowledge. + retrieved context back into the user-story prompt so stories are written + with cross-file knowledge instead of one file at a time. diff --git a/docs/MCP_SERVERS_USAGE.md b/docs/MCP_SERVERS_USAGE.md index 4376075..3976b0f 100644 --- a/docs/MCP_SERVERS_USAGE.md +++ b/docs/MCP_SERVERS_USAGE.md @@ -557,8 +557,10 @@ python src/mcp_servers/graph_db/neo4j_mcp_server.py **Location:** `src/mcp_servers/knowledge_base/knowledge_mcp_server.py` -Exposes the knowledge base that `archaeo --knowledge-base` builds (see -[KNOWLEDGE_MANAGEMENT.md](KNOWLEDGE_MANAGEMENT.md)). Backed by +The site archive, opened to agents. Exposes the knowledge base that +`archaeo --knowledge-base` builds (see [KNOWLEDGE_MANAGEMENT.md](KNOWLEDGE_MANAGEMENT.md)), +so an agent planning a migration can ask what the last dig found instead of +re-reading the code. Backed by [LightRAG](https://github.com/HKUDS/LightRAG): knowledge graph + vector indexes, file-based by default, Neo4j/Postgres/Qdrant/... via `KB_*` environment variables.