Build a document processing workflow in 30 minutes (or less!)

uvx mistralai-workflows-cli setup

.agents/skills/workflows
src/
	workflows/
		__init__.py
		hello.py
		start.py
	dev_worker.py
	discover.py
.gitignore
Makefile
pyproject.toml
README.md

dependencies = [
    "mistralai-workflows[mistralai]>=3.0.0,<4",
    "pydantic",
    "python-dotenv",
    "streamlit==1.55.0",
    "pymupdf>=1.23.0",
]

make installdeps

mkdir src/shared
touch src/shared/extraction_fields.py

"""
Fields to extract by document type.
Imported by workflow.py and app.py.
"""

COMMON_FIELDS: list[tuple[str, str]] = [
    ("full_name", "Full name of the patient (last name + first name)"),
    ("patient_address", "Full address of the patient"),
    ("social_security_number", "Social security number of the patient"),
]

SPECIFIC_FIELDS: dict[str, list[tuple[str, str]]] = {
    "prescription": [
        ("doctor_name", "Name of the doctor who made the prescription"),
        ("doctor_address", "Address of the doctor who made the prescription"),
        ("doctor_rpps", "RPPS number of the doctor"),
        ("medications", "List of prescribed medications (separated by commas)"),
        ("prescription_date", "Date of the prescription"),
    ],
    "medical_bill": [
        ("healthcare_professional_name", "Name of the healthcare professional"),
        ("healthcare_professional_address", "Address of the healthcare professional"),
        ("bill_amount", "Total billed amount (with currency)"),
        ("services", "Billed services or procedures"),
        ("care_date", "Date of care"),
    ],
    "hospitalization_report": [
        ("hospital_name", "Name of the hospital"),
        ("department", "Department or care unit"),
        ("responsible_doctor", "Responsible doctor"),
        ("admission_date", "Date of admission"),
        ("discharge_date", "Date of discharge"),
        ("primary_diagnosis", "Primary diagnosis"),
    ],
    "biological_analysis": [
        ("laboratory", "Name of the laboratory"),
        ("prescribing_doctor", "Prescribing doctor"),
        ("sample_date", "Date of the sample"),
        ("analyses", "Performed analyses"),
        ("abnormal_results", "Reported abnormal results"),
    ],
    "medical_imaging": [
        ("facility", "Name of the imaging facility"),
        ("radiologist_name", "Name of the radiologist"),
        ("exam_type", "Type of exam (MRI, CT, X-ray, etc.)"),
        ("anatomical_region", "Anatomical region examined"),
        ("exam_date", "Date of the exam"),
        ("conclusion", "Conclusion or main finding"),
    ],
    "medical_certificate": [
        ("doctor_name", "Name of the signing doctor"),
        ("doctor_address", "Address of the doctor"),
        ("certificate_subject", "Subject of the certificate"),
        ("certificate_date", "Date of the certificate"),
        ("sick_leave_duration", "Duration of sick leave if mentioned"),
    ],
    "mutual_reimbursement": [
        ("mutual_name", "Name of the mutual insurance"),
        ("member_number", "Member number"),
        ("reimbursement_amount", "Amount reimbursed by the mutual insurance"),
        ("reimbursement_date", "Reimbursement date"),
        ("reimbursed_acts", "Reimbursed acts"),
    ],
    "social_security_reimbursement": [
        ("fund_name", "Name of the social security fund"),
        ("reimbursement_amount", "Amount reimbursed by social security"),
        ("reimbursement_rate", "Applied reimbursement rate"),
        ("reimbursement_date", "Reimbursement date"),
        ("reimbursed_acts", "Reimbursed acts"),
    ],
    "consultation_report": [
        ("doctor_name", "Name of the consulting doctor"),
        ("doctor_specialty", "Specialty of the doctor"),
        ("consultation_date", "Date of the consultation"),
        ("reason", "Reason for the consultation"),
        ("diagnosis", "Diagnosis"),
        ("prescribed_treatment", "Prescribed treatment or follow-up"),
    ],
    "informed_consent": [
        ("doctor_name", "Name of the doctor"),
        ("facility", "Facility"),
        ("medical_procedure", "Medical procedure concerned"),
        ("signature_date", "Date of signature"),
    ],
    "other": [],
}

DOCUMENT_CATEGORIES: list[str] = list(SPECIFIC_FIELDS.keys())

CATEGORY_LABELS: dict[str, str] = {
    "prescription": "📋 Prescription",
    "medical_bill": "🧾 Medical Bill",
    "hospitalization_report": "🏥 Hospitalization Report",
    "biological_analysis": "🔬 Biological Analysis",
    "medical_imaging": "🩻 Medical Imaging",
    "medical_certificate": "📄 Medical Certificate",
    "mutual_reimbursement": "💳 Mutual Reimbursement",
    "social_security_reimbursement": "🏛️ Social Security Reimbursement",
    "consultation_report": "👨‍⚕️ Consultation Report",
    "informed_consent": "✍️ Informed Consent",
    "other": "❓ Other",
}

touch src/workflows/medical_doc_workflow.py

import asyncio
import logging
import os
from functools import lru_cache
from datetime import timedelta
from typing import Optional

import mistralai.workflows as workflows
import mistralai.workflows.plugins.mistralai as workflows_mistralai
from dotenv import load_dotenv
from pydantic import BaseModel, ConfigDict, Field, create_model

from shared.extraction_fields import COMMON_FIELDS, DOCUMENT_CATEGORIES, SPECIFIC_FIELDS

load_dotenv(override=True)

for name in ("mistralai_workflows", "httpx", "httpcore"):
    logging.getLogger(name).setLevel(logging.WARNING)

class ManualCategorySignal(BaseModel):
    category: str

class DocumentClassification(BaseModel):
    category: str = Field(description=f"One of: {', '.join(DOCUMENT_CATEGORIES)}")
    confidence: float = Field(ge=0.0, le=1.0)
    explanation: str

@lru_cache(maxsize=None)
def get_extraction_output_model(category: str) -> type[BaseModel]:
    common_model = create_model(
        "CommonExtractionFields",
        __config__=ConfigDict(extra="forbid"),
        **{key: (Optional[str], None) for key, _ in COMMON_FIELDS},
    )
    specific_model = create_model(
        f"SpecificExtractionFields_{category}",
        __config__=ConfigDict(extra="forbid"),
        **{key: (Optional[str], None) for key, _ in SPECIFIC_FIELDS.get(category, [])},
    )
    return create_model(
        f"ExtractionOutput_{category}",
        __config__=ConfigDict(extra="forbid"),
        common=(common_model, ...),
        specific=(specific_model, ...),
    )

# ── Activities ────────────────────────────────────────────────────────────────

# Resolve the uploaded file ID into a temporary signed URL for Document QnA.
@workflows.activity(start_to_close_timeout=timedelta(minutes=5), retry_policy_max_attempts=2)
async def get_document_signed_url(file_id: str) -> str:
    client = workflows_mistralai.get_mistral_client()
    signed_url = await client.files.get_signed_url_async(file_id=file_id)
    return signed_url.url

# Classify the document category directly from the document URL using structured output.
@workflows.activity(start_to_close_timeout=timedelta(minutes=2), retry_policy_max_attempts=2)
async def classify_document(document_url: str, filename: str) -> dict:
    client = workflows_mistralai.get_mistral_client()
    model = os.environ.get("MISTRAL_CLASSIFIER_MODEL", "mistral-medium-latest")
    response = await client.chat.parse_async(
        response_format=DocumentClassification,
        model=model,
        temperature=0.0,
        messages=[
            {
                "role": "system",
                "content": (
                    "You are an expert in medical document classification. "
                    "Return only valid JSON that matches the schema."
                ),
            },
            {
                "role": "user",
                "content": [
                    {
                        "type": "text",
                        "text": (
                            f"Classify the medical document '{filename}' into exactly one category from:\\n"
                            + "\\n".join(f"- {c}" for c in DOCUMENT_CATEGORIES)
                            + "\\n\\n"
                            "Return confidence between 0 and 1 and a short explanation."
                        ),
                    },
                    {
                        "type": "document_url",
                        "document_url": document_url,
                        "document_name": filename,
                    },
                ],
            },
        ],
    )
    parsed = response.choices[0].message.parsed if response.choices and response.choices[0].message else None
    if parsed is None:
        raise RuntimeError("Classification response could not be parsed.")
    return parsed.model_dump(mode="json")

# Extract common and category-specific fields from the document with schema-constrained JSON.
@workflows.activity(start_to_close_timeout=timedelta(minutes=2), retry_policy_max_attempts=2)
async def extract_patient_info(document_url: str, filename: str, category: str) -> dict:
    client = workflows_mistralai.get_mistral_client()
    model = os.environ.get("MISTRAL_EXTRACTOR_MODEL", "mistral-medium-latest")
    extraction_model = get_extraction_output_model(category)
    common_fields_text = "\\n".join(f"- {key}" for key, _ in COMMON_FIELDS)
    specific_fields_text = "\\n".join(f"- {key}" for key, _ in SPECIFIC_FIELDS.get(category, []))
    response = await client.chat.parse_async(
        response_format=extraction_model,
        model=model,
        temperature=0.0,
        messages=[
            {
                "role": "system",
                "content": (
                    "You extract administrative and medical information from a medical document. "
                    "Return only valid JSON that matches the schema. "
                    "If a field is missing, set it to null."
                ),
            },
            {
                "role": "user",
                "content": [
                    {
                        "type": "text",
                        "text": (
                            f"Extract fields from '{filename}' for category '{category}'.\\n\\n"
                            "Populate these common fields:\\n"
                            f"{common_fields_text}\\n\\n"
                            "Populate these category-specific fields:\\n"
                            f"{specific_fields_text if specific_fields_text else '- (none)'}\\n\\n"
                            "Return null for missing values."
                        ),
                    },
                    {
                        "type": "document_url",
                        "document_url": document_url,
                        "document_name": filename,
                    },
                ],
            },
        ],
    )
    parsed = response.choices[0].message.parsed if response.choices and response.choices[0].message else None
    if parsed is None:
        raise RuntimeError("Extraction response could not be parsed.")
    return parsed.model_dump(mode="json")

# ── Workflow ──────────────────────────────────────────────────────────────────

@workflows.workflow.define(name="pdf_ocr_workflow")
class PdfOcrWorkflow(workflows.InteractiveWorkflow):
    def __init__(self):
        # UI-facing step state exposed through the `get_steps` query.
        self.steps = {
            "ocr": {"status": "pending", "result": None},
            "classify": {"status": "pending", "result": None},
            "extract": {"status": "pending", "result": None},
        }
        # Set by the `manual_category` signal when a user overrides classification.
        self._manual_category = None

    @workflows.workflow.query(name="get_steps")
    def get_steps(self) -> dict:
        # Queries are synchronous/read-only: frontend polls this for progress.
        return self.steps

    @workflows.workflow.signal(name="manual_category")
    async def manual_category_signal(self, payload: ManualCategorySignal) -> None:
        # Signals are async/write: this mutates workflow state while run() is active.
        self._manual_category = payload.category

    @workflows.workflow.entrypoint
    async def run(
        self,
        file_id: str,
        filename: str,
        confidence_threshold: float = 0.9,
        manual_review_timeout_seconds: Optional[float] = None,
    ) -> workflows_mistralai.ChatAssistantWorkflowOutput:
        # TodoList items let clients display coarse workflow progress in addition to step data.
        ocr_item = workflows_mistralai.TodoListItem(
            title="Prepare document for Document QnA",
            description="Generate a signed URL so Mistral Document AI can read the document.",
        )
        classify_item = workflows_mistralai.TodoListItem(
            title="Classify document type",
            description="Predict the medical document category with confidence.",
        )
        extract_item = workflows_mistralai.TodoListItem(
            title="Extract structured fields",
            description="Extract patient and document-specific fields.",
        )

        async with workflows_mistralai.TodoList(items=[ocr_item, classify_item, extract_item]):
            self.steps["ocr"]["status"] = "running"
            async with ocr_item:
                signed_document_url = await get_document_signed_url(file_id)
            self.steps["ocr"] = {
                "status": "done",
                "result": "Document prepared for Document QnA (OCR handled by Mistral Document AI).",
            }

            self.steps["classify"]["status"] = "running"
            async with classify_item:
                classification = await classify_document(signed_document_url, filename)

                if classification.get("confidence", 0.0) < confidence_threshold:
                    # Low confidence enters a human-in-the-loop branch.
                    self.steps["classify"] = {"status": "waiting_human", "result": classification}
                    try:
                        # wait_condition pauses deterministically until signal or timeout.
                        await workflows.workflow.wait_condition(
                            lambda: self._manual_category is not None,
                            timeout=manual_review_timeout_seconds,
                            timeout_summary="manual_category_review",
                        )
                    except asyncio.TimeoutError:
                        # Non-interactive callers can set a timeout to avoid waiting forever.
                        classification["explanation"] = (
                            f"{classification.get('explanation', '').strip()} "
                            "Manual review timed out; using model-predicted category."
                        ).strip()
                    else:
                        # Manual override takes precedence when a signal arrives in time.
                        classification["category"] = self._manual_category
                        classification["confidence"] = 1.0
                        classification["explanation"] = f"Manually selected category: {self._manual_category}"

            self.steps["classify"] = {"status": "done", "result": classification}

            self.steps["extract"]["status"] = "running"
            async with extract_item:
                patient_info = await extract_patient_info(signed_document_url, filename, classification["category"])
            self.steps["extract"] = {"status": "done", "result": patient_info}

        return workflows_mistralai.ChatAssistantWorkflowOutput(
            # Return both human-readable text and structured payload for the UI.
            content=[workflows_mistralai.TextOutput(text=f"Processing complete for {filename}.")],
            structuredContent={
                "filename": filename,
                "ocr_text": "Document processed with Document QnA (no standalone OCR text payload).",
                "classification": classification,
                "patient_info": patient_info,
            },
        )
        
async def main() -> None:
    print("Worker ready — waiting for tasks...\\n")
    await workflows.run_worker([PdfOcrWorkflow])

if __name__ == "__main__":
    asyncio.run(main())