ai: Introduce Deep Search with JARVIS.

jarvis.py

@@ -29,8 +29,12 @@ os.system("apt-get update -q -y && apt-get install -q -y tesseract-ocr tesseract
 
 JARVIS_INIT = json.loads(os.getenv("HELLO", "[]"))
 
+DEEP_SEARCH_PROVIDER_HOST = os.getenv("DEEP_SEARCH_PROVIDER_HOST")
+DEEP_SEARCH_PROVIDER_KEY = os.getenv('DEEP_SEARCH_PROVIDER_KEY')
+DEEP_SEARCH_INSTRUCTIONS = os.getenv("DEEP_SEARCH_INSTRUCTIONS")
+
 INTERNAL_AI_GET_SERVER = os.getenv("INTERNAL_AI_GET_SERVER")
-INTERNAL_TRAINING_DATA = os.getenv("INTERNAL_TRAINING_DATA")
+INTERNAL_AI_INSTRUCTIONS = os.getenv("INTERNAL_TRAINING_DATA")
 
 SYSTEM_PROMPT_MAPPING = json.loads(os.getenv("SYSTEM_PROMPT_MAPPING", "{}"))
 SYSTEM_PROMPT_DEFAULT = os.getenv("DEFAULT_SYSTEM")
@@ -43,7 +47,7 @@ LINUX_SERVER_PROVIDER_KEYS_ATTEMPTS = {}
 
 LINUX_SERVER_ERRORS = set(map(int, os.getenv("LINUX_SERVER_ERROR", "").split(",")))
 
-AI_TYPES = {f"AI_TYPE_{i}": os.getenv(f"AI_TYPE_{i}") for i in range(1, 8)}
+AI_TYPES = {f"AI_TYPE_{i}": os.getenv(f"AI_TYPE_{i}") for i in range(1, 9)}
 
 RESPONSES = {f"RESPONSE_{i}": os.getenv(f"RESPONSE_{i}") for i in range(1, 11)}
 
@@ -235,7 +239,7 @@ async def fetch_response_stream_async(host, key, model, msgs, cfg, sid, stop_eve
         marked_item(key, LINUX_SERVER_PROVIDER_KEYS_MARKED, LINUX_SERVER_PROVIDER_KEYS_ATTEMPTS)
     return
 
-async def chat_with_model_async(history, user_input, model_display, sess, custom_prompt):
+async def chat_with_model_async(history, user_input, model_display, sess, custom_prompt, deep_search):
     ensure_stop_event(sess)
     sess.stop_event.clear()
     sess.cancel_token["cancelled"] = False
@@ -246,9 +250,37 @@ async def chat_with_model_async(history, user_input, model_display, sess, custom
         sess.session_id = str(uuid.uuid4())
     model_key = get_model_key(model_display)
     cfg = MODEL_CONFIG.get(model_key, DEFAULT_CONFIG)
-    msgs = [{"role": "user", "content": u} for u, _ in history] + [{"role": "assistant", "content": a} for _, a in history if a]
-    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))
-    msgs.insert(0, {"role": "system", "content": prompt})
+    msgs = []
+    if deep_search and model_display == MODEL_CHOICES[0]:
+        msgs.append({"role": "system", "content": DEEP_SEARCH_INSTRUCTIONS})
+        try:
+            async with httpx.AsyncClient() as client:
+                payload = {
+                    "query": user_input,
+                    "topic": "general",
+                    "search_depth": "basic",
+                    "chunks_per_source": 3,
+                    "max_results": 1,
+                    "time_range": None,
+                    "days": 7,
+                    "include_answer": True,
+                    "include_raw_content": False,
+                    "include_images": False,
+                    "include_image_descriptions": False,
+                    "include_domains": [],
+                    "exclude_domains": []
+                }
+                r = await client.post(DEEP_SEARCH_PROVIDER_HOST, headers={"Authorization": f"Bearer {DEEP_SEARCH_PROVIDER_KEY}"}, json=payload)
+                sr_json = r.json()
+                msgs.append({"role": "system", "content": json.dumps(sr_json)})
+        except:
+            pass
+        msgs.append({"role": "system", "content": INTERNAL_AI_INSTRUCTIONS})
+    elif model_display == MODEL_CHOICES[0]:
+        msgs.append({"role": "system", "content": INTERNAL_AI_INSTRUCTIONS})
+    else:
+        msgs.append({"role": "system", "content": custom_prompt or SYSTEM_PROMPT_MAPPING.get(model_key, SYSTEM_PROMPT_DEFAULT)})
+    msgs.extend([{"role": "user", "content": u} for u, _ in history] + [{"role": "assistant", "content": a} for _, a in history if a])
     msgs.append({"role": "user", "content": user_input})
     candidates = [(h, k) for h in LINUX_SERVER_HOSTS for k in LINUX_SERVER_PROVIDER_KEYS]
     random.shuffle(candidates)
@@ -264,7 +296,7 @@ async def chat_with_model_async(history, user_input, model_display, sess, custom
             return
     yield ("content", RESPONSES["RESPONSE_2"])
 
-async def respond_async(multi, history, model_display, sess, custom_prompt):
+async def respond_async(multi, history, model_display, sess, custom_prompt, deep_search):
     ensure_stop_event(sess)
     sess.stop_event.clear()
     sess.cancel_token["cancelled"] = False
@@ -286,7 +318,7 @@ async def respond_async(multi, history, model_display, sess, custom_prompt):
         responses = ""
         content_started = False
         ignore_reasoning = False
-        async for typ, chunk in chat_with_model_async(history, inp, model_display, sess, custom_prompt):
+        async for typ, chunk in chat_with_model_async(history, inp, model_display, sess, custom_prompt, deep_search):
             if sess.stop_event.is_set() or sess.cancel_token["cancelled"]:
                 break
             if typ == "reasoning":
@@ -332,9 +364,9 @@ async def respond_async(multi, history, model_display, sess, custom_prompt):
     yield history, gr.update(value="", interactive=True, submit_btn=True, stop_btn=False), sess
 
 def change_model(new):
-    visible = new != MODEL_CHOICES[0]
+    visible = new == MODEL_CHOICES[0]
     default = SYSTEM_PROMPT_MAPPING.get(get_model_key(new), SYSTEM_PROMPT_DEFAULT)
-    return [], create_session(), new, default
+    return [], create_session(), new, default, False, gr.update(visible=visible)
 
 def stop_response(history, sess):
     ensure_stop_event(sess)
@@ -350,16 +382,18 @@ with gr.Blocks(fill_height=True, fill_width=True, title=AI_TYPES["AI_TYPE_4"], h
     selected_model = gr.State(MODEL_CHOICES[0] if MODEL_CHOICES else "")
     J_A_R_V_I_S = gr.State("")
     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)
+    deep_search = gr.Checkbox(label=AI_TYPES["AI_TYPE_8"], value=False, visible=True)
     msg = gr.MultimodalTextbox(show_label=False, placeholder=RESPONSES["RESPONSE_5"], interactive=True, file_count="single", file_types=ALLOWED_EXTENSIONS)
     with gr.Sidebar(open=False):
         model_radio = gr.Radio(show_label=False, choices=MODEL_CHOICES, value=MODEL_CHOICES[0])
-    model_radio.change(fn=change_model, inputs=[model_radio], outputs=[user_history, user_session, selected_model, J_A_R_V_I_S])
+    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])
     def on_example_select(evt: gr.SelectData):
         return evt.value
-    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])
-    def clear_chat():
-        return [], create_session(), "", "", []
-    chatbot.clear(fn=clear_chat, inputs=[], outputs=[chatbot, user_session, J_A_R_V_I_S, msg, user_history])
-    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])
+    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])
+    def clear_chat(history, sess, prompt, model):
+        return [], create_session(), prompt, model, []
+    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])
+    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])
+    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)
     msg.stop(fn=stop_response, inputs=[user_history, user_session], outputs=[chatbot, msg, user_session])
 jarvis.queue(default_concurrency_limit=2).launch(max_file_size="1mb")