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:
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
|
||||
Reference in New Issue
Block a user