voiceagent/src/voice_assistant/agent/oci.py

"""
OCI Generative AI agent integration.

Provides parsing and streaming support for OCI Generative AI
chat inference using the GenericChatRequest format.
"""

from __future__ import annotations

import asyncio
import json
from typing import Any

import logging
from collections.abc import AsyncGenerator
from contextlib import aclosing

import oci

from voice_assistant.agent.adapter import StreamingAgentAdapter
from voice_assistant.agent.oci_stream import stream_sse_events
from voice_assistant.agent.oci_connection import open_oci_stream


def extract_text_delta(raw_data: str) -> str | None:
    """
    Extract incremental text from an OCI Generative AI SSE event.

    Returns:
        Text delta when available.
        None for empty or non-text events.

    Raises:
        ValueError for malformed JSON or invalid event structures.
    """
    if raw_data == "[DONE]":
        return None

    try:
        payload: Any = json.loads(raw_data)
    except json.JSONDecodeError as exc:
        raise ValueError(
            "Invalid JSON in OCI Generative AI SSE event."
        ) from exc

    if not isinstance(payload, dict):
        raise ValueError(
            "OCI Generative AI SSE event must be an object."
        )

    message = payload.get("message")

    if message is None:
        return None

    if not isinstance(message, dict):
        raise ValueError(
            "Invalid OCI Generative AI message structure."
        )

    content = message.get("content", [])

    if not isinstance(content, list):
        raise ValueError(
            "Invalid OCI Generative AI content structure."
        )

    deltas: list[str] = []

    for item in content:
        if not isinstance(item, dict):
            raise ValueError(
                "Invalid OCI Generative AI content item."
            )

        if item.get("type") != "TEXT":
            continue

        text = item.get("text")

        if not isinstance(text, str):
            raise ValueError(
                "Invalid OCI Generative AI text delta."
            )

        if text:
            deltas.append(text)

    return "".join(deltas) or None


logger = logging.getLogger(__name__)


class OCIAgentAdapter(StreamingAgentAdapter):
    """
    OCI Generative AI streaming Agent.

    Responsibilities:
        - Authenticate using Instance Principal.
        - Build OCI chat requests.
        - Execute inference outside the event loop.
        - Stream text deltas incrementally.
        - Close SSE resources after consumption.

    Non-responsibilities:
        - Managing semantic turn state.
        - Sending WebSocket events.
        - Performing TTS synthesis.
    """

    def __init__(
        self,
        region: str,
        compartment_id: str,
        model_id: str = "openai.gpt-4o-mini",
        temperature: float = 0.3,
        max_completion_tokens: int = 256,
        timeout_seconds: float = 60.0,
    ) -> None:
        if not region:
            raise ValueError("OCI region is required.")

        if not compartment_id:
            raise ValueError("OCI compartment_id is required.")

        if not model_id:
            raise ValueError("OCI model_id is required.")

        if timeout_seconds <= 0:
            raise ValueError("timeout_seconds must be positive.")

        self.region = region
        self.compartment_id = compartment_id
        self.model_id = model_id
        self.temperature = temperature
        self.max_completion_tokens = max_completion_tokens
        self.timeout_seconds = timeout_seconds
        self._closed = False

        signer = (
            oci.auth.signers
            .InstancePrincipalsSecurityTokenSigner()
        )

        self.client = (
            oci.generative_ai_inference
            .GenerativeAiInferenceClient(
                config={"region": region},
                signer=signer,
                timeout=(10, timeout_seconds),
            )
        )

    def _build_chat_details(
        self,
        text: str,
    ):
        models = oci.generative_ai_inference.models

        chat_request = models.GenericChatRequest(
            messages=[
                models.UserMessage(
                    content=[
                        models.TextContent(
                            text=text,
                        )
                    ]
                )
            ],
            is_stream=True,
            temperature=self.temperature,
            max_completion_tokens=self.max_completion_tokens,
        )

        return models.ChatDetails(
            compartment_id=self.compartment_id,
            serving_mode=models.OnDemandServingMode(
                model_id=self.model_id,
            ),
            chat_request=chat_request,
        )

    async def stream(
        self,
        session_id: str,
        turn_id: int,
        text: str,
    ) -> AsyncGenerator[str, None]:
        """
        Stream OCI Generative AI text deltas.

        The synchronous OCI request is executed outside the
        event loop. SSE events are consumed through the
        asynchronous bridge.
        """
        if self._closed:
            raise RuntimeError("OCI Agent is closed.")

        if not text.strip():
            raise ValueError("Agent input text cannot be empty.")

        details = self._build_chat_details(text)

        logger.debug(
            "Starting OCI Agent stream. "
            "session_id=%s turn_id=%s model=%s",
            session_id,
            turn_id,
            self.model_id,
        )

        response = await open_oci_stream(
            lambda: self.client.chat(details)
        )

        if response.status != 200:
            raise RuntimeError(
                f"OCI inference returned HTTP {response.status}."
            )

        sse = response.data

        async with aclosing(
            stream_sse_events(sse)
        ) as events:
            async for raw_data in events:
                delta = extract_text_delta(raw_data)

                if delta:
                    yield delta

    async def process(
        self,
        session_id: str,
        turn_id: int,
        text: str,
    ) -> str:
        """
        Collect the complete response for compatibility
        with the existing AgentAdapter contract.
        """
        parts: list[str] = []

        async with aclosing(
            self.stream(
                session_id=session_id,
                turn_id=turn_id,
                text=text,
            )
        ) as deltas:
            async for delta in deltas:
                parts.append(delta)

        response_text = "".join(parts).strip()

        if not response_text:
            raise RuntimeError(
                "OCI Agent returned an empty response."
            )

        return response_text


    async def aclose(self) -> None:
        """Release the OCI HTTP session after turn execution has stopped."""
        if self._closed:
            return
        await asyncio.to_thread(self.client.base_client.session.close)
        self._closed = True