hadadrjt commited on
Commit
5e49ce6
·
1 Parent(s): cf21eef

ai: Introduce Deep Search with JARVIS.

Browse files
Files changed (1) hide show
  1. jarvis.py +50 -16
jarvis.py CHANGED
@@ -29,8 +29,12 @@ os.system("apt-get update -q -y && apt-get install -q -y tesseract-ocr tesseract
29
 
30
  JARVIS_INIT = json.loads(os.getenv("HELLO", "[]"))
31
 
 
 
 
 
32
  INTERNAL_AI_GET_SERVER = os.getenv("INTERNAL_AI_GET_SERVER")
33
- INTERNAL_TRAINING_DATA = os.getenv("INTERNAL_TRAINING_DATA")
34
 
35
  SYSTEM_PROMPT_MAPPING = json.loads(os.getenv("SYSTEM_PROMPT_MAPPING", "{}"))
36
  SYSTEM_PROMPT_DEFAULT = os.getenv("DEFAULT_SYSTEM")
@@ -43,7 +47,7 @@ LINUX_SERVER_PROVIDER_KEYS_ATTEMPTS = {}
43
 
44
  LINUX_SERVER_ERRORS = set(map(int, os.getenv("LINUX_SERVER_ERROR", "").split(",")))
45
 
46
- AI_TYPES = {f"AI_TYPE_{i}": os.getenv(f"AI_TYPE_{i}") for i in range(1, 8)}
47
 
48
  RESPONSES = {f"RESPONSE_{i}": os.getenv(f"RESPONSE_{i}") for i in range(1, 11)}
49
 
@@ -235,7 +239,7 @@ async def fetch_response_stream_async(host, key, model, msgs, cfg, sid, stop_eve
235
  marked_item(key, LINUX_SERVER_PROVIDER_KEYS_MARKED, LINUX_SERVER_PROVIDER_KEYS_ATTEMPTS)
236
  return
237
 
238
- async def chat_with_model_async(history, user_input, model_display, sess, custom_prompt):
239
  ensure_stop_event(sess)
240
  sess.stop_event.clear()
241
  sess.cancel_token["cancelled"] = False
@@ -246,9 +250,37 @@ async def chat_with_model_async(history, user_input, model_display, sess, custom
246
  sess.session_id = str(uuid.uuid4())
247
  model_key = get_model_key(model_display)
248
  cfg = MODEL_CONFIG.get(model_key, DEFAULT_CONFIG)
249
- msgs = [{"role": "user", "content": u} for u, _ in history] + [{"role": "assistant", "content": a} for _, a in history if a]
250
- prompt = INTERNAL_TRAINING_DATA if model_key == DEFAULT_MODEL_KEY and INTERNAL_TRAINING_DATA else (custom_prompt or SYSTEM_PROMPT_MAPPING.get(model_key, SYSTEM_PROMPT_DEFAULT))
251
- msgs.insert(0, {"role": "system", "content": prompt})
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
252
  msgs.append({"role": "user", "content": user_input})
253
  candidates = [(h, k) for h in LINUX_SERVER_HOSTS for k in LINUX_SERVER_PROVIDER_KEYS]
254
  random.shuffle(candidates)
@@ -264,7 +296,7 @@ async def chat_with_model_async(history, user_input, model_display, sess, custom
264
  return
265
  yield ("content", RESPONSES["RESPONSE_2"])
266
 
267
- async def respond_async(multi, history, model_display, sess, custom_prompt):
268
  ensure_stop_event(sess)
269
  sess.stop_event.clear()
270
  sess.cancel_token["cancelled"] = False
@@ -286,7 +318,7 @@ async def respond_async(multi, history, model_display, sess, custom_prompt):
286
  responses = ""
287
  content_started = False
288
  ignore_reasoning = False
289
- async for typ, chunk in chat_with_model_async(history, inp, model_display, sess, custom_prompt):
290
  if sess.stop_event.is_set() or sess.cancel_token["cancelled"]:
291
  break
292
  if typ == "reasoning":
@@ -332,9 +364,9 @@ async def respond_async(multi, history, model_display, sess, custom_prompt):
332
  yield history, gr.update(value="", interactive=True, submit_btn=True, stop_btn=False), sess
333
 
334
  def change_model(new):
335
- visible = new != MODEL_CHOICES[0]
336
  default = SYSTEM_PROMPT_MAPPING.get(get_model_key(new), SYSTEM_PROMPT_DEFAULT)
337
- return [], create_session(), new, default
338
 
339
  def stop_response(history, sess):
340
  ensure_stop_event(sess)
@@ -350,16 +382,18 @@ with gr.Blocks(fill_height=True, fill_width=True, title=AI_TYPES["AI_TYPE_4"], h
350
  selected_model = gr.State(MODEL_CHOICES[0] if MODEL_CHOICES else "")
351
  J_A_R_V_I_S = gr.State("")
352
  chatbot = gr.Chatbot(label=AI_TYPES["AI_TYPE_1"], show_copy_button=True, scale=1, elem_id=AI_TYPES["AI_TYPE_2"], examples=JARVIS_INIT)
 
353
  msg = gr.MultimodalTextbox(show_label=False, placeholder=RESPONSES["RESPONSE_5"], interactive=True, file_count="single", file_types=ALLOWED_EXTENSIONS)
354
  with gr.Sidebar(open=False):
355
  model_radio = gr.Radio(show_label=False, choices=MODEL_CHOICES, value=MODEL_CHOICES[0])
356
- model_radio.change(fn=change_model, inputs=[model_radio], outputs=[user_history, user_session, selected_model, J_A_R_V_I_S])
357
  def on_example_select(evt: gr.SelectData):
358
  return evt.value
359
- chatbot.example_select(fn=on_example_select, inputs=[], outputs=[msg]).then(fn=respond_async, inputs=[msg, user_history, selected_model, user_session, J_A_R_V_I_S], outputs=[chatbot, msg, user_session])
360
- def clear_chat():
361
- return [], create_session(), "", "", []
362
- chatbot.clear(fn=clear_chat, inputs=[], outputs=[chatbot, user_session, J_A_R_V_I_S, msg, user_history])
363
- msg.submit(fn=respond_async, inputs=[msg, user_history, selected_model, user_session, J_A_R_V_I_S], outputs=[chatbot, msg, user_session], api_name=INTERNAL_AI_GET_SERVER, show_progress="full", show_progress_on=[chatbot, msg])
 
364
  msg.stop(fn=stop_response, inputs=[user_history, user_session], outputs=[chatbot, msg, user_session])
365
  jarvis.queue(default_concurrency_limit=2).launch(max_file_size="1mb")
 
29
 
30
  JARVIS_INIT = json.loads(os.getenv("HELLO", "[]"))
31
 
32
+ DEEP_SEARCH_PROVIDER_HOST = os.getenv("DEEP_SEARCH_PROVIDER_HOST")
33
+ DEEP_SEARCH_PROVIDER_KEY = os.getenv('DEEP_SEARCH_PROVIDER_KEY')
34
+ DEEP_SEARCH_INSTRUCTIONS = os.getenv("DEEP_SEARCH_INSTRUCTIONS")
35
+
36
  INTERNAL_AI_GET_SERVER = os.getenv("INTERNAL_AI_GET_SERVER")
37
+ INTERNAL_AI_INSTRUCTIONS = os.getenv("INTERNAL_TRAINING_DATA")
38
 
39
  SYSTEM_PROMPT_MAPPING = json.loads(os.getenv("SYSTEM_PROMPT_MAPPING", "{}"))
40
  SYSTEM_PROMPT_DEFAULT = os.getenv("DEFAULT_SYSTEM")
 
47
 
48
  LINUX_SERVER_ERRORS = set(map(int, os.getenv("LINUX_SERVER_ERROR", "").split(",")))
49
 
50
+ AI_TYPES = {f"AI_TYPE_{i}": os.getenv(f"AI_TYPE_{i}") for i in range(1, 9)}
51
 
52
  RESPONSES = {f"RESPONSE_{i}": os.getenv(f"RESPONSE_{i}") for i in range(1, 11)}
53
 
 
239
  marked_item(key, LINUX_SERVER_PROVIDER_KEYS_MARKED, LINUX_SERVER_PROVIDER_KEYS_ATTEMPTS)
240
  return
241
 
242
+ async def chat_with_model_async(history, user_input, model_display, sess, custom_prompt, deep_search):
243
  ensure_stop_event(sess)
244
  sess.stop_event.clear()
245
  sess.cancel_token["cancelled"] = False
 
250
  sess.session_id = str(uuid.uuid4())
251
  model_key = get_model_key(model_display)
252
  cfg = MODEL_CONFIG.get(model_key, DEFAULT_CONFIG)
253
+ msgs = []
254
+ if deep_search and model_display == MODEL_CHOICES[0]:
255
+ msgs.append({"role": "system", "content": DEEP_SEARCH_INSTRUCTIONS})
256
+ try:
257
+ async with httpx.AsyncClient() as client:
258
+ payload = {
259
+ "query": user_input,
260
+ "topic": "general",
261
+ "search_depth": "basic",
262
+ "chunks_per_source": 3,
263
+ "max_results": 1,
264
+ "time_range": None,
265
+ "days": 7,
266
+ "include_answer": True,
267
+ "include_raw_content": False,
268
+ "include_images": False,
269
+ "include_image_descriptions": False,
270
+ "include_domains": [],
271
+ "exclude_domains": []
272
+ }
273
+ r = await client.post(DEEP_SEARCH_PROVIDER_HOST, headers={"Authorization": f"Bearer {DEEP_SEARCH_PROVIDER_KEY}"}, json=payload)
274
+ sr_json = r.json()
275
+ msgs.append({"role": "system", "content": json.dumps(sr_json)})
276
+ except:
277
+ pass
278
+ msgs.append({"role": "system", "content": INTERNAL_AI_INSTRUCTIONS})
279
+ elif model_display == MODEL_CHOICES[0]:
280
+ msgs.append({"role": "system", "content": INTERNAL_AI_INSTRUCTIONS})
281
+ else:
282
+ msgs.append({"role": "system", "content": custom_prompt or SYSTEM_PROMPT_MAPPING.get(model_key, SYSTEM_PROMPT_DEFAULT)})
283
+ msgs.extend([{"role": "user", "content": u} for u, _ in history] + [{"role": "assistant", "content": a} for _, a in history if a])
284
  msgs.append({"role": "user", "content": user_input})
285
  candidates = [(h, k) for h in LINUX_SERVER_HOSTS for k in LINUX_SERVER_PROVIDER_KEYS]
286
  random.shuffle(candidates)
 
296
  return
297
  yield ("content", RESPONSES["RESPONSE_2"])
298
 
299
+ async def respond_async(multi, history, model_display, sess, custom_prompt, deep_search):
300
  ensure_stop_event(sess)
301
  sess.stop_event.clear()
302
  sess.cancel_token["cancelled"] = False
 
318
  responses = ""
319
  content_started = False
320
  ignore_reasoning = False
321
+ async for typ, chunk in chat_with_model_async(history, inp, model_display, sess, custom_prompt, deep_search):
322
  if sess.stop_event.is_set() or sess.cancel_token["cancelled"]:
323
  break
324
  if typ == "reasoning":
 
364
  yield history, gr.update(value="", interactive=True, submit_btn=True, stop_btn=False), sess
365
 
366
  def change_model(new):
367
+ visible = new == MODEL_CHOICES[0]
368
  default = SYSTEM_PROMPT_MAPPING.get(get_model_key(new), SYSTEM_PROMPT_DEFAULT)
369
+ return [], create_session(), new, default, False, gr.update(visible=visible)
370
 
371
  def stop_response(history, sess):
372
  ensure_stop_event(sess)
 
382
  selected_model = gr.State(MODEL_CHOICES[0] if MODEL_CHOICES else "")
383
  J_A_R_V_I_S = gr.State("")
384
  chatbot = gr.Chatbot(label=AI_TYPES["AI_TYPE_1"], show_copy_button=True, scale=1, elem_id=AI_TYPES["AI_TYPE_2"], examples=JARVIS_INIT)
385
+ deep_search = gr.Checkbox(label=AI_TYPES["AI_TYPE_8"], value=False, visible=True)
386
  msg = gr.MultimodalTextbox(show_label=False, placeholder=RESPONSES["RESPONSE_5"], interactive=True, file_count="single", file_types=ALLOWED_EXTENSIONS)
387
  with gr.Sidebar(open=False):
388
  model_radio = gr.Radio(show_label=False, choices=MODEL_CHOICES, value=MODEL_CHOICES[0])
389
+ model_radio.change(fn=change_model, inputs=[model_radio], outputs=[user_history, user_session, selected_model, J_A_R_V_I_S, deep_search, deep_search])
390
  def on_example_select(evt: gr.SelectData):
391
  return evt.value
392
+ chatbot.example_select(fn=on_example_select, inputs=[], outputs=[msg]).then(fn=respond_async, inputs=[msg, user_history, selected_model, user_session, J_A_R_V_I_S, deep_search], outputs=[chatbot, msg, user_session])
393
+ def clear_chat(history, sess, prompt, model):
394
+ return [], create_session(), prompt, model, []
395
+ deep_search.change(fn=clear_chat, inputs=[user_history, user_session, J_A_R_V_I_S, selected_model], outputs=[chatbot, user_session, J_A_R_V_I_S, selected_model, user_history])
396
+ chatbot.clear(fn=clear_chat, inputs=[user_history, user_session, J_A_R_V_I_S, selected_model], outputs=[chatbot, user_session, J_A_R_V_I_S, selected_model, user_history])
397
+ msg.submit(fn=respond_async, inputs=[msg, user_history, selected_model, user_session, J_A_R_V_I_S, deep_search], outputs=[chatbot, msg, user_session], api_name=INTERNAL_AI_GET_SERVER)
398
  msg.stop(fn=stop_response, inputs=[user_history, user_session], outputs=[chatbot, msg, user_session])
399
  jarvis.queue(default_concurrency_limit=2).launch(max_file_size="1mb")