aionboarding/ai-patient-onboarding/app/ai/sarvam_client.py

import base64
import time
import asyncio
import logging
import httpx2
from typing import Tuple, Optional
from app.config import SARVAM_API_KEY

logger = logging.getLogger("sarvam_client")

# Status Constants
STT_SUCCESS = "STT_SUCCESS"
STT_EMPTY_AUDIO = "STT_EMPTY_AUDIO"
STT_TIMEOUT = "STT_TIMEOUT"
STT_PROVIDER_ERROR = "STT_PROVIDER_ERROR"
STT_NETWORK_ERROR = "STT_NETWORK_ERROR"

TTS_SUCCESS = "TTS_SUCCESS"
TTS_TIMEOUT = "TTS_TIMEOUT"
TTS_PROVIDER_ERROR = "TTS_PROVIDER_ERROR"
TTS_NETWORK_ERROR = "TTS_NETWORK_ERROR"


class SarvamClient:

    _client = None

    def __init__(self):
        self.api_key = SARVAM_API_KEY
        self.base_url = "https://api.sarvam.ai"

    def _get_client(self) -> httpx2.AsyncClient:
        if SarvamClient._client is None:
            SarvamClient._client = httpx2.AsyncClient(
                timeout=8.0,
                limits=httpx2.Limits(max_keepalive_connections=5, max_connections=10),
            )
        return SarvamClient._client

    async def speech_to_text_with_status(
        self,
        audio_bytes: bytes,
        filename: str = "audio.webm",
        language_code: str = "ta-IN",
        session_id: Optional[str] = None,
        max_retries: int = 2,
    ) -> Tuple[str, str]:
        """
        Transcribes audio to text returning (transcript, status_code).
        Uses bounded retries for transient errors.
        """
        if not audio_bytes:
            return "", STT_EMPTY_AUDIO

        if not self.api_key or self.api_key in ["your_sarvam_key", ""]:
            return "", STT_PROVIDER_ERROR

        client = self._get_client()
        files = {"file": (filename, audio_bytes, "audio/webm")}
        data = {"model": "saaras:v2", "language_code": language_code}
        headers = {"api-subscription-key": self.api_key}

        attempt = 0
        last_status = STT_PROVIDER_ERROR

        while attempt <= max_retries:
            start_time = time.time()
            try:
                res = await client.post(
                    f"{self.base_url}/speech-to-text",
                    files=files,
                    data=data,
                    headers=headers,
                )
                duration = round(time.time() - start_time, 3)

                if res.status_code == 200:
                    transcript = res.json().get("transcript", "").strip()
                    logger.info(
                        "sarvam_stt_success",
                        extra={
                            "provider": "SarvamAI",
                            "operation": "STT",
                            "session_id": session_id,
                            "duration_sec": duration,
                            "attempt": attempt + 1,
                        },
                    )
                    return transcript, STT_SUCCESS
                elif 400 <= res.status_code < 500:
                    # Non-retryable 4xx client errors
                    logger.warning(
                        "sarvam_stt_client_error",
                        extra={
                            "provider": "SarvamAI",
                            "operation": "STT",
                            "session_id": session_id,
                            "http_status": res.status_code,
                        },
                    )
                    return "", STT_PROVIDER_ERROR
                else:
                    last_status = STT_PROVIDER_ERROR

            except (httpx2.TimeoutException, asyncio.TimeoutError):
                last_status = STT_TIMEOUT
                logger.warning("sarvam_stt_timeout", extra={"session_id": session_id, "attempt": attempt + 1})
            except (httpx2.NetworkError, Exception) as e:
                last_status = STT_NETWORK_ERROR
                logger.warning("sarvam_stt_network_error", extra={"session_id": session_id, "attempt": attempt + 1})

            attempt += 1
            if attempt <= max_retries:
                await asyncio.sleep(0.2 * attempt)

        return "", last_status

    async def speech_to_text(
        self,
        audio_bytes: bytes,
        filename: str = "audio.webm",
        language_code: str = "ta-IN",
    ) -> str:
        transcript, _ = await self.speech_to_text_with_status(audio_bytes, filename, language_code)
        return transcript

    async def text_to_speech_with_status(
        self,
        text: str,
        language_code: str = "ta-IN",
        session_id: Optional[str] = None,
        max_retries: int = 2,
    ) -> Tuple[Optional[bytes], str]:
        """
        Synthesizes text to speech returning (audio_bytes, status_code).
        Uses bounded retries for transient errors.
        """
        if not text or not text.strip():
            return None, TTS_PROVIDER_ERROR

        if not self.api_key or self.api_key in ["your_sarvam_key", ""]:
            return None, TTS_PROVIDER_ERROR

        client = self._get_client()
        headers = {
            "api-subscription-key": self.api_key,
            "Content-Type": "application/json",
        }
        payload = {
            "inputs": [text],
            "target_language_code": language_code,
            "speaker": "kavya",
            "model": "bulbul:v2",
        }

        attempt = 0
        last_status = TTS_PROVIDER_ERROR

        while attempt <= max_retries:
            start_time = time.time()
            try:
                res = await client.post(
                    f"{self.base_url}/text-to-speech",
                    json=payload,
                    headers=headers,
                )
                duration = round(time.time() - start_time, 3)

                if res.status_code == 200:
                    audios = res.json().get("audios", [])
                    if audios:
                        logger.info(
                            "sarvam_tts_success",
                            extra={
                                "provider": "SarvamAI",
                                "operation": "TTS",
                                "session_id": session_id,
                                "duration_sec": duration,
                                "attempt": attempt + 1,
                            },
                        )
                        return base64.b64decode(audios[0]), TTS_SUCCESS
                elif 400 <= res.status_code < 500:
                    return None, TTS_PROVIDER_ERROR
                else:
                    last_status = TTS_PROVIDER_ERROR

            except (httpx2.TimeoutException, asyncio.TimeoutError):
                last_status = TTS_TIMEOUT
            except (httpx2.NetworkError, Exception):
                last_status = TTS_NETWORK_ERROR

            attempt += 1
            if attempt <= max_retries:
                await asyncio.sleep(0.2 * attempt)

        return None, last_status

    async def text_to_speech(
        self,
        text: str,
        language_code: str = "ta-IN",
    ) -> Optional[bytes]:
        audio, _ = await self.text_to_speech_with_status(text, language_code)
        return audio