Update app.py

app.py

@@ -1,5 +1,12 @@
-import gradio as gr
 import time
+import json
+import uuid
+from fastapi import FastAPI, Request, HTTPException, Depends
+from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
+from starlette.responses import StreamingResponse
+from pydantic import BaseModel
+from typing import List, Optional
+
 from selenium import webdriver
 from selenium.webdriver.common.by import By
 from selenium.webdriver.common.keys import Keys
@@ -7,47 +14,70 @@ from selenium.webdriver.chrome.service import Service as ChromeService
 from selenium.webdriver.support.ui import WebDriverWait
 from selenium.webdriver.support import expected_conditions as EC
 
-def chat_with_sai_automation(prompt_text: str, progress=gr.Progress(track_tqdm=True)) -> str:
+# --- 1. FastAPI 应用和认证设置 ---
+
+app = FastAPI(
+    title="SAI-ChatBot OpenAI-Compatible API",
+    description="使用 Selenium 自动化在后台与 SAI-ChatBot 交互，并以 OpenAI API 格式返回结果。",
+    version="1.0.0"
+)
+
+# 一个简单的认证方案：只需要提供任何 Bearer Token 即可
+auth_scheme = HTTPBearer()
+
+def api_key_auth(credentials: HTTPAuthorizationCredentials = Depends(auth_scheme)):
+    """
+    我们不校验 key 的内容，只检查是否存在 Authorization header。
+    这满足了“key填什么都过”的要求。
+    """
+    if not credentials:
+        raise HTTPException(
+            status_code=401,
+            detail="Not authenticated",
+            headers={"WWW-Authenticate": "Bearer"},
+        )
+    # 这里可以添加日志，记录传入的 token，但我们不验证它
+    # print(f"Received token: {credentials.scheme} {credentials.token}")
+    return credentials.token
+
+# --- 2. OpenAI 格式的数据模型 (使用 Pydantic) ---
+
+class ChatMessage(BaseModel):
+    role: str
+    content: str
+
+class ChatCompletionRequest(BaseModel):
+    model: str  # 我们会忽略这个字段，但为了兼容性保留它
+    messages: List[ChatMessage]
+    stream: Optional[bool] = False
+
+# --- 3. Selenium 自动化核心函数 ---
+
+def get_sai_response(prompt_text: str):
     """
-    最终版：使用系统中通过 apt-get 安装的 chromium 和 chromedriver。
-    这是在 Hugging Face Spaces 上最稳定、最可靠的方案。
+    这是一个生成器函数，它会逐步返回（yield）AI的回复。
+    这使得我们可以实现流式API。
     """
-    full_response = ""
+    options = webdriver.ChromeOptions()
+    options.add_argument("--headless")
+    options.add_argument("--no-sandbox")
+    options.add_argument("--disable-dev-shm-usage")
+    options.add_argument("--disable-gpu")
+    options.binary_location = "/usr/bin/chromium"
+
+    service = ChromeService(executable_path='/usr/bin/chromedriver')
+    driver = None
     try:
-        # --- 1. 设置和初始化浏览器 ---
-        progress(0, desc="正在设置浏览器选项...")
-        options = webdriver.ChromeOptions()
-        options.add_argument("--headless")
-        options.add_argument("--no-sandbox")
-        options.add_argument("--disable-dev-shm-usage")
-        options.add_argument("--disable-gpu")
-        # 当通过 apt 安装时，chromium 通常在这里
-        options.binary_location = "/usr/bin/chromium"
-
-        progress(0.2, desc="正在指定系统驱动...")
-        # 直接指定由 apt 安装的 chromedriver 的路径
-        service = ChromeService(executable_path='/usr/bin/chromedriver')
-        
-        progress(0.3, desc="正在启动浏览器...")
         driver = webdriver.Chrome(service=service, options=options)
-        
-        # --- 2. 访问网页并等待输入框加载 ---
-        progress(0.4, desc="正在访问 sai.coludai.cn ...")
         driver.get("https://sai.coludai.cn/")
         
         wait = WebDriverWait(driver, 20)
         textarea_selector = 'textarea[placeholder="随时与未来对话，探索无限可能...."]'
-        progress(0.5, desc="等待页面输入框加载...")
         textarea = wait.until(EC.presence_of_element_located((By.CSS_SELECTOR, textarea_selector)))
         
-        # --- 3. 输入问题并发送 ---
-        progress(0.6, desc="输入问题并发送...")
         textarea.send_keys(prompt_text)
         textarea.send_keys(Keys.RETURN)
         
-        # --- 4. 等待并捕获 AI 回复 ---
-        progress(0.7, desc="等待 AI 回复...")
-        
         last_assistant_selector = "(.//div[@class='message-item' and @type='assistant'])[last()]"
         wait.until(EC.presence_of_element_located((By.XPATH, last_assistant_selector)))
         last_response_element = driver.find_element(By.XPATH, last_assistant_selector)
@@ -60,36 +90,90 @@ def chat_with_sai_automation(prompt_text: str, progress=gr.Progress(track_tqdm=T
             try:
                 markdown_body = last_response_element.find_element(By.CSS_SELECTOR, '.markdown-body')
                 current_text = markdown_body.text
-                progress(0.7 + 0.3 * (min(len(current_text), 500) / 500), desc=f"正在接收... (长度: {len(current_text)})")
-                if current_text == previous_text and current_text != "":
-                    full_response = current_text
-                    break
-                else:
+                
+                if current_text != previous_text:
+                    new_text_chunk = current_text[len(previous_text):]
+                    yield new_text_chunk
                     previous_text = current_text
+                
+                # 检查是否已经有一段时间没有新内容了，可以认为结束
+                if current_text == previous_text and current_text != "":
+                    # 再稍微等待一下，以防万一
                     time.sleep(1.5)
+                    final_text_check = markdown_body.text
+                    if final_text_check == previous_text:
+                         break # 确认结束，退出循环
+                
+                time.sleep(0.5)
+
             except Exception:
                 time.sleep(0.5)
-        
-        if not full_response:
-            full_response = previous_text if previous_text else "超时或未能捕获到完整回复。"
 
     except Exception as e:
-        full_response = f"自动化过程中发生错误: {e}\n\n请检查 Hugging Face Space 的日志获取详细信息。"
+        error_message = f"自动化过程中发生错误: {e}"
+        yield error_message
     finally:
-        if 'driver' in locals() and driver:
+        if driver:
             driver.quit()
-            
-    return full_response
-
-# --- 使用 Gradio 创建 Web 界面 ---
-with gr.Blocks() as iface:
-    gr.Markdown("# SAI-ChatBot 自动化工具")
-    gr.Markdown("在后台使用 Selenium 自动与 SAI-ChatBot 交互。请耐心等待，首次运行可能需要一些时间来构建环境。")
-    with gr.Row():
-        inp = gr.Textbox(lines=5, label="你的问题", placeholder="在这里输入你想问 SAI-ChatBot 的问题...")
-        out = gr.Textbox(lines=15, label="AI 的回复")
-    btn = gr.Button("提交")
-    btn.click(fn=chat_with_sai_automation, inputs=inp, outputs=out)
-
-if __name__ == "__main__":
-    iface.launch()
+
+# --- 4. API 端点定义 ---
+
+@app.post("/v1/chat/completions")
+async def chat_completions(request: ChatCompletionRequest, token: str = Depends(api_key_auth)):
+    # 从消息列表中提取最后一个用户的提问
+    last_user_message = next((msg.content for msg in reversed(request.messages) if msg.role == 'user'), None)
+    
+    if not last_user_message:
+        raise HTTPException(status_code=400, detail="No user message found")
+
+    response_id = f"chatcmpl-{uuid.uuid4()}"
+    created_timestamp = int(time.time())
+
+    # --- 流式响应 (Streaming Response) ---
+    if request.stream:
+        async def stream_generator():
+            for chunk in get_sai_response(last_user_message):
+                if not chunk:
+                    continue
+                # 构造符合 OpenAI SSE 格式的 JSON
+                response_chunk = {
+                    "id": response_id,
+                    "object": "chat.completion.chunk",
+                    "created": created_timestamp,
+                    "model": "sai-chatbot-l6",
+                    "choices": [{
+                        "index": 0,
+                        "delta": {"content": chunk},
+                        "finish_reason": None
+                    }]
+                }
+                yield f"data: {json.dumps(response_chunk)}\n\n"
+            # 发送结束标志
+            yield f"data: [DONE]\n\n"
+        
+        return StreamingResponse(stream_generator(), media_type="text/event-stream")
+
+    # --- 非流式响应 (Non-Streaming Response) ---
+    else:
+        full_content_chunks = [chunk for chunk in get_sai_response(last_user_message)]
+        full_content = "".join(full_content_chunks)
+
+        return {
+            "id": response_id,
+            "object": "chat.completion",
+            "created": created_timestamp,
+            "model": "sai-chatbot-l6",
+            "choices": [{
+                "index": 0,
+                "message": {
+                    "role": "assistant",
+                    "content": full_content,
+                },
+                "finish_reason": "stop"
+            }],
+            "usage": {
+                "prompt_tokens": len(last_user_message), # 伪造的 token 计数
+                "completion_tokens": len(full_content),
+                "total_tokens": len(last_user_message) + len(full_content)
+            }
+        }