Update main.py

main.py

@@ -1,37 +1,46 @@
-from fastapi import FastAPI, Request
-from fastapi.responses import PlainTextResponse
-from pydantic import BaseModel
-from typing import List, Optional, Literal
+from flask import Flask, request, jsonify
 from gradio_client import Client
-import uvicorn
-import time
 import uuid
+import time
 import logging
 
-# === Настройка логгера ===
-logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")
-logger = logging.getLogger(__name__)
-
-# === Задаём модель (используется глобально) ===
+# === Конфигурация ===
 AI_MODEL = "NousResearch/Nous-Hermes-2-Mixtral-8x7B-DPO"
+HF_SPACE_URL = "https://nymbo-serverless-textgen-hub.hf.space"
+
+# === Инициализация ===
+app = Flask(__name__)
+logging.basicConfig(level=logging.INFO)
 
-# === Подключаемся к Gradio Space по URL ===
 try:
-    gr_client = Client(src="https://nymbo-serverless-textgen-hub.hf.space")
-    logger.info("✅ Успешно подключено к Gradio Space")
+    gr_client = Client(src=HF_SPACE_URL)
 except Exception as e:
-    logger.error(f"❌ Ошибка подключения к Gradio Client: {e}")
+    logging.error(f"Ошибка подключения к HuggingFace Space: {e}")
     gr_client = None
 
-# === Функция вызова модели ===
-def ask(user_prompt, system_prompt):
-    if not gr_client:
-        return "[Ошибка: Gradio Client не инициализирован]"
+# === Корневой эндпоинт ===
+@app.route("/", methods=["GET"])
+def home():
+    return "Proxy free test", 200
 
+# === Основной OpenAI-подобный endpoint ===
+@app.route("/v1/chat/completions", methods=["POST"])
+def chat():
     try:
-        result = gr_client.predict(
-            history=[[user_prompt, None]],
-            system_msg=system_prompt,
+        data = request.get_json()
+        messages = data.get("messages", [])
+        user_msg = next((m["content"] for m in reversed(messages) if m["role"] == "user"), None)
+        system_msg = next((m["content"] for m in messages if m["role"] == "system"), "You are a helpful AI assistant.")
+    except Exception as e:
+        return jsonify({"error": f"Невалидный JSON: {e}"}), 400
+
+    if not user_msg or not gr_client:
+        return jsonify({"error": "Отсутствует сообщение пользователя или не подключена модель."}), 400
+
+    try:
+        output = gr_client.predict(
+            history=[[user_msg, None]],
+            system_msg=system_msg,
             max_tokens=512,
             temperature=0.7,
             top_p=0.95,
@@ -42,58 +51,10 @@ def ask(user_prompt, system_prompt):
             selected_model=AI_MODEL,
             api_name="/bot"
         )
-        return result
     except Exception as e:
-        logger.error(f"❌ Ошибка при вызове модели: {e}")
-        return f"[Ошибка: {str(e)}]"
-
-# === FastAPI приложение ===
-app = FastAPI()
-
-# === Модели запроса ===
-class Message(BaseModel):
-    role: Literal["user", "assistant", "system"]
-    content: str
-
-class ChatRequest(BaseModel):
-    model: str  # Заглушка — не используется
-    messages: List[Message]
-    temperature: Optional[float] = 0.7
-    top_p: Optional[float] = 0.95
-    max_tokens: Optional[int] = 512
+        return jsonify({"error": f"Ошибка при вызове модели: {e}"}), 500
 
-# === Корневой маршрут "/" ===
-@app.get("/", response_class=PlainTextResponse)
-async def root():
-    return "Proxy free test"
-
-# === Основной маршрут OpenAI-совместимый ===
-@app.post("/v1/chat/completions")
-async def chat_completion(request: Request):
-    headers = dict(request.headers)
-    body = await request.body()
-
-    logger.info("📥 Запрос получен")
-    logger.info(f"🔸 Заголовки: {headers}")
-    logger.info(f"🔸 Тело: {body.decode('utf-8')}")
-
-    try:
-        data = await request.json()
-        chat_request = ChatRequest(**data)
-    except Exception as e:
-        logger.error(f"❌ Ошибка разбора JSON: {e}")
-        return {"error": "Некорректный JSON"}
-
-    # Игнорируем model из запроса и используем глобальный AI_MODEL
-    user_msg = next((m.content for m in reversed(chat_request.messages) if m.role == "user"), None)
-    system_msg = next((m.content for m in chat_request.messages if m.role == "system"), "You are a helpful AI assistant.")
-
-    if not user_msg:
-        return {"error": "User message not found."}
-
-    assistant_reply = ask(user_msg, system_msg)
-
-    response = {
+    return jsonify({
         "id": f"chatcmpl-{uuid.uuid4().hex[:12]}",
         "object": "chat.completion",
         "created": int(time.time()),
@@ -103,7 +64,7 @@ async def chat_completion(request: Request):
                 "index": 0,
                 "message": {
                     "role": "assistant",
-                    "content": assistant_reply
+                    "content": output
                 },
                 "finish_reason": "stop"
             }
@@ -113,10 +74,8 @@ async def chat_completion(request: Request):
             "completion_tokens": 0,
             "total_tokens": 0
         }
-    }
-
-    return response
+    })
 
 # === Запуск сервера ===
 if __name__ == "__main__":
-    uvicorn.run("local_openai_server:app", host="0.0.0.0", port=7860, reload=True)
+    app.run(host="0.0.0.0", port=7860, debug=True)