Upload 9 files

app.py

@@ -2,37 +2,23 @@ import gradio as gr
 from fastapi import FastAPI, Request, HTTPException
 from fastapi.responses import JSONResponse, StreamingResponse
 from pydantic import BaseModel
-from typing import List, Optional
+from typing import List, Optional, Union
 import asyncio
 import random
 from datetime import datetime
 import threading
 import uvicorn
 import json
+import aiohttp
 
 from src.config import get_api_keys
 from src.model_tester import ModelTester
 from src.scheduler import Scheduler
-from src.utils import format_timestamp
 
 
 model_tester = ModelTester()
 scheduler = Scheduler(task_callback=lambda: model_tester.scan_all_models())
 
-
-class Message(BaseModel):
-    role: str
-    content: str
-
-
-class ChatCompletionRequest(BaseModel):
-    model: Optional[str] = None
-    messages: List[Message]
-    temperature: Optional[float] = 1.0
-    max_tokens: Optional[int] = None
-    stream: Optional[bool] = False
-
-
 fastapi_app = FastAPI(title="OpenRouter Free API")
 
 
@@ -54,48 +40,85 @@ async def list_models():
             "object": "model",
             "created": 1677610602,
             "owned_by": "openrouter",
-            "permission": [],
-            "root": model_id,
-            "parent": None,
             "free": is_free
         })
     
-    return {
-        "object": "list",
-        "data": models
+    return {"object": "list", "data": models}
+
+
+async def proxy_request(model_hint: Optional[str], messages: list, stream: bool):
+    """透传请求到OpenRouter，只修改model"""
+    api_keys = get_api_keys()
+    api_key = random.choice(api_keys)
+    
+    # 找到可用的模型ID
+    model_tester.refresh_model_list()
+    available_free = model_tester.get_all_free_models()
+    
+    target_model = None
+    
+    if model_hint:
+        # 尝试匹配用户指定的模型
+        for m in available_free:
+            model_name = m.replace(":free", "").split("/")[-1]
+            if model_hint.lower() in model_name.lower():
+                target_model = m
+                break
+        if not target_model:
+            target_model = f"{model_hint}:free" if ":free" not in model_hint else model_hint
+    
+    if not target_model and available_free:
+        target_model = random.choice(available_free[:5])
+    
+    if not target_model:
+        raise HTTPException(status_code=400, detail="No available model")
+    
+    # 构建请求
+    url = "https://openrouter.ai/api/v1/chat/completions"
+    headers = {
+        "Authorization": f"Bearer {api_key}",
+        "Content-Type": "application/json"
+    }
+    payload = {
+        "model": target_model,
+        "messages": messages,
+        "stream": stream
     }
+    
+    async with aiohttp.ClientSession() as session:
+        async with session.post(url, json=payload, headers=headers, timeout=aiohttp.ClientTimeout(total=120)) as response:
+            if stream:
+                async for chunk in response.content:
+                    yield chunk
+            else:
+                data = await response.json()
+                yield data
 
 
 @fastapi_app.post("/v1/chat/completions")
-async def chat_completions(request: ChatCompletionRequest):
-    prompt = request.messages[-1].content if request.messages else ""
-    model_hint = request.model
+async def chat_completions(request: Request):
+    body = await request.json()
+    stream = body.get("stream", False)
+    messages = body.get("messages", [])
     
-    if request.stream:
+    if stream:
         return StreamingResponse(
-            stream_chat(request.model, [{"role": m.role, "content": m.content} for m in request.messages]),
+            proxy_request(body.get("model"), messages, stream),
             media_type="text/event-stream"
         )
     
-    result = await model_tester.chat_completion(prompt, model_hint)
+    result = None
+    async for data in proxy_request(body.get("model"), messages, stream):
+        result = data
+        break
     
-    if not result.get("success"):
-        raise HTTPException(status_code=400, detail=result.get("error", "Request failed"))
+    if not result:
+        raise HTTPException(status_code=400, detail="Request failed")
     
-    # 直接返回原始响应
-    return result.get("response", {})
-
-
-async def stream_chat(model_hint: Optional[str], messages: list):
-    try:
-        async for chunk in model_tester.chat_completion_stream(model_hint, messages):
-            if isinstance(chunk, bytes):
-                yield chunk.decode('utf-8')
-            else:
-                yield chunk
-    except Exception as e:
-        yield f"data: {{\"error\": \"{str(e)}\"}}\n\n"
-        yield "data: [DONE]\n\n"
+    if "error" in result:
+        raise HTTPException(status_code=400, detail=result["error"])
+    
+    return result
 
 
 @fastapi_app.get("/health")
@@ -115,8 +138,7 @@ def get_scan_status():
     scan_result = model_tester.scan_result
     total = scan_result.get("total_available", 0)
     free = scan_result.get("free_available", 0)
-    timestamp = scan_result.get("timestamp", "Never")
-    return f"Free: {free} | Total: {total} | Last: {timestamp}"
+    return f"Free: {free} | Total: {total}"
 
 
 def format_model_list(models):
@@ -127,28 +149,11 @@ with gr.Blocks(title="OpenRouter Free API") as demo:
     gr.Markdown("# OpenRouter Free API")
     gr.Markdown("Standard OpenAI-compatible API with free model support")
     
-    gr.Markdown("## API Endpoints")
-    gr.Markdown("""
-    - `GET /v1/models` - List available models  
-    - `POST /v1/chat/completions` - Chat completion
-    """)
-    
-    gr.Markdown("---")
-    gr.Markdown(f"## Status: {get_scan_status()}")
+    gr.Markdown("## Status")
+    gr.Markdown(f"**{get_scan_status()}**")
     
     gr.Markdown("## Available Free Models")
-    free_display = gr.Textbox(
-        value=format_model_list(get_available_free_models()), 
-        lines=15, 
-        interactive=False
-    )
-    
-    gr.Markdown("## All Available Models")
-    all_display = gr.Textbox(
-        value=format_model_list(get_available_models()), 
-        lines=15, 
-        interactive=False
-    )
+    gr.Textbox(value=format_model_list(get_available_free_models()), lines=15, interactive=False)
 
 
 app = gr.mount_gradio_app(fastapi_app, demo, path="/")