Function bodies 33 total

def _extract_state_code(result: dict) -> str | None:
    for component in result.get("address_components", []):
        if "administrative_area_level_1" in component.get("types", []):
            short = component.get("short_name", "")
            if len(short) == 2 and short.isalpha():
                return short.upper()
    return None

def geocode_address(raw_input: str, config: dict) -> list[dict]:
    try:
        api_key = config.get("google_maps_api_key", "")
        resp = requests.get(
            "https://maps.googleapis.com/maps/api/geocode/json",
            params={"address": raw_input, "key": api_key},
            timeout=10,
        )
        resp.raise_for_status()
        data = resp.json()
        results = data.get("results", [])

        candidates = []
        for r in results:
            if not _is_us_result(r):
                continue
            loc = r.get("geometry", {}).get("location", {})
            state = _extract_state_code(r)
            if state is None:
                continue
            candidates.append({
                "formatted_address": r.get("formatted_address", ""),
                "lat": loc.get("lat"),
                "lng": loc.get("lng"),
                "state_code": state,
            })
            if len(candidates) >= _MAX_CANDIDATES:
                break

      

def pick_address(candidates: list[dict]) -> dict | None:
    if not candidates:
        return None

    console.print("\n[bold]Please select your address:[/bold]")
    for i, c in enumerate(candidates, 1):
        console.print(f"  {i}. {c['formatted_address']}")

    for attempt in range(_MAX_PICK_ATTEMPTS):
        choice = console.input("[bold]Enter number: [/bold]").strip()
        try:
            idx = int(choice) - 1
            if 0 <= idx < len(candidates):
                return candidates[idx]
        except ValueError:
            pass
        remaining = _MAX_PICK_ATTEMPTS - attempt - 1
        if remaining > 0:
            console.print(f"[red]Invalid choice. {remaining} attempt(s) remaining.[/red]")

    console.print("[yellow]Could not confirm address. Proceeding with raw input.[/yellow]")
    return None

def resolve_timezone(lat: float | None, lng: float | None) -> str:
    if lat is None or lng is None:
        return _DEFAULT_TIMEZONE
    try:
        tz = _get_tf().timezone_at(lat=lat, lng=lng)
        return tz if tz else _DEFAULT_TIMEZONE
    except Exception:
        logger.warning("Timezone lookup failed", exc_info=True)
        return _DEFAULT_TIMEZONE

def process_address(raw_input: str, config: dict) -> dict | None:
    """Geocode address. Returns dict on success, or None when geocoding fails."""
    candidates = geocode_address(raw_input, config)
    chosen = pick_address(candidates)

    if chosen:
        tz = resolve_timezone(chosen["lat"], chosen["lng"])
        return {
            "address_formatted": chosen["formatted_address"],
            "address_raw": raw_input,
            "timezone": tz,
        }

    return None

def get_available_slots(
    doctors: list[Doctor],
    reason_category: str,
    patient_timezone: str,
    now_utc: datetime,
) -> dict:
    """Return slots grouped by date string then doctor id.

    Structure: {
        "2025-06-02": {
            doctor_id: [{"start": datetime, "end": datetime}, ...],
            ...
        },
        ...
    }
    All datetimes are timezone-aware in the patient's timezone.
    """
    duration = DURATION_BY_CATEGORY.get(reason_category, 60)
    tz = ZoneInfo(patient_timezone)
    now_local = now_utc.astimezone(tz)
    earliest_start = now_local + timedelta(hours=1)

    start_date = now_local.date()
    end_date = start_date + timedelta(days=7)

    result: dict[str, dict[int, list[dict]]] = {}

    for doctor in doctors:
        slots = _generate_doctor_slots(doctor, duration, tz, earliest_start, start_date, end_date)
        for slot in slots:
            date_key = slot["start"].date().isoformat()
            if date_key not in result:
      

def _generate_doctor_slots(
    doctor: Doctor,
    duration: int,
    tz: ZoneInfo,
    earliest_start: datetime,
    start_date: date,
    end_date: date,
) -> list[dict]:
    weekly_hours = doctor.weekly_hours_json or []
    hours_by_day: dict[str, list[dict]] = {}
    for block in weekly_hours:
        day = block["day"]
        if day not in hours_by_day:
            hours_by_day[day] = []
        hours_by_day[day].append(block)

    slots = []
    current_date = start_date
    while current_date <= end_date:
        day_name = _DAY_NAMES[current_date.weekday()]
        for block in hours_by_day.get(day_name, []):
            block_start = _make_aware(current_date, _parse_time(block["start"]), tz)
            block_end = _make_aware(current_date, _parse_time(block["end"]), tz)

            candidate = _align_to_half_hour(max(block_start, earliest_start))
            while candidate + timedelta(minutes=duration) <= block_end:
                slots.append({
                    "star

def _align_to_half_hour(dt: datetime) -> datetime:
    if dt.minute <= 0 and dt.second == 0:
        return dt
    if dt.minute <= 30:
        aligned = dt.replace(minute=30, second=0, microsecond=0)
    else:
        aligned = (dt + timedelta(hours=1)).replace(minute=0, second=0, microsecond=0)
    return aligned

def _db_save(session: Session, obj) -> bool:
    """Try to commit; retry up to _DB_RETRY_MAX times. Returns True on success."""
    for attempt in range(_DB_RETRY_MAX + 1):
        try:
            session.add(obj)
            session.commit()
            return True
        except Exception:
            session.rollback()
            logger.warning("DB write failed (attempt %d)", attempt + 1, exc_info=True)
    return False

def _collect_demographics(console: Console, config: dict) -> dict | None:
    """Collect patient demographics. Returns dict or None if user backs past first field."""
    data: dict = {}
    states = list(_DEMO_STATES)
    idx = 0

    while idx < len(states):
        state = states[idx]

        if state == "first_name":
            result = ask_name("First name:")
        elif state == "last_name":
            result = ask_name("Last name:")
        elif state == "dob":
            result = ask_dob("Date of birth:", config=config)
        elif state == "payer":
            result = ask("Insurance payer name:")
            if result is not BACK and not result:
                console.print("[red]Payer name is required.[/red]")
                continue
        elif state == "insurance_id":
            result = ask_insurance_id("Insurance ID (leave empty if you want to enter it later):")
        else:
            break

        _check_quit(result)

        if result is BACK:
           

def _collect_address(config: dict, console: Console) -> dict | None:
    """Collect and geocode address. Returns dict or None if user backs."""
    last_raw = None  # track last unverified address for 'skip'

    while True:
        result = ask("Full address (or 'skip' to use previous unverified address):"
                      if last_raw else "Full address:")
        _check_quit(result)
        if result is BACK:
            return None

        # Allow 'skip' to accept the last unverified address as-is
        if result.lower() == "skip" and last_raw:
            return {
                "address_formatted": None,
                "address_raw": last_raw,
                "timezone": _DEFAULT_TIMEZONE,
            }

        if not validate_address(result):
            console.print(
                "[red]Please enter a valid address (at least 5 characters).[/red]"
            )
            continue

        addr_data = process_address(result, config)
        if addr_data is not None

def _ask_category(console: Console):
    """Present category menu. Returns category key string, or BACK/QUIT sentinel."""
    console.print("\n[bold]What type of appointment would you like?[/bold]")
    for i, (key, label, minutes) in enumerate(_CATEGORY_MENU, 1):
        console.print(f"  {i}. {label} ({minutes} min)")

    while True:
        result = ask("Enter a number (1-4):")
        _check_quit(result)
        if result is BACK:
            return BACK
        try:
            num = int(result)
            if 1 <= num <= len(_CATEGORY_MENU):
                return _CATEGORY_MENU[num - 1][0]
        except ValueError:
            pass
        console.print("[red]Please enter a number between 1 and 4.[/red]")

def _collect_medical(config: dict, console: Console) -> dict | None:
    """Collect appointment category, description, and LLM follow-ups.

    Returns dict with {issue_text, reason_category, transcript_raw} or None on back.
    """
    state = "category"
    category = None
    description = None
    llm_result = None
    follow_up_answers: list[str] = []
    fu_idx = 0

    while True:
        # -- Step 1: category selection --
        if state == "category":
            category = _ask_category(console)
            _check_quit(category)
            if category is BACK:
                return None  # back past medical -> address section
            state = "description"
            continue

        # -- Step 2: category-specific description --
        if state == "description":
            prompt = _CATEGORY_PROMPTS[category]
            description = ask(prompt)
            _check_quit(description)
            if description is BACK:
                state = "category"
             

def _collect_scheduling(
    session: Session,
    reason_category: str,
    patient_timezone: str,
    console: Console,
) -> dict | None:
    """Show availability and collect user's scheduling choice in-memory.

    Returns dict with {selection, preferred_time_text}, or None if user backs.
    """
    doctors = session.query(Doctor).all()
    doctors_by_id = {d.id: d.name for d in doctors}
    now_utc = datetime.now(UTC)

    slots = get_available_slots(doctors, reason_category, patient_timezone, now_utc)
    slot_index = build_slot_index(slots, doctors_by_id)

    if not slot_index:
        console.print("[yellow]No available slots in the next 7 days.[/yellow]")
        pref = prompt_preferred_time()
        _check_quit(pref)
        if pref is BACK:
            return None
        return {"selection": None, "preferred_time_text": pref}

    display_slots(slot_index)
    selection = prompt_slot_selection(slot_index)
    _check_quit(selection)
    if selection is BACK:
        return

def _show_review(data: dict, sched_data: dict, console: Console) -> bool:
    """Display summary and ask for confirmation. Returns True if confirmed."""
    console.print("\n[bold cyan]── Review Your Information ──[/bold cyan]")
    console.print(f"  Name:        {data['first_name']} {data['last_name']}")
    console.print(f"  DOB:         {data['dob']}")
    console.print(f"  Payer:       {data['payer']}")
    console.print(f"  Insurance:   {data['insurance_id']}")
    addr = data.get("address_formatted") or data.get("address_raw", "N/A")
    console.print(f"  Address:     {addr}")
    console.print(f"  Timezone:    {data['timezone']}")
    console.print(f"  Reason:      {data['reason_category']}")
    console.print(f"  Issue:       {data['issue_text']}")

    # Scheduling info
    selection = sched_data.get("selection")
    if selection:
        start_str = selection["start"].strftime("%B %d, %Y at %I:%M %p")
        console.print(f"  Appointment: {start_str} with {selection['doctor_

def run_intake(config: dict, session: Session) -> None:
    """Main intake state machine. Sections flow linearly with back support."""
    console.print("[bold green]Welcome to Assort Health Clinic![/bold green]")
    console.print("Type [bold]back[/bold] to go to the previous question, "
                  "or [bold]quit[/bold] to exit.\n")

    section = "demographics"

    demo_data: dict | None = None
    addr_data: dict | None = None
    med_data: dict | None = None
    sched_data: dict | None = None

    try:
        while True:
            # -- Demographics --
            if section == "demographics":
                demo_data = _collect_demographics(console, config)
                if demo_data is None:
                    console.print("[yellow]You're at the beginning of the form.[/yellow]")
                    continue
                section = "address"
                continue

            # -- Address --
            if section == "address":
                addr_data = _coll

def main():
    """Assort Health Clinic – Patient Intake Agent"""
    config = load_config()
    init_logging(config)

    engine = get_engine(config)
    init_db(engine)
    SessionFactory = get_session(engine)
    session = SessionFactory()

    seed_db(session)

    try:
        run_intake(config, session)
    finally:
        session.close()

def load_config(yaml_path: Path | None = None, env_path: Path | None = None) -> dict:
    if env_path is None:
        env_path = _PROJECT_ROOT / ".env"
    load_dotenv(env_path, override=True)

    if yaml_path is None:
        yaml_path = _PROJECT_ROOT / "config.yaml"
    with open(yaml_path) as f:
        cfg = yaml.safe_load(f)

    env_overrides = {
        "db_path": os.getenv("DB_PATH"),
        "retry_count": os.getenv("RETRY_COUNT"),
        "max_appointments_per_day": os.getenv("MAX_APPOINTMENTS_PER_DAY"),
        "log_level": os.getenv("LOG_LEVEL"),
        "openai_api_key": os.getenv("OPENAI_API_KEY"),
        "google_maps_api_key": os.getenv("GOOGLE_MAPS_API_KEY"),
    }
    for key, val in env_overrides.items():
        if val is not None:
            if key in ("retry_count", "max_appointments_per_day"):
                cfg[key] = int(val)
            else:
                cfg[key] = val

    return cfg

def parse_dob(text: str) -> date | None:
    cleaned = text.strip()
    # Deterministic parsing: try known formats
    for fmt in _DATE_FORMATS:
        try:
            dt = datetime.strptime(cleaned, fmt)
            return dt.date()
        except ValueError:
            continue
    return None

def _try_parse_iso(date_str: str) -> date | None:
    """Try to parse an ISO YYYY-MM-DD string from LLM response."""
    try:
        return datetime.strptime(date_str.strip(), "%Y-%m-%d").date()
    except (ValueError, AttributeError):
        return None

def ask(prompt: str) -> str:
    """Base prompt. Returns BACK or QUIT sentinel for navigation commands."""
    text = console.input(f"[bold]{prompt}[/bold] ").strip()
    if _is_quit(text):
        return QUIT
    if _is_back(text):
        return BACK
    return text

def ask_name(prompt: str) -> str:
    """Prompt for a name. Validates letters, hyphens, apostrophes only."""
    while True:
        result = ask(prompt)
        if result is QUIT:
            return QUIT
        if result is BACK:
            return BACK
        if validate_name(result):
            return result.strip()
        console.print(
            "[red]Invalid name. Only letters, hyphens, and apostrophes are allowed.[/red]"
        )

def ask_insurance_id(prompt: str) -> str:
    """Prompt for insurance ID. Empty string is allowed (skip)."""
    result = ask(prompt)
    if result is QUIT:
        return QUIT
    if result is BACK:
        return BACK
    return result

def ask_dob(prompt: str, config: dict | None = None) -> date:
    """Prompt for date of birth with LLM-assisted clarification.

    1. Try deterministic parsing first.
    2. If that fails, ask the LLM to parse or clarify.
    3. If LLM returns a date, accept it.
    4. If LLM returns a clarification, show it and re-prompt.
    5. Cap total attempts at 5; then tell user to call the office.
    """
    import assort_agent.llm as _llm_mod

    attempts = 0
    while attempts < _MAX_DOB_ATTEMPTS:
        result = ask(prompt)
        if result is QUIT:
            return QUIT
        if result is BACK:
            return BACK

        attempts += 1

        # Try deterministic parsing first
        parsed = parse_dob(result)
        if parsed is not None:
            return parsed

        # Fall back to LLM
        if config is not None:
            llm_result = _llm_mod.parse_dob_with_llm(result, config)
            if "date" in llm_result:
                parsed = _try_parse_iso(llm_res

def extract_followups(category: str, description: str, config: dict) -> dict:
    """Send category + description to LLM for follow-up questions.

    Returns {"follow_up_questions": [...], "issue_out_of_scope": bool}.
    On failure returns fallback with empty follow-ups and out_of_scope=False.
    """
    try:
        api_key = config.get("openai_api_key", "")
        system_msg = _SYSTEM_PROMPT.format(category=category)
        resp = requests.post(
            "https://api.openai.com/v1/chat/completions",
            headers={
                "Authorization": f"Bearer {api_key}",
                "Content-Type": "application/json",
            },
            json={
                "model": "gpt-4o-mini",
                "temperature": 0.2,
                "messages": [
                    {"role": "system", "content": system_msg},
                    {"role": "user", "content": description},
                ],
            },
            timeout=15,
        )
        resp.raise_for_st

def parse_dob_with_llm(text: str, config: dict) -> dict:
    """Ask LLM to parse or clarify a date of birth.

    Returns either {"date": "YYYY-MM-DD"} or {"clarification": "message"}.
    On failure returns {"clarification": "..."} with a generic message.
    """
    try:
        api_key = config.get("openai_api_key", "")
        resp = requests.post(
            "https://api.openai.com/v1/chat/completions",
            headers={
                "Authorization": f"Bearer {api_key}",
                "Content-Type": "application/json",
            },
            json={
                "model": "gpt-4o-mini",
                "temperature": 0.0,
                "messages": [
                    {"role": "system", "content": _DOB_SYSTEM_PROMPT},
                    {"role": "user", "content": text},
                ],
            },
            timeout=10,
        )
        resp.raise_for_status()
        content = resp.json()["choices"][0]["message"]["content"]
        parsed = json.loads

def _normalize(data: dict) -> dict:
    followups = data.get("follow_up_questions", [])
    if not isinstance(followups, list):
        followups = []
    followups = followups[:_MAX_FOLLOWUPS]

    out_of_scope = bool(data.get("issue_out_of_scope", False))

    return {
        "follow_up_questions": followups,
        "issue_out_of_scope": out_of_scope,
    }

def build_slot_index(slots: dict, doctors_by_id: dict) -> list[dict]:
    """Flatten grouped slots into a numbered list.

    Each entry: {number, date, doctor_id, doctor_name, start, end}
    """
    index = []
    num = 1
    for date_key in sorted(slots.keys()):
        doctor_slots = slots[date_key]
        for doctor_id in sorted(doctor_slots.keys()):
            doctor_name = doctors_by_id.get(doctor_id, f"Doctor {doctor_id}")
            for slot in doctor_slots[doctor_id]:
                index.append({
                    "number": num,
                    "date": date_key,
                    "doctor_id": doctor_id,
                    "doctor_name": doctor_name,
                    "start": slot["start"],
                    "end": slot["end"],
                })
                num += 1
    return index

def display_slots(slot_index: list[dict]) -> None:
    """Print slots grouped by date then doctor."""
    current_date = None
    current_doctor = None

    for entry in slot_index:
        if entry["date"] != current_date:
            current_date = entry["date"]
            current_doctor = None
            console.print(f"\n[bold cyan]{current_date}[/bold cyan]")

        if entry["doctor_name"] != current_doctor:
            current_doctor = entry["doctor_name"]
            console.print(f"  [bold]{current_doctor}[/bold]")

        start_str = entry["start"].strftime("%I:%M %p")
        end_str = entry["end"].strftime("%I:%M %p")
        console.print(f"    {entry['number']}. {start_str} - {end_str}")

def prompt_slot_selection(slot_index: list[dict]):
    """Prompt user to pick a slot number or 'N' for no times work.

    Returns the chosen slot dict, None (no preference), BACK, or QUIT.
    """
    console.print("\n[bold]Enter a slot number, or N if no times work.[/bold]")
    max_attempts = 3
    for attempt in range(max_attempts):
        choice = console.input("> ").strip()
        if _is_quit(choice):
            return QUIT
        if _is_back(choice):
            return BACK
        if choice.upper() == "N":
            return None
        try:
            num = int(choice)
            for entry in slot_index:
                if entry["number"] == num:
                    return entry
        except ValueError:
            pass
        remaining = max_attempts - attempt - 1
        if remaining > 0:
            console.print(f"[red]Invalid choice. {remaining} attempt(s) remaining.[/red]")

    console.print("[yellow]Could not confirm selection. Treating as no preference.[/yel

def prompt_preferred_time():
    """Ask for free-text preferred time. Returns text, BACK, or QUIT."""
    console.print("[bold]When would you prefer to be seen? (e.g. 'Tuesday morning')[/bold]")
    text = console.input("> ").strip()
    if _is_quit(text):
        return QUIT
    if _is_back(text):
        return BACK
    return text

def persist_selection(
    session: Session,
    appointment: Appointment,
    selection: dict | None,
    preferred_time_text: str | None,
    patient_timezone: str,
) -> None:
    """Update appointment based on user's scheduling choice."""
    if selection is not None:
        appointment.requested_start_local = selection["start"].replace(tzinfo=None)
        start_aware = selection["start"]
        if start_aware.tzinfo is None:
            start_aware = start_aware.replace(tzinfo=ZoneInfo(patient_timezone))
        appointment.requested_start_utc = start_aware.astimezone(ZoneInfo("UTC")).replace(tzinfo=None)
        appointment.requested_doctor_id = selection["doctor_id"]
    elif preferred_time_text:
        appointment.preferred_time_text = preferred_time_text

    appointment.status = "REQUESTED"
    session.commit()

def seed_db(session: Session) -> None:
    office = session.query(Office).filter_by(name=_OFFICE_NAME).first()
    if office is None:
        office = Office(name=_OFFICE_NAME)
        session.add(office)
        session.flush()

    for doc_data in _DOCTORS:
        exists = (
            session.query(Doctor)
            .filter_by(name=doc_data["name"], office_id=office.id)
            .first()
        )
        if exists is None:
            doctor = Doctor(
                office_id=office.id,
                name=doc_data["name"],
                weekly_hours_json=doc_data["weekly_hours"],
            )
            session.add(doctor)

    session.commit()