Implement script generation endpoint (Task 18)
- Added ScriptRequest, ScriptResponse, and ScriptSegment Pydantic models - Created /api/tts/script POST endpoint in main.py - Integrated with TTSService to generate script segments using AI - Added generate_script_segments method to TTSService - Returns segments with unique IDs for frontend tracking - Handles errors gracefully when AI segmentation fails - Falls back to pattern-based segmentation when Ollama unavailable - Added comprehensive tests for the endpoint - Validates: Requirements 15.1, 15.2
This commit is contained in:
104
backend/main.py
104
backend/main.py
@@ -34,6 +34,21 @@ class RolesResponse(BaseModel):
|
||||
"""Response model for role detection endpoint."""
|
||||
roles: List[RoleInfo] = Field(..., description="List of detected roles")
|
||||
|
||||
class ScriptSegment(BaseModel):
|
||||
"""Model representing a script segment with speaker assignment."""
|
||||
id: str = Field(..., description="Unique segment identifier")
|
||||
speaker: str = Field(..., description="Assigned speaker/role")
|
||||
text: str = Field(..., description="Text content for this segment")
|
||||
|
||||
class ScriptRequest(BaseModel):
|
||||
"""Request model for script generation endpoint."""
|
||||
text: str = Field(..., description="Text to segment and analyze")
|
||||
roles: List[str] = Field(..., description="List of detected role IDs")
|
||||
|
||||
class ScriptResponse(BaseModel):
|
||||
"""Response model for script generation endpoint."""
|
||||
segments: List[ScriptSegment] = Field(..., description="List of script segments with speaker assignments")
|
||||
|
||||
class RoleVoiceMapping(BaseModel):
|
||||
"""Model for role-to-voice mapping."""
|
||||
id: str = Field(..., description="Role identifier")
|
||||
@@ -129,6 +144,95 @@ async def detect_roles(request: RoleDetectionRequest):
|
||||
logger.error(f"Error during role detection: {str(e)}", exc_info=True)
|
||||
raise
|
||||
|
||||
@app.post("/api/tts/script", response_model=ScriptResponse)
|
||||
async def generate_script(request: ScriptRequest):
|
||||
"""
|
||||
Generates script preview with AI-segmented text and speaker assignments.
|
||||
|
||||
Uses AI (Ollama) to intelligently segment the text and assign speakers to each segment.
|
||||
Falls back to pattern-based segmentation if AI is unavailable.
|
||||
|
||||
Args:
|
||||
request: ScriptRequest containing text and list of detected roles
|
||||
|
||||
Returns:
|
||||
ScriptResponse with list of segments, each containing id, speaker, and text
|
||||
|
||||
Raises:
|
||||
400: If validation fails (empty text or roles)
|
||||
500: If script generation fails
|
||||
"""
|
||||
logger.info(
|
||||
f"Script generation requested for text of length {len(request.text)} "
|
||||
f"with {len(request.roles)} roles"
|
||||
)
|
||||
|
||||
# Validate text is not empty
|
||||
if not request.text or not request.text.strip():
|
||||
logger.error("Empty text provided for script generation")
|
||||
return JSONResponse(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
content={
|
||||
"error": "ValidationError",
|
||||
"message": "Text field cannot be empty",
|
||||
"details": {"field": "text"}
|
||||
}
|
||||
)
|
||||
|
||||
# Validate roles list is not empty
|
||||
if not request.roles:
|
||||
logger.error("No roles provided for script generation")
|
||||
return JSONResponse(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
content={
|
||||
"error": "ValidationError",
|
||||
"message": "At least one role is required",
|
||||
"details": {"field": "roles"}
|
||||
}
|
||||
)
|
||||
|
||||
try:
|
||||
# Generate script segments using TTS service
|
||||
segments = await tts_service.generate_script_segments(
|
||||
text=request.text,
|
||||
roles=request.roles
|
||||
)
|
||||
|
||||
if not segments:
|
||||
logger.warning("Script generation produced no segments")
|
||||
return JSONResponse(
|
||||
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
|
||||
content={
|
||||
"error": "ScriptGenerationError",
|
||||
"message": "Failed to generate script segments",
|
||||
"details": {"reason": "No segments produced"}
|
||||
}
|
||||
)
|
||||
|
||||
# Convert to response format with unique IDs
|
||||
script_segments = [
|
||||
ScriptSegment(
|
||||
id=f"seg_{i}",
|
||||
speaker=speaker,
|
||||
text=text
|
||||
)
|
||||
for i, (speaker, text) in enumerate(segments)
|
||||
]
|
||||
|
||||
logger.info(f"Generated script with {len(script_segments)} segments")
|
||||
return ScriptResponse(segments=script_segments)
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error during script generation: {str(e)}", exc_info=True)
|
||||
return JSONResponse(
|
||||
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
|
||||
content={
|
||||
"error": "ScriptGenerationError",
|
||||
"message": "Failed to generate script",
|
||||
"details": {"error": str(e)}
|
||||
}
|
||||
)
|
||||
|
||||
@app.post("/api/tts")
|
||||
async def generate_audiobook(request: TTSRequest):
|
||||
"""
|
||||
|
||||
@@ -288,6 +288,62 @@ class TTSService:
|
||||
|
||||
return concatenated
|
||||
|
||||
async def generate_script_segments(
|
||||
self,
|
||||
text: str,
|
||||
roles: List[str]
|
||||
) -> List[Tuple[str, str]]:
|
||||
"""
|
||||
Generates script segments with speaker assignments using AI.
|
||||
|
||||
This method is used for the script preview feature, allowing users to
|
||||
review and edit speaker assignments before generating audio.
|
||||
|
||||
Args:
|
||||
text: Input text to segment
|
||||
roles: List of detected role IDs/names
|
||||
|
||||
Returns:
|
||||
List of tuples: (speaker_name, text_content)
|
||||
|
||||
Raises:
|
||||
Exception: If both AI and pattern-based segmentation fail
|
||||
"""
|
||||
logger.info(f"Generating script segments for {len(text)} chars with {len(roles)} roles")
|
||||
|
||||
# Try AI segmentation first
|
||||
try:
|
||||
is_available = await self.ollama_service.is_available()
|
||||
if is_available:
|
||||
logger.info("Using AI for script segmentation")
|
||||
segments = await self.ollama_service.segment_text(text, roles)
|
||||
|
||||
if segments:
|
||||
logger.info(f"AI segmentation produced {len(segments)} segments")
|
||||
return segments
|
||||
else:
|
||||
logger.warning("AI segmentation returned empty result")
|
||||
else:
|
||||
logger.info("Ollama not available for script segmentation")
|
||||
except Exception as e:
|
||||
logger.warning(f"AI segmentation failed: {e}")
|
||||
|
||||
# Fall back to pattern-based segmentation
|
||||
logger.info("Using pattern-based segmentation as fallback")
|
||||
|
||||
# Create a dummy role_voice_map for pattern-based segmentation
|
||||
# We don't need actual voices here, just role assignments
|
||||
role_voice_map = {role: self.DEFAULT_VOICE for role in roles}
|
||||
|
||||
# Use existing segment_text method
|
||||
segments_with_voice = self.segment_text(text, role_voice_map)
|
||||
|
||||
# Convert from (role, voice, text) to (role, text)
|
||||
segments = [(role, text) for role, voice, text in segments_with_voice]
|
||||
|
||||
logger.info(f"Pattern-based segmentation produced {len(segments)} segments")
|
||||
return segments
|
||||
|
||||
async def generate_audiobook(
|
||||
self,
|
||||
text: str,
|
||||
|
||||
173
backend/test_script_endpoint.py
Normal file
173
backend/test_script_endpoint.py
Normal file
@@ -0,0 +1,173 @@
|
||||
"""
|
||||
Tests for the script generation endpoint.
|
||||
"""
|
||||
import pytest
|
||||
from fastapi.testclient import TestClient
|
||||
from main import app
|
||||
|
||||
client = TestClient(app)
|
||||
|
||||
|
||||
def test_script_endpoint_exists():
|
||||
"""Test that the /api/tts/script endpoint exists."""
|
||||
response = client.post(
|
||||
"/api/tts/script",
|
||||
json={
|
||||
"text": "Alice: Hello there!\nBob: Hi Alice!",
|
||||
"roles": ["Alice", "Bob"]
|
||||
}
|
||||
)
|
||||
# Should not return 404
|
||||
assert response.status_code != 404
|
||||
|
||||
|
||||
def test_script_generation_with_valid_input():
|
||||
"""Test script generation with valid text and roles."""
|
||||
response = client.post(
|
||||
"/api/tts/script",
|
||||
json={
|
||||
"text": "Alice: Hello there!\nBob: Hi Alice, how are you?",
|
||||
"roles": ["Alice", "Bob", "Narrator"]
|
||||
}
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
|
||||
# Check response structure
|
||||
assert "segments" in data
|
||||
assert isinstance(data["segments"], list)
|
||||
assert len(data["segments"]) > 0
|
||||
|
||||
# Check each segment has required fields
|
||||
for segment in data["segments"]:
|
||||
assert "id" in segment
|
||||
assert "speaker" in segment
|
||||
assert "text" in segment
|
||||
assert isinstance(segment["id"], str)
|
||||
assert isinstance(segment["speaker"], str)
|
||||
assert isinstance(segment["text"], str)
|
||||
assert len(segment["text"]) > 0
|
||||
|
||||
|
||||
def test_script_generation_with_empty_text():
|
||||
"""Test that empty text returns validation error."""
|
||||
response = client.post(
|
||||
"/api/tts/script",
|
||||
json={
|
||||
"text": "",
|
||||
"roles": ["Alice", "Bob"]
|
||||
}
|
||||
)
|
||||
|
||||
assert response.status_code == 400
|
||||
data = response.json()
|
||||
assert "error" in data
|
||||
assert data["error"] == "ValidationError"
|
||||
|
||||
|
||||
def test_script_generation_with_empty_roles():
|
||||
"""Test that empty roles list returns validation error."""
|
||||
response = client.post(
|
||||
"/api/tts/script",
|
||||
json={
|
||||
"text": "Some text here",
|
||||
"roles": []
|
||||
}
|
||||
)
|
||||
|
||||
assert response.status_code == 400
|
||||
data = response.json()
|
||||
assert "error" in data
|
||||
assert data["error"] == "ValidationError"
|
||||
|
||||
|
||||
def test_script_generation_missing_text_field():
|
||||
"""Test that missing text field returns validation error."""
|
||||
response = client.post(
|
||||
"/api/tts/script",
|
||||
json={
|
||||
"roles": ["Alice", "Bob"]
|
||||
}
|
||||
)
|
||||
|
||||
# Our error handler converts 422 to 400
|
||||
assert response.status_code == 400
|
||||
data = response.json()
|
||||
assert "error" in data
|
||||
assert data["error"] == "ValidationError"
|
||||
|
||||
|
||||
def test_script_generation_missing_roles_field():
|
||||
"""Test that missing roles field returns validation error."""
|
||||
response = client.post(
|
||||
"/api/tts/script",
|
||||
json={
|
||||
"text": "Some text here"
|
||||
}
|
||||
)
|
||||
|
||||
# Our error handler converts 422 to 400
|
||||
assert response.status_code == 400
|
||||
data = response.json()
|
||||
assert "error" in data
|
||||
assert data["error"] == "ValidationError"
|
||||
|
||||
|
||||
def test_script_segments_have_unique_ids():
|
||||
"""Test that each segment has a unique ID."""
|
||||
response = client.post(
|
||||
"/api/tts/script",
|
||||
json={
|
||||
"text": "Alice: Hello!\nBob: Hi!\nAlice: How are you?\nBob: Good!",
|
||||
"roles": ["Alice", "Bob"]
|
||||
}
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
|
||||
segment_ids = [seg["id"] for seg in data["segments"]]
|
||||
# Check all IDs are unique
|
||||
assert len(segment_ids) == len(set(segment_ids))
|
||||
|
||||
|
||||
def test_script_generation_with_narrator_text():
|
||||
"""Test script generation with narrative text (no role markers)."""
|
||||
response = client.post(
|
||||
"/api/tts/script",
|
||||
json={
|
||||
"text": "The sun was setting over the horizon. It was a beautiful evening.",
|
||||
"roles": ["Narrator"]
|
||||
}
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
|
||||
assert len(data["segments"]) > 0
|
||||
# Should assign to Narrator since no role markers
|
||||
assert any(seg["speaker"] == "Narrator" for seg in data["segments"])
|
||||
|
||||
|
||||
def test_script_generation_with_mixed_content():
|
||||
"""Test script generation with both dialogue and narrative."""
|
||||
response = client.post(
|
||||
"/api/tts/script",
|
||||
json={
|
||||
"text": "The wizard looked into his crystal ball.\nWizard: I see danger approaching.\nThe knight heard the warning.",
|
||||
"roles": ["Wizard", "Narrator"]
|
||||
}
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
|
||||
assert len(data["segments"]) >= 2
|
||||
speakers = [seg["speaker"] for seg in data["segments"]]
|
||||
# Should have both Wizard and Narrator
|
||||
assert "Wizard" in speakers or "Narrator" in speakers
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
pytest.main([__file__, "-v"])
|
||||
Reference in New Issue
Block a user