Paul Bailey
d18fddef57
- Added segment_text() method to OllamaService for AI-based text splitting - OllamaService now annotates each text segment with the appropriate speaker - Updated TTSService to use AI segmentation when available - Falls back to pattern-based segmentation when Ollama unavailable - AI intelligently assigns dialogue to speakers and narrative to Narrator - Added comprehensive test suite for AI segmentation - Verified AI correctly splits unmarked dialogue between multiple speakers This makes the tool much more powerful - users can now input plain narrative text without explicit role markers and the AI will intelligently assign voices to different characters and the narrator.
370 lines
13 KiB
Python
370 lines
13 KiB
Python
"""
|
|
Ollama integration service for AI-based speaker detection and text segmentation.
|
|
"""
|
|
import json
|
|
import logging
|
|
import os
|
|
from typing import List, Optional, Dict, Tuple
|
|
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 segment_text(self, text: str, speakers: List[str]) -> List[Tuple[str, str]]:
|
|
"""
|
|
Uses AI to segment text and assign each segment to a speaker.
|
|
|
|
Args:
|
|
text: Input text to segment
|
|
speakers: List of known speakers in the text
|
|
|
|
Returns:
|
|
List of tuples: (speaker_name, text_segment)
|
|
|
|
Raises:
|
|
Exception: If Ollama is unavailable or request fails
|
|
"""
|
|
try:
|
|
async with httpx.AsyncClient(timeout=60.0) as client: # Longer timeout for segmentation
|
|
prompt = self._build_segmentation_prompt(text, speakers)
|
|
|
|
logger.info(f"Requesting AI segmentation for {len(text)} chars with {len(speakers)} speakers")
|
|
|
|
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", "")
|
|
|
|
segments = self._parse_segmentation_response(response_text)
|
|
logger.info(f"AI segmentation produced {len(segments)} segments")
|
|
|
|
return segments
|
|
|
|
except httpx.TimeoutException:
|
|
logger.warning("Ollama segmentation request timed out")
|
|
raise Exception("Ollama segmentation timed out")
|
|
except httpx.HTTPError as e:
|
|
logger.warning(f"Ollama HTTP error during segmentation: {e}")
|
|
raise Exception(f"Ollama service error: {e}")
|
|
except Exception as e:
|
|
logger.warning(f"Ollama segmentation failed: {e}")
|
|
raise
|
|
|
|
def _build_segmentation_prompt(self, text: str, speakers: List[str]) -> str:
|
|
"""
|
|
Constructs prompt for text segmentation with speaker attribution.
|
|
|
|
Args:
|
|
text: Input text to segment
|
|
speakers: List of known speakers
|
|
|
|
Returns:
|
|
Formatted prompt string
|
|
"""
|
|
speakers_str = ", ".join(f'"{s}"' for s in speakers)
|
|
|
|
prompt = f"""You are analyzing a story or dialogue. Your task is to split the text into segments and assign each segment to the appropriate speaker.
|
|
|
|
Known speakers: {speakers_str}
|
|
|
|
Rules:
|
|
1. Split the text into logical segments (sentences or short paragraphs)
|
|
2. Assign each segment to the speaker who is speaking or being described
|
|
3. Use "Narrator" for narrative text, descriptions, or text that doesn't belong to a specific character
|
|
4. If a character is mentioned but not speaking, use "Narrator"
|
|
5. Only use speakers from the known speakers list
|
|
|
|
Return a JSON object with a "segments" array. Each segment should have:
|
|
- "speaker": the name of the speaker (must be from the known speakers list or "Narrator")
|
|
- "text": the text content for that segment
|
|
|
|
Example format:
|
|
{{
|
|
"segments": [
|
|
{{"speaker": "Narrator", "text": "The old wizard gazed into his crystal ball."}},
|
|
{{"speaker": "Wizard", "text": "I see danger approaching."}},
|
|
{{"speaker": "Narrator", "text": "A young knight heard the warning."}}
|
|
]
|
|
}}
|
|
|
|
Text to segment:
|
|
{text}
|
|
|
|
Return ONLY the JSON object with the segments array.
|
|
"""
|
|
return prompt
|
|
|
|
def _parse_segmentation_response(self, response: str) -> List[Tuple[str, str]]:
|
|
"""
|
|
Parses Ollama segmentation response to extract speaker-text pairs.
|
|
|
|
Args:
|
|
response: Response text from Ollama
|
|
|
|
Returns:
|
|
List of tuples: (speaker_name, text_content)
|
|
"""
|
|
if not response or not response.strip():
|
|
return []
|
|
|
|
try:
|
|
# Try to parse as JSON
|
|
data = json.loads(response.strip())
|
|
|
|
# Extract segments array
|
|
segments_data = []
|
|
if isinstance(data, dict):
|
|
if "segments" in data:
|
|
segments_data = data["segments"]
|
|
else:
|
|
# Try to find any array in the response
|
|
for value in data.values():
|
|
if isinstance(value, list):
|
|
segments_data = value
|
|
break
|
|
elif isinstance(data, list):
|
|
segments_data = data
|
|
|
|
# Convert to tuples
|
|
segments = []
|
|
for seg in segments_data:
|
|
if isinstance(seg, dict) and "speaker" in seg and "text" in seg:
|
|
speaker = str(seg["speaker"]).strip()
|
|
text = str(seg["text"]).strip()
|
|
if speaker and text:
|
|
segments.append((speaker, text))
|
|
|
|
return segments
|
|
|
|
except json.JSONDecodeError:
|
|
logger.warning(f"Failed to parse Ollama segmentation response as JSON: {response[:200]}")
|
|
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
|