diff --git a/.kiro/specs/text-to-audiobook/tasks.md b/.kiro/specs/text-to-audiobook/tasks.md index ff258cb..276fcde 100644 --- a/.kiro/specs/text-to-audiobook/tasks.md +++ b/.kiro/specs/text-to-audiobook/tasks.md @@ -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) diff --git a/backend/.env.example b/backend/.env.example index 8d653df..0abe249 100644 --- a/backend/.env.example +++ b/backend/.env.example @@ -1,2 +1,3 @@ PORT=8000 LOG_LEVEL=info +OLLAMA_BASE_URL=http://localhost:11434 diff --git a/backend/requirements.txt b/backend/requirements.txt index 1ef2e3f..98048ed 100644 --- a/backend/requirements.txt +++ b/backend/requirements.txt @@ -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 diff --git a/backend/services/ollama_service.py b/backend/services/ollama_service.py new file mode 100644 index 0000000..7037360 --- /dev/null +++ b/backend/services/ollama_service.py @@ -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 diff --git a/backend/test_ollama_service.py b/backend/test_ollama_service.py new file mode 100644 index 0000000..a01a52c --- /dev/null +++ b/backend/test_ollama_service.py @@ -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