vortex123 commited on
Commit
b4fa6b2
·
verified ·
1 Parent(s): cb994ee

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +106 -136
app.py CHANGED
@@ -13,297 +13,267 @@ GEMINI_API_KEY = "AIzaSyBoqoPX-9uzvXyxzse0gRwH8_P9xO6O3Bc"
13
  if not GEMINI_API_KEY:
14
  print("Error: GEMINI_API_KEY is not set.")
15
  exit()
16
-
17
  genai.configure(api_key=GEMINI_API_KEY)
18
 
 
19
  AVAILABLE_MODELS = [
20
  "gemini-1.5-flash",
21
  "gemini-1.5-pro",
22
- "gemini-2.0-flash-thinking-exp"
23
  ]
24
 
25
- # Создаём объекты-модели и храним их в словаре для быстрого доступа
26
  MODELS = {}
27
  for model_name in AVAILABLE_MODELS:
28
  try:
29
  MODELS[model_name] = genai.GenerativeModel(model_name=model_name)
30
  except Exception as e:
31
- print(f"Не удалось инициализировать модель {model_name}: {e}")
 
32
 
33
  ###############################################################################
34
- # 2. Функции для генерации ответов
35
  ###############################################################################
36
 
37
- def _convert_history_to_genai_format(history):
38
  """
39
- Gradio-хранит чат как список кортежей: [(user_msg, bot_msg), (user_msg, bot_msg), ...]
40
- Для Gemini же нужна история вида [{'role': 'user', 'content': ...}, {'role': 'assistant', 'content': ...}, ...].
41
-
42
- Возвращает список dict, где role = user/assistant, content = текст.
43
  """
44
  genai_history = []
45
  for user_text, bot_text in history:
46
- if user_text: # сообщение от пользователя
47
  genai_history.append({"role": "user", "content": user_text})
48
- if bot_text: # сообщение от ассистента
49
  genai_history.append({"role": "assistant", "content": bot_text})
50
  return genai_history
51
 
52
 
53
- def _convert_genai_to_chatbot_format(history):
54
- """
55
- Обратное преобразование:
56
- [{'role': 'user', 'content': ...}, {'role': 'assistant', 'content': ...}]
57
- -> [(user_msg, bot_msg), ...]
58
-
59
- Упрощённо полагаем, что роли будут идти поочерёдно user-assistant.
60
- Если вдруг несколько подряд от одного и того же role, объединяем по очереди.
61
- """
62
- display_history = []
63
- user_buffer = None
64
-
65
- for turn in history:
66
- role = turn["role"]
67
- content = turn["content"]
68
- if role == "user":
69
- # Запускаем новую пару
70
- user_buffer = content
71
- elif role == "assistant":
72
- if user_buffer is not None:
73
- display_history.append((user_buffer, content))
74
- user_buffer = None
75
- else:
76
- # На случай, если идёт assistant без user
77
- display_history.append((None, content))
78
-
79
- # Если вдруг остался user_buffer без ответа (редко бывает)
80
- if user_buffer is not None:
81
- display_history.append((user_buffer, None))
82
-
83
- return display_history
84
-
85
 
86
  async def _respond_stream(model_name, user_message, history):
87
  """
88
- Генерация ОТВЕТА С ПОСТЕПЕННЫМ стримингом (для моделей без "thinking").
89
- Возвращает генератор, чтобы Gradio мог пошагово обновлять вывод.
 
90
  """
91
  if model_name not in MODELS:
92
  yield "Ошибка: модель не найдена."
93
  return
94
 
95
  model = MODELS[model_name]
96
- genai_history = _convert_history_to_genai_format(history)
97
-
98
  try:
99
  chat = model.start_chat(history=genai_history)
100
- # Запрашиваем стриминг-ответ
101
  response_stream = chat.send_message(user_message, stream=True)
102
 
103
  partial_text = ""
104
  for chunk in response_stream:
105
  chunk_text = chunk.text or ""
106
  partial_text += chunk_text
107
- # Возвращаем накопленный ответ
108
  yield partial_text
 
 
109
  except Exception as e:
110
- yield f"Error during API call: {e}"
 
111
 
112
 
113
  async def _respond_thinking(model_name, user_message, history):
114
  """
115
- Генерация ответа для моделей с "thinking" (gemini-2.0-flash-thinking-exp).
116
- 1) Выдаём "Думаю..."
117
- 2) По завершении запроса возвращаем итоговый ответ
118
- 3) При желании можно вернуть размышления в отдельную переменную
 
119
  """
120
  if model_name not in MODELS:
121
- return "Ошибка: модель не найдена.", ""
 
 
122
 
123
  model = MODELS[model_name]
124
- genai_history = _convert_history_to_genai_format(history)
125
 
126
- # 1) "Думаю..."
127
- yield "Думаю...", "" # (ответ в чат, размышления)
128
 
129
  try:
130
  chat = model.start_chat(history=genai_history)
131
- response = chat.send_message(user_message, stream=False) # без stream
132
 
133
- # Для thinking-моделей ответ приходит в виде parts
134
  thinking_process_text = ""
135
  final_text = ""
136
  if response.candidates:
137
- for part in response.candidates[0].content.parts:
138
- # Если в part есть .thought == True, считаем это «размышлениями»
139
- if getattr(part, "thought", False):
140
- thinking_process_text += part.text or ""
141
  else:
142
- final_text += part.text or ""
143
 
144
- # Возвращаем итоговый ответ
145
  yield final_text, thinking_process_text
146
-
147
  except Exception as e:
148
- yield f"Error during API call: {e}", f"Error: {e}"
 
149
 
150
 
151
  ###############################################################################
152
- # 3. Основная функция (колбэк для Gradio), обрабатывающая один «шаг» диалога
153
  ###############################################################################
154
 
155
- async def user_send_message(user_message, history, model_name, thinking_store):
 
 
 
 
 
156
  """
157
- Эта функция вызывается при нажатии кнопки "Отправить" (или Enter).
158
  Параметры:
159
- user_message: текст от пользователя
160
- history: текущее состояние чата [(user_msg, bot_msg), ...]
161
  model_name: выбранная модель
162
- thinking_store: текущее значение поля "Размышления"
163
 
164
- Возвращаем:
165
- (updated_history, updated_thinking_store)
166
  """
 
167
  if not user_message.strip():
168
- # Пустая строка – ничего не делаем
169
- return history, thinking_store
170
 
171
- # Добавляем в историю шаг пользователя
172
- history.append((user_message, None)) # Пока ассистента нет
173
 
174
- # Если модель thinking
175
  if "thinking" in model_name.lower():
176
- # Асинхронный генератор: сначала "Думаю...", потом финальный ответ
177
- async for (assistant_text, thinking_text) in _respond_thinking(model_name, user_message, history):
178
- # Обновляем последний элемент истории (None -> assistant_text)
179
  history[-1] = (user_message, assistant_text)
180
- # Обновляем thinking_output
181
- yield history, thinking_text
 
182
  else:
183
- # Стриминговый ответ
184
  partial_answer = ""
185
- async for partial_chunk in _respond_stream(model_name, user_message, history):
186
- partial_answer = partial_chunk
187
  history[-1] = (user_message, partial_answer)
188
- # В режиме без thinking просто очищаем thinking_store
189
  yield history, ""
190
-
191
- # Когда генерация закончена, возвращаем итог
192
- return history, thinking_store
193
 
194
 
195
  ###############################################################################
196
- # 4. Колбэки для кнопки "Очистить"
197
  ###############################################################################
198
 
199
  def clear_all():
200
  """
201
- Полная очистка истории чата и поля «Размышления».
202
- Возвращает пустые значения для State и компонентов.
203
  """
204
- return [], "" # history, thinking_store
205
 
206
 
207
  ###############################################################################
208
- # 5. Определяем Gradio-интерфейс
209
  ###############################################################################
210
 
211
  with gr.Blocks() as demo:
212
- gr.Markdown("## Gemini Chatbot с сохранением истории и «thinking»-режимом")
213
-
214
- # Храним историю диалога (список кортежей) в gr.State
215
- history_state = gr.State([])
216
-
217
- # Храним «размышления» (thinking) в отдельном тексте
218
- thinking_store = gr.State("")
219
 
220
  with gr.Row():
221
  model_dropdown = gr.Dropdown(
222
  choices=AVAILABLE_MODELS,
223
  value="gemini-1.5-flash",
224
- label="Выберите модель"
225
  )
226
  clear_button = gr.Button("Очистить чат")
227
 
228
- # Сам чат: компонент Chatbot для отображения диалога
 
 
 
 
229
  chatbot = gr.Chatbot(label="Диалог с Gemini")
230
 
231
- # Поле ввода сообщения
232
  user_input = gr.Textbox(
233
- label="Ваш вопрос/сообщение",
234
- placeholder="Введите вопрос...",
235
  )
236
 
237
- # Поле для показа «Размышлений» (если модель thinking)
238
  thinking_output = gr.Textbox(
239
- label="Размышления модели (только для gemini-2.0-flash-thinking-exp)",
240
  interactive=False
241
  )
242
 
243
- # Кнопка «Отправить»
244
  send_btn = gr.Button("Отправить")
245
 
246
- ############################################################################
247
- # Связываем действия пользователя с колбэками
248
- ############################################################################
249
- # 1. При нажатии «Отправить» запускаем user_send_message
250
- # - обновляем историю и thinking_output
251
- # - одновременно обновляем отображение чата
252
  send_btn.click(
253
  fn=user_send_message,
254
  inputs=[user_input, history_state, model_dropdown, thinking_store],
255
  outputs=[history_state, thinking_store],
256
- queue=True,
257
  ).then(
258
- # После обновления history_state, пересчитываем отображение чата
259
  fn=lambda h: h,
260
  inputs=[history_state],
261
  outputs=[chatbot],
262
- preprocess=False,
263
- postprocess=True
264
  ).then(
265
- # После этого выводим текущее состояние thinking_store
266
  fn=lambda t: t,
267
  inputs=[thinking_store],
268
  outputs=[thinking_output],
 
269
  )
270
 
271
- # 2. При нажатии Enter в текстовом поле — та же логика, что и на кнопке
272
  user_input.submit(
273
  fn=user_send_message,
274
  inputs=[user_input, history_state, model_dropdown, thinking_store],
275
  outputs=[history_state, thinking_store],
276
- queue=True,
277
  ).then(
278
  fn=lambda h: h,
279
  inputs=[history_state],
280
  outputs=[chatbot],
281
- preprocess=False,
282
- postprocess=True
283
  ).then(
284
  fn=lambda t: t,
285
  inputs=[thinking_store],
286
  outputs=[thinking_output],
 
287
  )
288
 
289
- # 3. Кнопка «Очистить»
290
  clear_button.click(
291
  fn=clear_all,
292
  inputs=[],
293
  outputs=[history_state, thinking_store],
294
  queue=False
295
  ).then(
 
296
  fn=lambda h: h,
297
  inputs=[history_state],
298
- outputs=chatbot,
299
  ).then(
 
300
  fn=lambda _: "",
301
  inputs=[],
302
- outputs=thinking_output
303
  )
304
 
305
- ###############################################################################
306
- # 6. Запуск
307
- ###############################################################################
308
  if __name__ == "__main__":
309
- demo.launch()
 
13
  if not GEMINI_API_KEY:
14
  print("Error: GEMINI_API_KEY is not set.")
15
  exit()
 
16
  genai.configure(api_key=GEMINI_API_KEY)
17
 
18
+ # Доступные модели
19
  AVAILABLE_MODELS = [
20
  "gemini-1.5-flash",
21
  "gemini-1.5-pro",
22
+ "gemini-2.0-flash-thinking-exp",
23
  ]
24
 
25
+ # Пытаемся инициализировать все модели заранее и складываем в словарь
26
  MODELS = {}
27
  for model_name in AVAILABLE_MODELS:
28
  try:
29
  MODELS[model_name] = genai.GenerativeModel(model_name=model_name)
30
  except Exception as e:
31
+ print(f"[Предупреждение] Не удалось инициализировать модель {model_name}: {e}")
32
+
33
 
34
  ###############################################################################
35
+ # 2. Утилиты для преобразования истории Gradio <-> Gemini
36
  ###############################################################################
37
 
38
+ def _history_gradio_to_genai(history):
39
  """
40
+ Gradio хранит чат как [(user_msg, bot_msg), (user_msg, bot_msg), ...].
41
+ Для Gemini нужен формат [{'role': 'user', 'content': ...}, ...].
 
 
42
  """
43
  genai_history = []
44
  for user_text, bot_text in history:
45
+ if user_text:
46
  genai_history.append({"role": "user", "content": user_text})
47
+ if bot_text:
48
  genai_history.append({"role": "assistant", "content": bot_text})
49
  return genai_history
50
 
51
 
52
+ ###############################################################################
53
+ # 3. Функции-генераторы для запроса ответа от моделей (обычный/thinking)
54
+ ###############################################################################
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
55
 
56
  async def _respond_stream(model_name, user_message, history):
57
  """
58
+ Генерация ответа для Обычных моделей (без thinking) с помощью stream=True
59
+ - Возвращаем куски текста через yield
60
+ - В конце просто делаем return (без значения)
61
  """
62
  if model_name not in MODELS:
63
  yield "Ошибка: модель не найдена."
64
  return
65
 
66
  model = MODELS[model_name]
67
+ genai_history = _history_gradio_to_genai(history)
68
+
69
  try:
70
  chat = model.start_chat(history=genai_history)
 
71
  response_stream = chat.send_message(user_message, stream=True)
72
 
73
  partial_text = ""
74
  for chunk in response_stream:
75
  chunk_text = chunk.text or ""
76
  partial_text += chunk_text
77
+ # Возвращаем промежуточный вариант ответа
78
  yield partial_text
79
+
80
+ return # Завершить генератор без возвращения значения
81
  except Exception as e:
82
+ yield f"Ошибка при запросе к API: {e}"
83
+ return
84
 
85
 
86
  async def _respond_thinking(model_name, user_message, history):
87
  """
88
+ Генерация ответа для модели с "thinking" (например, gemini-2.0-flash-thinking-exp).
89
+ 1) Сначала выдаём "Думаю..."
90
+ 2) Затем, когда модель ответит (stream=False), выделяем thinking + финальный ответ.
91
+ 3) Возвращаем (полезный_ответ, размышления) в виде кортежа.
92
+ В Gradio это обычно [(assistant_text, thinking_text), ...].
93
  """
94
  if model_name not in MODELS:
95
+ # Выдаем ошибку через yield
96
+ yield "Ошибка: модель не найдена.", ""
97
+ return
98
 
99
  model = MODELS[model_name]
100
+ genai_history = _history_gradio_to_genai(history)
101
 
102
+ # Шаг 1: временно "Думаю..."
103
+ yield "Думаю...", ""
104
 
105
  try:
106
  chat = model.start_chat(history=genai_history)
107
+ response = chat.send_message(user_message, stream=False)
108
 
 
109
  thinking_process_text = ""
110
  final_text = ""
111
  if response.candidates:
112
+ parts = response.candidates[0].content.parts
113
+ for p in parts:
114
+ if getattr(p, "thought", False):
115
+ thinking_process_text += p.text or ""
116
  else:
117
+ final_text += p.text or ""
118
 
119
+ # Возвращаем готовый ответ и размышления
120
  yield final_text, thinking_process_text
121
+ return
122
  except Exception as e:
123
+ yield f"Ошибка при запросе к API: {e}", ""
124
+ return
125
 
126
 
127
  ###############################################################################
128
+ # 4. Основная асинхронная функция Gradio, обрабатывающая новый пользовательский ввод
129
  ###############################################################################
130
 
131
+ async def user_send_message(
132
+ user_message: str,
133
+ history: list[tuple[str, str]],
134
+ model_name: str,
135
+ thinking_text: str
136
+ ):
137
  """
 
138
  Параметры:
139
+ user_message: вход от пользователя (текущая реплика)
140
+ history: история [(user, assistant), ...]
141
  model_name: выбранная модель
142
+ thinking_text: текущее содержимое поля «Размышления»
143
 
144
+ Выход (через yield):
145
+ (обновлённая история, обновлённое thinking_text)
146
  """
147
+ # Если пользователь ничего не ввёл, просто возвращаем текущее состояние
148
  if not user_message.strip():
149
+ yield history, thinking_text
150
+ return
151
 
152
+ # Добавляем новую запись в историю: ассистент пока None
153
+ history.append((user_message, None))
154
 
155
+ # Проверяем, thinking-модель ли
156
  if "thinking" in model_name.lower():
157
+ # Обрабатываем через _respond_thinking
158
+ async for (assistant_text, thought_text) in _respond_thinking(model_name, user_message, history):
159
+ # Обновляем последнюю пару в истории
160
  history[-1] = (user_message, assistant_text)
161
+ # Обновляем thinking_text
162
+ yield history, thought_text
163
+ return
164
  else:
165
+ # Обычная модель (stream)
166
  partial_answer = ""
167
+ async for chunk in _respond_stream(model_name, user_message, history):
168
+ partial_answer = chunk
169
  history[-1] = (user_message, partial_answer)
170
+ # В обычном режиме thinking_text = ""
171
  yield history, ""
172
+ return
 
 
173
 
174
 
175
  ###############################################################################
176
+ # 5. Колбэки для очистки
177
  ###############################################################################
178
 
179
  def clear_all():
180
  """
181
+ Полная очистка чата и размышлений.
 
182
  """
183
+ return [], "" # (history, thinking_store)
184
 
185
 
186
  ###############################################################################
187
+ # 6. Определяем интерфейс Gradio
188
  ###############################################################################
189
 
190
  with gr.Blocks() as demo:
191
+ gr.Markdown("## Gemini Chatbot (с сохранением истории и thinking-режимом)")
 
 
 
 
 
 
192
 
193
  with gr.Row():
194
  model_dropdown = gr.Dropdown(
195
  choices=AVAILABLE_MODELS,
196
  value="gemini-1.5-flash",
197
+ label="Выберите модель",
198
  )
199
  clear_button = gr.Button("Очистить чат")
200
 
201
+ # Храним историю чата в gr.State
202
+ history_state = gr.State([]) # список кортежей (user, assistant)
203
+ # Храним «размышления» отдельно
204
+ thinking_store = gr.State("")
205
+
206
  chatbot = gr.Chatbot(label="Диалог с Gemini")
207
 
 
208
  user_input = gr.Textbox(
209
+ label="Ваш вопрос",
210
+ placeholder="Введите вопрос и нажмите Enter...",
211
  )
212
 
213
+ # Отдельный блок для показа «размышлений»
214
  thinking_output = gr.Textbox(
215
+ label="Размышления (только для gemini-2.0-flash-thinking-exp)",
216
  interactive=False
217
  )
218
 
 
219
  send_btn = gr.Button("Отправить")
220
 
221
+ # Связка: кнопка «Отправить» => user_send_message => обновление истории и размышлений
 
 
 
 
 
222
  send_btn.click(
223
  fn=user_send_message,
224
  inputs=[user_input, history_state, model_dropdown, thinking_store],
225
  outputs=[history_state, thinking_store],
226
+ queue=True
227
  ).then(
228
+ # После того как получили новую историю, обновляем чат
229
  fn=lambda h: h,
230
  inputs=[history_state],
231
  outputs=[chatbot],
232
+ queue=True
 
233
  ).then(
234
+ # После этого выводим thinking_store
235
  fn=lambda t: t,
236
  inputs=[thinking_store],
237
  outputs=[thinking_output],
238
+ queue=True
239
  )
240
 
241
+ # Аналогичное поведение при нажатии Enter в поле ввода
242
  user_input.submit(
243
  fn=user_send_message,
244
  inputs=[user_input, history_state, model_dropdown, thinking_store],
245
  outputs=[history_state, thinking_store],
246
+ queue=True
247
  ).then(
248
  fn=lambda h: h,
249
  inputs=[history_state],
250
  outputs=[chatbot],
251
+ queue=True
 
252
  ).then(
253
  fn=lambda t: t,
254
  inputs=[thinking_store],
255
  outputs=[thinking_output],
256
+ queue=True
257
  )
258
 
259
+ # Кнопка «Очистить» сбрасывает историю и thinking_store
260
  clear_button.click(
261
  fn=clear_all,
262
  inputs=[],
263
  outputs=[history_state, thinking_store],
264
  queue=False
265
  ).then(
266
+ # Затем обновляем чат
267
  fn=lambda h: h,
268
  inputs=[history_state],
269
+ outputs=[chatbot]
270
  ).then(
271
+ # И обнуляем thinking_output
272
  fn=lambda _: "",
273
  inputs=[],
274
+ outputs=[thinking_output]
275
  )
276
 
277
+ # Запуск Gradio
 
 
278
  if __name__ == "__main__":
279
+ demo.launch()