diff --git a/backend/services/ollama_service.py b/backend/services/ollama_service.py index 7037360..0377e82 100644 --- a/backend/services/ollama_service.py +++ b/backend/services/ollama_service.py @@ -34,37 +34,54 @@ class OllamaService: 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 + List of detected speaker names (deduplicated across all samples) Raises: Exception: If Ollama is unavailable or request fails """ - prompt = self._build_prompt(text) + # 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: - response = await client.post( - f"{self.base_url}/api/generate", - json={ - "model": self.model, - "prompt": prompt, - "stream": False, - "format": "json" - } - ) - response.raise_for_status() + # 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) - result = response.json() - response_text = result.get("response", "") - - speakers = self._parse_response(response_text) - logger.info(f"Ollama detected {len(speakers)} speakers: {speakers}") - return 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") @@ -76,6 +93,47 @@ class OllamaService: 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. @@ -83,7 +141,7 @@ class OllamaService: Instructs model to return JSON array of speaker names. Args: - text: Input text to analyze + text: Input text to analyze (should already be sampled/truncated) Returns: Formatted prompt string @@ -98,7 +156,7 @@ Examples: - For mixed content, include both characters and "Narrator" if appropriate Text to analyze: -{text[:2000]} +{text} Return format: {{"speakers": ["Speaker1", "Speaker2", ...]}} """ diff --git a/backend/test_ollama_service.py b/backend/test_ollama_service.py index a01a52c..1c26c1e 100644 --- a/backend/test_ollama_service.py +++ b/backend/test_ollama_service.py @@ -24,6 +24,44 @@ async def test_ollama_service_custom_base_url(): assert service.base_url == custom_url +@pytest.mark.asyncio +async def test_sample_text_short(): + """Test sampling short text returns full text.""" + service = OllamaService() + text = "Short text" + samples = service._sample_text(text, chunk_size=5000, num_samples=4) + + assert len(samples) == 1 + assert samples[0] == text + + +@pytest.mark.asyncio +async def test_sample_text_long(): + """Test sampling long text returns multiple samples.""" + service = OllamaService() + # Create text with position markers so we can verify different samples + text = "".join([f"Position{i:05d} " for i in range(2500)]) # ~25k characters + samples = service._sample_text(text, chunk_size=5000, num_samples=4) + + assert len(samples) == 4 + assert all(len(s) <= 5000 for s in samples) + # Samples should be from different parts - check they contain different position markers + assert "Position00000" in samples[0] # First sample should have early positions + assert "Position00000" not in samples[-1] # Last sample should have later positions + + +@pytest.mark.asyncio +async def test_sample_text_medium(): + """Test sampling medium text returns appropriate number of samples.""" + service = OllamaService() + text = "B" * 8000 # 8k characters + samples = service._sample_text(text, chunk_size=5000, num_samples=4) + + # Should get 2 samples (8000 / 5000 = 1.6, so 2 samples) + assert len(samples) >= 1 + assert all(len(s) <= 5000 for s in samples) + + @pytest.mark.asyncio async def test_build_prompt(): """Test prompt builder creates correct format.""" @@ -33,7 +71,7 @@ async def test_build_prompt(): assert "speakers" in prompt assert "JSON" in prompt - assert text in prompt or text[:2000] in prompt + assert text in prompt @pytest.mark.asyncio @@ -95,7 +133,7 @@ async def test_parse_response_invalid_json(): @pytest.mark.asyncio async def test_detect_speakers_success(): - """Test successful speaker detection.""" + """Test successful speaker detection with short text.""" service = OllamaService() text = "Alice: Hello!\nBob: Hi there!" @@ -109,7 +147,32 @@ async def test_detect_speakers_success(): mock_client.return_value.__aenter__.return_value.post = AsyncMock(return_value=mock_response) speakers = await service.detect_speakers(text) - assert speakers == ["Alice", "Bob"] + assert set(speakers) == {"Alice", "Bob"} + + +@pytest.mark.asyncio +async def test_detect_speakers_multiple_samples(): + """Test speaker detection with long text requiring multiple samples.""" + service = OllamaService() + text = "A" * 25000 # Long text requiring multiple samples + + # Mock responses for different samples + responses = [ + MagicMock(json=lambda: {"response": '{"speakers": ["Alice", "Bob"]}'}), + MagicMock(json=lambda: {"response": '{"speakers": ["Charlie", "Bob"]}'}), + MagicMock(json=lambda: {"response": '{"speakers": ["Alice", "David"]}'}), + MagicMock(json=lambda: {"response": '{"speakers": ["Eve"]}'}), + ] + + for r in responses: + r.raise_for_status = MagicMock() + + with patch('httpx.AsyncClient') as mock_client: + mock_client.return_value.__aenter__.return_value.post = AsyncMock(side_effect=responses) + + speakers = await service.detect_speakers(text) + # Should deduplicate across all samples + assert set(speakers) == {"Alice", "Bob", "Charlie", "David", "Eve"} @pytest.mark.asyncio