Update main.py

main.py

@@ -0,0 +1,84 @@
+from fastapi import FastAPI, Request
+from pydantic import BaseModel
+from typing import List, Optional, Literal
+from gradio_client import Client
+import uvicorn
+import time
+import uuid
+
+# ==== Инициализация Gradio Client ====
+gr_client = Client("Nymbo/Serverless-TextGen-Hub")
+
+# ==== Функция обращения к нейросети ====
+def ask(user_prompt, system_prompt, model):
+    result = gr_client.predict(
+        history=[[user_prompt, None]],
+        system_msg=system_prompt,
+        max_tokens=512,
+        temperature=0.7,
+        top_p=0.95,
+        freq_penalty=0,
+        seed=-1,
+        custom_model=model,
+        search_term="",
+        selected_model=model,
+        api_name="/bot"
+    )
+    return result
+
+# ==== 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
+    # остальные параметры можно добавить при необходимости
+
+@app.post("/v1/chat/completions")
+async def chat_completion(request: ChatRequest):
+    # Извлекаем последнее сообщение от пользователя
+    user_msg = next((m.content for m in reversed(request.messages) if m.role == "user"), None)
+    system_msg = next((m.content for m in 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, request.model)
+
+    # Формируем ответ в стиле OpenAI API
+    response = {
+        "id": f"chatcmpl-{uuid.uuid4().hex[:12]}",
+        "object": "chat.completion",
+        "created": int(time.time()),
+        "model": request.model,
+        "choices": [
+            {
+                "index": 0,
+                "message": {
+                    "role": "assistant",
+                    "content": assistant_reply
+                },
+                "finish_reason": "stop"
+            }
+        ],
+        "usage": {
+            "prompt_tokens": 0,   # Можно вычислить при необходимости
+            "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)