ydy0615/llm-in-text
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: LLM 应用网页开发及内联建议功能实现

Browse Source
This commit is contained in:
ydy0615
2026-04-05 13:42:29 +08:00
parent 9904b9bd78
commit 68ed783d6c
13 changed files with 800 additions and 513 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/llm.py
+7 -29
View File
@@ -6,44 +6,22 @@ from datetime import datetime
import ollama
from dotenv import load_dotenv
from prompts import get_vlm_ocr_prompt
load_dotenv()
OLLAMA_MODEL = os.getenv('OLLAMA_MODEL', 'gpt-oss:20b')
OLLAMA_HOST = os.getenv('OLLAMA_HOST', 'http://localhost:11434')
VLM_MODEL = os.getenv('VLM_MODEL', 'qwen3-vl:30b')
# Timeouts in seconds
COMPLETION_TIMEOUT = 30
OCR_TIMEOUT = 60
CONVERT_TIMEOUT = 30
# Timeouts in seconds (10 minutes for large model loading)
COMPLETION_TIMEOUT = 600
OCR_TIMEOUT = 120
CONVERT_TIMEOUT = 60
client = ollama.AsyncClient(host=OLLAMA_HOST)
logger = logging.getLogger("llm")
VLM_OCR_CONTEXT_PROMPT = """You are an OCR and visual-context extractor for markdown writing assistance.
Your output will be embedded inside an HTML comment as hidden context for a text-completion model.
Requirements:
- Keep output compact: maximum 120 words.
- Use plain text only (no markdown code fences).
- Never output <!-- or -->.
- Do not invent unreadable text; mark uncertain characters with ?.
- Preserve original script for recognized text (do not forcibly translate).
Return exactly this format:
TEXT:
<exact transcription of visible text; use " | " for line breaks; write "(none)" if no readable text>
KEY_DETAILS:
- <3-5 short factual bullets about relevant objects/layout>
LANGUAGE:
<dominant language(s) in visible text, e.g. English / Chinese / Mixed>
SUMMARY:
<one short sentence, <= 20 words>"""
def _extract_message(response) -> tuple[str, str]:
content = ""
@@ -166,7 +144,7 @@ async def call_vlm_ocr(image_bytes: bytes, language: str = 'auto') -> str:
model=VLM_MODEL,
messages=[{
'role': 'user',
'content': VLM_OCR_CONTEXT_PROMPT,
'content': get_vlm_ocr_prompt(),
'images': [image_bytes]
}],
stream=False,
backend/main.py
+11 -2
View File
@@ -26,6 +26,15 @@ logging.basicConfig(
)
logger = logging.getLogger("api")
_markitdown_instance = None
def _get_markitdown():
global _markitdown_instance
if _markitdown_instance is None:
_markitdown_instance = markitdown.MarkItDown()
return _markitdown_instance
app = FastAPI()
ACTIVE_COMPLETIONS: dict[str, asyncio.Task] = {}
@@ -310,8 +319,8 @@ async def convert_to_markdown(request: ConvertRequest, api_key: str = Security(g
try:
# Convert using MarkItDown
md = markitdown.MarkItDown()
result = md.convert(tmp_path)
md = _get_markitdown()
result = await asyncio.to_thread(md.convert, tmp_path)
markdown_text = _sanitize_converted_markdown(result.text_content)
logger.info(
backend/prompt.py
+12 -316
View File
@@ -2,6 +2,8 @@ from datetime import datetime, timedelta, timezone
import re
from typing import Protocol, Tuple, runtime_checkable
from prompts import get_language_guidance_map, get_system_prompt_template, get_inline_examples
@runtime_checkable
class UserPreferences(Protocol):
@@ -224,339 +226,33 @@ def _canonical_language_id(language_id: str) -> str:
_JS_LANGS = {"javascript", "typescript"}
_CODE_LANGS = {"python", "go", "rust", "java", "kotlin", "swift", "ruby", "php", "lua", "c", "cpp", "csharp", "r", "matlab", "dart"}
_LANG_GUIDANCE = {
"mermaid": """
Language-specific guidance (mermaid):
- Output valid Mermaid syntax only.
- Prefer concise, syntactically correct diagram statements.
- Avoid prose unless the user prompt explicitly requires it.""",
"latex": """
Language-specific guidance (latex):
- Output LaTeX math content only when completing LaTeX.
- If CURSOR_IN_FENCED_CODE_BLOCK=true and CURSOR_FENCE_LANGUAGE is latex/tex/katex:
- Output raw LaTeX lines only.
- Do not wrap with $ or $$.""",
"json": """
Language-specific guidance (json):
- Output strict JSON only (no comments, no trailing commas).
- Ensure valid quotes and braces.""",
"yaml": """
Language-specific guidance (yaml):
- Output valid YAML only.
- Use consistent indentation and avoid tabs.""",
"toml": """
Language-specific guidance (toml):
- Output valid TOML only.
- Keep key types consistent.""",
"ini": """
Language-specific guidance (ini):
- Output valid INI only.
- Keep section headers and key=value pairs consistent.""",
"sql": """
Language-specific guidance (sql):
- Output a single, valid SQL statement unless context requires multiple.
- Prefer ANSI SQL when dialect is unclear.""",
"bash": """
Language-specific guidance (bash):
- Output POSIX-compatible shell when possible.
- Avoid interactive prompts or destructive commands unless requested.""",
"powershell": """
Language-specific guidance (powershell):
- Output valid PowerShell commands.
- Avoid destructive commands unless explicitly requested.""",
"html": """
Language-specific guidance (html):
- Output valid HTML only.
- Keep markup minimal and well-formed.""",
"css": """
Language-specific guidance (css):
- Output valid CSS only.
- Use concise, readable selectors.""",
"diff": """
Language-specific guidance (diff):
- Output a unified diff only.
- Ensure @@ hunk headers and +/- lines are consistent.""",
"regex": """
Language-specific guidance (regex):
- Output the regex pattern only.
- Avoid delimiters unless explicitly requested.""",
"text": """
Language-specific guidance (text):
- Output plain text only.
- Avoid markdown formatting unless explicitly asked.""",
"xml": """
Language-specific guidance (xml):
- Output well-formed XML only.
- Ensure matching tags and proper escaping.""",
"dockerfile": """
Language-specific guidance (dockerfile):
- Output valid Dockerfile instructions only.
- Keep layers minimal and ordered logically.""",
"makefile": """
Language-specific guidance (makefile):
- Output valid Makefile syntax only.
- Use tabs for recipe lines.""",
}
_GENERIC_CODE = """
Language-specific guidance ({lang}):
- Output valid {lang} code.
- Avoid prose unless context clearly expects comments or docstrings."""
_JS_CODE = """
Language-specific guidance ({lang}):
- Output valid {lang} code.
- Prefer modern syntax and avoid prose unless comments are needed."""
def _language_guidance(language_id: str) -> str:
canonical = _canonical_language_id(language_id)
if canonical == "markdown":
return ""
guidance = _LANG_GUIDANCE.get(canonical)
guidance_map = get_language_guidance_map()
guidance = guidance_map.get(canonical)
if guidance:
return guidance
if canonical in _JS_LANGS:
return _JS_CODE.format(lang=canonical)
return guidance_map.get("_js_code", "").replace("{lang}", canonical)
if canonical in _CODE_LANGS:
return _GENERIC_CODE.format(lang=canonical)
return _GENERIC_CODE.format(lang=canonical)
return guidance_map.get("_generic_code", "").replace("{lang}", canonical)
return guidance_map.get("_generic_code", "").replace("{lang}", canonical)
def build_inline_system_prompt(language_id: str = "markdown") -> str:
safe_language_id = _canonical_language_id(language_id)
language_guidance = _language_guidance(safe_language_id)
system_prompt = f"""You are an inline completion engine for a {safe_language_id} editor with ghost-text suggestions.
Return only the insertion text that should be placed between PREFIX and SUFFIX.
CORE PRINCIPLE: Output insertion text only. No explanations, no meta labels, no wrapper quotes.
PRIORITY 1: CONTEXT AWARENESS (Read these flags from user prompt)
- CURSOR_IN_FENCED_CODE_BLOCK: Are you inside a code fence?
- CURSOR_FENCE_LANGUAGE: What language is the current fence?
- PREFIX_ENDS_WITH_NEWLINE: Does prefix end with newline?
- SUFFIX_STARTS_WITH_NEWLINE: Does suffix start with newline?
- MERMAID_CONTEXT: Is this a Mermaid diagram context?
PRIORITY 2: SPECIALIZED CONTENT RULES
2.1 Code Block Handling:
If CURSOR_IN_FENCED_CODE_BLOCK=true:
- You are inside a code fence
- Output code lines ONLY (no triple backticks)
- Use single \\n for code line separation
If CURSOR_IN_FENCED_CODE_BLOCK=false and code needed:
- Wrap code in fenced block with language tag:
```{{language}}
code here
```
- Never use inline backticks for code snippets
2.2 Math Formatting (KaTeX):
- Inline math: wrap with $...$
- Block math: wrap with $$...$$
- Never output bare formulas
- Exception: inside latex/tex/katex fence, output raw LaTeX
2.3 Mermaid Diagrams:
If CURSOR_FENCE_LANGUAGE=mermaid:
- Output Mermaid syntax ONLY
- No backticks, no explanations
If MERMAID_CONTEXT=true and outside fence:
- Output complete fenced block:
```mermaid
diagram syntax
```
PRIORITY 3: MARKDOWN STRUCTURE
3.1 Newline Semantics:
- Single \\n: soft break (same paragraph, renders as space or <br>)
- Double \\n\\n: hard break (new paragraph/block)
- Use \\n\\n for: new paragraphs, before headings, starting lists/tables
- Use \\n for: continuation within blocks (list items, table cells)
- Exception: inside code blocks, use \\n freely for code lines
3.2 Boundary Management:
Check PREFIX_ENDS_WITH_NEWLINE and SUFFIX_STARTS_WITH_NEWLINE:
- If PREFIX lacks needed newline: start OUTPUT with \\n
- If SUFFIX lacks needed newline: end OUTPUT with \\n
- Common cases requiring leading \\n:
* Starting a list after "Steps:"
* Creating new paragraph after text
* Adding heading after paragraph
- Common cases requiring trailing \\n:
* Before new heading
* End of section
3.3 Context Stitching:
- Never repeat text from SUFFIX beginning
- Match PREFIX tone, style, indentation
- Continue structures: lists, tables, quotes, headings
PRIORITY 4: HIDDEN CONTEXT
- OCR metadata like <OCR:...> is hidden context
- Never copy OCR tags to output
- Use OCR content as semantic hint only
"""
template = get_system_prompt_template()
system_prompt = template.replace("{language_id}", safe_language_id)
if language_guidance:
system_prompt = f"{system_prompt.rstrip()}\\n{language_guidance.strip()}"
system_prompt = f"{system_prompt.rstrip()}\n{language_guidance.strip()}"
return system_prompt.strip()
INLINE_EXAMPLES = """=== CATEGORY A: PROSE CONTINUATION ===
[EX01] Simple prose continuation
<PREFIX>The quick brown fox </PREFIX>
<SUFFIX>jumps over the lazy dog.</SUFFIX>
Expected OUTPUT:
moved quietly and then
[EX02] Avoid repeating suffix
<PREFIX>Our launch plan starts with </PREFIX>
<SUFFIX>phase one, followed by phase two.</SUFFIX>
Expected OUTPUT:
careful internal testing before
WRONG: phase one starts with (repeats suffix)
=== CATEGORY B: MARKDOWN STRUCTURES ===
[EX03] Continue checklist
<PREFIX>## TODO
- [ ] Buy milk
- [ ] </PREFIX>
<SUFFIX></SUFFIX>
Expected OUTPUT:
Write release notes and share draft with team
[EX04] Start list after header (PREFIX lacks newline)
PREFIX_ENDS_WITH_NEWLINE=false
<PREFIX>Deployment steps:</PREFIX>
<SUFFIX></SUFFIX>
Expected OUTPUT:
- Build artifact
- Deploy service
[EX05] Continue table row
<PREFIX>| Name | Score |
| --- | --- |
| Alice | 92 |
| Bob | </PREFIX>
<SUFFIX></SUFFIX>
Expected OUTPUT:
88 |
[EX06] Start new paragraph
<PREFIX>First paragraph ends.</PREFIX>
<SUFFIX></SUFFIX>
Expected OUTPUT:
Second paragraph starts.
WRONG: Second paragraph starts. (missing leading \\n\\n)
[EX07] Add newline before heading
PREFIX_ENDS_WITH_NEWLINE=false
<PREFIX>End of previous section.</PREFIX>
<SUFFIX>## Next Heading</SUFFIX>
Expected OUTPUT:
WRONG: (would join with heading without separation)
=== CATEGORY C: CODE BLOCKS ===
[EX08] Outside fence: wrap code in fence
CURSOR_IN_FENCED_CODE_BLOCK=false
<PREFIX>Parse this JSON payload in Python:</PREFIX>
<SUFFIX></SUFFIX>
Expected OUTPUT:
```python
import json
data = json.loads(payload)
```
WRONG: import json\\ndata = json.loads(payload) (no fence)
[EX09] Inside fence: output code only
CURSOR_IN_FENCED_CODE_BLOCK=true
<PREFIX>```python
def add(a, b):
return </PREFIX>
<SUFFIX>
```</SUFFIX>
Expected OUTPUT:
a + b
WRONG: ```python\\nreturn a + b\\n``` (duplicate fences)
[EX10] Code inside fence uses single newline
CURSOR_IN_FENCED_CODE_BLOCK=true
<PREFIX>```python
def hello():</PREFIX>
<SUFFIX>
```</SUFFIX>
Expected OUTPUT:
print("Hello")
return True
(Note: single \\n between code lines, no markdown rules)
=== CATEGORY D: MATH ===
[EX11] Inline math
<PREFIX>The derivative of x^2 is </PREFIX>
<SUFFIX>.</SUFFIX>
Expected OUTPUT:
$2x$
WRONG: 2x (bare formula)
[EX12] Block math
<PREFIX>We can write the Gaussian integral as:</PREFIX>
<SUFFIX></SUFFIX>
Expected OUTPUT:
$$
\\int_{-\\infty}^{\\infty} e^{-x^2}\\,dx = \\sqrt{\\pi}
$$
WRONG: \\int... (bare formula without $$)
=== CATEGORY E: MERMAID ===
[EX13] Inside mermaid fence
CURSOR_FENCE_LANGUAGE=mermaid
CURSOR_IN_FENCED_CODE_BLOCK=true
<PREFIX>```mermaid
flowchart TD
A[Start] --> </PREFIX>
<SUFFIX>
```</SUFFIX>
Expected OUTPUT:
B{Valid?}
B -->|Yes| C[Done]
WRONG: ```mermaid\\nB{Valid?}... (duplicate fence)
[EX14] Outside fence with mermaid context
CURSOR_IN_FENCED_CODE_BLOCK=false
MERMAID_CONTEXT=true
<PREFIX>Please provide a simple release pipeline diagram.</PREFIX>
<SUFFIX></SUFFIX>
Expected OUTPUT:
```mermaid
flowchart LR
Build --> Test --> Deploy
```
=== CATEGORY F: OCR METADATA ===
[EX15] Use OCR as context, never output
<PREFIX>![whiteboard](img.png) <OCR:equation y = mx + b>
The relationship is </PREFIX>
<SUFFIX>.</SUFFIX>
Expected OUTPUT:
$y = mx + b$
WRONG: <OCR:equation y = mx + b> (OCR tag in output)"""
_INLINE_EXAMPLES = get_inline_examples()
def build_completion_prompts(
@@ -636,7 +332,7 @@ Step 3: Choose newline type
- Do not repeat text from SUFFIX beginning
=== EXAMPLES BY CATEGORY ===
{INLINE_EXAMPLES}
{_INLINE_EXAMPLES}
=== NOW COMPLETE THE TASK ===
backend/prompts/__init__.py
+44
View File
@@ -0,0 +1,44 @@
import json
from pathlib import Path
from typing import Any
_PROMPTS_DIR = Path(__file__).parent
class PromptManager:
_instance = None
_data: dict[str, Any] = {}
def __new__(cls):
if cls._instance is None:
cls._instance = super().__new__(cls)
cls._instance._load_all()
return cls._instance
def _load_all(self):
for json_file in _PROMPTS_DIR.glob("*.json"):
key = json_file.stem
with open(json_file, "r", encoding="utf-8") as f:
self._data[key] = json.load(f)
def get(self, key: str, default: Any = None) -> Any:
return self._data.get(key, default)
_prompts = PromptManager()
def get_system_prompt_template() -> str:
return _prompts.get("system_prompt", {}).get("template", "")
def get_language_guidance_map() -> dict[str, str]:
return _prompts.get("language_guidance", {})
def get_inline_examples() -> str:
return _prompts.get("inline_examples", {}).get("content", "")
def get_vlm_ocr_prompt() -> str:
return _prompts.get("vlm_ocr", {}).get("prompt", "")
backend/prompts/inline_examples.json
+3
View File
@@ -0,0 +1,3 @@
{
"content": "=== CATEGORY A: PROSE CONTINUATION ===\n\n[EX01] Simple prose continuation\n<PREFIX>The quick brown fox </PREFIX>\n<SUFFIX>jumps over the lazy dog.</SUFFIX>\nExpected OUTPUT:\nmoved quietly and then\n\n[EX02] Avoid repeating suffix\n<PREFIX>Our launch plan starts with </PREFIX>\n<SUFFIX>phase one, followed by phase two.</SUFFIX>\nExpected OUTPUT:\ncareful internal testing before\nWRONG: phase one starts with (repeats suffix)\n\n=== CATEGORY B: MARKDOWN STRUCTURES ===\n\n[EX03] Continue checklist\n<PREFIX>## TODO\n- [ ] Buy milk\n- [ ] </PREFIX>\n<SUFFIX></SUFFIX>\nExpected OUTPUT:\nWrite release notes and share draft with team\n\n[EX04] Start list after header (PREFIX lacks newline)\nPREFIX_ENDS_WITH_NEWLINE=false\n<PREFIX>Deployment steps:</PREFIX>\n<SUFFIX></SUFFIX>\nExpected OUTPUT:\n\n- Build artifact\n- Deploy service\n\n[EX05] Continue table row\n<PREFIX>| Name | Score |\n| --- | --- |\n| Alice | 92 |\n| Bob | </PREFIX>\n<SUFFIX></SUFFIX>\nExpected OUTPUT:\n88 |\n\n[EX06] Start new paragraph\n<PREFIX>First paragraph ends.</PREFIX>\n<SUFFIX></SUFFIX>\nExpected OUTPUT:\n\nSecond paragraph starts.\nWRONG: Second paragraph starts. (missing leading \\n\\n)\n\n[EX07] Add newline before heading\nPREFIX_ENDS_WITH_NEWLINE=false\n<PREFIX>End of previous section.</PREFIX>\n<SUFFIX>## Next Heading</SUFFIX>\nExpected OUTPUT:\n\nWRONG: (would join with heading without separation)\n\n=== CATEGORY C: CODE BLOCKS ===\n\n[EX08] Outside fence: wrap code in fence\nCURSOR_IN_FENCED_CODE_BLOCK=false\n<PREFIX>Parse this JSON payload in Python:</PREFIX>\n<SUFFIX></SUFFIX>\nExpected OUTPUT:\n```python\nimport json\ndata = json.loads(payload)\n```\nWRONG: import json\\ndata = json.loads(payload) (no fence)\n\n[EX09] Inside fence: output code only\nCURSOR_IN_FENCED_CODE_BLOCK=true\n<PREFIX>```python\ndef add(a, b):\nreturn </PREFIX>\n<SUFFIX>\n```</SUFFIX>\nExpected OUTPUT:\na + b\nWRONG: ```python\\nreturn a + b\\n``` (duplicate fences)\n\n[EX10] Code inside fence uses single newline\nCURSOR_IN_FENCED_CODE_BLOCK=true\n<PREFIX>```python\ndef hello():</PREFIX>\n<SUFFIX>\n```</SUFFIX>\nExpected OUTPUT:\nprint(\"Hello\")\nreturn True\n(Note: single \\n between code lines, no markdown rules)\n\n=== CATEGORY D: MATH ===\n\n[EX11] Inline math\n<PREFIX>The derivative of x^2 is </PREFIX>\n<SUFFIX>.</SUFFIX>\nExpected OUTPUT:\n$2x$\nWRONG: 2x (bare formula)\n\n[EX12] Block math\n<PREFIX>We can write the Gaussian integral as:</PREFIX>\n<SUFFIX></SUFFIX>\nExpected OUTPUT:\n$$\n\\int_{-\\infty}^{\\infty} e^{-x^2}\\,dx = \\sqrt{\\pi}\n$$\nWRONG: \\int... (bare formula without $$)\n\n=== CATEGORY E: MERMAID ===\n\n[EX13] Inside mermaid fence\nCURSOR_FENCE_LANGUAGE=mermaid\nCURSOR_IN_FENCED_CODE_BLOCK=true\n<PREFIX>```mermaid\nflowchart TD\nA[Start] --> </PREFIX>\n<SUFFIX>\n```</SUFFIX>\nExpected OUTPUT:\nB{Valid?}\nB -->|Yes| C[Done]\nWRONG: ```mermaid\\nB{Valid?}... (duplicate fence)\n\n[EX14] Outside fence with mermaid context\nCURSOR_IN_FENCED_CODE_BLOCK=false\nMERMAID_CONTEXT=true\n<PREFIX>Please provide a simple release pipeline diagram.</PREFIX>\n<SUFFIX></SUFFIX>\nExpected OUTPUT:\n```mermaid\nflowchart LR\nBuild --> Test --> Deploy\n```\n\n=== CATEGORY F: OCR METADATA ===\n\n[EX15] Use OCR as context, never output\n<PREFIX>![whiteboard](img.png) <OCR:equation y = mx + b>\nThe relationship is </PREFIX>\n<SUFFIX>.</SUFFIX>\nExpected OUTPUT:\n$y = mx + b$\nWRONG: <OCR:equation y = mx + b> (OCR tag in output)"
}
backend/prompts/language_guidance.json
+21
View File
@@ -0,0 +1,21 @@
{
"mermaid": "\nLanguage-specific guidance (mermaid):\n- Output valid Mermaid syntax only.\n- Prefer concise, syntactically correct diagram statements.\n- Avoid prose unless the user prompt explicitly requires it.",
"latex": "\nLanguage-specific guidance (latex):\n- Output LaTeX math content only when completing LaTeX.\n- If CURSOR_IN_FENCED_CODE_BLOCK=true and CURSOR_FENCE_LANGUAGE is latex/tex/katex:\n- Output raw LaTeX lines only.\n- Do not wrap with $ or $$.",
"json": "\nLanguage-specific guidance (json):\n- Output strict JSON only (no comments, no trailing commas).\n- Ensure valid quotes and braces.",
"yaml": "\nLanguage-specific guidance (yaml):\n- Output valid YAML only.\n- Use consistent indentation and avoid tabs.",
"toml": "\nLanguage-specific guidance (toml):\n- Output valid TOML only.\n- Keep key types consistent.",
"ini": "\nLanguage-specific guidance (ini):\n- Output valid INI only.\n- Keep section headers and key=value pairs consistent.",
"sql": "\nLanguage-specific guidance (sql):\n- Output a single, valid SQL statement unless context requires multiple.\n- Prefer ANSI SQL when dialect is unclear.",
"bash": "\nLanguage-specific guidance (bash):\n- Output POSIX-compatible shell when possible.\n- Avoid interactive prompts or destructive commands unless requested.",
"powershell": "\nLanguage-specific guidance (powershell):\n- Output valid PowerShell commands.\n- Avoid destructive commands unless explicitly requested.",
"html": "\nLanguage-specific guidance (html):\n- Output valid HTML only.\n- Keep markup minimal and well-formed.",
"css": "\nLanguage-specific guidance (css):\n- Output valid CSS only.\n- Use concise, readable selectors.",
"diff": "\nLanguage-specific guidance (diff):\n- Output a unified diff only.\n- Ensure @@ hunk headers and +/- lines are consistent.",
"regex": "\nLanguage-specific guidance (regex):\n- Output the regex pattern only.\n- Avoid delimiters unless explicitly requested.",
"text": "\nLanguage-specific guidance (text):\n- Output plain text only.\n- Avoid markdown formatting unless explicitly asked.",
"xml": "\nLanguage-specific guidance (xml):\n- Output well-formed XML only.\n- Ensure matching tags and proper escaping.",
"dockerfile": "\nLanguage-specific guidance (dockerfile):\n- Output valid Dockerfile instructions only.\n- Keep layers minimal and ordered logically.",
"makefile": "\nLanguage-specific guidance (makefile):\n- Output valid Makefile syntax only.\n- Use tabs for recipe lines.",
"_generic_code": "\nLanguage-specific guidance ({lang}):\n- Output valid {lang} code.\n- Avoid prose unless context clearly expects comments or docstrings.",
"_js_code": "\nLanguage-specific guidance ({lang}):\n- Output valid {lang} code.\n- Prefer modern syntax and avoid prose unless comments are needed."
}
backend/prompts/system_prompt.json
+3
View File
@@ -0,0 +1,3 @@
{
"template": "You are an inline completion engine for a {language_id} editor with ghost-text suggestions.\n\nReturn only the insertion text that should be placed between PREFIX and SUFFIX.\n\nCORE PRINCIPLE: Output insertion text only. No explanations, no meta labels, no wrapper quotes.\n\nPRIORITY 1: CONTEXT AWARENESS (Read these flags from user prompt)\n- CURSOR_IN_FENCED_CODE_BLOCK: Are you inside a code fence?\n- CURSOR_FENCE_LANGUAGE: What language is the current fence?\n- PREFIX_ENDS_WITH_NEWLINE: Does prefix end with newline?\n- SUFFIX_STARTS_WITH_NEWLINE: Does suffix start with newline?\n- MERMAID_CONTEXT: Is this a Mermaid diagram context?\n\nPRIORITY 2: SPECIALIZED CONTENT RULES\n\n2.1 Code Block Handling:\nIf CURSOR_IN_FENCED_CODE_BLOCK=true:\n- You are inside a code fence\n- Output code lines ONLY (no triple backticks)\n- Use single \\n for code line separation\n\nIf CURSOR_IN_FENCED_CODE_BLOCK=false and code needed:\n- Wrap code in fenced block with language tag:\n```{language}\ncode here\n```\n- Never use inline backticks for code snippets\n\n2.2 Math Formatting (KaTeX):\n- Inline math: wrap with $...$\n- Block math: wrap with $$...$$\n- Never output bare formulas\n- Exception: inside latex/tex/katex fence, output raw LaTeX\n\n2.3 Mermaid Diagrams:\nIf CURSOR_FENCE_LANGUAGE=mermaid:\n- Output Mermaid syntax ONLY\n- No backticks, no explanations\n\nIf MERMAID_CONTEXT=true and outside fence:\n- Output complete fenced block:\n```mermaid\ndiagram syntax\n```\n\nPRIORITY 3: MARKDOWN STRUCTURE\n\n3.1 Newline Semantics:\n- Single \\n: soft break (same paragraph, renders as space or <br>)\n- Double \\n\\n: hard break (new paragraph/block)\n- Use \\n\\n for: new paragraphs, before headings, starting lists/tables\n- Use \\n for: continuation within blocks (list items, table cells)\n- Exception: inside code blocks, use \\n freely for code lines\n\n3.2 Boundary Management:\nCheck PREFIX_ENDS_WITH_NEWLINE and SUFFIX_STARTS_WITH_NEWLINE:\n- If PREFIX lacks needed newline: start OUTPUT with \\n\n- If SUFFIX lacks needed newline: end OUTPUT with \\n\n- Common cases requiring leading \\n:\n* Starting a list after \"Steps:\"\n* Creating new paragraph after text\n* Adding heading after paragraph\n- Common cases requiring trailing \\n:\n* Before new heading\n* End of section\n\n3.3 Context Stitching:\n- Never repeat text from SUFFIX beginning\n- Match PREFIX tone, style, indentation\n- Continue structures: lists, tables, quotes, headings\n\nPRIORITY 4: HIDDEN CONTEXT\n- OCR metadata like <OCR:...> is hidden context\n- Never copy OCR tags to output\n- Use OCR content as semantic hint only"
}
backend/prompts/vlm_ocr.json
+3
View File
@@ -0,0 +1,3 @@
{
"prompt": "You are an OCR and visual-context extractor for markdown writing assistance.\n\nYour output will be embedded inside an HTML comment as hidden context for a text-completion model.\n\nRequirements:\n- Keep output compact: maximum 120 words.\n- Use plain text only (no markdown code fences).\n- Never output <!-- or -->.\n- Do not invent unreadable text; mark uncertain characters with ?.\n- Preserve original script for recognized text (do not forcibly translate).\n\nReturn exactly this format:\n\nTEXT:\n<exact transcription of visible text; use \" | \" for line breaks; write \"(none)\" if no readable text>\n\nKEY_DETAILS:\n- <3-5 short factual bullets about relevant objects/layout>\n\nLANGUAGE:\n<dominant language(s) in visible text, e.g. English / Chinese / Mixed>\n\nSUMMARY:\n<one short sentence, <= 20 words>"
}
backend/requirements.txt
+1
View File
@@ -17,3 +17,4 @@ transformers
soundfile
numpy
accelerate
librosa
backend/tts_asr.py
+641 -100
View File
@@ -4,6 +4,9 @@ import base64
import logging
import os
import platform
import time
import traceback
from typing import Optional
from fastapi import APIRouter, HTTPException, Security
from pydantic import BaseModel
@@ -12,84 +15,447 @@ import numpy as np
router = APIRouter()
logger = logging.getLogger("tts_asr")
# Environment variables
TTS_ASR_DEVICE = os.environ.get("TTS_ASR_DEVICE", "auto")
TTS_ASR_WARMUP = os.environ.get("TTS_ASR_WARMUP", "true").lower() == "true"
TTS_ASR_WARMUP_TIMEOUT = int(os.environ.get("TTS_ASR_WARMUP_TIMEOUT", "120"))
TTS_ASR_IDLE_TIMEOUT = int(os.environ.get("TTS_ASR_IDLE_TIMEOUT", "0"))
# Warmup constants
TTS_WARMUP_TEXT = "你好，这是一个测试。"
ASR_WARMUP_AUDIO_SECONDS = 0.5
# Global state
_tts_pipeline = None
_asr_pipeline = None
_device = None
_device_tested = False
_tts_last_used = 0.0
_asr_last_used = 0.0
_tts_loading = False
_asr_loading = False
_tts_lock = asyncio.Lock()
_asr_lock = asyncio.Lock()
def _get_device():
global _device
if _device is not None:
def _test_device_capability(device_str: str) -> tuple[bool, str]:
"""
测试设备实际可用性
返回: (是否可用, 错误信息)
"""
try:
import torch
if device_str == "cpu":
return True, ""
if device_str == "mps":
if not hasattr(torch.backends, "mps") or not torch.backends.mps.is_available():
return False, "MPS 不可用"
if not torch.backends.mps.is_built():
return False, "MPS 未编译"
test_tensor = torch.randn(2, 2, device="mps")
_ = test_tensor @ test_tensor
del test_tensor
torch.mps.empty_cache()
return True, ""
if device_str.startswith("cuda"):
if not torch.cuda.is_available():
return False, "CUDA 不可用"
torch.cuda.empty_cache()
return True, ""
return False, f"未知设备类型: {device_str}"
except Exception as e:
return False, f"设备测试失败: {str(e)}"
def _get_device() -> str:
"""
获取最佳计算设备，支持环境变量覆盖和降级策略
"""
global _device, _device_tested
if _device is not None and _device_tested:
return _device
import torch
if platform.system() == "Darwin" and hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
_device = "mps"
logger.info("[Device] 使用 MPS 加速")
elif torch.cuda.is_available():
_device = "cuda"
logger.info("[Device] 使用 CUDA 加速")
else:
device_preference = []
if TTS_ASR_DEVICE == "cpu":
_device = "cpu"
logger.info("[Device] 使用 CPU")
_device_tested = True
logger.info("[Device] 强制使用 CPU (环境变量)")
return _device
elif TTS_ASR_DEVICE in ("mps", "cuda", "auto"):
if TTS_ASR_DEVICE != "auto":
device_preference = [TTS_ASR_DEVICE, "cpu"]
else:
if platform.system() == "Darwin":
device_preference = ["mps", "cpu"]
else:
device_preference = ["cuda", "cpu"]
else:
device_preference = ["mps", "cuda", "cpu"]
for dev in device_preference:
ok, err = _test_device_capability(dev)
if ok:
_device = dev
_device_tested = True
logger.info("[Device] 使用 %s 加速", dev.upper() if dev != "cpu" else "CPU")
return _device
else:
logger.warning("[Device] %s 不可用: %s", dev.upper() if dev != "cpu" else "CPU", err)
_device = "cpu"
_device_tested = True
logger.info("[Device] 降级使用 CPU")
return _device
def _device_arg():
def _device_arg() -> str:
device = _get_device()
if device == "cuda":
return "cuda:0"
return device
def _get_torch_dtype():
device = _get_device()
import torch
return torch.float16 if device != "cpu" else torch.float32
def _clear_cuda_cache():
try:
import torch
if _device and _device.startswith("cuda"):
torch.cuda.empty_cache()
except Exception:
pass
def _clear_mps_cache():
try:
import torch
if _device == "mps":
torch.mps.empty_cache()
except Exception:
pass
async def _load_tts_pipeline_with_retry(max_retries: int = 2) -> bool:
"""
加载TTS管道，支持重试和降级
"""
global _tts_pipeline, _tts_loading
async with _tts_lock:
if _tts_pipeline is not None:
return True
if _tts_loading:
return False
_tts_loading = True
try:
import torch
from transformers import pipeline
current_device = _get_device()
for attempt in range(max_retries):
try:
device_to_use = _device_arg()
torch_dtype = _get_torch_dtype()
logger.info("[TTS] 加载 Kokoro-82M 模型 (尝试 %d/%d, 设备: %s)...",
attempt + 1, max_retries, device_to_use)
_tts_pipeline = await asyncio.to_thread(
lambda: pipeline(
"text-to-speech",
model="hexgrad/Kokoro-82M",
trust_remote_code=True,
device=device_to_use,
torch_dtype=torch_dtype,
)
)
logger.info("[TTS] Kokoro-82M 模型加载完成")
return True
except RuntimeError as e:
error_str = str(e)
if "MPS" in error_str or "mps" in error_str:
logger.warning("[TTS] MPS 推理失败，尝试降级到 CPU: %s", error_str)
global _device
_device = "cpu"
_clear_mps_cache()
continue
elif "CUDA" in error_str or "cuda" in error_str:
logger.warning("[TTS] CUDA 推理失败，尝试降级到 CPU: %s", error_str)
_device = "cpu"
_clear_cuda_cache()
continue
else:
raise
except Exception as e:
logger.error("[TTS] 加载失败: %s", str(e))
if attempt == max_retries - 1:
raise
await asyncio.sleep(1)
return _tts_pipeline is not None
finally:
_tts_loading = False
async def _load_asr_pipeline_with_retry(max_retries: int = 2) -> bool:
"""
加载ASR管道，支持重试和降级
"""
global _asr_pipeline, _asr_loading
async with _asr_lock:
if _asr_pipeline is not None:
return True
if _asr_loading:
return False
_asr_loading = True
try:
import torch
from transformers import AutoModelForSpeechSeq2Seq, AutoProcessor, pipeline
for attempt in range(max_retries):
try:
device_to_use = _device_arg()
torch_dtype = _get_torch_dtype()
logger.info("[ASR] 加载 Whisper large-v3-turbo 模型 (尝试 %d/%d, 设备: %s)...",
attempt + 1, max_retries, device_to_use)
model_id = "openai/whisper-large-v3-turbo"
def load_model():
model = AutoModelForSpeechSeq2Seq.from_pretrained(
model_id,
torch_dtype=torch_dtype,
low_cpu_mem_usage=True,
use_safetensors=True,
)
processor = AutoProcessor.from_pretrained(model_id)
return pipeline(
"automatic-speech-recognition",
model=model,
tokenizer=processor.tokenizer,
feature_extractor=processor.feature_extractor,
torch_dtype=torch_dtype,
device=device_to_use,
)
_asr_pipeline = await asyncio.to_thread(load_model)
logger.info("[ASR] Whisper large-v3-turbo 模型加载完成")
return True
except RuntimeError as e:
error_str = str(e)
if "MPS" in error_str or "mps" in error_str:
logger.warning("[ASR] MPS 推理失败，尝试降级到 CPU: %s", error_str)
global _device
_device = "cpu"
_clear_mps_cache()
continue
elif "CUDA" in error_str or "cuda" in error_str:
logger.warning("[ASR] CUDA 推理失败，尝试降级到 CPU: %s", error_str)
_device = "cpu"
_clear_cuda_cache()
continue
else:
raise
except Exception as e:
logger.error("[ASR] 加载失败: %s", str(e))
if attempt == max_retries - 1:
raise
await asyncio.sleep(1)
return _asr_pipeline is not None
finally:
_asr_loading = False
def _get_tts_pipeline():
global _tts_pipeline
"""同步获取TTS管道（已弃用，保留兼容性）"""
if _tts_pipeline is not None:
return _tts_pipeline
import torch
from transformers import pipeline
logger.info("[TTS] 加载 Kokoro-82M 模型...")
_tts_pipeline = pipeline(
"text-to-speech",
model="hexgrad/Kokoro-82M",
trust_remote_code=True,
device=_device_arg(),
torch_dtype=torch.float16 if _get_device() != "cpu" else torch.float32,
)
logger.info("[TTS] Kokoro-82M 模型加载完成")
return _tts_pipeline
raise RuntimeError("TTS 管道未加载，请使用 _load_tts_pipeline_with_retry()")
def _get_asr_pipeline():
global _asr_pipeline
"""同步获取ASR管道（已弃用，保留兼容性）"""
if _asr_pipeline is not None:
return _asr_pipeline
raise RuntimeError("ASR 管道未加载，请使用 _load_asr_pipeline_with_retry()")
import torch
from transformers import AutoModelForSpeechSeq2Seq, AutoProcessor, pipeline
logger.info("[ASR] 加载 Whisper large-v3-turbo 模型...")
model_id = "openai/whisper-large-v3-turbo"
model = AutoModelForSpeechSeq2Seq.from_pretrained(
model_id,
torch_dtype=torch.float16 if _get_device() != "cpu" else torch.float32,
low_cpu_mem_usage=True,
use_safetensors=True,
)
processor = AutoProcessor.from_pretrained(model_id)
_asr_pipeline = pipeline(
"automatic-speech-recognition",
model=model,
tokenizer=processor.tokenizer,
feature_extractor=processor.feature_extractor,
torch_dtype=torch.float16 if _get_device() != "cpu" else torch.float32,
device=_device_arg(),
)
logger.info("[ASR] Whisper large-v3-turbo 模型加载完成")
return _asr_pipeline
async def _warmup_tts() -> bool:
"""
预热TTS模型，减少首次请求延迟
"""
global _tts_last_used
try:
logger.info("[TTS] 开始预热...")
if not await _load_tts_pipeline_with_retry():
logger.error("[TTS] 预热失败：无法加载管道")
return False
tts = _tts_pipeline
if tts is None:
return False
def warmup_inference():
try:
result = tts(TTS_WARMUP_TEXT, voice="af_bella")
if isinstance(result, dict):
audio = result.get("audio")
if hasattr(audio, "cpu"):
_ = audio.cpu()
return True
except Exception as e:
logger.warning("[TTS] 预热推理失败（可忽略）: %s", str(e))
return False
success = await asyncio.to_thread(warmup_inference)
_tts_last_used = time.time()
if success:
logger.info("[TTS] 预热完成")
return success
except Exception as e:
logger.error("[TTS] 预热异常: %s", str(e))
traceback.print_exc()
return False
async def _warmup_asr() -> bool:
"""
预热ASR模型，减少首次请求延迟
"""
global _asr_last_used
try:
logger.info("[ASR] 开始预热...")
if not await _load_asr_pipeline_with_retry():
logger.error("[ASR] 预热失败：无法加载管道")
return False
asr = _asr_pipeline
if asr is None:
return False
silence_samples = int(16000 * ASR_WARMUP_AUDIO_SECONDS)
silence_audio = np.zeros(silence_samples, dtype=np.float32)
def warmup_inference():
try:
result = asr(
silence_audio,
sampling_rate=16000,
generate_kwargs={"language": "zh", "task": "transcribe"},
)
return True
except Exception as e:
logger.warning("[ASR] 预热推理失败（可忽略）: %s", str(e))
return False
success = await asyncio.to_thread(warmup_inference)
_asr_last_used = time.time()
if success:
logger.info("[ASR] 预热完成")
return success
except Exception as e:
logger.error("[ASR] 预热异常: %s", str(e))
traceback.print_exc()
return False
async def _warmup_all() -> tuple[bool, bool]:
"""
预热所有模型
返回: (TTS预热结果, ASR预热结果)
"""
logger.info("[Warmup] 开始预热所有模型 (超时: %d秒)", TTS_ASR_WARMUP_TIMEOUT)
try:
tts_task = asyncio.create_task(_warmup_tts())
asr_task = asyncio.create_task(_warmup_asr())
done, pending = await asyncio.wait(
[tts_task, asr_task],
timeout=TTS_ASR_WARMUP_TIMEOUT,
return_when=asyncio.ALL_COMPLETED,
)
for task in pending:
task.cancel()
try:
await task
except asyncio.CancelledError:
pass
tts_result = tts_task.result() if tts_task in done else False
asr_result = asr_task.result() if asr_task in done else False
logger.info("[Warmup] 完成: TTS=%s, ASR=%s", tts_result, asr_result)
return tts_result, asr_result
except Exception as e:
logger.error("[Warmup] 异常: %s", str(e))
traceback.print_exc()
return False, False
def _check_and_unload_idle_models():
"""
检查并卸载空闲超过阈值的模型
"""
if TTS_ASR_IDLE_TIMEOUT <= 0:
return
global _tts_pipeline, _asr_pipeline
current_time = time.time()
if _tts_pipeline is not None:
idle_seconds = current_time - _tts_last_used
if idle_seconds > TTS_ASR_IDLE_TIMEOUT:
logger.info("[TTS] 空闲 %.0f 秒，卸载模型", idle_seconds)
_tts_pipeline = None
_clear_cuda_cache()
_clear_mps_cache()
if _asr_pipeline is not None:
idle_seconds = current_time - _asr_last_used
if idle_seconds > TTS_ASR_IDLE_TIMEOUT:
logger.info("[ASR] 空闲 %.0f 秒，卸载模型", idle_seconds)
_asr_pipeline = None
_clear_cuda_cache()
_clear_mps_cache()
def _save_audio_to_wav(audio_data: bytes, sample_rate: int = 16000) -> str:
@@ -105,74 +471,193 @@ def _save_audio_to_wav(audio_data: bytes, sample_rate: int = 16000) -> str:
return tmp.name
def _tts_sync(text: str, voice: str = "af_bella", rate: float = 1.0) -> tuple[bytes, int]:
tts = _get_tts_pipeline()
result = tts(text, voice=voice)
audio = None
async def _tts_sync_with_retry(text: str, voice: str = "af_bella", rate: float = 1.0, max_retries: int = 2) -> tuple[bytes, int]:
"""
TTS推理，支持重试和降级
"""
global _tts_last_used
_check_and_unload_idle_models()
if not await _load_tts_pipeline_with_retry():
raise RuntimeError("TTS 模型加载失败")
tts = _tts_pipeline
sample_rate = 24000
if isinstance(result, dict):
audio = result.get("audio")
sample_rate = int(result.get("sampling_rate", sample_rate))
elif isinstance(result, (list, tuple)) and result:
audio = result[0]
if audio is None:
raise RuntimeError("Kokoro 未返回音频数据")
for attempt in range(max_retries):
try:
def inference():
result = tts(text, voice=voice)
audio = None
sr = sample_rate
if hasattr(audio, "cpu"):
audio = audio.cpu().numpy()
if isinstance(result, dict):
audio = result.get("audio")
sr = int(result.get("sampling_rate", sr))
elif isinstance(result, (list, tuple)) and result:
audio = result[0]
duration_ms = int(len(audio) * 1000 / sample_rate)
if audio is None:
raise RuntimeError("Kokoro 未返回音频数据")
if audio.dtype != np.int16:
audio = (audio * 32767).astype(np.int16)
if hasattr(audio, "cpu"):
audio = audio.cpu().numpy()
import tempfile
import wave
if hasattr(audio, "numpy"):
audio = audio.numpy()
with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as tmp:
output_path = tmp.name
try:
with wave.open(output_path, "wb") as wf:
wf.setnchannels(1)
wf.setsampwidth(2)
wf.setframerate(sample_rate)
wf.writeframes(audio.tobytes())
with open(output_path, "rb") as f:
return f.read(), duration_ms
finally:
if os.path.exists(output_path):
os.unlink(output_path)
return audio, sr
audio, sample_rate = await asyncio.to_thread(inference)
duration_ms = int(len(audio) * 1000 / sample_rate)
if audio.dtype != np.int16:
audio = (audio * 32767).astype(np.int16)
import tempfile
import wave
with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as tmp:
output_path = tmp.name
try:
with wave.open(output_path, "wb") as wf:
wf.setnchannels(1)
wf.setsampwidth(2)
wf.setframerate(sample_rate)
wf.writeframes(audio.tobytes())
with open(output_path, "rb") as f:
audio_bytes = f.read()
_tts_last_used = time.time()
return audio_bytes, duration_ms
finally:
if os.path.exists(output_path):
os.unlink(output_path)
except RuntimeError as e:
error_str = str(e)
if "MPS" in error_str or "mps" in error_str:
logger.warning("[TTS] MPS 推理错误，尝试降级重试 (尝试 %d/%d): %s",
attempt + 1, max_retries, error_str)
global _device
_device = "cpu"
_clear_mps_cache()
if attempt < max_retries - 1:
continue
elif "CUDA" in error_str or "cuda" in error_str:
logger.warning("[TTS] CUDA 推理错误，尝试降级重试 (尝试 %d/%d): %s",
attempt + 1, max_retries, error_str)
_device = "cpu"
_clear_cuda_cache()
if attempt < max_retries - 1:
continue
raise
except Exception as e:
logger.error("[TTS] 推理失败: %s", str(e))
if attempt == max_retries - 1:
raise
await asyncio.sleep(0.5)
raise RuntimeError("TTS 推理失败")
async def _text_to_speech(text: str, voice: str = "af_bella", rate: float = 1.0) -> tuple[bytes, int]:
return await asyncio.to_thread(_tts_sync, text, voice, rate)
async def _asr_sync_with_retry(audio_data: bytes, language: str = "zh", max_retries: int = 2) -> str:
"""
ASR推理，支持重试和降级
"""
global _asr_last_used
_check_and_unload_idle_models()
def _asr_sync(audio_data: bytes, language: str = "zh") -> str:
import soundfile as sf
if not await _load_asr_pipeline_with_retry():
raise RuntimeError("ASR 模型加载失败")
asr = _get_asr_pipeline()
audio_path = _save_audio_to_wav(audio_data)
try:
audio_array, sample_rate = sf.read(audio_path)
result = asr(
audio_array,
sampling_rate=sample_rate,
generate_kwargs={"language": language, "task": "transcribe"},
)
if isinstance(result, dict):
return result.get("text", "").strip()
return str(result).strip()
import soundfile as sf
audio_array, sample_rate = await asyncio.to_thread(lambda: sf.read(audio_path))
if len(audio_array.shape) > 1:
audio_array = np.mean(audio_array, axis=1)
if sample_rate != 16000:
import librosa
audio_array = await asyncio.to_thread(
lambda: librosa.resample(audio_array, orig_sr=sample_rate, target_sr=16000)
)
sample_rate = 16000
audio_array = audio_array.astype(np.float32)
for attempt in range(max_retries):
try:
def inference():
asr = _asr_pipeline
result = asr(
audio_array,
sampling_rate=sample_rate,
generate_kwargs={"language": language, "task": "transcribe"},
)
if isinstance(result, dict):
return result.get("text", "").strip()
return str(result).strip()
text = await asyncio.to_thread(inference)
_asr_last_used = time.time()
return text
except RuntimeError as e:
error_str = str(e)
if "MPS" in error_str or "mps" in error_str:
logger.warning("[ASR] MPS 推理错误，尝试降级重试 (尝试 %d/%d): %s",
attempt + 1, max_retries, error_str)
global _device
_device = "cpu"
_clear_mps_cache()
if attempt < max_retries - 1:
continue
elif "CUDA" in error_str or "cuda" in error_str:
logger.warning("[ASR] CUDA 推理错误，尝试降级重试 (尝试 %d/%d): %s",
attempt + 1, max_retries, error_str)
_device = "cpu"
_clear_cuda_cache()
if attempt < max_retries - 1:
continue
raise
except Exception as e:
logger.error("[ASR] 推理失败: %s", str(e))
if attempt == max_retries - 1:
raise
await asyncio.sleep(0.5)
raise RuntimeError("ASR 推理失败")
finally:
if os.path.exists(audio_path):
os.unlink(audio_path)
# Legacy sync wrappers (for compatibility)
def _tts_sync(text: str, voice: str = "af_bella", rate: float = 1.0) -> tuple[bytes, int]:
raise RuntimeError("请使用 _tts_sync_with_retry()")
def _asr_sync(audio_data: bytes, language: str = "zh") -> str:
raise RuntimeError("请使用 _asr_sync_with_retry()")
async def _text_to_speech(text: str, voice: str = "af_bella", rate: float = 1.0) -> tuple[bytes, int]:
return await _tts_sync_with_retry(text, voice, rate)
async def _speech_to_text(audio_data: bytes, language: str = "zh") -> str:
return await asyncio.to_thread(_asr_sync, audio_data, language)
return await _asr_sync_with_retry(audio_data, language)
# Request/Response models
class TTSRequest(BaseModel):
text: str
voice: str = "af_bella"
@@ -196,21 +681,58 @@ class ASRResponse(BaseModel):
language: str
class ModelStatus(BaseModel):
tts_loaded: bool
asr_loaded: bool
device: str
tts_last_used: Optional[float] = None
asr_last_used: Optional[float] = None
def get_api_key(api_key: str):
import main
API_KEY = main.API_KEY
if api_key != API_KEY:
raise HTTPException(status_code=403, detail="API Key 无效")
return api_key
@router.get("/status", response_model=ModelStatus)
async def get_status(api_key: str = Security(get_api_key)):
"""
获取模型状态
"""
current_time = time.time()
return ModelStatus(
tts_loaded=_tts_pipeline is not None,
asr_loaded=_asr_pipeline is not None,
device=_get_device(),
tts_last_used=_tts_last_used if _tts_last_used > 0 else None,
asr_last_used=_asr_last_used if _asr_last_used > 0 else None,
)
@router.post("/warmup")
async def warmup_models(api_key: str = Security(get_api_key)):
"""
手动触发模型预热
"""
tts_result, asr_result = await _warmup_all()
return {
"tts_warmup": tts_result,
"asr_warmup": asr_result,
"device": _get_device(),
}
@router.post("/tts", response_model=TTSResponse)
async def text_to_speech(req: TTSRequest, api_key: str = Security(get_api_key)):
request_id = str(hash(req.text))[:8]
try:
logger.info("[TTS][%s] text_chars=%d voice=%s format=%s", request_id, len(req.text), req.voice, req.format)
logger.info("[TTS][%s] text_chars=%d voice=%s format=%s",
request_id, len(req.text), req.voice, req.format)
audio_data, duration_ms = await _text_to_speech(req.text, req.voice, req.rate)
if req.format.lower() == "mp3":
import subprocess
import tempfile
@@ -222,7 +744,9 @@ async def text_to_speech(req: TTSRequest, api_key: str = Security(get_api_key)):
output_path = tmp_out.name
try:
cmd = ["ffmpeg", "-i", input_path, "-acodec", "libmp3lame", "-ab", "128k", output_path]
result = await asyncio.to_thread(lambda: subprocess.run(cmd, capture_output=True, text=True, timeout=30))
result = await asyncio.to_thread(
lambda: subprocess.run(cmd, capture_output=True, text=True, timeout=30)
)
if result.returncode != 0:
raise RuntimeError(f"MP3 转换失败: {result.stderr}")
with open(output_path, "rb") as f:
@@ -231,8 +755,14 @@ async def text_to_speech(req: TTSRequest, api_key: str = Security(get_api_key)):
for path in [input_path, output_path]:
if os.path.exists(path):
os.unlink(path)
logger.info("[TTS][%s] success duration_ms=%d", request_id, duration_ms)
return TTSResponse(audio_base64=base64.b64encode(audio_data).decode(), format=req.format, duration_ms=duration_ms)
return TTSResponse(
audio_base64=base64.b64encode(audio_data).decode(),
format=req.format,
duration_ms=duration_ms,
)
except Exception as e:
logger.exception("[TTS] failed: %s", e)
raise HTTPException(status_code=500, detail=str(e))
@@ -242,15 +772,26 @@ async def text_to_speech(req: TTSRequest, api_key: str = Security(get_api_key)):
async def speech_to_text(req: ASRRequest, api_key: str = Security(get_api_key)):
request_id = str(hash(req.audio_base64))[:8]
try:
logger.info("[ASR][%s] audio_base64_chars=%d language=%s", request_id, len(req.audio_base64), req.language)
logger.info("[ASR][%s] audio_base64_chars=%d language=%s",
request_id, len(req.audio_base64), req.language)
audio_data = base64.b64decode(req.audio_base64)
text = await _speech_to_text(audio_data, req.language[:2])
logger.info("[ASR][%s] success text_chars=%d", request_id, len(text))
return ASRResponse(text=text, language=req.language)
except Exception as e:
logger.exception("[ASR] failed: %s", e)
raise HTTPException(status_code=500, detail=str(e))
def register_tts_asr_routes(app):
"""
注册TTS/ASR路由并可选执行预热
"""
app.include_router(router, prefix="/v1/tts-asr")
if TTS_ASR_WARMUP:
@app.on_event("startup")
async def warmup_on_startup():
logger.info("[Startup] 开始后台预热...")
asyncio.create_task(_warmup_all())