99 lines
3.8 KiB
Python
99 lines
3.8 KiB
Python
import json
|
|
import asyncio
|
|
import os
|
|
import httpx
|
|
from loguru import logger
|
|
from typing import List, Dict
|
|
|
|
from server.config import resolve_ollama_url
|
|
|
|
OLLAMA_URL = resolve_ollama_url(
|
|
os.environ.get("OLLAMA_URL", "http://127.0.0.1:11434/api/generate")
|
|
)
|
|
MODEL_NAME = os.environ.get("MODEL_NAME", "gemma4:e2b")
|
|
DATA_DIR = os.environ.get("DATA_DIR", "data")
|
|
INPUT_FILE = os.path.join(DATA_DIR, "output.json")
|
|
OUTPUT_FILE = os.path.join(DATA_DIR, "scripts.json")
|
|
|
|
async def generate_script(client: httpx.AsyncClient, headline: str, content: str, attempt: int = 1) -> str:
|
|
system_prompt = (
|
|
"You are an expert newscaster script writer. Your sole purpose is to convert articles into scripts for a TTS engine. "
|
|
"CRITICAL RULE: Provide ONLY the literal words to be spoken by the news anchor. "
|
|
"STRICTLY FORBIDDEN: Do not include peripheral annotations, stage directions, or non-spoken text. "
|
|
"Specifically, do NOT include:\n"
|
|
"- Bracketed text (e.g., [Music], [Pause], [Cut to B-roll])\n"
|
|
"- Parenthetical notes (e.g., (smiling), (slowly))\n"
|
|
"- Speaker labels or names (e.g., 'Anchor:', 'Reporter:', 'Host:')\n"
|
|
"- Scene headings, timestamps, or markdown formatting.\n"
|
|
"The output must be a clean stream of spoken text ready for a teleprompter."
|
|
)
|
|
prompt = f"Write a detailed, lengthy newscaster's script based on the following text:\n\n```text\n{content}```"
|
|
payload = {
|
|
"model": MODEL_NAME,
|
|
"prompt": prompt,
|
|
"system": system_prompt,
|
|
"stream": False
|
|
}
|
|
|
|
try:
|
|
logger.info(f"Generating script for: {headline} (Attempt {attempt})")
|
|
# High timeout for slow local LLM
|
|
response = await client.post(OLLAMA_URL, json=payload, timeout=600.0)
|
|
response.raise_for_status()
|
|
result = response.json()
|
|
return result.get("response", "")
|
|
except (httpx.HTTPStatusError, httpx.RequestError) as e:
|
|
if attempt < 3:
|
|
wait = attempt * 10
|
|
logger.warning(f"Error occurred for {headline}: {e}. Retrying in {wait}s...")
|
|
await asyncio.sleep(wait)
|
|
return await generate_script(client, headline, content, attempt + 1)
|
|
logger.error(f"Max retries reached for {headline}. Final error: {e}")
|
|
return ""
|
|
except Exception:
|
|
logger.exception(f"Unexpected error occurred while requesting script for {headline}")
|
|
return ""
|
|
|
|
async def main():
|
|
try:
|
|
with open(INPUT_FILE, "r", encoding="utf-8") as f:
|
|
data = json.load(f)
|
|
except (FileNotFoundError, json.JSONDecodeError) as e:
|
|
logger.error(f"Failed to read {INPUT_FILE}: {e}")
|
|
return
|
|
|
|
articles = data.get("articles", [])
|
|
if not articles:
|
|
logger.warning("No articles found in the input file.")
|
|
return
|
|
|
|
results = []
|
|
async with httpx.AsyncClient() as client:
|
|
for i, article in enumerate(articles):
|
|
headline = article.get("headline", "Unknown Headline")
|
|
body = article.get("body", "")
|
|
|
|
if not body:
|
|
logger.warning(f"Skipping article {headline} because body is empty.")
|
|
continue
|
|
|
|
script = await generate_script(client, headline, body)
|
|
if script:
|
|
results.append({
|
|
"headline": headline,
|
|
"script": script
|
|
})
|
|
# Incremental save
|
|
with open(OUTPUT_FILE, "w", encoding="utf-8") as f:
|
|
json.dump(results, f, indent=4, ensure_ascii=False)
|
|
else:
|
|
logger.warning(f"Failed to generate script for {headline}")
|
|
|
|
# Avoid overloading Ollama
|
|
await asyncio.sleep(2)
|
|
|
|
logger.info(f"Completed processing. Total scripts written: {len(results)}")
|
|
|
|
if __name__ == "__main__":
|
|
asyncio.run(main())
|