EmoCube commited on
Commit
8b04581
·
verified ·
1 Parent(s): 0c0641a

Update main.py

Browse files
Files changed (1) hide show
  1. main.py +31 -23
main.py CHANGED
@@ -1,4 +1,5 @@
1
  from fastapi import FastAPI, Request
 
2
  from pydantic import BaseModel
3
  from typing import List, Optional, Literal
4
  from gradio_client import Client
@@ -6,24 +7,27 @@ import uvicorn
6
  import time
7
  import uuid
8
  import logging
9
- import json
10
 
11
  # === Настройка логгера ===
12
  logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")
13
  logger = logging.getLogger(__name__)
14
 
15
- # === Подключаемся к Gradio Space напрямую по URL ===
 
 
 
16
  try:
17
- gr_client = Client("Nymbo/Serverless-TextGen-Hub")
 
18
  except Exception as e:
19
- logger.error(f"Ошибка при подключении к Gradio Client: {e}")
20
  gr_client = None
21
 
22
- # === Вызов нейросети ===
23
- def ask(user_prompt, system_prompt, model):
24
  if not gr_client:
25
  return "[Ошибка: Gradio Client не инициализирован]"
26
-
27
  try:
28
  result = gr_client.predict(
29
  history=[[user_prompt, None]],
@@ -33,63 +37,67 @@ def ask(user_prompt, system_prompt, model):
33
  top_p=0.95,
34
  freq_penalty=0,
35
  seed=-1,
36
- custom_model=model,
37
  search_term="",
38
- selected_model=model,
39
  api_name="/bot"
40
  )
41
  return result
42
  except Exception as e:
43
- logger.error(f"Ошибка при вызове Gradio predict: {e}")
44
  return f"[Ошибка: {str(e)}]"
45
 
46
- # === Инициализация FastAPI ===
47
  app = FastAPI()
48
 
49
- # === Pydantic модели ===
50
  class Message(BaseModel):
51
  role: Literal["user", "assistant", "system"]
52
  content: str
53
 
54
  class ChatRequest(BaseModel):
55
- model: str
56
  messages: List[Message]
57
  temperature: Optional[float] = 0.7
58
  top_p: Optional[float] = 0.95
59
  max_tokens: Optional[int] = 512
60
 
61
- # === Основной маршрут ===
 
 
 
 
 
62
  @app.post("/v1/chat/completions")
63
  async def chat_completion(request: Request):
64
- # Логгируем заголовки и тело запроса
65
  headers = dict(request.headers)
66
  body = await request.body()
67
- logger.info("== Входящий запрос ==")
68
- logger.info(f"Заголовки: {headers}")
69
- logger.info(f"Тело: {body.decode('utf-8')}")
 
70
 
71
  try:
72
  data = await request.json()
73
  chat_request = ChatRequest(**data)
74
  except Exception as e:
75
- logger.error(f"Ошибка парсинга запроса: {e}")
76
  return {"error": "Некорректный JSON"}
77
 
78
- # Извлекаем сообщения
79
  user_msg = next((m.content for m in reversed(chat_request.messages) if m.role == "user"), None)
80
  system_msg = next((m.content for m in chat_request.messages if m.role == "system"), "You are a helpful AI assistant.")
81
 
82
  if not user_msg:
83
  return {"error": "User message not found."}
84
 
85
- # Ответ от модели
86
- assistant_reply = ask(user_msg, system_msg, chat_request.model)
87
 
88
  response = {
89
  "id": f"chatcmpl-{uuid.uuid4().hex[:12]}",
90
  "object": "chat.completion",
91
  "created": int(time.time()),
92
- "model": chat_request.model,
93
  "choices": [
94
  {
95
  "index": 0,
 
1
  from fastapi import FastAPI, Request
2
+ from fastapi.responses import PlainTextResponse
3
  from pydantic import BaseModel
4
  from typing import List, Optional, Literal
5
  from gradio_client import Client
 
7
  import time
8
  import uuid
9
  import logging
 
10
 
11
  # === Настройка логгера ===
12
  logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")
13
  logger = logging.getLogger(__name__)
14
 
15
+ # === Задаём модель (используется глобально) ===
16
+ AI_MODEL = "NousResearch/Nous-Hermes-2-Mixtral-8x7B-DPO"
17
+
18
+ # === Подключаемся к Gradio Space по URL ===
19
  try:
20
+ gr_client = Client("https://nymbo-serverless-textgen-hub.hf.space")
21
+ logger.info("✅ Успешно подключено к Gradio Space")
22
  except Exception as e:
23
+ logger.error(f"Ошибка подключения к Gradio Client: {e}")
24
  gr_client = None
25
 
26
+ # === Функция вызова модели ===
27
+ def ask(user_prompt, system_prompt):
28
  if not gr_client:
29
  return "[Ошибка: Gradio Client не инициализирован]"
30
+
31
  try:
32
  result = gr_client.predict(
33
  history=[[user_prompt, None]],
 
37
  top_p=0.95,
38
  freq_penalty=0,
39
  seed=-1,
40
+ custom_model=AI_MODEL,
41
  search_term="",
42
+ selected_model=AI_MODEL,
43
  api_name="/bot"
44
  )
45
  return result
46
  except Exception as e:
47
+ logger.error(f"Ошибка при вызове модели: {e}")
48
  return f"[Ошибка: {str(e)}]"
49
 
50
+ # === FastAPI приложение ===
51
  app = FastAPI()
52
 
53
+ # === Модели запроса ===
54
  class Message(BaseModel):
55
  role: Literal["user", "assistant", "system"]
56
  content: str
57
 
58
  class ChatRequest(BaseModel):
59
+ model: str # Заглушка — не используется
60
  messages: List[Message]
61
  temperature: Optional[float] = 0.7
62
  top_p: Optional[float] = 0.95
63
  max_tokens: Optional[int] = 512
64
 
65
+ # === Корневой маршрут "/" ===
66
+ @app.get("/", response_class=PlainTextResponse)
67
+ async def root():
68
+ return "Proxy free test"
69
+
70
+ # === Основной маршрут OpenAI-совместимый ===
71
  @app.post("/v1/chat/completions")
72
  async def chat_completion(request: Request):
 
73
  headers = dict(request.headers)
74
  body = await request.body()
75
+
76
+ logger.info("📥 Запрос получен")
77
+ logger.info(f"🔸 Заголовки: {headers}")
78
+ logger.info(f"🔸 Тело: {body.decode('utf-8')}")
79
 
80
  try:
81
  data = await request.json()
82
  chat_request = ChatRequest(**data)
83
  except Exception as e:
84
+ logger.error(f"Ошибка разбора JSON: {e}")
85
  return {"error": "Некорректный JSON"}
86
 
87
+ # Игнорируем model из запроса и используем глобальный AI_MODEL
88
  user_msg = next((m.content for m in reversed(chat_request.messages) if m.role == "user"), None)
89
  system_msg = next((m.content for m in chat_request.messages if m.role == "system"), "You are a helpful AI assistant.")
90
 
91
  if not user_msg:
92
  return {"error": "User message not found."}
93
 
94
+ assistant_reply = ask(user_msg, system_msg)
 
95
 
96
  response = {
97
  "id": f"chatcmpl-{uuid.uuid4().hex[:12]}",
98
  "object": "chat.completion",
99
  "created": int(time.time()),
100
+ "model": AI_MODEL,
101
  "choices": [
102
  {
103
  "index": 0,