FewShotPromptTemplate in LangChain

# Zero-shot: just instructions
"Extract the drug and dose from this text."

# Few-shot: instructions + examples
"""Extract drug name and dose from clinical text.

Text: "Patient was started on warfarin 5mg daily"
Drug: warfarin
Dose: 5mg daily

Text: "Initiated metformin 500mg BID with meals"
Drug: metformin
Dose: 500mg twice daily

Text: {input}
Drug:"""

from langchain_core.prompts import (
    FewShotPromptTemplate,
    PromptTemplate,
    FewShotChatMessagePromptTemplate,
    ChatPromptTemplate,
)

# Define examples
drug_extraction_examples = [
    {
        "input": "Patient was started on warfarin 5mg daily for AFib.",
        "drug": "warfarin",
        "dose": "5mg",
        "frequency": "daily",
        "indication": "atrial fibrillation",
    },
    {
        "input": "Initiated metformin 500mg BID with meals for T2DM.",
        "drug": "metformin",
        "dose": "500mg",
        "frequency": "twice daily (BID)",
        "indication": "type 2 diabetes",
    },
    {
        "input": "Lisinopril 10mg once daily prescribed for hypertension.",
        "drug": "lisinopril",
        "dose": "10mg",
        "frequency": "once daily",
        "indication": "hypertension",
    },
]

# Template for each example
example_prompt = PromptTemplate(
    input_variables=["input", "drug", "dose", "frequency", "indication"],
    template=(
        "Text: {input}\n"
        "Drug: {drug}\n"
        "Dose: {dose}\n"
        "Frequency: {frequency}\n"
        "Indication: {indication}"
    ),
)

# FewShotPromptTemplate wraps examples + suffix for the actual query
few_shot_prompt = FewShotPromptTemplate(
    examples=drug_extraction_examples,
    example_prompt=example_prompt,
    prefix="Extract drug information from clinical text:\n",
    suffix="Text: {input}\nDrug:",    # Suffix with the actual query variable
    input_variables=["input"],
    example_separator="\n\n",
)

# Format with a new query
formatted = few_shot_prompt.format(
    input="Atorvastatin 40mg at bedtime started for hyperlipidemia."
)
print(formatted)

from langchain_openai import ChatOpenAI
from langchain_core.output_parsers import StrOutputParser

# Chat-specific few-shot template
chat_examples = [
    {
        "input": "What is the mechanism of warfarin?",
        "output": "Warfarin inhibits vitamin K epoxide reductase (VKORC1), preventing regeneration of vitamin K, which is required for synthesis of clotting factors II, VII, IX, and X.",
    },
    {
        "input": "What is the mechanism of metformin?",
        "output": "Metformin activates AMP-activated protein kinase (AMPK) in the liver, reducing hepatic glucose output (gluconeogenesis) and improving insulin sensitivity.",
    },
]

example_prompt = ChatPromptTemplate.from_messages([
    ("human", "{input}"),
    ("ai", "{output}"),
])

few_shot_chat = FewShotChatMessagePromptTemplate(
    example_prompt=example_prompt,
    examples=chat_examples,
)

# Wrap in full system prompt + few-shot examples + current query
final_prompt = ChatPromptTemplate.from_messages([
    ("system", "You are a clinical pharmacologist. Explain drug mechanisms precisely, including molecular targets."),
    few_shot_chat,   # Inserts the few-shot examples here
    ("human", "{input}"),
])

model = ChatOpenAI(model="gpt-4o", temperature=0)
parser = StrOutputParser()

chain = final_prompt | model | parser
result = chain.invoke({"input": "What is the mechanism of aspirin?"})
print(result)
# Model mimics the precise, molecular-level style shown in examples

from langchain_core.example_selectors import LengthBasedExampleSelector

# Large pool of examples
all_examples = [
    {"input": "warfarin 5mg daily", "output": "Drug: warfarin, Dose: 5mg, Frequency: daily"},
    {"input": "metformin 500mg BID", "output": "Drug: metformin, Dose: 500mg, Frequency: BID"},
    {"input": "aspirin 81mg daily", "output": "Drug: aspirin, Dose: 81mg, Frequency: daily"},
    {"input": "lisinopril 10mg QD", "output": "Drug: lisinopril, Dose: 10mg, Frequency: daily"},
    {"input": "atorvastatin 40mg QHS", "output": "Drug: atorvastatin, Dose: 40mg, Frequency: nightly"},
    # ... more examples
]

# Select examples to keep total prompt under a token budget
length_selector = LengthBasedExampleSelector(
    examples=all_examples,
    example_prompt=PromptTemplate(
        input_variables=["input", "output"],
        template="Input: {input}\nOutput: {output}",
    ),
    max_length=200,   # Max word count for all examples combined
)

dynamic_few_shot = FewShotPromptTemplate(
    example_selector=length_selector,
    example_prompt=PromptTemplate(
        input_variables=["input", "output"],
        template="Input: {input}\nOutput: {output}",
    ),
    prefix="Extract drug information:",
    suffix="Input: {input}\nOutput:",
    input_variables=["input"],
)

# Will select as many examples as fit within 200 words
formatted = dynamic_few_shot.format(input="omeprazole 20mg daily")

from langchain_core.example_selectors import SemanticSimilarityExampleSelector
from langchain_openai import OpenAIEmbeddings
from langchain_chroma import Chroma

# Larger pool with diverse examples
interaction_examples = [
    {
        "query": "Can warfarin be used with aspirin?",
        "answer": "Warfarin + aspirin: Major interaction. Both increase bleeding risk. Co-administration significantly raises risk of serious hemorrhage. Monitor INR closely if unavoidable.",
    },
    {
        "query": "Is metformin safe with alcohol?",
        "answer": "Metformin + alcohol: Moderate interaction. Alcohol increases lactic acidosis risk. Advise patients to limit alcohol. Monitor for signs of lactic acidosis.",
    },
    {
        "query": "Can SSRIs be combined with MAOIs?",
        "answer": "SSRIs + MAOIs: Contraindicated. Risk of serotonin syndrome (life-threatening). Require 14-day washout between agents.",
    },
    {
        "query": "Does ciprofloxacin interact with warfarin?",
        "answer": "Ciprofloxacin + warfarin: Major. Ciprofloxacin inhibits CYP1A2 and CYP3A4, increasing warfarin exposure. INR can double — reduce warfarin dose and monitor INR closely.",
    },
    {
        "query": "What about amiodarone and warfarin?",
        "answer": "Amiodarone + warfarin: Major. Amiodarone inhibits CYP2C9, drastically increasing warfarin levels. INR may triple. Reduce warfarin dose by 30-50% and monitor frequently.",
    },
]

# Create semantic selector backed by Chroma
embeddings = OpenAIEmbeddings(model="text-embedding-3-small")
selector = SemanticSimilarityExampleSelector.from_examples(
    examples=interaction_examples,
    embeddings=embeddings,
    vectorstore_cls=Chroma,
    k=2,   # Select 2 most similar examples
)

# FewShotChatMessagePromptTemplate with semantic selection
similar_example_prompt = ChatPromptTemplate.from_messages([
    ("human", "{query}"),
    ("ai", "{answer}"),
])

semantic_few_shot = FewShotChatMessagePromptTemplate(
    example_selector=selector,
    example_prompt=similar_example_prompt,
    input_variables=["query"],
)

final_prompt_with_semantic = ChatPromptTemplate.from_messages([
    ("system", "You are a clinical pharmacist specializing in drug interactions. Classify severity as: Contraindicated / Major / Moderate / Minor."),
    semantic_few_shot,   # 2 most relevant examples will be inserted here
    ("human", "{query}"),
])

chain = final_prompt_with_semantic | ChatOpenAI(model="gpt-4o") | StrOutputParser()

# Query: similar to "warfarin + aspirin" example → that example will be selected
result = chain.invoke({"query": "Can I give clopidogrel with warfarin?"})
print(result)

from dataclasses import dataclass
from typing import Optional

@dataclass
class FewShotExample:
    """A single few-shot example with metadata."""
    input: str
    output: str
    category: str
    quality_score: float = 1.0
    verified: bool = False

class FewShotLibrary:
    """Manage and version-control few-shot examples."""

    def __init__(self):
        self.examples: list[FewShotExample] = []

    def add(self, example: FewShotExample) -> None:
        self.examples.append(example)

    def get_verified(self, category: Optional[str] = None) -> list[dict]:
        """Return only verified examples for production use."""
        filtered = [e for e in self.examples if e.verified]
        if category:
            filtered = [e for e in filtered if e.category == category]
        return [{"input": e.input, "output": e.output} for e in filtered]

    def build_few_shot_prompt(
        self,
        category: str,
        max_examples: int = 3,
    ) -> FewShotChatMessagePromptTemplate:
        examples = self.get_verified(category)[:max_examples]
        return FewShotChatMessagePromptTemplate(
            examples=examples,
            example_prompt=ChatPromptTemplate.from_messages([
                ("human", "{input}"),
                ("ai", "{output}"),
            ]),
        )


# Use the library
library = FewShotLibrary()
library.add(FewShotExample(
    input="warfarin 5mg daily",
    output='{"drug": "warfarin", "dose": "5mg", "frequency": "daily"}',
    category="drug_extraction",
    verified=True,
))

few_shot = library.build_few_shot_prompt("drug_extraction", max_examples=3)