Improve Ollama speaker detection with multi-sample strategy
- Sample up to 4 chunks of 5000 characters from different text positions - Distribute samples evenly across beginning, middle, and end of text - Deduplicate speakers across all samples for comprehensive detection - Log individual sample results for debugging - Handle short texts efficiently (single sample) - Updated tests to verify multi-sample behavior This ensures speakers introduced anywhere in long texts are detected, not just those appearing in the first 2000 characters.
This commit is contained in:
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user