src/agents/IntakeAgent.py

import pandas as pd
import json
from langchain_openai import ChatOpenAI
from langchain_core.prompts import PromptTemplate
from src.agents.CallState import CallState

class IntakeAgent:
    def __init__(self):
        self.llm = ChatOpenAI(model="gpt-4o", temperature=0)
        
    def __call__(self, state: CallState) -> CallState:
        """Validates input formats, cleans up profanity, extracts meta."""
        file_path = state["file_path"]
        file_type = state["file_type"]
        
        state["metadata"] = {"file_type": file_type, "file_path": file_path}
        
        if file_type == "csv":
            try:
                df = pd.read_csv(file_path)
                normalized = {str(c).strip().lower(): c for c in df.columns}
                required = {"id", "transcript"}
                missing = sorted(required - set(normalized.keys()))
                if missing:
                    state["metadata"]["intake_error"] = (
                        f"CSV is missing required headers: {', '.join(missing)}"
                    )
                    state["content"] = ""
                    state["clean_content"] = ""
                    return state

                id_col = normalized["id"]
                transcript_col = normalized["transcript"]
                transcripts = df[transcript_col].fillna("").astype(str).tolist()
                state["content"] = " ".join(transcripts).strip()
                state["metadata"]["row_count"] = int(len(df))
                state["metadata"]["ids_preview"] = (
                    df[id_col].fillna("").astype(str).head(20).tolist()
                )
            except Exception as e:
                state["metadata"]["intake_error"] = str(e)
                state["content"] = ""
                state["clean_content"] = ""
                return state

        if file_type == "json":
            try:
                with open(file_path, "r", encoding="utf-8") as f:
                    raw = f.read()

                try:
                    payload = json.loads(raw)
                except json.JSONDecodeError:
                    # Common failure mode: users paste multi-line transcripts with raw control characters
                    # (e.g., literal newlines) inside JSON strings. JSON requires these to be escaped.
                    def _sanitize_control_chars_in_strings(text: str) -> str:
                        out: list[str] = []
                        in_string = False
                        escape = False
                        for ch in text:
                            if not in_string:
                                out.append(ch)
                                if ch == '"':
                                    in_string = True
                                continue

                            # Inside string
                            if escape:
                                out.append(ch)
                                escape = False
                                continue
                            if ch == "\\":
                                out.append(ch)
                                escape = True
                                continue
                            if ch == '"':
                                out.append(ch)
                                in_string = False
                                continue

                            code = ord(ch)
                            if ch == "\n":
                                out.append("\\n")
                            elif ch == "\r":
                                out.append("\\r")
                            elif ch == "\t":
                                out.append("\\t")
                            elif code < 0x20:
                                out.append(f"\\u{code:04x}")
                            else:
                                out.append(ch)
                        return "".join(out)

                    payload = json.loads(_sanitize_control_chars_in_strings(raw))

                transcripts: list[str] = []
                ids_preview: list[str] = []

                def add_item(item: object) -> None:
                    if not isinstance(item, dict):
                        return
                    transcript = item.get("transcript")
                    if isinstance(transcript, str) and transcript.strip():
                        transcripts.append(transcript.strip())
                    item_id = item.get("id")
                    if item_id is not None and len(ids_preview) < 20:
                        ids_preview.append(str(item_id))

                if isinstance(payload, list):
                    for item in payload:
                        add_item(item)
                elif isinstance(payload, dict):
                    if isinstance(payload.get("transcript"), str):
                        add_item(payload)
                    elif isinstance(payload.get("transcripts"), list):
                        for item in payload["transcripts"]:
                            add_item(item)
                    elif isinstance(payload.get("calls"), list):
                        for item in payload["calls"]:
                            add_item(item)
                    else:
                        state["metadata"]["intake_error"] = (
                            "JSON must be either a list of {id, transcript} objects, "
                            "a single {id, transcript} object, or a dict with 'transcripts'/'calls' arrays."
                        )
                        state["content"] = ""
                        state["clean_content"] = ""
                        return state
                else:
                    state["metadata"]["intake_error"] = "Unsupported JSON root type."
                    state["content"] = ""
                    state["clean_content"] = ""
                    return state

                if not transcripts:
                    state["metadata"]["intake_error"] = "JSON contained no non-empty 'transcript' fields."
                    state["content"] = ""
                    state["clean_content"] = ""
                    return state

                state["content"] = " ".join(transcripts).strip()
                state["metadata"]["row_count"] = int(len(transcripts))
                if ids_preview:
                    state["metadata"]["ids_preview"] = ids_preview
            except Exception as e:
                state["metadata"]["intake_error"] = str(e)
                state["content"] = ""
                state["clean_content"] = ""
                return state

        if state.get("content"):
            prompt = PromptTemplate.from_template(
                "Clean the following text of any profanity and fix basic grammatical errors. Return only the clean text:\n{text}"
            )
            chain = prompt | self.llm
            clean_text = chain.invoke({"text": state["content"]}).content
            state["clean_content"] = clean_text
            
        return state