""" 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