From f282785c1924097b4ceeae7073743b8bd1feae1f Mon Sep 17 00:00:00 2001
From: yuetsh <517252939@qq.com>
Date: Mon, 6 Apr 2026 05:27:13 -0600
Subject: [PATCH] fix

---
 prompt/llm.py | 41 ++++++++++++++++++++++++++++++++++++-----
 1 file changed, 36 insertions(+), 5 deletions(-)

diff --git a/prompt/llm.py b/prompt/llm.py
index 0bfbb8a..20d2c5e 100644
--- a/prompt/llm.py
+++ b/prompt/llm.py
@@ -11,7 +11,25 @@ SYSTEM_PROMPT = """你是一个网页生成助手。根据用户的需求描述
 3. CSS 和 JS 可以为空，但仍然需要返回空的代码块
 4. 用中文回复，先简要说明你做了什么，然后给出代码
 5. 在已有代码基础上修改时，返回完整的修改后代码，不要只返回片段
-6. 由于任何外部链接都被屏蔽，使用纯 HTML、CSS 和 JS 实现功能，不要依赖外部库"""
+6. 由于任何外部链接都被屏蔽，使用纯 HTML、CSS 和 JS 实现功能，不要依赖外部库
+
+输出格式示例（必须严格遵守，三个代码块缺一不可）：
+
+好的，我为你创建了一个点击按钮变色的示例。
+
+```html
+<button id="btn">点击我</button>
+```
+
+```css
+button { padding: 8px 16px; }
+```
+
+```js
+document.getElementById('btn').onclick = function() {
+  this.style.background = 'red';
+};
+```"""
 
 DEFAULT_MODEL = "deepseek-chat"
 
@@ -66,12 +84,25 @@ async def stream_chat(history: list[dict], model: str = ""):
 def extract_code(text: str) -> dict:
     """Extract HTML, CSS, JS code blocks from AI response text."""
     result = {"html": None, "css": None, "js": None}
-    pattern = r"```(html|css|js|javascript)\s*\n(.*?)```"
-    matches = re.findall(pattern, text, re.DOTALL)
+    pattern = r"```(html|css|js|javascript|typescript|ts|jsx|tsx)\s*\n(.*?)```"
+    matches = re.findall(pattern, text, re.DOTALL | re.IGNORECASE)
     for lang, code in matches:
         lang = lang.lower()
-        if lang == "javascript":
+        if lang in ("javascript", "typescript", "ts", "jsx", "tsx"):
             lang = "js"
-        if lang in result:
+        if lang in result and result[lang] is None:
             result[lang] = code.strip()
+
+    # Fallback: extract <style> and <script> from HTML block
+    if result["html"]:
+        if result["css"] is None:
+            style_match = re.search(r"<style[^>]*>(.*?)</style>", result["html"], re.DOTALL | re.IGNORECASE)
+            if style_match:
+                result["css"] = style_match.group(1).strip()
+
+        if result["js"] is None:
+            script_match = re.search(r"<script[^>]*>(.*?)</script>", result["html"], re.DOTALL | re.IGNORECASE)
+            if script_match:
+                result["js"] = script_match.group(1).strip()
+
     return result