AI Knowledge Base Pipeline

A fully serverless RAG (Retrieval-Augmented Generation) pipeline built with AWS Step Functions, Amazon Bedrock, and S3 Vectors. Upload gym documents, ask natural-language questions, get accurate answers grounded in your data.

What It Does

• Ingests PDF documents via an event-driven pipeline
• Extracts text with Amazon Textract (async, parallel)
• Enriches documents with AI-generated metadata using Nova Pro
• Indexes everything into a Bedrock Knowledge Base backed by S3 Vectors
• Exposes a conversational HTTP API (POST /ask) with multi-turn memory

Stack: Go · Terraform · AWS Step Functions · Amazon Bedrock · S3 Vectors · API Gateway

---

Architecture

PDF documents (S3)
        │
        ▼
[ EventBridge ] ← triggered by manifest.json upload
        │
        ▼
[ Step Functions Workflow ]
        │
        ├─ DistributedMap (per document, up to 10 parallel)
        │       ├─ StartTextractJob      (direct SDK call)
        │       ├─ StoreTaskTokenAndWait ← ⏸ WAIT_FOR_TASK_TOKEN
        │       └─ FetchTextractResults  (Lambda)
        │
        ├─ EnrichDocuments (Lambda → Nova Pro via Converse API)
        │       └─ writes .metadata.json files to S3
        │
        └─ StartKBIngestionJob (direct SDK call → Bedrock KB)
                └─ chunks + embeds into S3 Vectors

[ API Gateway + ask-gym Lambda ]
        │
        ▼
POST /ask → RetrieveAndGenerate → FitBot answer

Textract Async Coordination

Textract is asynchronous — it fires an SNS notification when done. Instead of polling, the pipeline uses WAIT_FOR_TASK_TOKEN:

SFN stores task token in DynamoDB (keyed by Textract JobId)
        │
        ⏸ FROZEN (zero cost, no polling)
        │
Textract → SNS → textract-callback Lambda
        │
DynamoDB.GetItem(JobId) → states.SendTaskSuccess(token)
        │
        ▶ SFN resumes

---

Project Structure

├── main.tf                          # Root Terraform — wires modules together
├── variables.tf
├── output.tf
├── modules/
│   ├── knowledge_base/              # S3 buckets, Bedrock KB, S3 Vectors, data source
│   └── pipeline/
│       ├── sfn.tf                   # Step Functions state machine + IAM
│       ├── lambdas.tf               # Lambda functions + IAM policies
│       ├── api.tf                   # API Gateway + ask-gym Lambda
│       ├── eventbridge.tf           # EventBridge rule (manifest.json trigger)
│       ├── dynamodb.tf              # Task token table
│       └── sns.tf                   # Textract completion notifications
└── lambdas/
    ├── textract-callback/           # Receives SNS, calls SendTaskSuccess
    ├── textract-results/            # Fetches OCR pages, writes raw text to S3
    ├── enrich-documents/            # Calls Nova Pro, writes .metadata.json
    ├── ask-gym/                     # FitBot HTTP handler (RetrieveAndGenerate)
    └── shared/                      # Shared types and DynamoDB token store

---

Prerequisites

• Terraform >= 1.0
• Go >= 1.21
• AWS CLI configured with appropriate permissions
• Amazon Bedrock model access enabled for:
  • amazon.titan-embed-text-v2:0 (embedding)
  • eu.amazon.nova-pro-v1:0 (enrichment)
  • eu.amazon.nova-lite-v1:0 (query answering)

---

Deploy

# 1. Build all Lambda binaries
make build

# 2. Initialize Terraform
make init

# 3. Preview changes
make plan

# 4. Deploy
make apply

Terraform will output:

• manifest_bucket_name — where to upload documents and the manifest
• ask_gym_api_url — the FitBot API endpoint
• test_command — ready-to-run command to trigger the pipeline

---

Running the Pipeline

1. Upload your PDFs to S3

Documents must follow the path convention docs/{gym_id}/filename.pdf. The gym_id segment is used as the multi-tenancy key — it's never trusted from the model.

aws s3 cp my-pricing.pdf s3://<manifest_bucket_name>/docs/my-gym/my-pricing.pdf

2. Upload a manifest to trigger the pipeline

manifest.json lists the files to process:

[
  { "filename": "docs/my-gym/pricing.pdf" },
  { "filename": "docs/my-gym/schedule.pdf" }
]

Upload it (this triggers EventBridge → Step Functions):

aws s3 cp manifest.json s3://<manifest_bucket_name>/manifest.json
# or use the ready-made output:
$(terraform output -raw test_command)

3. Wait for the pipeline to complete

Monitor the execution in the AWS Step Functions console. When the state machine completes, Bedrock KB ingestion runs in the background (typically 1–3 minutes for small datasets).

---

Querying the API

curl -X POST <ask_gym_api_url> \
  -H "Content-Type: application/json" \
  -H "X-Gym-Id: my-gym" \
  -d '{"question": "What are the membership prices?"}'

Response:

{
  "answer": "Hey! Here are the membership options...",
  "session_id": "0e0818d8-6825-43e1-99f8-619280ddffa7"
}

Multi-turn conversation

Pass the session_id from the previous response as X-Session-Id to maintain conversation context — Bedrock stores the history server-side, no storage required on your end:

curl -X POST <ask_gym_api_url> \
  -H "Content-Type: application/json" \
  -H "X-Gym-Id: my-gym" \
  -H "X-Session-Id: 0e0818d8-6825-43e1-99f8-619280ddffa7" \
  -d '{"question": "Which option did you just mention is cheaper?"}'

---

Key Design Decisions

Decision	Why
S3 Vectors instead of OpenSearch	No idle cluster, pay per query, native Bedrock KB integration
WAIT_FOR_TASK_TOKEN for Textract	Zero cost while waiting, no Lambda polling loops
gym_id from S3 path, not model	Models can hallucinate names — paths can't
Nova Pro for enrichment, Nova Lite for queries	Match model capability to latency/cost requirements
FIXED_SIZE chunking at 200 tokens	S3 Vectors has a 2MB per-document limit; smaller chunks also improve retrieval precision
Fire-and-forget on StartIngestionJob	No downstream step depends on completion; Bedrock handles it async

---

License

MIT