Paul Bailey
f56f8086e5
- Updated RoleDetector to use OllamaService for AI-based speaker detection - Implemented _detect_with_ollama() method with availability check - Added fallback logic: tries Ollama first, then pattern matching - Implemented role merging and deduplication across both methods - Handles Ollama unavailable gracefully with logging - Logs detection method used (AI vs pattern) - Updated main.py endpoint to await async detect_roles - Added comprehensive integration tests Requirements: 14.1, 14.4, 14.5
157 lines
5.2 KiB
Python
157 lines
5.2 KiB
Python
"""
|
|
Role detection service for identifying distinct speakers in text.
|
|
"""
|
|
import re
|
|
import logging
|
|
from typing import List, Set, Dict, Optional
|
|
from .ollama_service import OllamaService
|
|
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class RoleDetector:
|
|
"""
|
|
Service for detecting distinct speaker roles in text.
|
|
|
|
Uses AI-based detection (Ollama) as primary method, with fallback to pattern matching.
|
|
|
|
Supports two pattern types:
|
|
1. Colon pattern: "RoleName:" at the start of a line
|
|
2. Bracket pattern: "[RoleName]" preceding dialogue
|
|
"""
|
|
|
|
# Regex patterns for role detection
|
|
COLON_PATTERN = r'^([A-Z][a-zA-Z\s]+?):\s'
|
|
BRACKET_PATTERN = r'\[([A-Z][a-zA-Z\s]+?)\]'
|
|
DEFAULT_ROLE = "Narrator"
|
|
|
|
def __init__(self, ollama_service: Optional[OllamaService] = None):
|
|
"""
|
|
Initializes RoleDetector with optional OllamaService.
|
|
|
|
Args:
|
|
ollama_service: Optional OllamaService instance for AI-based detection.
|
|
If None, creates a new instance.
|
|
"""
|
|
self.ollama_service = ollama_service or OllamaService()
|
|
|
|
async def detect_roles(self, text: str) -> List[Dict[str, str]]:
|
|
"""
|
|
Analyzes text to identify distinct speaker roles.
|
|
|
|
Uses AI-based detection (Ollama) first, then falls back to pattern matching.
|
|
Merges and deduplicates results from both methods.
|
|
|
|
Args:
|
|
text: Input text to analyze
|
|
|
|
Returns:
|
|
List of role dictionaries with 'id' and 'label' fields.
|
|
Returns default "Narrator" role if no patterns found.
|
|
"""
|
|
if not text or not text.strip():
|
|
return [{"id": self.DEFAULT_ROLE, "label": self.DEFAULT_ROLE}]
|
|
|
|
all_roles = set()
|
|
|
|
# Try AI-based detection first
|
|
ollama_roles = await self._detect_with_ollama(text)
|
|
if ollama_roles:
|
|
all_roles.update(ollama_roles)
|
|
logger.info(f"AI detection found {len(ollama_roles)} roles: {sorted(ollama_roles)}")
|
|
|
|
# Always run pattern matching as well (for marked text or as fallback)
|
|
colon_roles = self._extract_colon_pattern(text)
|
|
bracket_roles = self._extract_bracket_pattern(text)
|
|
pattern_roles = colon_roles.union(bracket_roles)
|
|
|
|
if pattern_roles:
|
|
all_roles.update(pattern_roles)
|
|
logger.info(f"Pattern detection found {len(pattern_roles)} roles: {sorted(pattern_roles)}")
|
|
|
|
# If no roles found by either method, return default
|
|
if not all_roles:
|
|
logger.info("No roles detected, using default Narrator")
|
|
return [{"id": self.DEFAULT_ROLE, "label": self.DEFAULT_ROLE}]
|
|
|
|
# Convert to role dictionaries and sort for consistent ordering
|
|
role_list = [
|
|
{"id": role.strip(), "label": role.strip()}
|
|
for role in sorted(all_roles)
|
|
]
|
|
|
|
logger.info(f"Total unique roles after merging: {len(role_list)}")
|
|
return role_list
|
|
|
|
async def _detect_with_ollama(self, text: str) -> Set[str]:
|
|
"""
|
|
Uses Ollama with llama3:8B to intelligently detect speakers.
|
|
|
|
Args:
|
|
text: Input text to analyze
|
|
|
|
Returns:
|
|
Set of speaker names or empty set if Ollama unavailable/fails.
|
|
"""
|
|
try:
|
|
# Check if Ollama is available
|
|
is_available = await self.ollama_service.is_available()
|
|
if not is_available:
|
|
logger.info("Ollama service not available, skipping AI detection")
|
|
return set()
|
|
|
|
# Detect speakers using Ollama
|
|
speakers = await self.ollama_service.detect_speakers(text)
|
|
|
|
if speakers:
|
|
logger.info(f"Ollama detected speakers: {speakers}")
|
|
return set(speakers)
|
|
else:
|
|
logger.info("Ollama returned no speakers")
|
|
return set()
|
|
|
|
except Exception as e:
|
|
logger.warning(f"Ollama detection failed: {e}, falling back to pattern matching")
|
|
return set()
|
|
|
|
def _extract_colon_pattern(self, text: str) -> Set[str]:
|
|
"""
|
|
Extracts roles from "RoleName:" pattern at line starts.
|
|
|
|
Args:
|
|
text: Input text to analyze
|
|
|
|
Returns:
|
|
Set of role names found
|
|
"""
|
|
roles = set()
|
|
pattern = re.compile(self.COLON_PATTERN, re.MULTILINE)
|
|
|
|
for match in pattern.finditer(text):
|
|
role_name = match.group(1).strip()
|
|
if role_name:
|
|
roles.add(role_name)
|
|
|
|
return roles
|
|
|
|
def _extract_bracket_pattern(self, text: str) -> Set[str]:
|
|
"""
|
|
Extracts roles from "[RoleName]" pattern.
|
|
|
|
Args:
|
|
text: Input text to analyze
|
|
|
|
Returns:
|
|
Set of role names found
|
|
"""
|
|
roles = set()
|
|
pattern = re.compile(self.BRACKET_PATTERN)
|
|
|
|
for match in pattern.finditer(text):
|
|
role_name = match.group(1).strip()
|
|
if role_name:
|
|
roles.add(role_name)
|
|
|
|
return roles
|