src/agents/input_agent.py

import logging

from langchain_core.messages import AIMessage, HumanMessage, SystemMessage
from langgraph.types import interrupt

from src.workflow.state import GovBridgeState
from src.integrations.llm_factory import get_llm

_log = logging.getLogger(__name__)

_SYSTEM = """You are the intake agent for GovBridge, an app that helps citizens write formal emails to their government.

Your job is to determine TWO things from the conversation:
1. CATEGORY — what kind of request this is
2. LOCATION — the city, state, or region the email is about

## CATEGORY
- idea: a suggestion, proposal, or improvement the citizen wants to pitch
- complaint: a problem, grievance, or broken service they want to report
Classify immediately if clear. Do NOT ask "Is that correct?"

## LOCATION
Extract the location from what the user has written. Look for:
- A city or region name ("Chennai", "Ottawa", "London")
- A government official that implies a location ("CM Vijay" → Tamil Nadu, "NYC Mayor" → New York City)
- A state or country name
If you cannot find any location in the conversation, ask the user: "Which city or region is your email about?"
Ask this ONLY if there is genuinely no location clue anywhere in the conversation.

## OUTPUT TAGS — include on separate lines at the end of your message:
  CATEGORY: idea          (or complaint)
  LOCATION: [city/region]

Include both tags when you have both. Include only the ones you have found.
Never ask more than 1 question per response."""


def input_agent(state: GovBridgeState) -> dict:
    llm = get_llm("small")
    messages = state.get("messages", [])
    _log.info("starting — messages=%d", len(messages))

    system = SystemMessage(content=_SYSTEM)
    response = llm.invoke([system] + messages)

    content = response.content
    category = None
    location = None
    clean_lines = []
    for line in content.splitlines():
        stripped = line.strip()
        if stripped.startswith("CATEGORY:"):
            category = stripped.split(":", 1)[1].strip().lower()
        elif stripped.startswith("LOCATION:"):
            location = stripped.split(":", 1)[1].strip()
        else:
            clean_lines.append(line)

    display_text = "\n".join(clean_lines).strip()

    if location:
        _log.info("location extracted: %s", location)
    if category:
        _log.info("category detected: %s", category)

    if category:
        result = {
            "messages": [AIMessage(content=display_text)],
            "category": category,
            "collected_details": {},
            "email_cc": [],
            "conversation_complete": False,
            "draft_approved": False,
        }
        if location:
            result["location"] = location
        return result

    _log.info("no category yet — interrupting for user reply")
    partial = {}
    if location:
        partial["location"] = location
    user_reply = interrupt(display_text)
    return {
        "messages": [AIMessage(content=display_text), HumanMessage(content=user_reply)],
        **partial,
    }