Composing Complex Prompts from Parts

from langchain_core.prompts import ChatPromptTemplate
from langchain_openai import ChatOpenAI
from langchain_core.output_parsers import StrOutputParser
import datetime

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

# Base prompt with all variables
base_prompt = ChatPromptTemplate.from_messages([
    ("system",
     "You are a {specialty} at {institution}. "
     "Today is {date}. "
     "Follow {guidelines_standard} guidelines."),
    ("human", "{question}"),
])

# Create specialized variants via partial()
clinical_pharmacist_prompt = base_prompt.partial(
    specialty="clinical pharmacist",
    institution="General Hospital",
    date=datetime.date.today().isoformat(),
    guidelines_standard="ASHP",
)

icu_pharmacist_prompt = base_prompt.partial(
    specialty="ICU clinical pharmacist",
    institution="Medical Center ICU",
    date=datetime.date.today().isoformat(),
    guidelines_standard="SCCM",
)

# Each variant only needs "question" at runtime
clinical_chain = clinical_pharmacist_prompt | model | parser
icu_chain = icu_pharmacist_prompt | model | parser

clinical_answer = clinical_chain.invoke({"question": "What is the warfarin dose for AFib?"})
icu_answer = icu_chain.invoke({"question": "What is the heparin protocol for DVT prophylaxis?"})

from langchain_core.prompts import ChatPromptTemplate
from langchain_core.messages import SystemMessage

# Reusable system instruction blocks
SAFETY_DISCLAIMER = (
    "Clinical decisions must be individualized. "
    "This information is for educational purposes. "
    "Professional consultation is required for patient care."
)

CITATION_REQUIREMENT = (
    "When making clinical claims, indicate the evidence level: "
    "(A) RCT evidence, (B) observational evidence, (C) expert consensus."
)

RESPONSE_FORMAT = (
    "Structure your response with:\n"
    "1. Direct answer\n"
    "2. Clinical rationale\n"
    "3. Monitoring/follow-up\n"
)

def build_clinical_prompt(
    role: str,
    include_safety: bool = True,
    include_citations: bool = True,
    include_format: bool = True,
) -> ChatPromptTemplate:
    """Build a clinical prompt from modular components."""
    system_parts = [f"You are a {role}."]

    if include_safety:
        system_parts.append(SAFETY_DISCLAIMER)
    if include_citations:
        system_parts.append(CITATION_REQUIREMENT)
    if include_format:
        system_parts.append(RESPONSE_FORMAT)

    system_content = " ".join(system_parts)

    return ChatPromptTemplate.from_messages([
        ("system", system_content),
        ("human", "{question}"),
    ])

# Different configurations
research_prompt = build_clinical_prompt(
    "clinical researcher",
    include_safety=False,
    include_citations=True,
)

patient_prompt = build_clinical_prompt(
    "pharmacist counseling patients",
    include_citations=False,
    include_format=False,
)

clinician_prompt = build_clinical_prompt(
    "clinical pharmacist",
    include_safety=True,
    include_citations=True,
    include_format=True,
)

from langchain_core.prompts import MessagesPlaceholder, FewShotChatMessagePromptTemplate

# Compose from discrete placeholder sections
def build_full_clinical_prompt(
    role: str,
    examples: list[dict] = None,
    with_history: bool = True,
) -> ChatPromptTemplate:
    """Build a complete clinical prompt with optional few-shot and history."""
    messages = [
        ("system", f"You are a {role}. Answer with clinical precision."),
    ]

    # Inject few-shot examples if provided
    if examples:
        example_prompt = ChatPromptTemplate.from_messages([
            ("human", "{input}"),
            ("ai", "{output}"),
        ])
        few_shot_section = FewShotChatMessagePromptTemplate(
            examples=examples,
            example_prompt=example_prompt,
        )
        messages.append(few_shot_section)

    # Chat history placeholder
    if with_history:
        messages.append(MessagesPlaceholder("chat_history", optional=True))

    # Current question
    messages.append(("human", "{question}"))

    return ChatPromptTemplate.from_messages(messages)


# Demonstrate usage
examples = [
    {
        "input": "What is the mechanism of warfarin?",
        "output": "Warfarin inhibits VKORC1, blocking vitamin K recycling. Result: reduced synthesis of clotting factors II, VII, IX, X. (Evidence Level A)",
    },
]

full_prompt = build_full_clinical_prompt(
    role="clinical pharmacologist",
    examples=examples,
    with_history=True,
)

from langchain_core.messages import HumanMessage, AIMessage

chain = full_prompt | model | parser
result = chain.invoke({
    "question": "What is the mechanism of rivaroxaban?",
    "chat_history": [
        HumanMessage(content="What anticoagulants do you know?"),
        AIMessage(content="Warfarin, rivaroxaban, apixaban, dabigatran are major anticoagulants."),
    ],
})

from langchain_core.runnables import RunnablePassthrough, RunnableLambda

# Pattern: inject retrieved context into prompt at runtime
CONTEXT_BLOCK_TEMPLATE = """Relevant clinical information:
{context}

---
Based on the above information, answer the following question.
If the information is insufficient, say so explicitly."""

rag_prompt = ChatPromptTemplate.from_messages([
    ("system", "You are a clinical pharmacist. Answer only from provided context."),
    ("human",
     f"{CONTEXT_BLOCK_TEMPLATE}\n\nQuestion: {{question}}"),
])

# The context variable gets filled with retrieved documents
def format_retrieved_docs(docs: list) -> str:
    return "\n\n".join([
        f"[Source {i+1}: {d.get('title', 'Unknown')}]\n{d['content']}"
        for i, d in enumerate(docs)
    ])

# RAG chain with dynamic context injection
rag_chain = (
    RunnablePassthrough.assign(
        context=lambda inputs: format_retrieved_docs(inputs.get("retrieved_docs", []))
    )
    | rag_prompt
    | model
    | parser
)

result = rag_chain.invoke({
    "question": "What monitoring is needed for warfarin?",
    "retrieved_docs": [
        {"title": "Warfarin FDA Label", "content": "Monitor INR every 1-4 weeks..."},
        {"title": "Lexicomp", "content": "Target INR for AFib: 2.0-3.0..."},
    ],
})

from dataclasses import dataclass
from typing import Callable

@dataclass
class PromptConfig:
    """Configuration for a named prompt."""
    name: str
    description: str
    builder: Callable[[], ChatPromptTemplate]
    required_vars: list[str]
    optional_vars: list[str]


class PromptRegistry:
    """Central registry of prompt configurations."""

    def __init__(self):
        self._registry: dict[str, PromptConfig] = {}

    def register(self, config: PromptConfig) -> None:
        self._registry[config.name] = config

    def get(self, name: str) -> ChatPromptTemplate:
        if name not in self._registry:
            raise KeyError(f"Prompt '{name}' not in registry")
        return self._registry[name].builder()

    def list_prompts(self) -> list[str]:
        return list(self._registry.keys())


# Build the registry
registry = PromptRegistry()

registry.register(PromptConfig(
    name="clinical_pharmacist",
    description="Standard clinical pharmacist prompt",
    builder=lambda: build_clinical_prompt("clinical pharmacist"),
    required_vars=["question"],
    optional_vars=["chat_history"],
))

registry.register(PromptConfig(
    name="drug_interaction_specialist",
    description="Specialized for drug interaction queries",
    builder=lambda: ChatPromptTemplate.from_messages([
        ("system", "You are a clinical pharmacologist specializing in drug interactions. "
                   "Classify interactions as: Contraindicated / Major / Moderate / Minor. "
                   "Explain the mechanism for each interaction."),
        ("human", "{question}"),
    ]),
    required_vars=["question"],
    optional_vars=[],
))

registry.register(PromptConfig(
    name="patient_counselor",
    description="Patient-friendly explanations",
    builder=lambda: ChatPromptTemplate.from_messages([
        ("system", "You are a pharmacist counseling patients. Use plain language. "
                   "Avoid jargon. Focus on what the patient needs to know and do."),
        ("human", "{question}"),
    ]),
    required_vars=["question"],
    optional_vars=[],
))

# Use registry to build chains dynamically
def build_chain_for_prompt(prompt_name: str) -> "Chain":
    prompt = registry.get(prompt_name)
    return prompt | model | parser

# Select prompt based on context
def smart_chain(question: str, user_type: str) -> str:
    prompt_map = {
        "clinician": "clinical_pharmacist",
        "specialist": "drug_interaction_specialist",
        "patient": "patient_counselor",
    }
    prompt_name = prompt_map.get(user_type, "clinical_pharmacist")
    chain = build_chain_for_prompt(prompt_name)
    return chain.invoke({"question": question})