kamidso-back/kamidso/native_connectors/result_store.py

"""Save raw connector responses where code_interpreter can read them later."""

from __future__ import annotations

import json
import os
from typing import Any, Dict, List, Optional, Sequence

from kamidso.native_connectors.models import ConnectorFunction, NativeProvider

_PREVIEW_RECORDS = 10
_CELL_CHARS = 120
_MAX_OBS_CHARS = 20_000


def get_nested(payload: Any, path: Optional[str]) -> Any:
    if not path or not isinstance(payload, dict):
        return None
    current: Any = payload
    for part in path.split("."):
        if not isinstance(current, dict) or part not in current:
            return None
        current = current[part]
    return current


def _json_default(value: Any) -> str:
    if isinstance(value, (bytes, bytearray)):
        return value.decode("utf-8", errors="replace")
    return str(value)


def _cell(value: Any) -> str:
    text = str(value).replace("\n", " ")
    if len(text) > _CELL_CHARS:
        return text[:_CELL_CHARS] + "…"
    return text


def materialize_connector_result(
    react_state: Dict[str, Any],
    response_json: Any,
    *,
    work_dir: Optional[str] = None,
) -> Dict[str, Any]:
    """Write the raw API JSON under ``connector_results/``.

    Returns metadata. The file itself stays the untouched API payload.
    """
    from kamidso.engine.app.openapi.api_v1.tools.sql_query import session_work_dir

    folder_root = work_dir or session_work_dir(react_state.get("conv_id"))
    out_dir = os.path.join(folder_root, "connector_results")
    os.makedirs(out_dir, exist_ok=True)
    seq = int(react_state.get("connector_result_seq") or 0) + 1
    react_state["connector_result_seq"] = seq
    name = f"c_{seq:03d}.json"
    payload = json.dumps(response_json, ensure_ascii=False, default=_json_default)
    path = os.path.join(out_dir, name)
    with open(path, "w", encoding="utf-8") as handle:
        handle.write(payload)
    with open(os.path.join(out_dir, "latest.json"), "w", encoding="utf-8") as handle:
        handle.write(payload)
    relative = f"connector_results/{name}"
    react_state["latest_connector_result"] = relative
    return {
        "result_id": f"connector_{seq:03d}",
        "relative_path": relative,
        "size_bytes": len(payload.encode("utf-8")),
    }


def _preview_line(index: int, record: Any) -> str:
    if isinstance(record, dict):
        preferred = [
            key
            for key in ("Title", "Name", "FullName", "Id", "UserId")
            if key in record
        ]
        if not preferred:
            preferred = list(record.keys())[:3]
        bits = [f"{key}={_cell(record.get(key))}" for key in preferred[:3]]
        return f"{index}. " + ", ".join(bits)
    return f"{index}. {_cell(record)}"


def _pagination_lines(
    function: ConnectorFunction,
    body: Dict[str, Any],
    payload: Any,
    records: Optional[Sequence[Any]],
) -> List[str]:
    pagination = function.pagination or {}
    if pagination.get("type") != "offset_limit":
        return []
    offset_field = str(pagination.get("offset_field") or "From")
    limit_field = str(pagination.get("limit_field") or "Limit")
    returned = len(records) if isinstance(records, list) else 0
    lines = [
        f"Current page: {offset_field}={body.get(offset_field)} {limit_field}={body.get(limit_field)}",
        f"Returned: {returned}",
    ]
    total = get_nested(payload, f"{function.response_root}.TotalCount")
    if isinstance(total, int):
        lines.append(f"TotalCount: {total}")
        offset = body.get(offset_field)
        if isinstance(offset, int) and offset + returned < total:
            lines.append("More data exists.")
    return lines


def render_success_observation(
    *,
    provider: NativeProvider,
    function: ConnectorFunction,
    body: Dict[str, Any],
    status_code: int,
    payload: Any,
    relative_path: str,
) -> str:
    lines = [
        f"{provider.display_name} / {function.name} completed successfully.",
        "",
        f"HTTP: {status_code}",
        "",
        "Summary:",
    ]
    summaries = False
    for path in function.summary_fields:
        value = get_nested(payload, path)
        if value is None:
            continue
        summaries = True
        label = path.split(".")[-1]
        lines.append(f"- {label}: {value}")
    if not summaries:
        lines.append("- (no summary fields)")
    if function.records_path and function.records_path.endswith(".List") and function.summary_fields:
        lines.append(
            "Response.List is only the current page. "
            "Use the aggregate counts above for totals."
        )

    records = get_nested(payload, function.records_path) if function.records_path else None
    if isinstance(records, list):
        lines.extend(["", f"Records in current page: {len(records)}", "", "Preview:"])
        if not records:
            lines.append("(empty)")
        for index, record in enumerate(records[:_PREVIEW_RECORDS], start=1):
            lines.append(_preview_line(index, record))
        if len(records) > _PREVIEW_RECORDS:
            lines.append(f"(Showing first {_PREVIEW_RECORDS} of {len(records)} records.)")
    else:
        root = get_nested(payload, function.response_root)
        preview = json.dumps(root if root is not None else payload, ensure_ascii=False, default=_json_default)
        if len(preview) > 1000:
            preview = preview[:1000] + "…"
        lines.extend(["", "Object preview:", preview])

    page_lines = _pagination_lines(function, body, payload, records if isinstance(records, list) else None)
    if page_lines:
        lines.extend([""] + page_lines)

    lines.extend(
        [
            "",
            f"[Saved full result → PLOT_DIR/{relative_path}]",
            "",
            "Use code_interpreter only if computation, dataframe operations, joining, "
            "statistics, charts or export are needed.",
        ]
    )
    text = "\n".join(lines)
    if len(text) > _MAX_OBS_CHARS:
        text = (
            text[:_MAX_OBS_CHARS]
            + f"\n\n... [Output truncated at {_MAX_OBS_CHARS} chars. "
            f"Full result at PLOT_DIR/{relative_path}]"
        )
    return text