Building a medical literature Q&A agent in LangGraph

Building a retrieval-augmented system for clinical text demands strict control over context and traceability. This tutorial constructs a medical literature qa agent langgraph application that ingests abstracts, retrieves relevant passages, and generates cited answers using a compiled state graph.

Prerequisites

• Python 3.10 or newer
• Familiarity with LangChain primitives (Document, Retriever)
• API key for an OpenAI-compatible chat model
• Install dependencies:

pip install langgraph langchain langchain-community langchain-openai chromadb tiktoken

Set your key in the environment:

export OPENAI_API_KEY="sk-..."

If you prefer a single endpoint that fronts many models, the base_url swap shown later works without code changes.

Build the corpus

We use three short synthetic abstracts. In production you would load PubMed XML or PDFs; the graph does not care about source format as long as you emit Document objects with page_content and metadata.

from langchain_core.documents import Document

docs = [
    Document(
        page_content="Metformin is a biguanide used for type 2 diabetes. Common adverse effects include gastrointestinal upset, diarrhea, and nausea. Lactic acidosis is rare but serious.",
        metadata={"source": "PMID-001", "title": "Metformin overview"},
    ),
    Document(
        page_content="SGLT2 inhibitors reduce cardiovascular mortality in heart failure patients. Genital mycotic infections occur more frequently than with placebo.",
        metadata={"source": "PMID-002", "title": "SGLT2 inhibitors in HF"},
    ),
    Document(
        page_content="Warfarin dosing is guided by INR targets. Bleeding risk increases with age and concurrent aspirin use. Pharmacogenetic testing can inform initial dose.",
        metadata={"source": "PMID-003", "title": "Warfarin management"},
    ),
]

Embed and index

We embed with text-embedding-3-small and store in an in-memory Chroma collection.

from langchain_openai import OpenAIEmbeddings
from langchain_community.vectorstores import Chroma

embeddings = OpenAIEmbeddings(model="text-embedding-3-small")
vectordb = Chroma.from_documents(
    documents=docs,
    embedding=embeddings,
    collection_name="medlit",
)
retriever = vectordb.as_retriever(search_kwargs={"k": 2})

Define graph state and nodes

The medical literature qa agent langgraph design uses a typed state. Each node returns a partial dict that LangGraph merges.

from typing import TypedDict, List
from langchain_core.documents import Document

class State(TypedDict):
    question: str
    context: List[Document]
    answer: str

Retrieval node

def retrieve(state: State) -> dict:
    retrieved = retriever.invoke(state["question"])
    return {"context": retrieved}

Generation node

We force the model to ground on retrieved context only. The prompt is deliberately strict.

from langchain_openai import ChatOpenAI

chat = ChatOpenAI(model="gpt-4o-mini", temperature=0)

def generate(state: State) -> dict:
    ctx = "\n\n".join(
        f"[{d.metadata['source']}] {d.page_content}" for d in state["context"]
    )
    prompt = (
        "You are a clinical assistant. Answer the question using ONLY the context. "
        "Cite the source tag (e.g., [PMID-001]) after each claim.\n\n"
        f"Context:\n{ctx}\n\nQuestion: {state['question']}"
    )
    resp = chat.invoke(prompt)
    return {"answer": resp.content}

Compile the graph

from langgraph.graph import StateGraph, END

builder = StateGraph(State)
builder.add_node("retrieve", retrieve)
builder.add_node("generate", generate)
builder.set_entry_point("retrieve")
builder.add_edge("retrieve", "generate")
builder.add_edge("generate", END)

graph = builder.compile()

Run and verify

result = graph.invoke(
    {"question": "What are common side effects of metformin?"}
)
print(result["answer"])

Expected output (wording may vary):

Common side effects of metformin include gastrointestinal upset, diarrhea, and nausea [PMID-001]. 
Lactic acidosis is a rare but serious adverse effect [PMID-001].

The context field contains the two retrieved documents. Inspect it:

for d in result["context"]:
    print(d.metadata["source"])

PMID-001
PMID-003

Note that the retriever pulled PMID-003 as a secondary match; the generator correctly ignored warfarin content because the prompt constrained it to the question.

Pointing at an OpenAI-compatible gateway

If you want one endpoint that addresses 240+ models and automatically fails over when a provider is rate-limited, swap the ChatOpenAI constructor. The medical literature qa agent langgraph code stays identical:

chat = ChatOpenAI(
    model="gpt-4o-mini",
    temperature=0,
    base_url="https://api.n4n.ai/v1",
    api_key="your-n4n-key",
)

n4n.ai forwards provider cache-control hints and meters per token, so you can add extra_headers={"x-cache": "true"} to enable prompt caching where the underlying provider supports it.

Add a relevance gate

A real deployment should reject low-relevance context. Extend the state and add a grading node that scores the top document.

from langgraph.graph import END

def grade(state: State) -> str:
    if not state["context"]:
        return "no_context"
    # Simple heuristic: require the question token overlap
    q_tokens = set(state["question"].lower().split())
    top = state["context"][0].page_content.lower()
    overlap = any(t in top for t in q_tokens if len(t) > 4)
    return "generate" if overlap else "no_context"

builder.add_node("grade", lambda s: s)  # passthrough, logic in edge
builder.add_conditional_edges(
    "retrieve",
    grade,
    {"generate": "generate", "no_context": END},
)

Now if retrieval fails, the graph ends without calling the model, saving tokens and avoiding hallucination.

Running with the gate

graph = builder.compile()
out = graph.invoke({"question": "How do I fly a kite?"})
# out == {"question": "...", "context": [], "answer": ""}

The gate routed to END because no medical document matched.

Closing the loop with citations

For auditability in healthcare, emit the metadata alongside the answer. Modify generate to return both:

def generate(state: State) -> dict:
    ctx = "\n\n".join(
        f"[{d.metadata['source']}] {d.page_content}" for d in state["context"]
    )
    prompt = (
        "Answer using ONLY context. Cite source tags. "
        "Then list the sources used.\n\n"
        f"Context:\n{ctx}\n\nQuestion: {state['question']}"
    )
    resp = chat.invoke(prompt)
    sources = [d.metadata["source"] for d in state["context"]]
    return {"answer": resp.content, "sources": sources}

Add "sources": List[str] to State. The medical literature qa agent langgraph now produces a trail that satisfies most institutional review requirements.

Deployment notes

• Run the graph inside a FastAPI route; LangGraph state is JSON-serializable except for Document objects. Return only answer and sources to clients.
• Set retriever k based on context window. With gpt-4o-mini (128k), k=4 is safe for short abstracts.
• If you batch questions, compile the graph once and reuse the instance; node functions are stateless.

The pattern above is the minimal viable core. From here you can add multi-hop retrieval, re-ranking, or human-in-the-loop approval before the generate node—all as additional edges in the same compiled graph.