Paul Bailey
73dc5695ca
- Sample up to 4 chunks of 5000 characters from different text positions - Distribute samples evenly across beginning, middle, and end of text - Deduplicate speakers across all samples for comprehensive detection - Log individual sample results for debugging - Handle short texts efficiently (single sample) - Updated tests to verify multi-sample behavior This ensures speakers introduced anywhere in long texts are detected, not just those appearing in the first 2000 characters.
231 lines
8.1 KiB
Python
231 lines
8.1 KiB
Python
"""
|
|
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.
|
|
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
|