From f56f8086e550b49be7c2e516584d43263dd876e2 Mon Sep 17 00:00:00 2001 From: Paul Bailey Date: Tue, 9 Dec 2025 10:49:37 +0000 Subject: [PATCH] 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 --- backend/main.py | 6 +- backend/services/role_detector.py | 75 ++++++++++++- backend/test_ollama_integration.py | 175 +++++++++++++++++++++++++++++ 3 files changed, 248 insertions(+), 8 deletions(-) create mode 100644 backend/test_ollama_integration.py diff --git a/backend/main.py b/backend/main.py index 3c453db..bc1c11a 100644 --- a/backend/main.py +++ b/backend/main.py @@ -62,7 +62,7 @@ app = FastAPI( # Configure CORS app.add_middleware( CORSMiddleware, - allow_origins=["http://localhost:5173", "http://localhost:3000"], + allow_origins=["http://localhost:5173", "http://localhost:5174", "http://localhost:3000"], allow_credentials=True, allow_methods=["*"], allow_headers=["*"], @@ -111,6 +111,8 @@ async def detect_roles(request: RoleDetectionRequest): """ Analyzes text and returns detected speaker roles. + Uses AI-based detection (Ollama) first, then falls back to pattern matching. + Supports two pattern types: - Colon pattern: "RoleName:" at line start - 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)}") 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]}") return RolesResponse(roles=roles) except Exception as e: diff --git a/backend/services/role_detector.py b/backend/services/role_detector.py index 5c569fa..e134d6b 100644 --- a/backend/services/role_detector.py +++ b/backend/services/role_detector.py @@ -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. diff --git a/backend/test_ollama_integration.py b/backend/test_ollama_integration.py new file mode 100644 index 0000000..3766eb6 --- /dev/null +++ b/backend/test_ollama_integration.py @@ -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)