Skip to main content
Glama
PShinde630

nexla-document-qa-mcp

by PShinde630
README.md
# Nexla Document Q&A MCP Server

A locally runnable Model Context Protocol (MCP) server that answers natural
language questions over a small collection of research papers. Answers are
grounded in retrieved PDF passages and include validated document and page
citations.

The project was built for the Nexla Software Engineer take-home assignment. It
focuses on a clear local implementation, observable behavior, and honest
trade-offs rather than production infrastructure.

## What It Does

- Parses and indexes five PDF research papers.
- Preserves document names and page numbers throughout the pipeline.
- Creates local semantic embeddings with `BAAI/bge-small-en-v1.5`.
- Stores 191 cleaned chunks and their metadata in persistent ChromaDB.
- Retrieves relevant passages for natural-language questions.
- Generates grounded answers with Gemini.
- Handles focused, cross-document, and insufficient-evidence questions.
- Exposes the required `query_documents(question: str)` MCP tool.
- Includes a Python MCP client for end-to-end verification without requiring
  Claude Desktop, Gemini CLI, Node.js, or a paid MCP client.

## Architecture

```text
PDF files in data/
    |
    v
pypdf page extraction
    |
    v
paragraph-aware, page-bounded token chunking
    |
    v
BGE embeddings (384 dimensions)
    |
    v
persistent ChromaDB index

Question
    |
    v
BGE query embedding
    |
    v
cosine similarity retrieval
    |
    +--> relevance threshold
    +--> duplicate removal
    +--> bibliography filtering
    +--> comparison-aware document diversity
    |
    v
Gemini grounded generation
    |
    v
source ID validation and trusted metadata mapping
    |
    v
MCP structured answer with PDF names, pages, and excerpts
```

The MCP transport is intentionally thin. `server.py` delegates to
`QueryService`; retrieval, prompting, generation, and citation mapping remain
separate and independently testable.

## Technology Choices

| Component | Choice | Reason |
|---|---|---|
| Language | Python 3.11 or 3.12 | Strong PDF, ML, vector-store, and MCP ecosystem |
| MCP framework | Official Python MCP SDK with FastMCP | Small standards-compliant tool wrapper |
| PDF parser | `pypdf` | Lightweight and sufficient for the selected text PDFs |
| Embeddings | `BAAI/bge-small-en-v1.5` | Local embeddings, 512-token input limit, 384-dimensional vectors |
| Vector store | ChromaDB | Simple persistent local cosine-similarity search |
| Answer model | `gemini-3.1-flash-lite` | Fast structured output with an accessible API tier |
| Configuration | `.env` with `python-dotenv` | Keeps runtime configuration outside source code |
| Tests | `pytest` | Fast deterministic unit and integration tests |

## Selected Documents

The `data/` directory contains the five provided research papers:

```text
2020.acl-main.408.pdf
D18-1360.pdf
D19-1539.pdf
N19-1240.pdf
P19-1598.pdf
```

The current index contains:

```text
Documents: 5
Pages with extractable text: 62
Cleaned chunks: 191
Embedding dimensions: 384
```

## Setup

### 1. Prerequisites

- Python 3.11 or 3.12
- A Gemini API key
- The five PDFs in `data/`

### 2. Create a virtual environment

Windows PowerShell:

```powershell
python -m venv .venv
.\.venv\Scripts\Activate.ps1
```

macOS or Linux:

```bash
python3 -m venv .venv
source .venv/bin/activate
```

### 3. Install the project

Runtime and test dependencies:

```powershell
python -m pip install --upgrade pip
python -m pip install -e ".[dev]"
```

### 4. Configure the environment

Windows PowerShell:

```powershell
Copy-Item .env.example .env
```

macOS or Linux:

```bash
cp .env.example .env
```

Add the Gemini key to `.env`:

```dotenv
LLM_PROVIDER=gemini
GEMINI_API_KEY=replace_with_your_key
GEMINI_MODEL=gemini-3.1-flash-lite

DATA_DIR=./data
CHROMA_DIR=./.chroma
CHROMA_COLLECTION=document_qa
EMBEDDING_MODEL=BAAI/bge-small-en-v1.5
HF_HOME=./.model-cache

RETRIEVAL_CANDIDATE_K=12
RETRIEVAL_CONTEXT_K=6
RETRIEVAL_MIN_SCORE=0.60
CHUNK_SIZE_TOKENS=450
CHUNK_OVERLAP_TOKENS=60
```

## Build the Index

Place the PDFs in `data/`, then run:

```powershell
python -m document_qa.ingest
```

Ingestion parses every PDF page, creates page-bounded chunks, embeds them, and
rebuilds the managed ChromaDB collection. Re-ingestion replaces the previous
collection instead of accumulating duplicate records.

Expected summary for the current corpus:

```text
Documents: 5
Pages: 62
Chunks: 191
Embeddings: 191
Stored records: 191
Collection: document_qa
```

## Ask Questions Directly

The terminal Q&A command is useful for debugging without MCP:

```powershell
python -m document_qa.answer "What is the main purpose of the KGLM?" --offline
```

`--offline` applies only to loading the cached embedding model. Gemini answer
generation still requires internet access.

## Run the MCP Server

Start the stdio server:

```powershell
python -m document_qa.server
```

The command waits silently for MCP protocol messages. It is not an interactive
terminal prompt. Stop it with `Ctrl+C`.

The server loads the embedding model before starting the FastMCP event loop.
This is important on Windows because loading SentenceTransformer/PyTorch inside
the first asynchronous tool call caused the original MCP verification request
to block.

## Verify MCP with Python

The included Python client starts the server, initializes an MCP session, lists
the tools, calls `query_documents`, prints the structured result, and closes the
server:

```powershell
python -m document_qa.mcp_client "What is the main purpose of the KGLM?" --offline
```

The client has a configurable MCP response timeout:

```powershell
python -m document_qa.mcp_client "Your question" --offline --timeout 180
```

No Node.js, Gemini CLI, Claude Desktop, or ChatGPT connector is required for
this verification path.

## MCP Tool

### `query_documents`

Answers a question using only indexed PDF evidence.

Input:

```json
{
  "question": "Compare the KGLM and ERASER approaches."
}
```

Output:

```json
{
  "answer": "KGLM improves factual generation ... [S1]. ERASER evaluates rationales ... [S4].",
  "sources": [
    {
      "source_id": "S1",
      "document": "P19-1598.pdf",
      "page": 9,
      "excerpt": "..."
    },
    {
      "source_id": "S4",
      "document": "2020.acl-main.408.pdf",
      "page": 1,
      "excerpt": "..."
    }
  ]
}
```

Source IDs returned by Gemini are validated against the retrieved passages.
Document names and pages are mapped from trusted ChromaDB metadata, not copied
from model-generated filenames or page numbers.

## Retrieval and Grounding

### Parsing and chunking

- PDFs are extracted one page at a time.
- Paragraph boundaries are preserved where extraction allows.
- A chunk never combines content from different pages.
- Chunks contain at most 450 BGE content tokens.
- Consecutive chunks reuse up to 60 tokens for context continuity.
- Common PDF ligatures and tokenizer spacing artifacts are normalized.

### Semantic retrieval

- Every question is embedded using the same BGE model used for documents.
- ChromaDB returns cosine-similarity candidates.
- Exact duplicate chunks and duplicate normalized text are removed.
- Candidates below the configurable `0.60` relevance floor are rejected.
- Citation-dense bibliography chunks are filtered before context selection.
- Focused questions preserve global similarity order.
- Explicit comparison questions encourage evidence from multiple relevant
  documents.
- Broad collection questions inspect more candidates and allow a larger,
  document-balanced context.

Document diversity is not forced for every question. A focused KGLM question
should cite only the KGLM paper if it contains the strongest evidence.

### Answer generation

Gemini receives the question and complete retrieved passages. The prompt
requires it to:

- use only supplied evidence;
- cite supported claims using source labels;
- synthesize independently described approaches for comparison questions;
- return insufficient evidence when an essential part is unsupported; and
- avoid inventing source IDs, documents, pages, or facts.

Generation uses structured JSON, `temperature=0`, bounded HTTP timeouts, and
limited transient retries. Hosted model wording may still vary slightly.

### Insufficient evidence

Unrelated questions are rejected before Gemini when no candidate meets the
relevance threshold:

```text
The indexed documents do not contain enough information to answer this question.
```

No sources are returned in this case.

## Example Interactions

These examples were executed through the Python MCP client.

### 1. Focused document question

Question:

```text
What is the main purpose of the KGLM?
```

Answer returned:

```text
The primary goal of the Knowledge Graph Language Model (KGLM) is to enable a
neural language model to generate facts and entities from a knowledge graph.
By maintaining a local knowledge graph of entities mentioned in the context,
the model can produce factual completions and specific tokens, such as dates
and numbers, rather than generic words.
```

Sources included `P19-1598.pdf`, pages 2, 7, and 9.

### 2. Cross-document comparison

Question:

```text
Compare the KGLM paper's approach with the ERASER benchmark's approach.
```

Answer returned:

```text
The KGLM paper and the ERASER benchmark serve fundamentally different purposes
in NLP research. KGLM is a specific neural language model designed to improve
factual accuracy in text generation by accessing an external knowledge graph.
In contrast, ERASER is a standardized benchmark designed to evaluate the
interpretability of existing NLP models through rationales and supporting
evidence.
```

Sources included `P19-1598.pdf`, pages 7 and 9, and
`2020.acl-main.408.pdf`, page 1.

### 3. Insufficient evidence

Question:

```text
What medical treatment do these papers recommend for diabetes?
```

Answer:

```text
The indexed documents do not contain enough information to answer this question.
```

Sources: none.

## Testing

Run the complete deterministic suite:

```powershell
python -m pytest
```

Current result:

```text
84 passed
```

The automated suite does not require a Gemini key or network access. It covers:

- PDF discovery, normalization, metadata, and empty-text errors;
- page-bounded token chunking and deterministic chunk IDs;
- embedding validation;
- ChromaDB persistence and repeat ingestion;
- semantic retrieval, deduplication, relevance gating, reference filtering,
  comparison diversity, and broad-query balancing;
- Gemini response parsing and citation validation using fake clients;
- trusted citation mapping and insufficient-evidence behavior;
- MCP tool discovery, invocation, structured output, and safe errors; and
- Python MCP client validation and formatting.

Manual checks performed during development included focused questions,
cross-document comparisons, unrelated questions, a real Gemini connectivity
check, a complete MCP stdio call, and installation from a built wheel in a
temporary virtual environment.

## Diagnostics

Useful inspection commands:

```powershell
python -m document_qa.inspect_chunks
python -m document_qa.inspect_embeddings
python -m document_qa.inspect_store
python -m document_qa.search "Your retrieval question" --offline
python -m document_qa.check_gemini
```

The MCP client and server log each major stage to stderr without displaying API
keys, environment contents, vectors, or complete document passages.

## Limitations and Trade-offs

- Image-only or scanned PDFs require OCR. The selected corpus contains
  extractable text, so OCR was kept outside the assignment scope.
- Tables, equations, and multi-column layouts can lose visual structure during
  PDF text extraction.
- The bibliography filter is heuristic and intentionally conservative.
- A relevance threshold calibrated on this small corpus may need adjustment for
  a different document collection.
- Citation validation proves that source IDs are retrieved and attached; it
  cannot mathematically prove claim-level semantic entailment.
- Very broad questions may still summarize only the papers with relevant
  retrieved evidence rather than mentioning every indexed document.
- Gemini is a remote dependency and may return rate-limit, availability, or
  timeout errors.
- `temperature=0` improves consistency but does not guarantee identical wording
  from a hosted model.
- The local index is designed for five documents. A much larger corpus would
  benefit from hybrid retrieval, reranking, ingestion versioning, and more
  formal evaluation.

## Vibe-Coding Reflection

I used three AI-assisted tools during this project, but for different purposes.

**Drive-integrated Gemini** helped me get an initial overview of the larger
document pool and narrow it down to five related papers. This was useful for
understanding themes quickly, but I still manually reviewed the selected PDFs
before treating them as the final corpus.

**Claude Sonnet** was used as a second opinion on the project brief and early
implementation plan. It helped challenge technology choices and check whether
the planned components covered the assignment. I did not treat the plan as
final simply because an AI approved it; I compared its suggestions with the
assignment and the behavior of the actual libraries.

**OpenAI Codex** was used for repository inspection, implementation, tests,
terminal diagnostics, and iterative debugging. I directed it in small stages:
PDF parsing, chunking, embeddings, ChromaDB, retrieval, Gemini generation,
citation mapping, MCP integration, and final quality checks. After each stage I
inspected real terminal output instead of moving directly to the next feature.

Several AI-generated or AI-supported decisions needed correction:

- The first plan used MiniLM with chunks that could exceed its effective token
  limit. After inspecting tokenizer behavior, I moved to
  `BAAI/bge-small-en-v1.5` and token-aware 450-token chunks.
- The initially selected Gemini model was no longer available to new users. I
  queried the available model catalog, tested alternatives, and selected
  `gemini-3.1-flash-lite`.
- MCP verification initially timed out. Generic timeout messages were not
  enough, so I added logs at client connection, service startup, embedding
  loading, retrieval, Gemini generation, citation mapping, and shutdown. The
  logs showed that SentenceTransformer was being loaded lazily inside the first
  asynchronous MCP tool request on Windows. Loading it before starting the MCP
  event loop fixed the issue.
- Ordinary top-k retrieval failed an explicit KGLM-versus-ERASER comparison
  because six KGLM chunks occupied the full context and ERASER ranked seventh.
  I added comparison-aware document diversity, then tested the exact failed
  question again.
- A first citation check rejected valid grouped citations such as `[S1, S2]`.
  The validator was changed to parse source IDs inside citation groups.
- Some broad queries retrieved bibliography entries. A relevance floor and
  citation-density filter were added and tested against real corpus rankings.

The most useful part of AI assistance was speed: it helped generate focused
tests, inspect library APIs, and propose debugging paths. The risky part was
that plausible-looking code could still be wrong at runtime. The MCP timeout
and multi-document retrieval failures were only found because I ran real
questions, challenged unexpected answers, and asked for stage-level evidence.

My view is that AI works best as an implementation partner and reviewer, not as
the final decision-maker. I relied on it for acceleration and alternative
ideas, while keeping acceptance criteria, security decisions, trade-offs, and
final validation under human control.

## Repository Hygiene

Secrets and generated local artifacts are excluded from version control, including
`.env`, virtual environments, model caches, ChromaDB data, build outputs, and
temporary test directories. Use `.env.example` to configure the application
without committing credentials.