sample-okf-llm-wiki

# Data Wiki: AI-powered knowledge bundles for your AWS data, served to agents over MCP **Data Wiki** turns AWS Glue databases into portable knowledge bundles and serves them to AI agents over the Model Context Protocol (MCP). It reads a Glue catalog, has an LLM author a set of markdown documents describing each dataset — its tables, joins, metrics, and known issues — keeps those documents in sync as the catalog changes, and exposes them to agents through an MCP server. The result: agents that understand *what your data means*, not just what columns exist. The bundles are Open Knowledge Format (OKF) bundles — a bundle is a directory of markdown files with YAML frontmatter. Data Wiki maps OKF's model onto AWS as `data domain → dataset (Glue database) → table`. The `okf`/`OKF_` prefix on code identifiers, resource names, and env vars refers to the format, not the product. **Note:** This is a sample implementation provided for educational and demonstration purposes. You should review, test, and customize this code for your specific requirements before deploying to any environment. --- ## Quick Links - 📚 [Using the Console](./quick-start-guide/using-the-console.md) — register a dataset, harvest, and browse a bundle - 🔌 [Connect an Agent (MCP)](./quick-start-guide/connect-an-agent.md) — mint credentials and query bundles from your agents - 🚀 [Getting Started](#getting-started) — prerequisites and one-command deploy - 🏗️ [Architecture](./docs/ARCHITECTURE.md) — how the pieces fit and the reasoning behind them - 📐 [Conventions](./docs/CONVENTIONS.md) — the contract between services (S3 layout, DynamoDB shapes, env vars) - 📘 [API Reference](./docs/API_REFERENCE.md) — external API shapes the code was written against --- ## Features - **Automated Induction** - An AI agent reads your Glue catalog, samples real data, and authors a knowledge bundle describing tables, joins, metrics, and known issues. - **Grounded in Your Docs** - Attach your own source documents (DDL, data dictionaries, runbooks, specs) so the bundle reflects your organization's knowledge, not just the raw schema. - **Self-Healing Freshness** - Glue catalog changes automatically trigger scoped re-harvests, and the semantic search index stays in sync with the bundle on every write. - **Served Over MCP** - Any MCP-capable agent can discover, read, and semantically search your bundles with `list_domains`, `list_directory`, `read_page`, `glob`, `grep`, `get_backlinks`, and `semantic_search`. - **Machine Credentials** - Mint scoped OAuth2 credentials so applications and agents can call the MCP server programmatically. - **Web Console** - Register datasets, upload context docs, run and watch harvests, browse bundles, and explore the link graph from a React UI. --- ## Architecture ### Solution Architecture <p align="center"> <img alt="Data Wiki solution architecture" src="./assets/architecture.png" width="900"> </p> **AWS Services Used:** - Amazon S3 (bundle storage) and S3 Vectors (semantic index) - AWS Glue (source catalog) and Amazon Athena (data sampling) - Amazon Bedrock (Claude for induction, Titan V2 for embeddings) - Amazon Bedrock AgentCore Runtime (harvest + consumption MCP) - AWS Lambda (Control API, reindex, incremental) - Amazon API Gateway (HTTP API) and Amazon Cognito (auth) - Amazon EventBridge + Amazon SQS (event-driven freshness) - Amazon DynamoDB (registry, freshness, credentials) - Amazon CloudFront (SPA delivery) ### Repository Layout ``` data-wiki/ ├── services/ │ ├── okf_core/ # pure-Python OKF primitives (no AWS or agent deps) │ ├── okf_aws/ # shared boto3 helpers: Titan embed, S3 Vectors, S3 keys │ ├── harvest/ # induction: deepagents agent on AgentCore │ ├── reindex/ # freshness: S3 events → S3 Vectors Lambda │ ├── incremental/ # Glue change → scoped re-harvest + nightly reconcile │ ├── control_api/ # Cognito-authed REST (API GW + Lambda) │ └── consumption_mcp/ # streamable-HTTP MCP server on AgentCore ├── infra/ │ ├── durable/ # S3 buckets, S3 Vectors index, Cognito, DynamoDB │ ├── compute/ # Lambdas, API GW, AgentCore runtimes + IAM, EventBridge/SQS, CloudFront │ └── modules/lambda/ # reusable zip-Lambda + least-privilege role ├── ui/ # React SPA (JS), shadcn/ui + Tailwind, Cognito OIDC ├── okf-mcp/ # Claude Code plugin: connect an agent to the MCP server ├── scripts/ # run_tests.sh, build_lambdas.sh, deploy.sh ├── quick-start-guide/ # task-oriented guides for using the console └── docs/ # ARCHITECTURE.md, CONVENTIONS.md, API_REFERENCE.md ``` See [`docs/ARCHITECTURE.md`](./docs/ARCHITECTURE.md) for how the pieces fit together and the reasoning behind the main decisions. --- ## Getting Started ### Prerequisites **Required Tools:** The following tools must be installed on your local machine: - [AWS CLI](https://docs.aws.amazon.com/cli/latest/userguide/getting-started-install.html) configured with [appropriate credentials](https://docs.aws.amazon.com/cli/latest/userguide/cli-configure-files.html) - [Terraform CLI](https://developer.hashicorp.com/terraform/install) - [Docker](https://docs.docker.com/engine/install/) running (with `buildx` for ARM64 image builds) - [Node.js](https://nodejs.org/en/download) and npm - [Python](https://www.python.org/downloads/) (v3.12 or later) - [jq](https://jqlang.org/download/) **AWS Setup:** - A **pre-created, versioned S3 bucket** for Terraform state (the deploy uses split state and won't create this for you). - **Amazon Bedrock model access** enabled in your region for the Claude induction model and Amazon Titan Text Embeddings V2. Follow the instructions [here](https://docs.aws.amazon.com/bedrock/latest/userguide/model-access.html); without access you'll get an `AccessDeniedException` during a harvest. - At least one **AWS Glue database** with tables you want to harvest. > **Note:** S3 Vectors and Bedrock AgentCore Runtime are used by this stack. Verify they are available in your target region before deploying. ### Installation and Deployment 1. **Clone the Repository** ```bash git clone <your-repo-url> cd sample-okf-llm-wiki ``` 2. **Authenticate to AWS** ```bash # Option I: Export AWS temporary credentials export AWS_ACCESS_KEY_ID="your_temp_access_key" export AWS_SECRET_ACCESS_KEY="your_temp_secret_key" export AWS_SESSION_TOKEN="your_temp_session_token" export AWS_DEFAULT_REGION="your_region" # Option II: Export AWS Profile export AWS_PROFILE="your_profile_name" ``` 3. **Run the deployment** ```bash ./scripts/deploy.sh # full pipeline (prompts on first run) ``` On the **first run**, you'll be prompted for: - The AWS **region** - The **Terraform state bucket** name (pre-created, versioned) - The **initial admin user** (email and name) These are saved to `scripts/.deployment.config` and reused on subsequent runs. The deploy then runs all five stages: push the container images, apply both Terraform stacks, wire the CloudFront URL into Cognito, and sync the SPA. > **Note:** The admin user is created in the Cognito User Pool from your inputs; Cognito emails a temporary password to that address. ### Accessing the Console After a successful deploy, print the deployed URLs: ```bash ./scripts/deploy.sh summary ``` Open the CloudFront URL, sign in with the admin credentials from your email, and follow [Using the Console](./quick-start-guide/using-the-console.md) to register your first dataset and run a harvest. --- ## Configuration Options ### Deploy Stages `deploy.sh` orchestrates the full pipeline, but each stage can be run on its own — useful for iterating on one layer without re-running everything: ```bash ./scripts/deploy.sh <durable|images|compute|cognito-urls|ui|summary|destroy> ``` | Stage | What it does | | -------------- | ---------------------------------------------------------------------------- | | `durable` | S3 buckets, S3 Vectors index, Cognito, DynamoDB (the long-lived stack). | | `images` | Build and push the harvest + consumption container images to ECR. | | `compute` | Lambdas, API Gateway, AgentCore runtimes, EventBridge/SQS, CloudFront. | | `cognito-urls` | Re-apply the durable stack to inject the CloudFront URL into Cognito. | | `ui` | Build the React SPA and sync it to S3 / CloudFront. | | `summary` | Print the deployed URLs. | | `destroy` | Tear everything down (prompts for confirmation). | A few things the deploy handles that would otherwise be manual: - The harvest runtime's **VPC** (2 private subnets + NAT) is auto-provisioned, since S3 Files mounts require VPC networking. Bring your own with `TF_VAR_harvest_vpc_subnet_ids` / `TF_VAR_harvest_vpc_security_group_ids`. - The **S3 Files** file system, mount target, and access point are managed by Terraform and mounted at `/mnt/data`. - The Cognito **callback/logout URLs** get the CloudFront URL through the `cognito-urls` re-apply, not the console. - Image URIs and the CloudFront URL flow between stages via `.deployment.config` and Terraform outputs. ### Connecting Agents (MCP) Once bundles are harvested, agents consume them over the MCP server. Mint a machine credential in the console's **Credentials** view, then connect: - **From Claude Code** — use the bundled [`okf-mcp`](./okf-mcp/) plugin, which handles token refresh automatically. - **From any MCP client** — exchange your credential for a bearer token at the Cognito token endpoint and call the runtime directly. See [Connect an Agent (MCP)](./quick-start-guide/connect-an-agent.md) for both paths. ### Local UI Development To run the console against your deployed backend on `localhost`: ```bash cd ui npm run dev:env # writes ui/.env.local from the compute stack outputs npm install # first run only npm run dev # http://localhost:5173 ``` Port `5173` is whitelisted in Cognito for the OIDC redirect, so local dev works against the real deployed Cognito and Control API. See [`ui/README.md`](./ui/README.md) for details. --- ## Running the Tests The suite runs entirely offline: `moto` mocks S3 and DynamoDB, and the remaining clients (s3vectors, bedrock-runtime, glue, athena, agentcore) are injected fakes. There are no live AWS calls. ```bash python3 -m venv .venv && source .venv/bin/activate pip install -e services/okf_core -e services/okf_aws # shared libs pip install -e services/harvest -e services/reindex -e services/incremental \ -e services/control_api -e services/consumption_mcp --no-deps pip install pytest "moto[s3,dynamodb]" ./scripts/run_tests.sh # unit tests + the offline end-to-end harvest test cd ui && npm ci && npm run build cd ../infra/durable && terraform init -backend=false && terraform validate cd ../compute && terraform init -backend=false && terraform validate ``` `tests/test_e2e_harvest_offline.py` drives the non-LLM half of the harvest pipeline end to end (Glue source → guard engine → link-graph impact analysis → `finalize_bundle`) against a fake F1-shaped Glue/Athena source and asserts the output is a valid OKF bundle. Actual Bedrock calls, S3 Vectors queries, and AgentCore hosting need a real account and aren't exercised here. --- ## Clean Up To tear down everything the deploy created: 1. **Export AWS credentials** (same as the deploy step above). 2. **Run the destroy stage:** ```bash ./scripts/deploy.sh destroy ``` This tears down both Terraform stacks (durable and compute). The Terraform **state bucket** you pre-created is not managed by the deploy — remove it manually if you no longer need it. > **Warning:** This permanently deletes the harvested bundles, the vector index, and all registry/credential state. The underlying Glue databases and their data are not touched. --- ## Contributing See [CONTRIBUTING](./CONTRIBUTING.md) for more information. ## License This project is licensed under the MIT License. See the [LICENSE](./LICENSE) file. ## Authors The following authors have contributed this sample within AWS and prepared it for open-source release: - [Edvin Hallvaxhiu](https://github.com/edvinhallvaxhiu) - [Aubrey Oosthuizen](https://github.com/a13zen) - [Rahul Shaurya](https://github.com/shauryarahul) - [Sindi Cali](https://github.com/scali123) - [Shukhrat Khodjaev](https://github.com/khodjaevsh) - [Ravi Kiran Ganji](https://github.com/RaviGanjiAWS) Special thanks to [Bruno Dhefto](https://github.com/dhefto) for his thorough review and suggestions that helped improve the solution.