vortex123 commited on
Commit
cfd7ca0
·
verified ·
1 Parent(s): e94650b

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +36 -165
app.py CHANGED
@@ -8,18 +8,17 @@ import asyncio
8
  ###############################################################################
9
 
10
  # Подставьте свой ключ или берите из окружения
11
- GEMINI_API_KEY = "AIzaSyBoqoPX-9uzvXyxzse0gRwH8_P9xO6O3Bc"
12
  if not GEMINI_API_KEY:
13
- print("Error: GEMINI_API_KEY is not set.")
14
  exit()
15
 
16
  genai.configure(api_key=GEMINI_API_KEY)
17
 
18
  # Выберите доступные модели (пример)
19
  AVAILABLE_MODELS = [
20
- "gemini-2.0-flash-exp",
21
- "gemini-exp-1206",
22
- "gemini-2.0-flash-thinking-exp-1219",
23
  ]
24
 
25
  MODELS = {}
@@ -35,19 +34,11 @@ for model_name in AVAILABLE_MODELS:
35
  # Когда пользователь переключается на модель, мы добавляем это сообщение в историю.
36
 
37
  DEFAULT_DEVELOPER_PROMPTS = {
38
- "gemini-2.0-flash-exp": (
39
- "You are a normal model (developer role). "
40
- "Provide direct answers with no JSON wrapping."
41
- ),
42
- "gemini-exp-1206": (
43
- "You are an experimental normal model (developer role). "
44
- "Provide direct answers with no JSON wrapping."
45
- ),
46
- "gemini-2.0-flash-thinking-exp-1219": (
47
- "You are a thinking model (developer role). "
48
- "Please provide your final answer in the format {output: ...}. "
49
- "You may use internal thoughts but do not show them directly to the user."
50
  ),
 
51
  }
52
 
53
  ###############################################################################
@@ -58,9 +49,7 @@ def _assistant_role(model_name: str) -> str:
58
  """
59
  Некоторые новые модели не принимают 'assistant', а требуют 'model'.
60
  """
61
- # Допустим "gemini-exp-1206" и "gemini-2.0-flash-thinking-exp-1219" хотят "model"
62
- if model_name in ["gemini-exp-1206", "gemini-2.0-flash-thinking-exp-1219"]:
63
- return "model"
64
  return "assistant"
65
 
66
  ###############################################################################
@@ -78,8 +67,11 @@ def _history_to_genai(history, model_name):
78
 
79
  for user_text, bot_text in history:
80
  if user_text:
81
- # Пользовательское сообщение
82
- genai_history.append({"role": "user", "parts": user_text})
 
 
 
83
  if bot_text:
84
  # Ответ ассистента (или model)
85
  genai_history.append({"role": asst_role, "parts": bot_text})
@@ -91,7 +83,7 @@ def _history_to_genai(history, model_name):
91
 
92
  async def _respond_stream(model_name, user_message, history):
93
  """
94
- Стриминговый ответ для обычных моделей:
95
  - Кусочек за кусочком (partial_text).
96
  """
97
  if model_name not in MODELS:
@@ -103,10 +95,10 @@ async def _respond_stream(model_name, user_message, history):
103
 
104
  try:
105
  chat = model.start_chat(history=genai_history)
106
- stream = chat.send_message(user_message, stream=True)
107
 
108
  partial_text = ""
109
- for chunk in stream:
110
  partial_text += (chunk.text or "")
111
  yield partial_text
112
 
@@ -115,11 +107,10 @@ async def _respond_stream(model_name, user_message, history):
115
  yield f"Ошибка при запросе к API: {e}"
116
  return
117
 
118
-
119
  async def _respond_thinking(model_name, user_message, history):
120
  """
121
  Для thinking-моделей:
122
- 1) Выводим "Думаю..."
123
  2) После завершения — финальный ответ в формате {output: ...} + размышления.
124
  """
125
  if model_name not in MODELS:
@@ -142,10 +133,14 @@ async def _respond_thinking(model_name, user_message, history):
142
  if response.candidates:
143
  parts = response.candidates[0].content.parts
144
  for p in parts:
145
- if getattr(p, "thought", False):
 
 
146
  thinking_process_text += p.text or ""
 
 
147
  else:
148
- final_text += p.text or ""
149
 
150
  # Для thinking-моделей просили итоговый ответ в {output: ...}
151
  final_text_formatted = f"{{output: {final_text}}}"
@@ -168,7 +163,7 @@ async def user_send_message(
168
  thinking_text: str
169
  ):
170
  """
171
- Колбэк, когда пользователь отправляет запрос.
172
  Добавляем в history новый (user_msg, None), затем генерируем ответ.
173
  """
174
  # Пустой ввод
@@ -179,9 +174,9 @@ async def user_send_message(
179
  # Добавляем (user_message, None)
180
  history.append((user_message, None))
181
 
182
- # Если модель — thinking
183
  if "thinking" in model_name.lower():
184
- async for (assistant_text, thought_text) in _respond_thinking(model_name, user_message, history):
185
  history[-1] = (user_message, assistant_text)
186
  yield history, thought_text
187
  return
@@ -208,43 +203,24 @@ def clear_all():
208
 
209
  def on_model_change(selected_model, history, thinking_text):
210
  """
211
- При переключении модели добавляем в историю developer-сообщение,
212
  + добавляем дефолтный промпт этой модели (тоже developer).
213
  """
214
  new_history = history.copy()
215
 
216
  # Cообщаем модели, что переключились (developer role)
217
  new_history.append((
218
- # developer-сообщение выглядит как "role=developer" → user_text
219
- # => для удобства считаем, что (user_text, None) => потом в _history_to_genai
220
- # оно будет "role=developer", а не user.
221
- # Но проще завести отдельное поле?
222
- # Для упрощения – сделаем pseudo-user c developer-меткой.
223
- # Или сделаем костыль: user_msg="[developer]Switched to X" →
224
- # и в _history_to_genai, если строка начинается с "[developer]",
225
- # ставим role="developer"?
226
- #
227
- # Более аккуратно — используем 3-элементный кортеж?
228
- # Но Gradio Chatbot ждёт (str, str).
229
- #
230
- # Проще при конвертации _history_to_genai:
231
- # если user_text.startswith("<developer>"),
232
- # ставим role="developer".
233
- #
234
- # По заданию: "замени роль system на developer".
235
- # Ок, будем распознавать префикс "<developer>: ".
236
- #
237
- "<developer>: Switched to model: " + selected_model,
238
  None
239
  ))
240
 
241
  # Добавляем дефолтный промпт (developer role)
242
  default_prompt = DEFAULT_DEVELOPER_PROMPTS.get(
243
- selected_model,
244
  "No default prompt for this model."
245
  )
246
  new_history.append((
247
- "<developer>: " + default_prompt,
248
  None
249
  ))
250
 
@@ -254,112 +230,7 @@ def on_model_change(selected_model, history, thinking_text):
254
  # 9. Функция конвертации истории с учётом developer role
255
  ###############################################################################
256
 
257
- def _history_to_genai_enhanced(history, model_name):
258
- """
259
- Улучшенная версия, отличающая developer-сообщения
260
- (префикс "<developer>: ") от user-сообщений.
261
- """
262
- asst_role = _assistant_role(model_name)
263
- genai_history = []
264
-
265
- for user_text, bot_text in history:
266
- if user_text:
267
- if user_text.startswith("<developer>: "):
268
- # Считаем это developer role
269
- dev_content = user_text.replace("<developer>: ", "", 1)
270
- genai_history.append({"role": "developer", "parts": dev_content})
271
- else:
272
- # Обычный пользователь
273
- genai_history.append({"role": "user", "parts": user_text})
274
-
275
- if bot_text:
276
- # Ответ ассистента / модель
277
- genai_history.append({"role": asst_role, "parts": bot_text})
278
-
279
- return genai_history
280
-
281
- # Обновим наши ге��ераторы, чтобы использовать _history_to_genai_enhanced:
282
- async def _respond_stream_enh(model_name, user_message, history):
283
- if model_name not in MODELS:
284
- yield "Ошибка: модель не найдена."
285
- return
286
-
287
- model = MODELS[model_name]
288
- genai_history = _history_to_genai_enhanced(history, model_name)
289
-
290
- try:
291
- chat = model.start_chat(history=genai_history)
292
- stream = chat.send_message(user_message, stream=True)
293
-
294
- partial_text = ""
295
- for chunk in stream:
296
- partial_text += (chunk.text or "")
297
- yield partial_text
298
-
299
- return
300
- except Exception as e:
301
- yield f"Ошибка при запросе к API: {e}"
302
- return
303
-
304
- async def _respond_thinking_enh(model_name, user_message, history):
305
- if model_name not in MODELS:
306
- yield "Ошибка: модель не найдена.", ""
307
- return
308
-
309
- model = MODELS[model_name]
310
- genai_history = _history_to_genai_enhanced(history, model_name)
311
-
312
- # "Думаю..."
313
- yield "Думаю...", ""
314
-
315
- try:
316
- chat = model.start_chat(history=genai_history)
317
- response = chat.send_message(user_message, stream=False)
318
-
319
- thinking_process_text = ""
320
- final_text = ""
321
-
322
- if response.candidates:
323
- parts = response.candidates[0].content.parts
324
- for p in parts:
325
- if getattr(p, "thought", False):
326
- thinking_process_text += p.text or ""
327
- else:
328
- final_text += p.text or ""
329
-
330
- # JSON-обёртка
331
- final_text_formatted = f"{{output: {final_text}}}"
332
-
333
- yield final_text_formatted, thinking_process_text
334
- return
335
- except Exception as e:
336
- yield f"Ошибка при запросе к API: {e}", ""
337
- return
338
-
339
- async def user_send_message_enh(
340
- user_message: str,
341
- history: list[tuple[str, str]],
342
- model_name: str,
343
- thinking_text: str
344
- ):
345
- if not user_message.strip():
346
- yield history, thinking_text
347
- return
348
-
349
- history.append((user_message, None))
350
-
351
- if "thinking" in model_name.lower():
352
- async for (assistant_text, thought_text) in _respond_thinking_enh(model_name, user_message, history):
353
- history[-1] = (user_message, assistant_text)
354
- yield history, thought_text
355
- return
356
- else:
357
- partial_answer = ""
358
- async for chunk in _respond_stream_enh(model_name, user_message, history):
359
- partial_answer = chunk
360
- history[-1] = (user_message, partial_answer)
361
- yield history, ""
362
- return
363
 
364
  ###############################################################################
365
  # 10. Построение интерфейса Gradio
@@ -371,7 +242,7 @@ with gr.Blocks() as demo:
371
  with gr.Row():
372
  model_dropdown = gr.Dropdown(
373
  choices=AVAILABLE_MODELS,
374
- value="gemini-2.0-flash-exp",
375
  label="Выберите модель"
376
  )
377
  clear_button = gr.Button("Очистить чат")
@@ -409,7 +280,7 @@ with gr.Blocks() as demo:
409
  # (B) При нажатии «Отправить»
410
  ################################################
411
  send_chain = send_btn.click(
412
- fn=user_send_message_enh,
413
  inputs=[user_input, history_state, model_dropdown, thinking_store],
414
  outputs=[history_state, thinking_store],
415
  queue=True
@@ -438,7 +309,7 @@ with gr.Blocks() as demo:
438
  # (C) При нажатии Enter в textbox
439
  ################################################
440
  submit_chain = user_input.submit(
441
- fn=user_send_message_enh,
442
  inputs=[user_input, history_state, model_dropdown, thinking_store],
443
  outputs=[history_state, thinking_store],
444
  queue=True
@@ -490,4 +361,4 @@ with gr.Blocks() as demo:
490
 
491
  # Запуск
492
  if __name__ == "__main__":
493
- demo.launch()
 
8
  ###############################################################################
9
 
10
  # Подставьте свой ключ или берите из окружения
11
+ GEMINI_API_KEY = os.environ.get("GEMINI_API_KEY")
12
  if not GEMINI_API_KEY:
13
+ print("Error: GEMINI_API_KEY is not set. Please set the environment variable.")
14
  exit()
15
 
16
  genai.configure(api_key=GEMINI_API_KEY)
17
 
18
  # Выберите доступные модели (пример)
19
  AVAILABLE_MODELS = [
20
+ "gemini-1.5-flash",
21
+ # Add other available models if needed
 
22
  ]
23
 
24
  MODELS = {}
 
34
  # Когда пользователь переключается на модель, мы добавляем это сообщение в историю.
35
 
36
  DEFAULT_DEVELOPER_PROMPTS = {
37
+ "gemini-1.5-flash": (
38
+ "You are a helpful assistant (developer role). "
39
+ "Provide direct answers."
 
 
 
 
 
 
 
 
 
40
  ),
41
+ # Add default prompts for other models if needed
42
  }
43
 
44
  ###############################################################################
 
49
  """
50
  Некоторые новые модели не принимают 'assistant', а требуют 'model'.
51
  """
52
+ # Adjust based on model requirements
 
 
53
  return "assistant"
54
 
55
  ###############################################################################
 
67
 
68
  for user_text, bot_text in history:
69
  if user_text:
70
+ if user_text.startswith("<developer>: "):
71
+ genai_history.append({"role": "system", "parts": user_text.replace("<developer>: ", "", 1)})
72
+ else:
73
+ # Пользовательское сообщение
74
+ genai_history.append({"role": "user", "parts": user_text})
75
  if bot_text:
76
  # Ответ ассистента (или model)
77
  genai_history.append({"role": asst_role, "parts": bot_text})
 
83
 
84
  async def _respond_stream(model_name, user_message, history):
85
  """
86
+ Стриминговый ответ для обычных моделей:
87
  - Кусочек за кусочком (partial_text).
88
  """
89
  if model_name not in MODELS:
 
95
 
96
  try:
97
  chat = model.start_chat(history=genai_history)
98
+ response = chat.send_message(user_message, stream=True)
99
 
100
  partial_text = ""
101
+ async for chunk in response:
102
  partial_text += (chunk.text or "")
103
  yield partial_text
104
 
 
107
  yield f"Ошибка при запросе к API: {e}"
108
  return
109
 
 
110
  async def _respond_thinking(model_name, user_message, history):
111
  """
112
  Для thinking-моделей:
113
+ 1) Выводим "Думаю..."
114
  2) После завершения — финальный ответ в формате {output: ...} + размышления.
115
  """
116
  if model_name not in MODELS:
 
133
  if response.candidates:
134
  parts = response.candidates[0].content.parts
135
  for p in parts:
136
+ # Assuming a specific structure for "thinking" models' output
137
+ # This part might need adjustments based on actual model output
138
+ if isinstance(p, dict) and "output" not in p: # Heuristic for thought
139
  thinking_process_text += p.text or ""
140
+ elif isinstance(p, dict) and "output" in p:
141
+ final_text += p["output"] # Assuming output is directly in the dict
142
  else:
143
+ final_text += p.text or "" # Fallback
144
 
145
  # Для thinking-моделей просили итоговый ответ в {output: ...}
146
  final_text_formatted = f"{{output: {final_text}}}"
 
163
  thinking_text: str
164
  ):
165
  """
166
+ Колбэк, когда пользователь отправляет запрос.
167
  Добавляем в history новый (user_msg, None), затем генерируем ответ.
168
  """
169
  # Пустой ввод
 
174
  # Добавляем (user_message, None)
175
  history.append((user_message, None))
176
 
177
+ # Если модель — thinking (adjust logic based on actual model names)
178
  if "thinking" in model_name.lower():
179
+ async for assistant_text, thought_text in _respond_thinking(model_name, user_message, history):
180
  history[-1] = (user_message, assistant_text)
181
  yield history, thought_text
182
  return
 
203
 
204
  def on_model_change(selected_model, history, thinking_text):
205
  """
206
+ При переключении модели добавляем в историю developer-сообщение,
207
  + добавляем дефолтный промпт этой модели (тоже developer).
208
  """
209
  new_history = history.copy()
210
 
211
  # Cообщаем модели, что переключились (developer role)
212
  new_history.append((
213
+ f"<developer>: Switched to model: {selected_model}",
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
214
  None
215
  ))
216
 
217
  # Добавляем дефолтный промпт (developer role)
218
  default_prompt = DEFAULT_DEVELOPER_PROMPTS.get(
219
+ selected_model,
220
  "No default prompt for this model."
221
  )
222
  new_history.append((
223
+ f"<developer>: {default_prompt}",
224
  None
225
  ))
226
 
 
230
  # 9. Функция конвертации истории с учётом developer role
231
  ###############################################################################
232
 
233
+ # This part is already implemented in _history_to_genai
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
234
 
235
  ###############################################################################
236
  # 10. Построение интерфейса Gradio
 
242
  with gr.Row():
243
  model_dropdown = gr.Dropdown(
244
  choices=AVAILABLE_MODELS,
245
+ value=AVAILABLE_MODELS[0] if AVAILABLE_MODELS else None,
246
  label="Выберите модель"
247
  )
248
  clear_button = gr.Button("Очистить чат")
 
280
  # (B) При нажатии «Отправить»
281
  ################################################
282
  send_chain = send_btn.click(
283
+ fn=user_send_message,
284
  inputs=[user_input, history_state, model_dropdown, thinking_store],
285
  outputs=[history_state, thinking_store],
286
  queue=True
 
309
  # (C) При нажатии Enter в textbox
310
  ################################################
311
  submit_chain = user_input.submit(
312
+ fn=user_send_message,
313
  inputs=[user_input, history_state, model_dropdown, thinking_store],
314
  outputs=[history_state, thinking_store],
315
  queue=True
 
361
 
362
  # Запуск
363
  if __name__ == "__main__":
364
+ demo.launch()