""" 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:latest" # Using llama3:latest (8B model) self.timeout = 30.0 # 30 seconds timeout async def detect_speakers(self, text: str) -> List[str]: """ Sends text to Ollama for speaker detection. Samples up to 4 chunks of 5000 characters from different parts of the text to ensure comprehensive speaker detection in long texts. Args: text: Input text to analyze for speakers Returns: List of detected speaker names (deduplicated across all samples) Raises: Exception: If Ollama is unavailable or request fails """ # Get text samples from different parts of the text samples = self._sample_text(text, chunk_size=5000, num_samples=4) all_speakers = set() try: async with httpx.AsyncClient(timeout=self.timeout) as client: # Process each sample for i, sample in enumerate(samples): prompt = self._build_prompt(sample) logger.debug(f"Sending sample {i+1}/{len(samples)} to Ollama ({len(sample)} chars)") 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"Sample {i+1} detected {len(speakers)} speakers: {speakers}") # Add to our set (automatically deduplicates) all_speakers.update(speakers) speaker_list = sorted(list(all_speakers)) logger.info(f"Ollama detected {len(speaker_list)} total unique speakers: {speaker_list}") return speaker_list 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 _sample_text(self, text: str, chunk_size: int = 5000, num_samples: int = 4) -> List[str]: """ Samples text from different parts to ensure comprehensive speaker detection. For short texts, returns the full text. For long texts, samples from beginning, middle sections, and end. Args: text: Input text to sample chunk_size: Size of each sample in characters num_samples: Number of samples to extract Returns: List of text samples """ text_len = len(text) # If text is shorter than one chunk, return it as-is if text_len <= chunk_size: return [text] # If text is shorter than what we'd sample, return fewer samples max_possible_samples = (text_len // chunk_size) + 1 actual_samples = min(num_samples, max_possible_samples) samples = [] if actual_samples == 1: # Just take the beginning samples.append(text[:chunk_size]) else: # Distribute samples evenly across the text step = text_len // actual_samples for i in range(actual_samples): start = i * step end = min(start + chunk_size, text_len) samples.append(text[start:end]) return samples 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 (should already be sampled/truncated) 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} 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