extraction.py

"""
Core extraction logic: takes raw, messy order text (as pasted from social media /
WhatsApp channels) and returns structured order data using Cohere's LLM.
"""

import json
import os
import cohere

EXTRACTION_PROMPT = """You are an assistant helping a medical logistics fulfillment team \
(Zipline-style drone delivery) convert messy, informally-written order requests into \
structured data.

Orders are pasted directly from social media / WhatsApp channels. They may include:
- A sender name or facility name
- A list of items with quantities, often using local abbreviations \
(e.g. "PCM syrp" = Paracetamol Syrup, "Scalvein" = IV Cannula / Scalp Vein set, \
"AD Sy" = Auto-Disable Syringe, "Penta" = Pentavalent Vaccine, "MR" = Measles-Rubella)
- Sometimes a requested date and urgency note (e.g. "tomorrow", "please consider...")
- Quantities may be in various units: doses, packs, pieces, bottles, sachets, vials, \
or just bare numbers

Your job: extract a clean JSON object with this exact structure:

{{
  "sender": "<name or facility name as written, or null>",
  "facility": "<facility name if mentioned, or null>",
  "requested_date": "<date in YYYY-MM-DD format if mentioned, or null>",
  "urgency": "<one of: routine, urgent, or null if not indicated>",
  "notes": "<any extra free-text instructions/notes that don't fit elsewhere, or null>",
  "items": [
    {{
      "raw_text": "<the original text for this line item, verbatim>",
      "item_name": "<your best-guess normalized/expanded product name>",
      "quantity": <number, your best guess>,
      "unit": "<unit if specified e.g. doses, pkt, pieces, or null>"
    }}
  ]
}}

Rules:
- Today's date is 2026-06-11. Resolve relative dates ("tomorrow", "today") accordingly.
- If urgency is implied by phrasing like "please", "urgent", "ASAP", "today" or a \
near-term date, mark "urgent". Otherwise "routine".
- Expand abbreviations to full product names where you can confidently infer them.
- Keep "raw_text" exactly as written for each item, for traceability.
- If a single line covers multiple distinct sizes/variants (e.g. "Scalvein 10 mixed \
23g and 21g"), split into separate item entries if you can infer the split, \
otherwise keep as one item with a note.
- Output ONLY the JSON object. No preamble, no markdown code fences, no explanation.

Order text to process:
---
{order_text}
---
"""


def get_cohere_client(api_key: str | None = None) -> cohere.ClientV2:
    key = api_key or os.environ.get("COHERE_API_KEY")
    if not key:
        raise ValueError(
            "No Cohere API key provided. Set COHERE_API_KEY env var or pass api_key."
        )
    return cohere.ClientV2(api_key=key)


def extract_order(order_text: str, api_key: str | None = None, model: str = "command-r7b-12-2024") -> dict:
    """
    Send raw order text to Cohere and return a parsed structured dict.

    Raises ValueError if the model output isn't valid JSON.
    """
    if not order_text or not order_text.strip():
        raise ValueError("Order text is empty.")

    client = get_cohere_client(api_key)
    prompt = EXTRACTION_PROMPT.format(order_text=order_text.strip())

    response = client.chat(
        model=model,
        messages=[{"role": "user", "content": prompt}],
        temperature=0.1,
    )

    raw_output = response.message.content[0].text.strip()

    # Defensive cleanup in case the model wraps in code fences despite instructions
    if raw_output.startswith("```"):
        raw_output = raw_output.strip("`")
        if raw_output.startswith("json"):
            raw_output = raw_output[4:]
        raw_output = raw_output.strip()

    try:
        data = json.loads(raw_output)
    except json.JSONDecodeError as e:
        raise ValueError(f"Model did not return valid JSON: {e}\n\nRaw output:\n{raw_output}")

    # Basic shape validation / defaults
    data.setdefault("sender", None)
    data.setdefault("facility", None)
    data.setdefault("requested_date", None)
    data.setdefault("urgency", None)
    data.setdefault("notes", None)
    data.setdefault("items", [])

    for item in data["items"]:
        item.setdefault("raw_text", "")
        item.setdefault("item_name", "")
        item.setdefault("quantity", None)
        item.setdefault("unit", None)

    return data