fix api key

app.py

@@ -5,22 +5,17 @@ import os
 import re
 import string
 import time
+from typing import Optional, Any
 
 import gradio as gr
 import openai
 import google.generativeai as genai
 
 
-openai_key = os.environ.get('OPEN_AI_KEY')
-gpt_client = openai.OpenAI(api_key=openai_key)
-
-gemini_key = os.environ.get('GEMINI_API_KEY')
-genai.configure(api_key=gemini_key)
-
-
 def query_gpt_model(
     prompt: str,
     llm: str = 'gpt-3.5-turbo-1106',
+    client: Optional[Any] = None,
     temperature: float = 0.0,
     max_decode_steps: int = 512,
     seconds_to_reset_tokens: float = 30.0,
@@ -28,7 +23,7 @@ def query_gpt_model(
 
   while True:
     try:
-      raw_response = gpt_client.chat.completions.with_raw_response.create(
+      raw_response = client.chat.completions.with_raw_response.create(
         model=llm,
         max_tokens=max_decode_steps,
         temperature=temperature,
@@ -56,8 +51,10 @@ safety_settings=[
 def query_gemini_model(
     prompt: str,
     llm: str = 'gemini-pro',
+    client: Optional[Any] = None,
     retries: int = 10,
 ) -> str:
+  del client
   model = genai.GenerativeModel(llm)
   while True and retries > 0:
     try:
@@ -74,12 +71,13 @@ def query_gemini_model(
 def query_model(
     prompt: str,
     model_name: str = 'gemini-pro',
+    client: Optional[Any] = None,
 ) -> str:
   model_type = model_name.split('-')[0]
   if model_type == "gpt":
-    return query_gpt_model(prompt, llm=model_name)
+    return query_gpt_model(prompt, llm=model_name, client=client)
   elif model_type == "gemini":
-    return query_gemini_model(prompt, llm=model_name)
+    return query_gemini_model(prompt, llm=model_name, client=client)
   else:
     raise ValueError('Unexpected model_name: ', model_name)
 
@@ -201,6 +199,7 @@ def parse_pause_point(text):
 
 def quality_pagination(example,
                        model_name='gemini-pro',
+                       client=None,
                        word_limit=600,
                        start_threshold=280,
                        max_retires=10,
@@ -232,7 +231,7 @@ def quality_pagination(example,
       pause_point = len(paragraphs)
     else:
       prompt = prompt_pagination_template.format(preceding, '\n'.join(passage), end_tag)
-      response = query_model(prompt=prompt, model_name=model_name).strip()
+      response = query_model(prompt=prompt, model_name=model_name, client=client).strip()
       pause_point = parse_pause_point(response)
       if pause_point and (pause_point <= i or pause_point > j):
         # process += f"prompt:\n{prompt},\nresponse:\n{response}\n"
@@ -264,7 +263,7 @@ Passage:
 
 """
 
-def quality_gisting(example, pages, model_name, word_limit=600, start_threshold=280, verbose=True):
+def quality_gisting(example, pages, model_name, client=None, word_limit=600, start_threshold=280, verbose=True):
   article = example['article']
   title = example['title']
   word_count = count_words(article)
@@ -273,7 +272,7 @@ def quality_gisting(example, pages, model_name, word_limit=600, start_threshold=
   shortened_pages = []
   for i, page in enumerate(pages):
     prompt = prompt_shorten_template.format('\n'.join(page))
-    response = query_model(prompt, model_name)
+    response = query_model(prompt, model_name, client)
     shortened_text = response.strip()
     shortened_pages.append(shortened_text)
     text_output += "[gist] page {}: {}\n\n".format(i, shortened_text)
@@ -323,7 +322,7 @@ Question:
 
 """
 
-def quality_parallel_lookup(example, verbose=True):
+def quality_parallel_lookup(example, model_name, client, verbose=True):
   preprocessed_pages = example['pages']
   article = example['article']
   title = example['title']
@@ -360,7 +359,7 @@ def quality_parallel_lookup(example, verbose=True):
 
     page_ids = []
 
-    response = query_model(prompt=prompt_lookup).strip()
+    response = query_model(prompt=prompt_lookup, model_name=model_name, client=client).strip()
 
     try: start = response.index('[')
     except ValueError: start = len(response)
@@ -391,7 +390,7 @@ def quality_parallel_lookup(example, verbose=True):
     prompt_answer = prompt_answer_template.format(expanded_shortened_article, q, '\n'.join(options_i))
 
     model_choice = None
-    response = query_model(prompt=prompt_answer)
+    response = query_model(prompt=prompt_answer, model_name=model_name, client=client)
     response = response.strip()
     for j, choice in enumerate(choices):
       if response.startswith(f"Answer: {choice}") or response.startswith(f"Answer: {choice[1]}"):
@@ -408,14 +407,25 @@ def quality_parallel_lookup(example, verbose=True):
 
 def query_model_with_quality(
       index: int,
-      model_name: str = 'gemini-pro'
+      model_name: str = 'gemini-pro',
+      api_key: Optional[str] = None,
 ):
+   # setup api key first
+   client = None
+   model_type = model_name.split('-')[0]
+   if model_type == "gpt":
+    # api_key = os.environ.get('OPEN_AI_KEY')
+    client = openai.OpenAI(api_key=api_key)
+   elif model_type == "gemini":
+     # api_key = os.environ.get('GEMINI_API_KEY')
+     genai.configure(api_key=api_key)
+
    example = quality_dev[index]
-   pages, pagination = quality_pagination(example, model_name)
+   pages, pagination = quality_pagination(example, model_name, client)
    print('Finish Pagination.')
-   example_with_gists, gisting = quality_gisting(example, pages, model_name)
+   example_with_gists, gisting = quality_gisting(example, pages, model_name, client)
    print('Finish Gisting.')
-   answers = quality_parallel_lookup(example_with_gists)
+   answers = quality_parallel_lookup(example_with_gists, model_name, client)
    return prompt_pagination_template, pagination, prompt_shorten_template, gisting, prompt_lookup_template, '\n\n'.join(answers)
 
 
@@ -428,6 +438,11 @@ with gr.Blocks() as demo:
     """)
     with gr.Tab('ReadAgent (QuALITY)'):
         llm_options = gr.Radio(llm_api_options, label="Backend LLM API", value='gemini-pro')
+        llm_api_key = gr.Textbox(
+          label="Paste your OpenAI API key (sk-...) or Gemini API key",
+          lines=1,
+          type="password",
+        )
         with gr.Row():
             with gr.Column():
                 index = gr.Dropdown(list(range(len(quality_dev))), value=13, label="QuALITY Index",)
@@ -443,7 +458,8 @@ with gr.Blocks() as demo:
         fn=query_model_with_quality,
         inputs=[
            index,
-           llm_options
+           llm_options,
+           llm_api_key,
         ],
         outputs=[
           prompt_pagination, pagination_results,