Ver Fonte

修改cache逻辑

sequoia00 há 2 semanas atrás
pai
commit
dd694ba282
3 ficheiros alterados com 272 adições e 5 exclusões
  1. 1 0
      config.py
  2. 270 4
      main_server.py
  3. 1 1
      start.sh

+ 1 - 0
config.py

@@ -9,5 +9,6 @@ MYSQL_DATABASE = os.getenv("READER_PRO_DB_NAME", "reader_pro")
 
 # TTS upstream service configuration
 # "http://141.140.15.30:8028"
+#http://127.0.0.1:8028
 TTS_API_BASE_URL = os.getenv("READER_PRO_TTS_API_BASE_URL", "http://127.0.0.1:8028")
 TTS_GENERATE_ENDPOINT = os.getenv("READER_PRO_TTS_GENERATE_ENDPOINT", "/generate")

+ 270 - 4
main_server.py

@@ -1,7 +1,7 @@
 from fastapi import FastAPI, File, UploadFile, HTTPException, Form, Response, Cookie
 from fastapi.encoders import jsonable_encoder
 from fastapi.middleware.cors import CORSMiddleware
-from fastapi.responses import JSONResponse, RedirectResponse, StreamingResponse, HTMLResponse
+from fastapi.responses import FileResponse, JSONResponse, RedirectResponse, StreamingResponse, HTMLResponse
 from fastapi.staticfiles import StaticFiles
 from pydantic import BaseModel
 
@@ -15,6 +15,7 @@ import io
 import logging
 import base64
 import json
+import wave
 from datetime import datetime, timedelta, timezone
 import secrets
 import pymysql
@@ -58,6 +59,7 @@ app.mount("/static", StaticFiles(directory="static"), name="static")
 # Audio cache directory
 CACHE_DIR = "audio_cache"
 os.makedirs(CACHE_DIR, exist_ok=True)
+MAX_CACHE_ENTRIES = 100
 
 SESSION_COOKIE = "reader_pro_session"
 SESSION_TTL_DAYS = 1
@@ -65,6 +67,128 @@ SESSION_TTL_DAYS_REMEMBER = 30
 TTS_GENERATE_URL = f"{TTS_API_BASE_URL.rstrip('/')}/{TTS_GENERATE_ENDPOINT.lstrip('/')}"
 
 
+def make_tts_cache_key(text: str, voice: str, speed: float, suffix: str = "") -> str:
+    cache_input = json.dumps(
+        {"text": text, "voice": voice, "speed": speed},
+        ensure_ascii=False,
+        sort_keys=True,
+        separators=(",", ":"),
+    )
+    return hashlib.md5(cache_input.encode("utf-8")).hexdigest() + suffix
+
+
+def cache_paths_for_tts(text: str, voice: str, speed: float) -> tuple[str, str, str]:
+    cache_id = make_tts_cache_key(text, voice, speed)
+    return (
+        cache_id,
+        os.path.join(CACHE_DIR, f"{cache_id}.ndjson"),
+        os.path.join(CACHE_DIR, f"{cache_id}.wav"),
+    )
+
+
+def wav_bytes_from_segments(audio_segments: list[bytes]) -> bytes:
+    if not audio_segments:
+        return b""
+    if len(audio_segments) == 1:
+        return audio_segments[0]
+
+    output = io.BytesIO()
+    writer = None
+    try:
+        for segment in audio_segments:
+            with wave.open(io.BytesIO(segment), "rb") as reader:
+                params = reader.getparams()
+                frames = reader.readframes(reader.getnframes())
+                if writer is None:
+                    writer = wave.open(output, "wb")
+                    writer.setparams(params)
+                elif reader.getparams()[:3] != writer.getparams()[:3]:
+                    raise ValueError("WAV segment format mismatch")
+                writer.writeframes(frames)
+    except Exception:
+        logger.warning("failed to merge WAV segments; falling back to raw concatenation", exc_info=True)
+        return b"".join(audio_segments)
+    finally:
+        if writer is not None:
+            writer.close()
+
+    return output.getvalue()
+
+
+def audio_segments_from_ndjson_file(path: str) -> list[bytes]:
+    segments: list[bytes] = []
+    with open(path, "rb") as f:
+        for raw_line in f:
+            line = raw_line.strip()
+            if not line:
+                continue
+            try:
+                data = json.loads(line)
+            except json.JSONDecodeError:
+                continue
+            audio_b64 = data.get("audio")
+            if audio_b64:
+                segments.append(base64.b64decode(audio_b64))
+    return segments
+
+
+def write_wav_cache_from_ndjson(ndjson_path: str, wav_path: str) -> None:
+    audio_segments = audio_segments_from_ndjson_file(ndjson_path)
+    wav_bytes = wav_bytes_from_segments(audio_segments)
+    if wav_bytes:
+        tmp_path = f"{wav_path}.tmp.{os.getpid()}"
+        with open(tmp_path, "wb") as f:
+            f.write(wav_bytes)
+        os.replace(tmp_path, wav_path)
+
+
+def trim_audio_cache(max_entries: int = MAX_CACHE_ENTRIES) -> None:
+    entries: dict[str, dict[str, object]] = {}
+    try:
+        filenames = os.listdir(CACHE_DIR)
+    except OSError as e:
+        logger.warning("failed to list audio cache: %s", e)
+        return
+
+    for filename in filenames:
+        path = os.path.join(CACHE_DIR, filename)
+        if not os.path.isfile(path) or ".tmp." in filename:
+            continue
+
+        stem, ext = os.path.splitext(filename)
+        if ext not in {".ndjson", ".wav", ".mp3"}:
+            continue
+
+        try:
+            mtime = os.path.getmtime(path)
+        except OSError:
+            continue
+
+        entry = entries.setdefault(stem, {"mtime": mtime, "paths": []})
+        entry["mtime"] = min(float(entry["mtime"]), mtime)
+        paths = entry["paths"]
+        if isinstance(paths, list):
+            paths.append(path)
+
+    overflow = len(entries) - max_entries
+    if overflow <= 0:
+        return
+
+    old_entries = sorted(entries.items(), key=lambda item: float(item[1]["mtime"]))[:overflow]
+    for cache_id, entry in old_entries:
+        paths = entry["paths"]
+        if not isinstance(paths, list):
+            continue
+        for path in paths:
+            try:
+                os.remove(path)
+            except FileNotFoundError:
+                pass
+            except OSError as e:
+                logger.warning("failed to remove cache file %s: %s", path, e)
+        logger.info("audio cache trimmed: %s", cache_id)
+
+
 def db_conn():
     return pymysql.connect(
         host=MYSQL_HOST,
@@ -258,7 +382,7 @@ def login_page(session_token: Optional[str] = Cookie(default=None, alias=SESSION
 <head>
   <meta charset="utf-8"/>
   <meta name="viewport" content="width=device-width, initial-scale=1"/>
-  <title>VoiceFlow AI Reader 【小满TTS英文听书】 - 登录</title>
+  <title>VoiceFlow AI Reader</title>
   <style>
     *{box-sizing:border-box}
     body{margin:0;min-height:100vh;display:grid;place-items:center;background:linear-gradient(145deg,#f2f6fb,#e7eef8);font-family:Arial,sans-serif;color:#1f2937}
@@ -793,7 +917,30 @@ async def generate_proxy(request: TextToSpeechRequest):
     if not user_input:
         raise HTTPException(status_code=400, detail="输入文本为空")
 
+    cache_id, cache_path, wav_path = cache_paths_for_tts(user_input, request.voice, request.speed)
+
+    if os.path.exists(cache_path) and os.path.getsize(cache_path) > 0:
+        logger.info("generate cache hit: %s", cache_id)
+        if not os.path.exists(wav_path) or os.path.getsize(wav_path) == 0:
+            write_wav_cache_from_ndjson(cache_path, wav_path)
+
+        async def cached_stream() -> AsyncGenerator[bytes, None]:
+            with open(cache_path, "rb") as f:
+                while True:
+                    chunk = f.read(64 * 1024)
+                    if not chunk:
+                        break
+                    yield chunk
+                    await asyncio.sleep(0)
+
+        return StreamingResponse(cached_stream(), media_type="application/x-ndjson")
+
     async def stream_generator() -> AsyncGenerator[bytes, None]:
+        tmp_path = f"{cache_path}.tmp.{os.getpid()}.{id(request)}"
+        wav_tmp_path = f"{wav_path}.tmp.{os.getpid()}.{id(request)}"
+        cache_complete = False
+        audio_segments: list[bytes] = []
+        text_buffer = ""
         try:
             async with aiohttp.ClientSession() as session:
                 async with session.post(
@@ -804,15 +951,130 @@ async def generate_proxy(request: TextToSpeechRequest):
                     if response.status != 200:
                         raise HTTPException(status_code=500, detail="TTS API 请求失败")
 
-                    async for chunk in response.content.iter_any():
-                        yield chunk
+                    with open(tmp_path, "wb") as cache_file:
+                        async for chunk in response.content.iter_any():
+                            cache_file.write(chunk)
+                            text_buffer += chunk.decode("utf-8")
+                            lines = text_buffer.split("\n")
+                            text_buffer = lines[-1]
+                            for line in lines[:-1]:
+                                if not line.strip():
+                                    continue
+                                try:
+                                    data = json.loads(line)
+                                except json.JSONDecodeError:
+                                    continue
+                                audio_b64 = data.get("audio")
+                                if audio_b64:
+                                    audio_segments.append(base64.b64decode(audio_b64))
+                            yield chunk
+                    if text_buffer.strip():
+                        try:
+                            data = json.loads(text_buffer)
+                            audio_b64 = data.get("audio")
+                            if audio_b64:
+                                audio_segments.append(base64.b64decode(audio_b64))
+                        except json.JSONDecodeError:
+                            pass
+                    cache_complete = True
         except Exception as e:
             logger.error(f"generate proxy error: {str(e)}")
             raise HTTPException(status_code=500, detail=str(e))
+        finally:
+            if cache_complete and os.path.exists(tmp_path) and os.path.getsize(tmp_path) > 0:
+                os.replace(tmp_path, cache_path)
+                wav_bytes = wav_bytes_from_segments(audio_segments)
+                if wav_bytes:
+                    with open(wav_tmp_path, "wb") as wav_file:
+                        wav_file.write(wav_bytes)
+                    os.replace(wav_tmp_path, wav_path)
+                logger.info("generate cache stored: %s", cache_id)
+                trim_audio_cache()
+            elif os.path.exists(tmp_path):
+                os.remove(tmp_path)
+            if os.path.exists(wav_tmp_path):
+                os.remove(wav_tmp_path)
 
     return StreamingResponse(stream_generator(), media_type="application/x-ndjson")
 
 
+@app.post("/generate-audio")
+async def generate_audio_file(request: TextToSpeechRequest):
+    user_input = request.user_input.strip()
+    if not user_input:
+        raise HTTPException(status_code=400, detail="输入文本为空")
+
+    cache_id, ndjson_path, wav_path = cache_paths_for_tts(user_input, request.voice, request.speed)
+    if os.path.exists(wav_path) and os.path.getsize(wav_path) > 0:
+        return FileResponse(wav_path, media_type="audio/wav", filename=f"{cache_id}.wav")
+
+    if os.path.exists(ndjson_path) and os.path.getsize(ndjson_path) > 0:
+        write_wav_cache_from_ndjson(ndjson_path, wav_path)
+        if os.path.exists(wav_path) and os.path.getsize(wav_path) > 0:
+            return FileResponse(wav_path, media_type="audio/wav", filename=f"{cache_id}.wav")
+
+    ndjson_tmp_path = f"{ndjson_path}.tmp.{os.getpid()}.{id(request)}"
+    wav_tmp_path = f"{wav_path}.tmp.{os.getpid()}.{id(request)}"
+    audio_segments: list[bytes] = []
+    text_buffer = ""
+    try:
+        async with aiohttp.ClientSession() as session:
+            async with session.post(
+                TTS_GENERATE_URL,
+                headers={"Content-Type": "application/json"},
+                json={"text": user_input, "voice": request.voice, "speed": request.speed},
+            ) as response:
+                if response.status != 200:
+                    raise HTTPException(status_code=500, detail="TTS API 请求失败")
+
+                with open(ndjson_tmp_path, "wb") as ndjson_file:
+                    async for chunk in response.content.iter_any():
+                        ndjson_file.write(chunk)
+                        text_buffer += chunk.decode("utf-8")
+                        lines = text_buffer.split("\n")
+                        text_buffer = lines[-1]
+                        for line in lines[:-1]:
+                            if not line.strip():
+                                continue
+                            try:
+                                data = json.loads(line)
+                            except json.JSONDecodeError:
+                                continue
+                            audio_b64 = data.get("audio")
+                            if audio_b64:
+                                audio_segments.append(base64.b64decode(audio_b64))
+
+        if text_buffer.strip():
+            try:
+                data = json.loads(text_buffer)
+                audio_b64 = data.get("audio")
+                if audio_b64:
+                    audio_segments.append(base64.b64decode(audio_b64))
+            except json.JSONDecodeError:
+                pass
+
+        wav_bytes = wav_bytes_from_segments(audio_segments)
+        if not wav_bytes:
+            raise HTTPException(status_code=500, detail="TTS API 未返回音频")
+
+        with open(wav_tmp_path, "wb") as wav_file:
+            wav_file.write(wav_bytes)
+        os.replace(ndjson_tmp_path, ndjson_path)
+        os.replace(wav_tmp_path, wav_path)
+        logger.info("generate audio cache stored: %s", cache_id)
+        trim_audio_cache()
+        return FileResponse(wav_path, media_type="audio/wav", filename=f"{cache_id}.wav")
+    except HTTPException:
+        raise
+    except Exception as e:
+        logger.error(f"generate audio file error: {str(e)}")
+        raise HTTPException(status_code=500, detail=str(e))
+    finally:
+        for tmp_path in (ndjson_tmp_path, wav_tmp_path):
+            if os.path.exists(tmp_path):
+                os.remove(tmp_path)
+
+
 @app.post("/text-to-speech/")
 async def text_to_speech(request: TextToSpeechRequest):
     user_input = request.user_input.strip()
@@ -873,6 +1135,7 @@ async def text_to_speech(request: TextToSpeechRequest):
                     full_audio.seek(0)
                     with open(audio_path, "wb") as f:
                         f.write(full_audio.getvalue())
+                    trim_audio_cache()
 
         except Exception as e:
             logger.error(f"TTS error: {str(e)}")
@@ -944,6 +1207,7 @@ async def generate_api_audio(chunk: str, voice: str, speed: float) -> AsyncGener
                                     yield audio_bytes
                                     with open(audio_path, "wb") as f:
                                         f.write(audio_bytes)
+                                    trim_audio_cache()
                             except json.JSONDecodeError as e:
                                 logger.error(f"JSON decode error: {str(e)}")
                                 continue
@@ -956,6 +1220,7 @@ async def generate_api_audio(chunk: str, voice: str, speed: float) -> AsyncGener
                                 yield audio_bytes
                                 with open(audio_path, "wb") as f:
                                     f.write(audio_bytes)
+                                trim_audio_cache()
                         except json.JSONDecodeError:
                             pass
 
@@ -988,6 +1253,7 @@ async def page_to_speech(request: TextToSpeechRequest):
         full_audio_buffer.seek(0)
         with open(full_audio_path, "wb") as f:
             f.write(full_audio_buffer.getvalue())
+        trim_audio_cache()
 
     return StreamingResponse(audio_generator(), media_type="audio/wav")
 

+ 1 - 1
start.sh

@@ -1 +1 @@
-nohup uvicorn main_server:app --host 0.0.0.0 --port 8005  &
+nohup uvicorn main_server:app --host 0.0.0.0 --port 8005 --reload &