Upload 9 files

app.py

@@ -66,6 +66,31 @@ async def list_models():
     }
 
 
+def parse_openrouter_response(response_data: dict) -> str:
+    """从OpenRouter响应中提取内容"""
+    content = ""
+    
+    # 标准OpenAI格式
+    if "choices" in response_data and response_data["choices"]:
+        choices = response_data["choices"]
+        if choices:
+            msg = choices[0].get("message", {})
+            content = msg.get("content", "")
+            if not content:
+                # 可能是delta格式
+                delta = choices[0].get("delta", {})
+                content = delta.get("content", "")
+    
+    # 直接返回的情况
+    if not content and "message" in response_data:
+        content = response_data.get("message", {}).get("content", "")
+    
+    if not content and "content" in response_data:
+        content = response_data.get("content", "")
+    
+    return content
+
+
 @fastapi_app.post("/v1/chat/completions")
 async def chat_completions(request: ChatCompletionRequest):
     prompt = request.messages[-1].content if request.messages else ""
@@ -73,7 +98,7 @@ async def chat_completions(request: ChatCompletionRequest):
     
     if request.stream:
         return StreamingResponse(
-            stream_chat(request.model, prompt, request.messages),
+            stream_chat(request.model, [{"role": m.role, "content": m.content} for m in request.messages]),
             media_type="text/event-stream"
         )
     
@@ -83,10 +108,7 @@ async def chat_completions(request: ChatCompletionRequest):
         raise HTTPException(status_code=400, detail=result.get("error", "Request failed"))
     
     response_data = result.get("response", {})
-    
-    content = ""
-    if "choices" in response_data and response_data["choices"]:
-        content = response_data["choices"][0].get("message", {}).get("content", "")
+    content = parse_openrouter_response(response_data)
     
     return {
         "id": response_data.get("id", f"chatcmpl-{random.randint(100000, 999999)}"),
@@ -111,59 +133,19 @@ async def chat_completions(request: ChatCompletionRequest):
     }
 
 
-async def stream_chat(model_hint: Optional[str], prompt: str, messages: list):
-    model_hint = model_hint or ""
-    
-    result = await model_tester.chat_completion(prompt, model_hint)
+async def stream_chat(model_hint: Optional[str], messages: list):
+    # 直接代理OpenRouter的流式响应
+    stream, used_model = await model_tester.chat_completion_stream(model_hint, messages)
     
-    if not result.get("success"):
-        yield f'data: {{"error": "{result.get("error", "Request failed")}"}}\n\n'
+    if not stream:
+        yield f'data: {{"error": "No available model found"}}\n\n'
         yield "data: [DONE]\n\n"
         return
     
-    response_data = result.get("response", {})
-    content = ""
-    if "choices" in response_data and response_data["choices"]:
-        content = response_data["choices"][0].get("message", {}).get("content", "")
-    
-    model_id = result.get("model", model_hint or "unknown")
-    completion_id = f"chatcmpl-{random.randint(100000, 999999)}"
-    created = int(datetime.now().timestamp())
-    
-    # 流式输出每个字
-    for i, char in enumerate(content):
-        chunk = {
-            "id": completion_id,
-            "object": "chat.completion.chunk",
-            "created": created,
-            "model": model_id,
-            "choices": [
-                {
-                    "index": 0,
-                    "delta": {
-                        "content": char
-                    },
-                    "finish_reason": None
-                }
-            ]
-        }
-        yield f"data: {json.dumps(chunk)}\n\n"
-    
-    # 发送完成信号
-    final_chunk = {
-        "id": completion_id,
-        "object": "chat.completion.chunk",
-        "created": created,
-        "model": model_id,
-        "choices": [
-            {
-                "index": 0,
-                "delta": {},
-                "finish_reason": "stop"
-            }
-        ]
-    }
-    yield f"data: {json.dumps(final_chunk)}\n\n"
+    # 直接转发流式数据
+    async for chunk in stream:
+        yield chunk.decode() if isinstance(chunk, bytes) else chunk
+    
     yield "data: [DONE]\n\n"
 
 

src/model_tester.py

@@ -139,6 +139,47 @@ class ModelTester:
         """Get all free models from API list (not tested)"""
         return self._free_models
 
+    async def try_model_direct_stream(
+        self,
+        session: aiohttp.ClientSession,
+        model_id: str,
+        api_key: str,
+        messages: List[Dict[str, str]]
+    ) -> Optional[Dict[str, Any]]:
+        """发送流式请求到OpenRouter"""
+        url = "https://openrouter.ai/api/v1/chat/completions"
+        payload = {
+            "model": model_id,
+            "messages": messages,
+            "max_tokens": 2048,
+            "stream": True
+        }
+        headers = {
+            "Authorization": f"Bearer {api_key}",
+            "Content-Type": "application/json"
+        }
+
+        try:
+            timeout = aiohttp.ClientTimeout(total=config.get_request_timeout())
+            async with session.post(url, json=payload, headers=headers, timeout=timeout) as response:
+                if response.status == 200:
+                    return {
+                        "success": True,
+                        "model": model_id,
+                        "stream": response.content,
+                        "method": "direct"
+                    }
+                else:
+                    body = await response.text()
+                    return {
+                        "success": False,
+                        "model": model_id,
+                        "error": f"HTTP {response.status}: {body[:100]}",
+                        "method": "direct"
+                    }
+        except Exception as e:
+            return {"success": False, "model": model_id, "error": str(e), "method": "direct"}
+
     async def try_model_direct(
         self, 
         session: aiohttp.ClientSession, 
@@ -356,3 +397,39 @@ class ModelTester:
     def test_all_models(self) -> Dict[str, Any]:
         """Legacy sync method - use scan_all_models instead"""
         return self.scan_all_models()
+
+    async def chat_completion_stream(self, model_hint: Optional[str], messages: List[Dict[str, str]]):
+        """流式聊天 - 返回流式响应对象"""
+        api_keys = config.get_api_keys()
+        api_key = random.choice(api_keys)
+
+        # 方案1：尝试用户指定的模型
+        if model_hint:
+            full_model_id = self.find_model_in_list(model_hint)
+            if full_model_id:
+                async with aiohttp.ClientSession() as session:
+                    result = await self.try_model_direct_stream(session, full_model_id, api_key, messages)
+                    if result and result.get("success"):
+                        return result.get("stream"), result.get("model")
+
+        # 方案2：从列表中找到可用模型
+        self.refresh_model_list()
+        available_free = self.get_all_free_models()
+
+        candidates = []
+        if model_hint and available_free:
+            for m in available_free:
+                model_name = m.replace(":free", "").split("/")[-1]
+                if model_hint.lower() in model_name.lower():
+                    candidates.append(m)
+
+        if not candidates and available_free:
+            candidates = available_free[:10]
+
+        async with aiohttp.ClientSession() as session:
+            for model_id in candidates:
+                result = await self.try_model_direct_stream(session, model_id, api_key, messages)
+                if result and result.get("success"):
+                    return result.get("stream"), result.get("model")
+
+        return None, None