ai: Implementing server-side streaming responses.

* Say Hi! to fast responses.
* Say No! to slow responses.

jarvis.py

@@ -43,11 +43,13 @@ 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)}
-RESPONSES = {f"RESPONSE_{i}": os.getenv(f"RESPONSE_{i}") for i in range(1, 10)}
+
+RESPONSES = {f"RESPONSE_{i}": os.getenv(f"RESPONSE_{i}") for i in range(1, 11)}
 
 MODEL_MAPPING = json.loads(os.getenv("MODEL_MAPPING", "{}"))
 MODEL_CONFIG = json.loads(os.getenv("MODEL_CONFIG", "{}"))
 MODEL_CHOICES = list(MODEL_MAPPING.values())
+
 DEFAULT_CONFIG = json.loads(os.getenv("DEFAULT_CONFIG", "{}"))
 DEFAULT_MODEL_KEY = list(MODEL_MAPPING.keys())[0] if MODEL_MAPPING else None
 
@@ -199,30 +201,36 @@ def extract_file_content(fp):
         except Exception as e:
             return f"{fp}: {e}"
 
-async def fetch_response_async(host, key, model, msgs, cfg, sid):
+async def fetch_response_stream_async(host, key, model, msgs, cfg, sid):
     for t in [1, 2]:
         try:
             async with httpx.AsyncClient(timeout=t) as client:
-                r = await client.post(host, json={"model": model, "messages": msgs, **cfg, "session_id": sid}, headers={"Authorization": f"Bearer {key}"})
-                if r.status_code in LINUX_SERVER_ERRORS:
-                    marked_item(key, LINUX_SERVER_PROVIDER_KEYS_MARKED, LINUX_SERVER_PROVIDER_KEYS_ATTEMPTS)
-                    return None
-                r.raise_for_status()
-                j = r.json()
-                if isinstance(j, dict) and j.get("choices"):
-                    ch = j["choices"][0]
-                    if ch.get("message") and isinstance(ch["message"].get("content"), str):
-                        return ch["message"]["content"]
-                return None
+                async with client.stream("POST", host, json={**{"model": model, "messages": msgs, "session_id": sid, "stream": True}, **cfg}, headers={"Authorization": f"Bearer {key}"}) as response:
+                    async for line in response.aiter_lines():
+                        if not line:
+                            continue
+                        if line.startswith("data: "):
+                            data = line[6:]
+                            if data.strip() == RESPONSES["RESPONSE_10"]:
+                                return
+                            try:
+                                j = json.loads(data)
+                                if isinstance(j, dict) and j.get("choices"):
+                                    ch = j["choices"][0]
+                                    if ch.get("delta") and isinstance(ch["delta"].get("content"), str):
+                                        yield ch["delta"]["content"]
+                            except:
+                                continue
         except:
             continue
-    marked_item(key, LINUX_SERVER_PROVIDER_KEYS_MARKED, LINUX_SERVER_PROVIDER_KEYS_ATTEMPTS)
-    return None
+        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):
     ensure_stop_event(sess)
     if not get_available_items(LINUX_SERVER_PROVIDER_KEYS, LINUX_SERVER_PROVIDER_KEYS_MARKED) or not get_available_items(LINUX_SERVER_HOSTS, LINUX_SERVER_HOSTS_ATTEMPTS):
-        return RESPONSES["RESPONSE_3"]
+        yield RESPONSES["RESPONSE_3"]
+        return
     if not hasattr(sess, "session_id") or not sess.session_id:
         sess.session_id = str(uuid.uuid4())
         sess.stop_event = asyncio.Event()
@@ -235,25 +243,27 @@ async def chat_with_model_async(history, user_input, model_display, sess, custom
     msgs.insert(0, {"role": "system", "content": prompt})
     msgs.append({"role": "user", "content": user_input})
     if sess.active_candidate:
-        res = await fetch_response_async(sess.active_candidate[0], sess.active_candidate[1], model_key, msgs, cfg, sess.session_id)
-        if res:
-            return res
-        sess.active_candidate = None
+        async for chunk in fetch_response_stream_async(sess.active_candidate[0], sess.active_candidate[1], model_key, msgs, cfg, sess.session_id):
+            yield chunk
+        return
     keys = get_available_items(LINUX_SERVER_PROVIDER_KEYS, LINUX_SERVER_PROVIDER_KEYS_MARKED)
     hosts = get_available_items(LINUX_SERVER_HOSTS, LINUX_SERVER_HOSTS_ATTEMPTS)
     random.shuffle(keys)
     random.shuffle(hosts)
     for k in keys:
         for h in hosts:
-            task = asyncio.create_task(fetch_response_async(h, k, model_key, msgs, cfg, sess.session_id))
-            done, _ = await asyncio.wait({task}, return_when=asyncio.FIRST_COMPLETED)
-            if task in done:
-                result = task.result()
-                if result:
+            stream_gen = fetch_response_stream_async(h, k, model_key, msgs, cfg, sess.session_id)
+            full_text = ""
+            got_any = False
+            async for chunk in stream_gen:
+                if not got_any:
+                    got_any = True
                     sess.active_candidate = (h, k)
-                    return result
-            task.cancel()
-    return RESPONSES["RESPONSE_2"]
+                full_text += chunk
+                yield chunk
+            if got_any and full_text:
+                return
+    yield RESPONSES["RESPONSE_2"]
 
 async def respond_async(multi, history, model_display, sess, custom_prompt):
     ensure_stop_event(sess)
@@ -270,37 +280,42 @@ async def respond_async(multi, history, model_display, sess, custom_prompt):
         inp += msg_input["text"]
     history.append([inp, RESPONSES["RESPONSE_8"]])
     yield history, gr.update(interactive=False, submit_btn=False, stop_btn=True), sess
-    task = asyncio.create_task(chat_with_model_async(history, inp, model_display, sess, custom_prompt))
+    queue = asyncio.Queue()
+    async def background():
+        full = ""
+        async for chunk in chat_with_model_async(history, inp, model_display, sess, custom_prompt):
+            full += chunk
+            await queue.put(chunk)
+        await queue.put(None)
+        return full
+    bg_task = asyncio.create_task(background())
     stop_task = asyncio.create_task(sess.stop_event.wait())
-    done, pending = await asyncio.wait({task, stop_task}, return_when=asyncio.FIRST_COMPLETED)
-    if stop_task in done:
-        task.cancel()
-        history[-1][1] = RESPONSES["RESPONSE_1"]
-        yield history, gr.update(value="", interactive=True, submit_btn=True, stop_btn=False), sess
-        sess.stop_event.clear()
-        return
-    stop_task.cancel()
-    ai = task.result()
-    history[-1][1] = ""
-    buffer = []
-    last_update = asyncio.get_event_loop().time()
-    for char in ai:
-        if sess.stop_event.is_set():
-            history[-1][1] = RESPONSES["RESPONSE_1"]
-            yield history, gr.update(value="", interactive=True, submit_btn=True, stop_btn=False), sess
-            sess.stop_event.clear()
-            return
-        buffer.append(char)
-        current_time = asyncio.get_event_loop().time()
-        if len(buffer) >= 4 or (current_time - last_update) > 0.001:
-            history[-1][1] += "".join(buffer)
-            buffer.clear()
-            last_update = current_time
-            yield history, gr.update(interactive=False, submit_btn=False, stop_btn=True), sess
-            await asyncio.sleep(0.003)
-    if buffer:
-        history[-1][1] += "".join(buffer)
-        yield history, gr.update(interactive=False, submit_btn=False, stop_btn=True), sess
+    first_meaningful_chunk_found = False
+    try:
+        while True:
+            done, _ = await asyncio.wait({stop_task, asyncio.create_task(queue.get())}, return_when=asyncio.FIRST_COMPLETED)
+            if stop_task in done:
+                bg_task.cancel()
+                history[-1][1] = RESPONSES["RESPONSE_1"]
+                yield history, gr.update(value="", interactive=True, submit_btn=True, stop_btn=False), sess
+                sess.stop_event.clear()
+                return
+            for d in done:
+                chunk = d.result()
+                if chunk is None:
+                    raise StopAsyncIteration
+                if not first_meaningful_chunk_found:
+                    if chunk.strip():
+                        history[-1][1] = chunk
+                        first_meaningful_chunk_found = True
+                else:
+                    history[-1][1] += chunk
+                yield history, gr.update(interactive=False, submit_btn=False, stop_btn=True), sess
+    except StopAsyncIteration:
+        pass
+    finally:
+        stop_task.cancel()
+    full_response = await bg_task
     yield history, gr.update(value="", interactive=True, submit_btn=True, stop_btn=False), sess
 
 def change_model(new):