feat: OllamaClient.identify_speakers() and summarize() for diarization pipeline
This commit is contained in:
@@ -35,3 +35,47 @@ async def test_list_models_returns_list():
|
||||
client = OllamaClient(base_url="http://localhost:11434")
|
||||
models = await client.list_models()
|
||||
assert "gemma3:12b" in models
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_identify_speakers_returns_dict():
|
||||
import respx, httpx, json
|
||||
from llm import OllamaClient
|
||||
client = OllamaClient()
|
||||
mapping = {"SPEAKER_00": "Thomas", "SPEAKER_01": "Möller"}
|
||||
transcript_excerpt = "SPEAKER_00: Gut, Herr Möller.\nSPEAKER_01: Danke, Thomas."
|
||||
|
||||
with respx.mock:
|
||||
respx.post("http://localhost:11434/api/generate").mock(
|
||||
return_value=httpx.Response(200, json={"response": json.dumps(mapping)})
|
||||
)
|
||||
result = await client.identify_speakers(transcript_excerpt)
|
||||
assert result == {"SPEAKER_00": "Thomas", "SPEAKER_01": "Möller"}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_identify_speakers_returns_empty_on_parse_failure():
|
||||
import respx, httpx
|
||||
from llm import OllamaClient
|
||||
client = OllamaClient()
|
||||
|
||||
with respx.mock:
|
||||
respx.post("http://localhost:11434/api/generate").mock(
|
||||
return_value=httpx.Response(200, json={"response": "kein json hier"})
|
||||
)
|
||||
result = await client.identify_speakers("irgendwas")
|
||||
assert result == {}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_summarize_returns_string():
|
||||
import respx, httpx
|
||||
from llm import OllamaClient
|
||||
client = OllamaClient()
|
||||
|
||||
with respx.mock:
|
||||
respx.post("http://localhost:11434/api/generate").mock(
|
||||
return_value=httpx.Response(200, json={"response": "# Zusammenfassung\n\nKurzer Text."})
|
||||
)
|
||||
result = await client.summarize("Thomas: Hallo.\nMöller: Hi.", model="gemma3:12b")
|
||||
assert "Zusammenfassung" in result
|
||||
|
||||
Reference in New Issue
Block a user