ALMOST SO ABCK


app.py

@@ -3,6 +3,7 @@ import requests
 import json
 import os
 import datetime
+import time
 from requests.exceptions import RequestException
 
 api_keys_env = os.environ.get('API_KEYS')
@@ -10,6 +11,7 @@ if api_keys_env:
     API_KEYS = [key.strip() for key in api_keys_env.strip().split('\n') if key.strip()]
 else:
     raise ValueError("all keez ded go kys")
+
 API_URL = os.environ.get('API_URL')
 
 DEFAULT_PARAMS = {
@@ -79,60 +81,118 @@ def predict(message, history, system_prompt, temperature, top_p, top_k, min_p, t
 
     api_key_index = 0
     retries = 0
-    max_retries = len(API_KEYS)
+    max_retries = len(API_KEYS) * 5
+    partial_message = ""
+    processing_count = 0
+    max_processing_attempts = 10
 
     while retries < max_retries:
-        current_api_key = API_KEYS[api_key_index]
+        if stop_flag[0]:
+            return
 
+        current_api_key = API_KEYS[api_key_index]
         headers = {
             "Authorization": f"Bearer {current_api_key}",
             "Content-Type": "application/json"
         }
 
         try:
-            with requests.post(API_URL, headers=headers, data=json.dumps(data), stream=True) as response:
-                if response.status_code == 200:
-                    partial_message = ""
-                    for line in response.iter_lines():
-                        if stop_flag[0]:
-                            response.close()
+            response = requests.post(API_URL, headers=headers, data=json.dumps(data), stream=True)
+            
+            if response.status_code == 200:
+                processing_count = 0
+                for line in response.iter_lines():
+                    if stop_flag[0]:
+                        response.close()
+                        return
+
+                    if not line:
+                        continue
+
+                    line = line.decode('utf-8')
+                    if RESPONSE_LOGGING_ENABLED:
+                        print(f"API Response: {line}")
+
+                    if ": OPENROUTER PROCESSING" in line:
+                        processing_count += 1
+                        if processing_count >= max_processing_attempts:
+                            print("Too many processing attempts, cycling to next key...")
                             break
-                        if line:
-                            line = line.decode('utf-8')
-                            if RESPONSE_LOGGING_ENABLED:
-                                print(f"API Response: {line}")
-                            if line.startswith("data: "):
-                                if line.strip() == "data: [DONE]":
-                                    break
+                        continue
+
+                    if not line.startswith("data: "):
+                        continue
+
+                    if line.strip() == "data: [DONE]":
+                        response.close()
+                        if partial_message:
+                            yield partial_message
+                        return
+
+                    try:
+                        json_data = json.loads(line[6:])
+                        
+                        # Check for rate limit error
+                        if 'error' in json_data:
+                            error_msg = json_data.get('error', {}).get('message', '')
+                            if isinstance(error_msg, str):
                                 try:
-                                    json_data = json.loads(line[6:])
-                                    if 'choices' in json_data and json_data['choices']:
-                                        delta = json_data['choices'][0]['delta']
-                                        content = delta.get('content', '')
-                                        if content:
-                                            partial_message += content
-                                            yield partial_message
+                                    error_obj = json.loads(error_msg)
+                                    if error_obj.get('error', {}).get('type') == 'rate_limit_exceeded':
+                                        print("Rate limit hit in streaming response, cycling keys...")
+                                        break
                                 except json.JSONDecodeError:
-                                    continue  
-                    if partial_message:
-                        yield partial_message
-                    # successful response, break out of retry loop
-                    break
-                elif response.status_code == 429:
-                    print("rate limit hit, cycling keys...")
-                    retries += 1
-                    api_key_index = (api_key_index + 1) % len(API_KEYS)
-                    continue
+                                    pass
+                            continue
+
+                        if 'choices' in json_data and json_data['choices']:
+                            delta = json_data['choices'][0]['delta']
+                            content = delta.get('content', '')
+                            if content:
+                                partial_message += content
+                                yield partial_message
+
+                    except json.JSONDecodeError as e:
+                        print(f"JSON decode error: {e}")
+                        continue
+
+                response.close()
+                if partial_message:
+                    return
+                
+            elif response.status_code == 429:
+                print("Rate limit hit from status code, cycling keys...")
+                time.sleep(1)
+            elif response.status_code == 401:
+                print(f"Invalid API key {api_key_index}, cycling to next...")
+                api_key_index = (api_key_index + 1) % len(API_KEYS)
+                retries += 1
+                continue
+            else:
+                error_message = f"Error: Received status code {response.status_code} - {response.text}"
+                print(error_message)
+                if partial_message:
+                    yield partial_message + f"\n[Error: {error_message}]"
                 else:
-                    # theothershits
-                    error_message = f"Error: Received status code {response.status_code} - {response.text}"
-                    print(error_message)
                     yield f"An error occurred: {error_message}"
-                    break
+                return
+
         except RequestException as e:
             print(f"Request error: {e}")
-            yield f"An error occurred: {str(e)}"
-            break
+            if partial_message:
+                yield partial_message + f"\n[Error: {str(e)}]"
+            else:
+                yield f"An error occurred: {str(e)}"
+            return
+
+        retries += 1
+        api_key_index = (api_key_index + 1) % len(API_KEYS)
+        time.sleep(1)
+
+    if partial_message:
+        yield partial_message + "\n[Error: Maximum retries reached]"
+    else:
+        yield "Error: Maximum retries reached. Please try again later."
 
 def import_chat(custom_format_string):
     try: