audiobook/backend/services/ollama_service.py

"""
Ollama integration service for AI-based speaker detection.
"""
import json
import logging
import os
from typing import List, Optional
import httpx


logger = logging.getLogger(__name__)


class OllamaService:
    """
    Service for integrating with Ollama API for AI-based speaker detection.
    
    Uses llama3:8b model to intelligently detect speakers in text without
    explicit role markers.
    """
    
    def __init__(self, base_url: Optional[str] = None):
        """
        Initializes Ollama service with base URL.
        
        Args:
            base_url: Ollama API base URL. Defaults to OLLAMA_BASE_URL env var
                     or http://localhost:11434 if not set.
        """
        self.base_url = base_url or os.getenv("OLLAMA_BASE_URL", "http://localhost:11434")
        self.model = "llama3:8b"
        self.timeout = 30.0  # 30 seconds timeout
        
    async def detect_speakers(self, text: str) -> List[str]:
        """
        Sends text to Ollama for speaker detection.
        
        Args:
            text: Input text to analyze for speakers
            
        Returns:
            List of detected speaker names
            
        Raises:
            Exception: If Ollama is unavailable or request fails
        """
        prompt = self._build_prompt(text)
        
        try:
            async with httpx.AsyncClient(timeout=self.timeout) as client:
                response = await client.post(
                    f"{self.base_url}/api/generate",
                    json={
                        "model": self.model,
                        "prompt": prompt,
                        "stream": False,
                        "format": "json"
                    }
                )
                response.raise_for_status()
                
                result = response.json()
                response_text = result.get("response", "")
                
                speakers = self._parse_response(response_text)
                logger.info(f"Ollama detected {len(speakers)} speakers: {speakers}")
                return speakers
                
        except httpx.TimeoutException:
            logger.warning("Ollama request timed out after 30 seconds")
            raise Exception("Ollama request timed out")
        except httpx.HTTPError as e:
            logger.warning(f"Ollama HTTP error: {e}")
            raise Exception(f"Ollama service error: {e}")
        except Exception as e:
            logger.warning(f"Ollama detection failed: {e}")
            raise
    
    def _build_prompt(self, text: str) -> str:
        """
        Constructs prompt for speaker detection.
        
        Instructs model to return JSON array of speaker names.
        
        Args:
            text: Input text to analyze
            
        Returns:
            Formatted prompt string
        """
        prompt = f"""Analyze the following text and identify all distinct speakers or characters.
Return ONLY a JSON object with a "speakers" array containing the names of speakers.
If there are no clear speakers, return {{"speakers": ["Narrator"]}}.

Examples:
- For dialogue with character names, list those characters
- For narrative text, use "Narrator"
- For mixed content, include both characters and "Narrator" if appropriate

Text to analyze:
{text[:2000]}

Return format: {{"speakers": ["Speaker1", "Speaker2", ...]}}
"""
        return prompt
    
    def _parse_response(self, response: str) -> List[str]:
        """
        Parses Ollama response to extract speaker names.
        
        Handles both JSON and plain text responses.
        
        Args:
            response: Response text from Ollama
            
        Returns:
            List of speaker names
        """
        if not response or not response.strip():
            return []
        
        try:
            # Try to parse as JSON
            data = json.loads(response.strip())
            
            # Handle different possible JSON structures
            if isinstance(data, dict):
                if "speakers" in data:
                    speakers = data["speakers"]
                elif "characters" in data:
                    speakers = data["characters"]
                elif "roles" in data:
                    speakers = data["roles"]
                else:
                    # Try to find any array in the response
                    for value in data.values():
                        if isinstance(value, list):
                            speakers = value
                            break
                    else:
                        return []
            elif isinstance(data, list):
                speakers = data
            else:
                return []
            
            # Ensure we have a list of strings
            if isinstance(speakers, list):
                return [str(s).strip() for s in speakers if s]
            
            return []
            
        except json.JSONDecodeError:
            logger.warning(f"Failed to parse Ollama response as JSON: {response[:100]}")
            # Try to extract speaker names from plain text
            # Look for common patterns like lists or quoted names
            return []
    
    async def is_available(self) -> bool:
        """
        Checks if Ollama service is available.
        
        Returns:
            True if service responds to health check, False otherwise
        """
        try:
            async with httpx.AsyncClient(timeout=5.0) as client:
                response = await client.get(f"{self.base_url}/api/tags")
                return response.status_code == 200
        except Exception as e:
            logger.debug(f"Ollama availability check failed: {e}")
            return False