Paul Bailey
63bdf57593
Updated the speaker detection prompt to explicitly instruct the AI to: - Identify narrative text and assign it to 'Narrator' - Detect character dialogue and list character names - Include BOTH Narrator and characters when story has mixed content This ensures stories with narrative descriptions get proper voice assignments for both the narrator and speaking characters. Example: Story with 'Once upon a time...' narrative and dialogue from Lily, Benny, Sammy now correctly detects all 4 roles.
377 lines
14 KiB
Python
377 lines
14 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 who need voice assignments.
|
|
|
|
CRITICAL: Look for TWO types of content:
|
|
1. NARRATIVE TEXT: Descriptions, scene-setting, "he said", "she thought", etc. → This needs "Narrator"
|
|
2. CHARACTER DIALOGUE: Direct speech by characters → These need character names
|
|
|
|
You MUST include "Narrator" if there is ANY narrative text, descriptions, or non-dialogue content.
|
|
|
|
Return ONLY a JSON object with a "speakers" array.
|
|
|
|
Examples:
|
|
- Text: "Alice said hello. Bob waved back." → {{"speakers": ["Narrator", "Alice", "Bob"]}}
|
|
- Text: "Once upon a time, there was a princess." → {{"speakers": ["Narrator"]}}
|
|
- Text: "The wizard spoke. 'I see danger,' he warned." → {{"speakers": ["Narrator"]}}
|
|
- Text: "Alice: Hello! Bob: Hi there!" → {{"speakers": ["Alice", "Bob"]}}
|
|
|
|
Text to analyze:
|
|
{text}
|
|
|
|
Return format: {{"speakers": ["Narrator", "Character1", "Character2", ...]}}
|
|
"""
|
|
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
|