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:
@@ -62,7 +62,7 @@ app = FastAPI(
|
|||||||
# Configure CORS
|
# Configure CORS
|
||||||
app.add_middleware(
|
app.add_middleware(
|
||||||
CORSMiddleware,
|
CORSMiddleware,
|
||||||
allow_origins=["http://localhost:5173", "http://localhost:3000"],
|
allow_origins=["http://localhost:5173", "http://localhost:5174", "http://localhost:3000"],
|
||||||
allow_credentials=True,
|
allow_credentials=True,
|
||||||
allow_methods=["*"],
|
allow_methods=["*"],
|
||||||
allow_headers=["*"],
|
allow_headers=["*"],
|
||||||
@@ -111,6 +111,8 @@ async def detect_roles(request: RoleDetectionRequest):
|
|||||||
"""
|
"""
|
||||||
Analyzes text and returns detected speaker roles.
|
Analyzes text and returns detected speaker roles.
|
||||||
|
|
||||||
|
Uses AI-based detection (Ollama) first, then falls back to pattern matching.
|
||||||
|
|
||||||
Supports two pattern types:
|
Supports two pattern types:
|
||||||
- Colon pattern: "RoleName:" at line start
|
- Colon pattern: "RoleName:" at line start
|
||||||
- Bracket pattern: "[RoleName]" preceding dialogue
|
- Bracket pattern: "[RoleName]" preceding dialogue
|
||||||
@@ -120,7 +122,7 @@ async def detect_roles(request: RoleDetectionRequest):
|
|||||||
logger.info(f"Role detection requested for text of length {len(request.text)}")
|
logger.info(f"Role detection requested for text of length {len(request.text)}")
|
||||||
|
|
||||||
try:
|
try:
|
||||||
roles = role_detector.detect_roles(request.text)
|
roles = await role_detector.detect_roles(request.text)
|
||||||
logger.info(f"Detected {len(roles)} roles: {[r['label'] for r in roles]}")
|
logger.info(f"Detected {len(roles)} roles: {[r['label'] for r in roles]}")
|
||||||
return RolesResponse(roles=roles)
|
return RolesResponse(roles=roles)
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
|
|||||||
@@ -2,13 +2,20 @@
|
|||||||
Role detection service for identifying distinct speakers in text.
|
Role detection service for identifying distinct speakers in text.
|
||||||
"""
|
"""
|
||||||
import re
|
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:
|
class RoleDetector:
|
||||||
"""
|
"""
|
||||||
Service for detecting distinct speaker roles in text.
|
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:
|
Supports two pattern types:
|
||||||
1. Colon pattern: "RoleName:" at the start of a line
|
1. Colon pattern: "RoleName:" at the start of a line
|
||||||
2. Bracket pattern: "[RoleName]" preceding dialogue
|
2. Bracket pattern: "[RoleName]" preceding dialogue
|
||||||
@@ -19,10 +26,23 @@ class RoleDetector:
|
|||||||
BRACKET_PATTERN = r'\[([A-Z][a-zA-Z\s]+?)\]'
|
BRACKET_PATTERN = r'\[([A-Z][a-zA-Z\s]+?)\]'
|
||||||
DEFAULT_ROLE = "Narrator"
|
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.
|
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:
|
Args:
|
||||||
text: Input text to analyze
|
text: Input text to analyze
|
||||||
|
|
||||||
@@ -33,15 +53,26 @@ class RoleDetector:
|
|||||||
if not text or not text.strip():
|
if not text or not text.strip():
|
||||||
return [{"id": self.DEFAULT_ROLE, "label": self.DEFAULT_ROLE}]
|
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)
|
colon_roles = self._extract_colon_pattern(text)
|
||||||
bracket_roles = self._extract_bracket_pattern(text)
|
bracket_roles = self._extract_bracket_pattern(text)
|
||||||
|
pattern_roles = colon_roles.union(bracket_roles)
|
||||||
|
|
||||||
# Combine and deduplicate
|
if pattern_roles:
|
||||||
all_roles = colon_roles.union(bracket_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:
|
if not all_roles:
|
||||||
|
logger.info("No roles detected, using default Narrator")
|
||||||
return [{"id": self.DEFAULT_ROLE, "label": self.DEFAULT_ROLE}]
|
return [{"id": self.DEFAULT_ROLE, "label": self.DEFAULT_ROLE}]
|
||||||
|
|
||||||
# Convert to role dictionaries and sort for consistent ordering
|
# Convert to role dictionaries and sort for consistent ordering
|
||||||
@@ -50,8 +81,40 @@ class RoleDetector:
|
|||||||
for role in sorted(all_roles)
|
for role in sorted(all_roles)
|
||||||
]
|
]
|
||||||
|
|
||||||
|
logger.info(f"Total unique roles after merging: {len(role_list)}")
|
||||||
return 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]:
|
def _extract_colon_pattern(self, text: str) -> Set[str]:
|
||||||
"""
|
"""
|
||||||
Extracts roles from "RoleName:" pattern at line starts.
|
Extracts roles from "RoleName:" pattern at line starts.
|
||||||
|
|||||||
175
backend/test_ollama_integration.py
Normal file
175
backend/test_ollama_integration.py
Normal file
@@ -0,0 +1,175 @@
|
|||||||
|
#!/usr/bin/env python3
|
||||||
|
"""
|
||||||
|
Unit tests for Ollama integration in RoleDetector.
|
||||||
|
Tests the integration between RoleDetector and OllamaService.
|
||||||
|
"""
|
||||||
|
import pytest
|
||||||
|
import asyncio
|
||||||
|
from unittest.mock import AsyncMock, MagicMock, patch
|
||||||
|
from services.role_detector import RoleDetector
|
||||||
|
from services.ollama_service import OllamaService
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_detect_roles_with_ollama_available():
|
||||||
|
"""Test role detection when Ollama is available and returns speakers."""
|
||||||
|
# Create mock OllamaService
|
||||||
|
mock_ollama = AsyncMock(spec=OllamaService)
|
||||||
|
mock_ollama.is_available.return_value = True
|
||||||
|
mock_ollama.detect_speakers.return_value = ["Alice", "Bob", "Charlie"]
|
||||||
|
|
||||||
|
# Create RoleDetector with mock
|
||||||
|
detector = RoleDetector(ollama_service=mock_ollama)
|
||||||
|
|
||||||
|
# Test with text that has no pattern markers
|
||||||
|
text = "This is a conversation between Alice and Bob. Charlie joins later."
|
||||||
|
roles = await detector.detect_roles(text)
|
||||||
|
|
||||||
|
# Verify Ollama was called
|
||||||
|
mock_ollama.is_available.assert_called_once()
|
||||||
|
mock_ollama.detect_speakers.assert_called_once_with(text)
|
||||||
|
|
||||||
|
# Verify roles were detected
|
||||||
|
role_labels = [r['label'] for r in roles]
|
||||||
|
assert "Alice" in role_labels
|
||||||
|
assert "Bob" in role_labels
|
||||||
|
assert "Charlie" in role_labels
|
||||||
|
print("✓ PASS: Ollama detection works when available")
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_detect_roles_with_ollama_unavailable():
|
||||||
|
"""Test role detection falls back to pattern matching when Ollama unavailable."""
|
||||||
|
# Create mock OllamaService that's unavailable
|
||||||
|
mock_ollama = AsyncMock(spec=OllamaService)
|
||||||
|
mock_ollama.is_available.return_value = False
|
||||||
|
|
||||||
|
# Create RoleDetector with mock
|
||||||
|
detector = RoleDetector(ollama_service=mock_ollama)
|
||||||
|
|
||||||
|
# Test with text that has pattern markers
|
||||||
|
text = "Alice: Hello there!\nBob: Hi Alice!"
|
||||||
|
roles = await detector.detect_roles(text)
|
||||||
|
|
||||||
|
# Verify Ollama availability was checked
|
||||||
|
mock_ollama.is_available.assert_called_once()
|
||||||
|
# Verify detect_speakers was NOT called since unavailable
|
||||||
|
mock_ollama.detect_speakers.assert_not_called()
|
||||||
|
|
||||||
|
# Verify pattern matching still worked
|
||||||
|
role_labels = [r['label'] for r in roles]
|
||||||
|
assert "Alice" in role_labels
|
||||||
|
assert "Bob" in role_labels
|
||||||
|
print("✓ PASS: Falls back to pattern matching when Ollama unavailable")
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_detect_roles_with_ollama_error():
|
||||||
|
"""Test role detection handles Ollama errors gracefully."""
|
||||||
|
# Create mock OllamaService that throws error
|
||||||
|
mock_ollama = AsyncMock(spec=OllamaService)
|
||||||
|
mock_ollama.is_available.return_value = True
|
||||||
|
mock_ollama.detect_speakers.side_effect = Exception("Ollama connection failed")
|
||||||
|
|
||||||
|
# Create RoleDetector with mock
|
||||||
|
detector = RoleDetector(ollama_service=mock_ollama)
|
||||||
|
|
||||||
|
# Test with text that has pattern markers
|
||||||
|
text = "Alice: Hello there!\nBob: Hi Alice!"
|
||||||
|
roles = await detector.detect_roles(text)
|
||||||
|
|
||||||
|
# Verify Ollama was attempted
|
||||||
|
mock_ollama.is_available.assert_called_once()
|
||||||
|
mock_ollama.detect_speakers.assert_called_once()
|
||||||
|
|
||||||
|
# Verify pattern matching still worked as fallback
|
||||||
|
role_labels = [r['label'] for r in roles]
|
||||||
|
assert "Alice" in role_labels
|
||||||
|
assert "Bob" in role_labels
|
||||||
|
print("✓ PASS: Handles Ollama errors gracefully and falls back to patterns")
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_detect_roles_merges_ollama_and_pattern_results():
|
||||||
|
"""Test that roles from both Ollama and pattern matching are merged."""
|
||||||
|
# Create mock OllamaService
|
||||||
|
mock_ollama = AsyncMock(spec=OllamaService)
|
||||||
|
mock_ollama.is_available.return_value = True
|
||||||
|
mock_ollama.detect_speakers.return_value = ["Alice", "Charlie"]
|
||||||
|
|
||||||
|
# Create RoleDetector with mock
|
||||||
|
detector = RoleDetector(ollama_service=mock_ollama)
|
||||||
|
|
||||||
|
# Test with text that has pattern markers for different roles
|
||||||
|
text = "Alice: Hello!\nBob: Hi there!\nCharlie appears later."
|
||||||
|
roles = await detector.detect_roles(text)
|
||||||
|
|
||||||
|
# Verify both methods were used
|
||||||
|
mock_ollama.detect_speakers.assert_called_once()
|
||||||
|
|
||||||
|
# Verify roles from both sources are present
|
||||||
|
role_labels = [r['label'] for r in roles]
|
||||||
|
assert "Alice" in role_labels # From both Ollama and pattern
|
||||||
|
assert "Bob" in role_labels # From pattern only
|
||||||
|
assert "Charlie" in role_labels # From Ollama only
|
||||||
|
print("✓ PASS: Merges results from both Ollama and pattern matching")
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_detect_roles_deduplicates_across_methods():
|
||||||
|
"""Test that duplicate roles from Ollama and patterns are deduplicated."""
|
||||||
|
# Create mock OllamaService
|
||||||
|
mock_ollama = AsyncMock(spec=OllamaService)
|
||||||
|
mock_ollama.is_available.return_value = True
|
||||||
|
mock_ollama.detect_speakers.return_value = ["Alice", "Bob"]
|
||||||
|
|
||||||
|
# Create RoleDetector with mock
|
||||||
|
detector = RoleDetector(ollama_service=mock_ollama)
|
||||||
|
|
||||||
|
# Test with text where patterns also detect Alice and Bob
|
||||||
|
text = "Alice: Hello!\nBob: Hi there!"
|
||||||
|
roles = await detector.detect_roles(text)
|
||||||
|
|
||||||
|
# Verify roles are deduplicated
|
||||||
|
role_labels = [r['label'] for r in roles]
|
||||||
|
assert role_labels.count("Alice") == 1
|
||||||
|
assert role_labels.count("Bob") == 1
|
||||||
|
assert len(role_labels) == 2
|
||||||
|
print("✓ PASS: Deduplicates roles across both detection methods")
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_detect_roles_returns_narrator_when_no_roles_found():
|
||||||
|
"""Test that default Narrator is returned when no roles detected by either method."""
|
||||||
|
# Create mock OllamaService that returns empty
|
||||||
|
mock_ollama = AsyncMock(spec=OllamaService)
|
||||||
|
mock_ollama.is_available.return_value = True
|
||||||
|
mock_ollama.detect_speakers.return_value = []
|
||||||
|
|
||||||
|
# Create RoleDetector with mock
|
||||||
|
detector = RoleDetector(ollama_service=mock_ollama)
|
||||||
|
|
||||||
|
# Test with text that has no role markers
|
||||||
|
text = "This is just plain text without any speakers."
|
||||||
|
roles = await detector.detect_roles(text)
|
||||||
|
|
||||||
|
# Verify default Narrator is returned
|
||||||
|
assert len(roles) == 1
|
||||||
|
assert roles[0]['label'] == "Narrator"
|
||||||
|
print("✓ PASS: Returns default Narrator when no roles found")
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == '__main__':
|
||||||
|
# Run tests
|
||||||
|
print("Testing Ollama Integration in RoleDetector...\n")
|
||||||
|
|
||||||
|
asyncio.run(test_detect_roles_with_ollama_available())
|
||||||
|
asyncio.run(test_detect_roles_with_ollama_unavailable())
|
||||||
|
asyncio.run(test_detect_roles_with_ollama_error())
|
||||||
|
asyncio.run(test_detect_roles_merges_ollama_and_pattern_results())
|
||||||
|
asyncio.run(test_detect_roles_deduplicates_across_methods())
|
||||||
|
asyncio.run(test_detect_roles_returns_narrator_when_no_roles_found())
|
||||||
|
|
||||||
|
print("\n" + "="*50)
|
||||||
|
print("All integration tests passed!")
|
||||||
|
print("="*50)
|
||||||
Reference in New Issue
Block a user