Integrate Ollama into role detection with fallback to pattern matching

- 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
This commit is contained in:
2025-12-09 10:49:37 +00:00
parent 73dc5695ca
commit f56f8086e5
3 changed files with 248 additions and 8 deletions

View File

@@ -2,13 +2,20 @@
Role detection service for identifying distinct speakers in text.
"""
import re
from typing import List, Set, Dict
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
@@ -19,10 +26,23 @@ class RoleDetector:
BRACKET_PATTERN = r'\[([A-Z][a-zA-Z\s]+?)\]'
DEFAULT_ROLE = "Narrator"
def detect_roles(self, text: str) -> List[Dict[str, str]]:
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
@@ -33,15 +53,26 @@ class RoleDetector:
if not text or not text.strip():
return [{"id": self.DEFAULT_ROLE, "label": self.DEFAULT_ROLE}]
# Extract roles from both patterns
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)
# Combine and deduplicate
all_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, return default
# 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
@@ -50,8 +81,40 @@ class RoleDetector:
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.