Update app.py

app.py

@@ -4,21 +4,15 @@ import io
 import base64
 import requests
 import json
-import threading
-import uuid
 import os
-from datetime import datetime
 import urllib.parse
-
 from dotenv import load_dotenv
 
 load_dotenv()
 hugging_face_api_key = os.getenv("HUGGING_FACE_API_KEY")
 openai_key = os.getenv("OPENAI_API_KEY")
 
-# openai_key = openai_api_key
-
-user_prompt_examples =[
+user_prompt_examples = [
     "Find me the stars in the hyades cluster.",
     "Find me the stars in the orion cluster.",
     "Find me the stars in the andromeda cluster.",
@@ -37,7 +31,7 @@ user_prompt_examples =[
     "Find stars in the Scorpius constellation with a magnitude brighter than 5.",
     "Search for stars with high proper motion in the Ursa Major group."
 ]
-    
+
 def talk_to_llm(user_prompt):
     headers = {
         "Content-Type": "application/json",
@@ -48,11 +42,9 @@ def talk_to_llm(user_prompt):
     As StarGateVR, your role is specialized in customizing ADQL (Astronomical Data Query Language) 
     queries for astronomers. Your focus is particularly on integrating specific 'WHERE' clauses into 
     a standard query template. We will put your WHERE clause into the completed query template. 
-
     The query includes essential SELECT fields like source_id, positional data (ra, dec), 
     motion data (pmra, pmdec), and light parameters. Note that any
     fields used in the WHERE clause must also be added to the SELECT clause.
-
     Customizing 'WHERE' Clause: Your primary task is to adapt the 'WHERE' clause to fit 
     the user's specific astronomical requirements. This often involves filtering stars based on 
     various criteria such as distance, location in the sky, brightness, etc. 
@@ -71,25 +63,21 @@ def talk_to_llm(user_prompt):
     phot_g_mean_mag,
     bp_rp as bp_rp_mag,
     ```
-    
+
     Note that the WHERE clause must reference variables by the field name and not the "AS" name. 
     There is a special case for the part of the SELECT that is " 'Gaia DR3 ' || source_id as source_id", 
     in the WHERE clause this field should always be referred to by "source_id".
-
     Here is an example of the WHERE clause:
     ```
     WHERE (parallax >= 11.11 AND parallax_over_error>=20 AND
     astrometric_excess_noise<=2)
     ```
-
     Here is the preferred structure for the FROM clause:
     ```
     FROM gaiadr3.gaia_source
     ```
-
     Bounds on Parallax: Always include bounds on parallax in the 'WHERE' clause. This is 
     important as it helps in retrieving stars within a specified 3D region of space.
-
     The json structure to return is
     {{
         "reasoning": "<Think through what the user is asking for, and what you know about the GAIA DB 
@@ -99,9 +87,7 @@ def talk_to_llm(user_prompt):
         the downstream logic.>",
         "the_query": "<a properly formatted ADQL query that will return the stars the 
         user is asking for>"
-        
     }}
-
     The users prompt is "{user_prompt}"
     '''
 
@@ -114,38 +100,28 @@ def talk_to_llm(user_prompt):
         "temperature": .2
     }
 
-    print ("PROMPT TEXT:\n")
-    print (prompt_text)
     response = requests.post("https://api.openai.com/v1/chat/completions", headers=headers, json=data)
     response_json = response.json()
-    print(response_json)
-    output_json = json.loads(response_json['choices'][0]['message']['content'])
-    
+
     try:
+        output_json = json.loads(response_json['choices'][0]['message']['content'])
         the_query = output_json['the_query']
         reasoning = output_json['reasoning']
-        print(f"dialog: {reasoning}")
-        print(f"query: {the_query}")
-        # Update theme on the server
         return reasoning, the_query
     except KeyError as e:
         print(f"Key error: {e}")
-        return "Failed to generate query."
+        return "Failed to generate query.", ""
     except json.JSONDecodeError:
         print("JSON decoding failed")
-        return "Invalid response query."
+        return "Invalid response query.", ""
 
 def complete_query(partial_query):
     query_template = f'''
     SELECT TOP 300000
-
     -- IMPORTANT NOTE: Parameters that are in units of Magnitude must have an "as" name that ends in "_mag"
-
     --Required parameters
-
     -- ID - force a leading hash symbol to stop Excel from reading the ID number as a float
     'Gaia DR3 ' || source_id as source_id,
-
     -- Measured Position
     ra,
     dec,
@@ -154,11 +130,9 @@ def complete_query(partial_query):
     pmra,
     pmdec,
     radial_velocity as rv,
-
     --Key source light params for HR diagram
     phot_g_mean_mag,
     bp_rp as bp_rp_mag,
-
     --Optional plot parameters (you can add anything you want here, just give good unique "as" names - it will show up in the .cvs and hence in StarGate
     phot_rp_mean_mag,
     phot_bp_mean_mag,
@@ -166,9 +140,7 @@ def complete_query(partial_query):
     bp_g as bp_g_mag,
     radial_velocity_error as rv_error,
     parallax_error,
-
     -- Additional parameters that appear in the WHERE clause should be added here
-
     -- Note: No comma after this last SELECT item
     parallax_over_error
     -- Use DR3
@@ -176,10 +148,6 @@ def complete_query(partial_query):
     {partial_query}
     '''
     
-    print ("completed query")
-    print (query_template)
-    print ("  ")
-    
     return query_template
 
 def download_url_from_query(query, user_prompt):
@@ -187,9 +155,7 @@ def download_url_from_query(query, user_prompt):
     tap_url = "https://gea.esac.esa.int/tap-server/tap/sync?REQUEST=doQuery&LANG=ADQL&FORMAT=csv&QUERY="
     query = urllib.parse.quote_plus(query)
     tap_url = tap_url + query
-    print  ("downloading from tap")
     filename = user_prompt.replace(' ', '_').replace('.', '') + ".csv"
-    print (filename)
     download_from_tap(tap_url, filename)
 
 def create_markdown_url_from_query(query):
@@ -197,8 +163,6 @@ def create_markdown_url_from_query(query):
     tap_url = "https://gea.esac.esa.int/tap-server/tap/sync?REQUEST=doQuery&LANG=ADQL&FORMAT=csv&QUERY="
     query = urllib.parse.quote_plus(query)
     tap_url = tap_url + query
-    print  (tap_url)
-    download_from_tap(tap_url, "output_file.csv")
     markdown_link = download_data(tap_url)
     return markdown_link
 
@@ -221,10 +185,9 @@ def download_data(tap_url):
 def process_queries():
     results = {}
     for prompt in user_prompt_examples:
-        print(f"Processing prompt: {prompt}")
         reasoning, the_query = talk_to_llm(prompt)
-        query = the_query
-        download_url_from_query(query, prompt)
+        if the_query:
+            download_url_from_query(the_query, prompt)
     return results
 
 with gr.Blocks() as demo:
@@ -237,10 +200,8 @@ with gr.Blocks() as demo:
 
     create_tap_url_bt = gr.Button("Create TAP URL")
     download_data_output = gr.Markdown()
-    create_tap_url_bt.click(fn=create_markdown_url_from_query, inputs=the_query_output, 
-                            outputs=[download_data_output])
+    create_tap_url_bt.click(fn=create_markdown_url_from_query, inputs=the_query_output, outputs=[download_data_output])
 
-#demo.launch(server_name="0.0.0.0", server_port=7861, share=True, debug=True)
 demo.launch()