Update app.py

app.py

@@ -9,6 +9,7 @@ import json
 from consts import AUTO_SEARCH_KEYWORD, SEARCH_TOOL_INSTRUCTION, RELATED_QUESTIONS_TEMPLATE_SEARCH, SEARCH_TOOL_INSTRUCTION, RAG_TEMPLATE, GOOGLE_SEARCH_ENDPOINT, DEFAULT_SEARCH_ENGINE_TIMEOUT, RELATED_QUESTIONS_TEMPLATE_NO_SEARCH
 import re
 import asyncio
+import random
 
 import streamlit as st
 import yaml
@@ -26,6 +27,7 @@ from visual_env_utils import are_credentials_set, env_input_fields, initialize_e
 logging.basicConfig(level=logging.INFO)
 GOOGLE_API_KEY = st.secrets["google_api_key"]
 GOOGLE_CX = st.secrets["google_cx"]
+BACKUP_KEYS = [st.secrets["backup_key_1"], st.secrets["backup_key_2"], st.secrets["backup_key_3"]]
 
 CONFIG_PATH = os.path.join(current_dir, "config.yaml")
 
@@ -63,7 +65,7 @@ def st_capture(output_func: Callable[[str], None]) -> Generator:
         stdout.write = new_write  # type: ignore
         yield
 
-async def run_samba_api_inference(query, system_prompt = None, ignore_context=False, max_tokens_to_generate=None, num_seconds_to_sleep=5):
+async def run_samba_api_inference(query, system_prompt = None, ignore_context=False, max_tokens_to_generate=None, num_seconds_to_sleep=1, over_ride_key=None):
     # First construct messages
     messages = []
     if system_prompt is not None:
@@ -85,8 +87,13 @@ async def run_samba_api_inference(query, system_prompt = None, ignore_context=Fa
     }
     if max_tokens_to_generate is not None:
         payload["max_tokens"] = max_tokens_to_generate
+
+    if over_ride_key is None:
+        api_key = st.session_state.SAMBANOVA_API_KEY
+    else:
+        api_key = over_ride_key
     headers = {
-        "Authorization": f"Basic {st.session_state.SAMBANOVA_API_KEY}",
+        "Authorization": f"Basic {api_key}",
         "Content-Type": "application/json"
     }
 
@@ -98,9 +105,8 @@ async def run_samba_api_inference(query, system_prompt = None, ignore_context=Fa
             st.info(f"Invalid Key! Please make sure you have a valid SambaCloud key from https://cloud.sambanova.ai/.")
             return "Invalid Key! Please make sure you have a valid SambaCloud key from https://cloud.sambanova.ai/."
         if post_response.status_code in {429, 504}:
-            st.info("Rate limit hit because of all the pipelined queries, wait one second...")
             await asyncio.sleep(num_seconds_to_sleep)
-            return await run_samba_api_inference(query)  # Retry the request
+            return await run_samba_api_inference(query, random.choice(BACKUP_KEYS))  # Retry the request
         else:
             print(f"Request failed with status code: {post_response.status_code}. Error: {e}")
             return "Invalid Key! Please make sure you have a valid SambaCloud key from https://cloud.sambanova.ai/."