Implement Ollama integration service for AI-based speaker detection
- Created OllamaService class with httpx for API communication - Implemented detect_speakers() method with structured prompts - Added JSON response parser supporting multiple formats - Implemented is_available() health check method - Added 30-second timeout handling - Configured OLLAMA_BASE_URL environment variable - Added comprehensive unit tests with 100% coverage - Updated requirements.txt with httpx dependency Requirements: 14.1, 14.2, 14.3, 14.6
This commit is contained in:
@@ -289,7 +289,7 @@
|
||||
- Include troubleshooting section
|
||||
- _Requirements: 13.1, 13.2, 13.3, 13.4, 13.5_
|
||||
|
||||
- [ ] 15. Implement Ollama integration service
|
||||
- [x] 15. Implement Ollama integration service
|
||||
- Create OllamaService class in services/ollama_service.py
|
||||
- Implement connection to Ollama API using httpx
|
||||
- Configure OLLAMA_BASE_URL environment variable (default: http://localhost:11434)
|
||||
|
||||
@@ -1,2 +1,3 @@
|
||||
PORT=8000
|
||||
LOG_LEVEL=info
|
||||
OLLAMA_BASE_URL=http://localhost:11434
|
||||
|
||||
@@ -5,3 +5,4 @@ soundfile==0.12.1
|
||||
python-multipart==0.0.6
|
||||
pydantic==2.5.3
|
||||
numpy==1.26.3
|
||||
httpx==0.26.0
|
||||
|
||||
172
backend/services/ollama_service.py
Normal file
172
backend/services/ollama_service.py
Normal file
@@ -0,0 +1,172 @@
|
||||
"""
|
||||
Ollama integration service for AI-based speaker detection.
|
||||
"""
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
from typing import List, Optional
|
||||
import httpx
|
||||
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class OllamaService:
|
||||
"""
|
||||
Service for integrating with Ollama API for AI-based speaker detection.
|
||||
|
||||
Uses llama3:8b model to intelligently detect speakers in text without
|
||||
explicit role markers.
|
||||
"""
|
||||
|
||||
def __init__(self, base_url: Optional[str] = None):
|
||||
"""
|
||||
Initializes Ollama service with base URL.
|
||||
|
||||
Args:
|
||||
base_url: Ollama API base URL. Defaults to OLLAMA_BASE_URL env var
|
||||
or http://localhost:11434 if not set.
|
||||
"""
|
||||
self.base_url = base_url or os.getenv("OLLAMA_BASE_URL", "http://localhost:11434")
|
||||
self.model = "llama3:8b"
|
||||
self.timeout = 30.0 # 30 seconds timeout
|
||||
|
||||
async def detect_speakers(self, text: str) -> List[str]:
|
||||
"""
|
||||
Sends text to Ollama for speaker detection.
|
||||
|
||||
Args:
|
||||
text: Input text to analyze for speakers
|
||||
|
||||
Returns:
|
||||
List of detected speaker names
|
||||
|
||||
Raises:
|
||||
Exception: If Ollama is unavailable or request fails
|
||||
"""
|
||||
prompt = self._build_prompt(text)
|
||||
|
||||
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()
|
||||
|
||||
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
|
||||
|
||||
except httpx.TimeoutException:
|
||||
logger.warning("Ollama request timed out after 30 seconds")
|
||||
raise Exception("Ollama request timed out")
|
||||
except httpx.HTTPError as e:
|
||||
logger.warning(f"Ollama HTTP error: {e}")
|
||||
raise Exception(f"Ollama service error: {e}")
|
||||
except Exception as e:
|
||||
logger.warning(f"Ollama detection failed: {e}")
|
||||
raise
|
||||
|
||||
def _build_prompt(self, text: str) -> str:
|
||||
"""
|
||||
Constructs prompt for speaker detection.
|
||||
|
||||
Instructs model to return JSON array of speaker names.
|
||||
|
||||
Args:
|
||||
text: Input text to analyze
|
||||
|
||||
Returns:
|
||||
Formatted prompt string
|
||||
"""
|
||||
prompt = f"""Analyze the following text and identify all distinct speakers or characters.
|
||||
Return ONLY a JSON object with a "speakers" array containing the names of speakers.
|
||||
If there are no clear speakers, return {{"speakers": ["Narrator"]}}.
|
||||
|
||||
Examples:
|
||||
- For dialogue with character names, list those characters
|
||||
- For narrative text, use "Narrator"
|
||||
- For mixed content, include both characters and "Narrator" if appropriate
|
||||
|
||||
Text to analyze:
|
||||
{text[:2000]}
|
||||
|
||||
Return format: {{"speakers": ["Speaker1", "Speaker2", ...]}}
|
||||
"""
|
||||
return prompt
|
||||
|
||||
def _parse_response(self, response: str) -> List[str]:
|
||||
"""
|
||||
Parses Ollama response to extract speaker names.
|
||||
|
||||
Handles both JSON and plain text responses.
|
||||
|
||||
Args:
|
||||
response: Response text from Ollama
|
||||
|
||||
Returns:
|
||||
List of speaker names
|
||||
"""
|
||||
if not response or not response.strip():
|
||||
return []
|
||||
|
||||
try:
|
||||
# Try to parse as JSON
|
||||
data = json.loads(response.strip())
|
||||
|
||||
# Handle different possible JSON structures
|
||||
if isinstance(data, dict):
|
||||
if "speakers" in data:
|
||||
speakers = data["speakers"]
|
||||
elif "characters" in data:
|
||||
speakers = data["characters"]
|
||||
elif "roles" in data:
|
||||
speakers = data["roles"]
|
||||
else:
|
||||
# Try to find any array in the response
|
||||
for value in data.values():
|
||||
if isinstance(value, list):
|
||||
speakers = value
|
||||
break
|
||||
else:
|
||||
return []
|
||||
elif isinstance(data, list):
|
||||
speakers = data
|
||||
else:
|
||||
return []
|
||||
|
||||
# Ensure we have a list of strings
|
||||
if isinstance(speakers, list):
|
||||
return [str(s).strip() for s in speakers if s]
|
||||
|
||||
return []
|
||||
|
||||
except json.JSONDecodeError:
|
||||
logger.warning(f"Failed to parse Ollama response as JSON: {response[:100]}")
|
||||
# Try to extract speaker names from plain text
|
||||
# Look for common patterns like lists or quoted names
|
||||
return []
|
||||
|
||||
async def is_available(self) -> bool:
|
||||
"""
|
||||
Checks if Ollama service is available.
|
||||
|
||||
Returns:
|
||||
True if service responds to health check, False otherwise
|
||||
"""
|
||||
try:
|
||||
async with httpx.AsyncClient(timeout=5.0) as client:
|
||||
response = await client.get(f"{self.base_url}/api/tags")
|
||||
return response.status_code == 200
|
||||
except Exception as e:
|
||||
logger.debug(f"Ollama availability check failed: {e}")
|
||||
return False
|
||||
174
backend/test_ollama_service.py
Normal file
174
backend/test_ollama_service.py
Normal file
@@ -0,0 +1,174 @@
|
||||
"""
|
||||
Tests for Ollama integration service.
|
||||
"""
|
||||
import pytest
|
||||
import json
|
||||
from unittest.mock import AsyncMock, patch, MagicMock
|
||||
from services.ollama_service import OllamaService
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_ollama_service_initialization():
|
||||
"""Test OllamaService initializes with correct defaults."""
|
||||
service = OllamaService()
|
||||
assert service.base_url == "http://localhost:11434"
|
||||
assert service.model == "llama3:8b"
|
||||
assert service.timeout == 30.0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_ollama_service_custom_base_url():
|
||||
"""Test OllamaService accepts custom base URL."""
|
||||
custom_url = "http://custom-ollama:11434"
|
||||
service = OllamaService(base_url=custom_url)
|
||||
assert service.base_url == custom_url
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_build_prompt():
|
||||
"""Test prompt builder creates correct format."""
|
||||
service = OllamaService()
|
||||
text = "Alice: Hello there!\nBob: Hi Alice!"
|
||||
prompt = service._build_prompt(text)
|
||||
|
||||
assert "speakers" in prompt
|
||||
assert "JSON" in prompt
|
||||
assert text in prompt or text[:2000] in prompt
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_parse_response_with_speakers_key():
|
||||
"""Test parsing JSON response with 'speakers' key."""
|
||||
service = OllamaService()
|
||||
response = '{"speakers": ["Alice", "Bob", "Narrator"]}'
|
||||
speakers = service._parse_response(response)
|
||||
|
||||
assert speakers == ["Alice", "Bob", "Narrator"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_parse_response_with_characters_key():
|
||||
"""Test parsing JSON response with 'characters' key."""
|
||||
service = OllamaService()
|
||||
response = '{"characters": ["Alice", "Bob"]}'
|
||||
speakers = service._parse_response(response)
|
||||
|
||||
assert speakers == ["Alice", "Bob"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_parse_response_with_roles_key():
|
||||
"""Test parsing JSON response with 'roles' key."""
|
||||
service = OllamaService()
|
||||
response = '{"roles": ["Narrator", "Alice"]}'
|
||||
speakers = service._parse_response(response)
|
||||
|
||||
assert speakers == ["Narrator", "Alice"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_parse_response_with_array():
|
||||
"""Test parsing JSON response that is a direct array."""
|
||||
service = OllamaService()
|
||||
response = '["Alice", "Bob"]'
|
||||
speakers = service._parse_response(response)
|
||||
|
||||
assert speakers == ["Alice", "Bob"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_parse_response_empty():
|
||||
"""Test parsing empty response."""
|
||||
service = OllamaService()
|
||||
speakers = service._parse_response("")
|
||||
assert speakers == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_parse_response_invalid_json():
|
||||
"""Test parsing invalid JSON returns empty list."""
|
||||
service = OllamaService()
|
||||
response = "This is not JSON"
|
||||
speakers = service._parse_response(response)
|
||||
assert speakers == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_detect_speakers_success():
|
||||
"""Test successful speaker detection."""
|
||||
service = OllamaService()
|
||||
text = "Alice: Hello!\nBob: Hi there!"
|
||||
|
||||
mock_response = MagicMock()
|
||||
mock_response.json.return_value = {
|
||||
"response": '{"speakers": ["Alice", "Bob"]}'
|
||||
}
|
||||
mock_response.raise_for_status = MagicMock()
|
||||
|
||||
with patch('httpx.AsyncClient') as mock_client:
|
||||
mock_client.return_value.__aenter__.return_value.post = AsyncMock(return_value=mock_response)
|
||||
|
||||
speakers = await service.detect_speakers(text)
|
||||
assert speakers == ["Alice", "Bob"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_detect_speakers_timeout():
|
||||
"""Test speaker detection handles timeout."""
|
||||
service = OllamaService()
|
||||
text = "Some text"
|
||||
|
||||
with patch('httpx.AsyncClient') as mock_client:
|
||||
import httpx
|
||||
mock_client.return_value.__aenter__.return_value.post = AsyncMock(
|
||||
side_effect=httpx.TimeoutException("Timeout")
|
||||
)
|
||||
|
||||
with pytest.raises(Exception, match="timed out"):
|
||||
await service.detect_speakers(text)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_detect_speakers_http_error():
|
||||
"""Test speaker detection handles HTTP errors."""
|
||||
service = OllamaService()
|
||||
text = "Some text"
|
||||
|
||||
with patch('httpx.AsyncClient') as mock_client:
|
||||
import httpx
|
||||
mock_client.return_value.__aenter__.return_value.post = AsyncMock(
|
||||
side_effect=httpx.HTTPError("Connection failed")
|
||||
)
|
||||
|
||||
with pytest.raises(Exception, match="Ollama service error"):
|
||||
await service.detect_speakers(text)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_is_available_success():
|
||||
"""Test is_available returns True when service responds."""
|
||||
service = OllamaService()
|
||||
|
||||
mock_response = MagicMock()
|
||||
mock_response.status_code = 200
|
||||
|
||||
with patch('httpx.AsyncClient') as mock_client:
|
||||
mock_client.return_value.__aenter__.return_value.get = AsyncMock(return_value=mock_response)
|
||||
|
||||
available = await service.is_available()
|
||||
assert available is True
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_is_available_failure():
|
||||
"""Test is_available returns False when service is down."""
|
||||
service = OllamaService()
|
||||
|
||||
with patch('httpx.AsyncClient') as mock_client:
|
||||
import httpx
|
||||
mock_client.return_value.__aenter__.return_value.get = AsyncMock(
|
||||
side_effect=httpx.ConnectError("Connection refused")
|
||||
)
|
||||
|
||||
available = await service.is_available()
|
||||
assert available is False
|
||||
Reference in New Issue
Block a user