Update app.py

app.py

@@ -8,6 +8,7 @@ import requests
 import ast
 
 
+
 api_base = "https://api.endpoints.anyscale.com/v1"
 token = os.environ["OPENAI_API_KEY"]
 url = f"{api_base}/chat/completions"
@@ -54,7 +55,7 @@ def remove_leading_whitespace(func_str):
     return "\n".join(new_lines)
 
 
-def main(fxn: str, examples: str = "", temperature: float = 0.7):
+def main(fxn: str, openai_api_key, examples: str = "", temperature: float = 0.7):
     """Requires Anyscale Endpoints Alpha API access.
 
     If examples is not a empty string, it will be formatted into
@@ -63,7 +64,7 @@ def main(fxn: str, examples: str = "", temperature: float = 0.7):
 
     s = requests.Session()
     api_base = os.environ["OPENAI_API_BASE"]
-    token = os.environ["OPENAI_API_KEY"]
+    token = openai_api_key
     url = f"{api_base}/chat/completions"
 
     message = "Write me a test of this function\n{}".format(fxn)
@@ -234,30 +235,18 @@ def is_semipositive_definite(matrix):
 
     # Check if all eigenvalues are non-negative
     return all(val >= 0 for val in eigenvalues)
-""",
-"""
-def key_exists(data, key_to_find):
-    # Base case: If data is not a dictionary, return False
-    if not isinstance(data, dict):
-        return False
-
-    # If the key is in the dictionary, return True
-    if key_to_find in data:
-        return True
-
-    # Otherwise, recurse into each nested dictionary
-    for key, value in data.items():
-        if key_exists(value, key_to_find):
-            return True
-
-    # If the key was not found in this dictionary or any sub-dictionary, return False
-    return False
 """
     ]
 example = examples[0]
 
 with gr.Blocks() as demo:
     gr.Markdown("<h1><center>Llama_test: generate unit test for your Python code</center></h1>")
+    openai_api_key = gr.Textbox(
+                show_label=False,
+                label="Set your Anyscale API key here.",
+                lines=1,
+                type="password"
+            )
     code_input = gr.Code(example, language="python", label="Provide the code of the function you want to test")
     gr.Examples(
     examples=examples,
@@ -268,6 +257,6 @@ with gr.Blocks() as demo:
         code_output = gr.Code(language="python", label="Passed tests")
         code_output2 = gr.Code(language="python", label="Failed tests")
 
-    generate_btn.click(main, inputs=code_input, outputs=[code_output, code_output2])
+    generate_btn.click(main, inputs=[code_input, openai_api_key], outputs=[code_output, code_output2])
 if __name__ == "__main__":
     demo.launch()