wyndham-ARR/agent_integration/service.py

"""Conversation-oriented service built on top of the low-level Agent API client."""

from __future__ import annotations

import hashlib
import json
import uuid
from typing import Any, Dict, Iterator, Optional

from agent_integration.client import OpenAgentAPIClient, OpenAgentEvent
from agent_integration.sessions import SessionRecord, SessionStore


class SessionNotFoundError(LookupError):
    """Raised when a local conversation has no saved Agent session."""


class AgentResponseError(RuntimeError):
    """Raised when the Agent API omits a required response field."""


class OpenAgentService:
    """Maintain stable Agent sessions for external conversations."""

    def __init__(self, client: OpenAgentAPIClient, session_store: SessionStore) -> None:
        self.client = client
        self.session_store = session_store

    def ensure_session(
        self,
        conversation_id: str,
        *,
        external_subject_id: Optional[str] = None,
        metadata: Optional[Dict[str, Any]] = None,
    ) -> SessionRecord:
        conversation_id = self._required("conversation_id", conversation_id)
        existing = self.session_store.get(conversation_id)
        if existing is not None:
            return existing

        subject_id = self._required(
            "external_subject_id",
            external_subject_id if external_subject_id is not None else conversation_id,
        )
        metadata_value = self._metadata(metadata)
        response = self.client.create_session(
            external_subject_id=subject_id,
            idempotency_key=self.session_idempotency_key(conversation_id),
            metadata=metadata_value,
        )
        session_id = response.get("session_id")
        if not isinstance(session_id, str) or not session_id.strip():
            raise AgentResponseError("create-session response is missing session_id")

        return self.session_store.put(
            conversation_id=conversation_id,
            session_id=session_id,
            external_subject_id=subject_id,
            metadata=metadata_value,
        )

    def stream_message(
        self,
        conversation_id: str,
        message: str,
        *,
        message_id: Optional[str] = None,
        external_subject_id: Optional[str] = None,
        metadata: Optional[Dict[str, Any]] = None,
        include_trace: bool = False,
    ) -> Iterator[OpenAgentEvent]:
        conversation_id = self._required("conversation_id", conversation_id)
        message = self._message(message)
        session = self.ensure_session(
            conversation_id,
            external_subject_id=external_subject_id,
            metadata=metadata,
        )
        idempotency_key = self.message_idempotency_key(conversation_id, message_id)
        yield from self.client.stream_message(
            session.session_id,
            message,
            idempotency_key=idempotency_key,
            metadata=self._metadata(metadata),
            include_trace=include_trace,
        )

    def send_message(
        self,
        conversation_id: str,
        message: str,
        *,
        message_id: Optional[str] = None,
        external_subject_id: Optional[str] = None,
        metadata: Optional[Dict[str, Any]] = None,
    ) -> Dict[str, Any]:
        conversation_id = self._required("conversation_id", conversation_id)
        message = self._message(message)
        session = self.ensure_session(
            conversation_id,
            external_subject_id=external_subject_id,
            metadata=metadata,
        )
        return self.client.send_message(
            session.session_id,
            message,
            idempotency_key=self.message_idempotency_key(conversation_id, message_id),
            metadata=self._metadata(metadata),
        )

    def get_run(self, conversation_id: str, run_id: str) -> Dict[str, Any]:
        session = self.get_session(conversation_id)
        return self.client.get_run(session.session_id, self._required("run_id", run_id))

    def cancel_run(self, conversation_id: str, run_id: str) -> Dict[str, Any]:
        session = self.get_session(conversation_id)
        return self.client.cancel_run(session.session_id, self._required("run_id", run_id))

    def get_session(self, conversation_id: str) -> SessionRecord:
        conversation_id = self._required("conversation_id", conversation_id)
        session = self.session_store.get(conversation_id)
        if session is None:
            raise SessionNotFoundError(
                f"no Agent session is stored for conversation_id={conversation_id!r}"
            )
        return session

    def forget_session(self, conversation_id: str) -> bool:
        return self.session_store.delete(self._required("conversation_id", conversation_id))

    @classmethod
    def session_idempotency_key(cls, conversation_id: str) -> str:
        return cls._stable_key("sess", cls._required("conversation_id", conversation_id))

    @classmethod
    def message_idempotency_key(
        cls,
        conversation_id: str,
        message_id: Optional[str],
    ) -> str:
        message_token = message_id if message_id is not None else uuid.uuid4().hex
        return cls._stable_key(
            "msg",
            cls._required("conversation_id", conversation_id),
            cls._required("message_id", message_token),
        )

    @staticmethod
    def _stable_key(prefix: str, *parts: str) -> str:
        digest = hashlib.sha256("\x1f".join(parts).encode("utf-8")).hexdigest()
        return f"{prefix}_{digest[:48]}"

    @staticmethod
    def _required(name: str, value: str) -> str:
        normalized = str(value).strip()
        if not normalized:
            raise ValueError(f"{name} must not be empty")
        return normalized

    @staticmethod
    def _message(value: str) -> str:
        message = str(value)
        if not message.strip():
            raise ValueError("message must not be empty")
        return message

    @staticmethod
    def _metadata(value: Optional[Dict[str, Any]]) -> Dict[str, Any]:
        if value is None:
            return {}
        if not isinstance(value, dict):
            raise ValueError("metadata must be a JSON object")
        metadata = dict(value)
        try:
            json.dumps(metadata, ensure_ascii=False, allow_nan=False)
        except (TypeError, ValueError) as exc:
            raise ValueError("metadata must be strict JSON-serializable data") from exc
        return metadata