Update app.py

app.py

@@ -19,21 +19,23 @@ import json, openai, os, time
 
 from datetime import date
 from openai import OpenAI
+from tavily import TavilyClient
 from typing import List
 from utils import function_to_schema, show_json
 
-client = OpenAI(api_key=os.environ.get("OPENAI_API_KEY"))
+openai_client = OpenAI(api_key=os.environ.get("OPENAI_API_KEY"))
+tavily_client = TavilyClient(api_key=os.environ.get("TAVILY_API_KEY"))
 
 assistant_id = "asst_YWSKwQEiYVGApJdxxrTH0Cy8"
 
 assistant, thread = None, None
 
-def today_tool() -> str:
+def today_tool() -> today: str:
     """Returns today's date. Use this function for any questions related to knowing today's date. 
        There should be no input. This function always returns today's date."""
     return str(date.today())
 
-def yf_download_tool(tickers: List[str], start_date: date, end_date: date) -> pd.DataFrame:
+def yf_download_tool(tickers: List[str], start_date: date, end_date: date) -> stock_data: pd.DataFrame:
     """Returns historical stock data for a list of given tickers from start date to end date 
        using the yfinance library download function. 
        Use this function for any questions related to getting historical stock data. 
@@ -41,13 +43,21 @@ def yf_download_tool(tickers: List[str], start_date: date, end_date: date) -> pd
        This function always returns a pandas DataFrame."""
     return yf.download(tickers, start=start_date, end=end_date)
 
+def tavily_search_tool(query: str) -> answer: str:
+    """Searches the web for a given query and returns an answer, "
+       ready for use as context in a RAG application, using the Tavily API. 
+       Use this function for any questions requiring knowledge not available to the model. 
+       The input should be the query string. This function always returns an answer string."""
+    return tavily_client.get_search_context(query=query, max_results=5)
+
 tools = {
     "today_tool": today_tool,
     "yf_download_tool": yf_download_tool,
+    "tavily_search_tool": tavily_search_tool,
 }
 
-def create_assistant(client):
-    assistant = client.beta.assistants.create(
+def create_assistant(openai_client):
+    assistant = openai_client.beta.assistants.create(
         name="Python Code Generator",
         instructions=(
              "You are a Python programming language expert that "
@@ -59,6 +69,7 @@ def create_assistant(client):
             {"type": "code_interpreter"},
             {"type": "function", "function": function_to_schema(today_tool)},
             {"type": "function", "function": function_to_schema(yf_download_tool)},
+            {"type": "function", "function": function_to_schema(tavily_search_tool)},
         ],
     )
     
@@ -66,22 +77,22 @@ def create_assistant(client):
     
     return assistant
 
-def load_assistant(client):   
-    assistant = client.beta.assistants.retrieve(assistant_id)
+def load_assistant(openai_client):   
+    assistant = openai_client.beta.assistants.retrieve(assistant_id)
 
     show_json("assistant", assistant)
     
     return assistant
 
-def create_thread(client):
-    thread = client.beta.threads.create()
+def create_thread(openai_client):
+    thread = openai_client.beta.threads.create()
     
     show_json("thread", thread)
     
     return thread
 
-def create_message(client, thread, msg):        
-    message = client.beta.threads.messages.create(
+def create_message(openai_client, thread, msg):        
+    message = openai_client.beta.threads.messages.create(
         role="user",
         thread_id=thread.id,
         content=msg,
@@ -91,8 +102,8 @@ def create_message(client, thread, msg):
     
     return message
 
-def create_run(client, assistant, thread):
-    run = client.beta.threads.runs.create(
+def create_run(openai_client, assistant, thread):
+    run = openai_client.beta.threads.runs.create(
         assistant_id=assistant.id,
         thread_id=thread.id,
         parallel_tool_calls=False,
@@ -102,14 +113,12 @@ def create_run(client, assistant, thread):
     
     return run
 
-def wait_on_run(client, thread, run):
+def wait_on_run(openai_client, thread, run):
     while run.status == "queued" or run.status == "in_progress":
-        run = client.beta.threads.runs.retrieve(
+        run = openai_client.beta.threads.runs.retrieve(
             thread_id=thread.id,
             run_id=run.id,
         )
-
-        show_json("run", run)
             
         time.sleep(1)
     
@@ -120,8 +129,8 @@ def wait_on_run(client, thread, run):
 
     return run
 
-def get_run_steps(client, thread, run):
-    run_steps = client.beta.threads.runs.steps.list(
+def get_run_steps(openai_client, thread, run):
+    run_steps = openai_client.beta.threads.runs.steps.list(
         thread_id=thread.id,
         run_id=run.id,
         order="asc",
@@ -161,8 +170,8 @@ def execute_tool_calls(run_steps):
 
     return tool_call_ids, tool_call_results
 
-def get_messages(client, thread):
-    messages = client.beta.threads.messages.list(
+def get_messages(openai_client, thread):
+    messages = openai_client.beta.threads.messages.list(
         thread_id=thread.id
     )
     
@@ -224,19 +233,19 @@ def chat(message, history):
     global assistant, thread     
     
     if assistant == None:
-        #assistant = create_assistant(client) # on first run, create assistant and update assistant_id
-                                              # see https://platform.openai.com/playground/assistants
-        assistant = load_assistant(client) # on subsequent runs, load assistant
+        assistant = create_assistant(openai_client) # on first run, create assistant and update assistant_id
+                                                     # see https://platform.openai.com/playground/assistants
+        #assistant = load_assistant(openai_client) # on subsequent runs, load assistant
     
     if thread == None or len(history) == 0:
-        thread = create_thread(client)
+        thread = create_thread(openai_client)
         
-    create_message(client, thread, message)
+    create_message(openai_client, thread, message)
 
-    run = create_run(client, assistant, thread)
+    run = create_run(openai_client, assistant, thread)
 
-    run = wait_on_run(client, thread, run)
-    run_steps = get_run_steps(client, thread, run)
+    run = wait_on_run(openai_client, thread, run)
+    run_steps = get_run_steps(openai_client, thread, run)
 
     ### TODO
     tool_call_ids, tool_call_results = execute_tool_calls(run_steps)
@@ -264,14 +273,14 @@ def chat(message, history):
                 "output": tool_call_results[0]
             }
         
-        run = client.beta.threads.runs.submit_tool_outputs(
+        run = openai_client.beta.threads.runs.submit_tool_outputs(
             thread_id=thread.id,
             run_id=run.id,
             tool_outputs=[tool_output]
         )
     
-        run = wait_on_run(client, thread, run)
-        run_steps = get_run_steps(client, thread, run)
+        run = wait_on_run(openai_client, thread, run)
+        run_steps = get_run_steps(openai_client, thread, run)
     ###
         tool_call_ids, tool_call_results = execute_tool_calls(run_steps)
             
@@ -298,17 +307,17 @@ def chat(message, history):
                     "output": tool_call_results[1]
                 }
             
-            run = client.beta.threads.runs.submit_tool_outputs(
+            run = openai_client.beta.threads.runs.submit_tool_outputs(
                 thread_id=thread.id,
                 run_id=run.id,
                 tool_outputs=[tool_output]
             )
         
-            run = wait_on_run(client, thread, run)
-            run_steps = get_run_steps(client, thread, run)    
+            run = wait_on_run(openai_client, thread, run)
+            run_steps = get_run_steps(openai_client, thread, run)    
     ###
     
-    messages = get_messages(client, thread)
+    messages = get_messages(openai_client, thread)
 
     text_values, image_values = extract_content_values(messages)
 
@@ -328,8 +337,9 @@ gr.ChatInterface(
         description=(
             "The assistant can **generate, explain, fix, optimize, document, and test code**. "
             "It can also **execute code**. "
-            "It has access to <b>today tool</b> (get current date) and "
-            "to **yfinance download tool** (get stock data)."
+            "It has access to <b>today tool</b> (get current date), to "
+            "**yfinance download tool** (get stock data), and to "
+            "**tavily search tool** (search web)."
         ),
         clear_btn="Clear",
         retry_btn=None,
@@ -340,7 +350,8 @@ gr.ChatInterface(
                   ["Fix: x = [5, 2, 1, 3, 4]; print(x.sort())"],
                   ["Optimize: x = []; for i in range(0, 10000): x.append(i)"],
                   ["Execute: First 25 Fibbonaci numbers"],
-                  ["Execute with tools: Create a plot showing stock gain QTD for NVDA and AMD, x-axis is \"Day\" and y-axis is \"Gain %\""]
+                  ["Execute with tools: Create a plot showing stock gain QTD for NVDA and AMD, x-axis is \"Day\" and y-axis is \"Gain %\""],
+                  ["Execute with tools: TODO"]
                  ],
         cache_examples=False,
     ).launch()