app.py

import gradio as gr
from PIL import Image, ImageDraw, ImageFont
import io
import base64
import requests
import json
import threading
import uuid
import os
from datetime import datetime
import urllib.parse


openai_key = "sk-proj-HmBsYhzg5xvirTFW4HS2T3BlbkFJWf90z4T9oXhMgepP3LQL"

user_prompt_examples =[
    "Find me the stars in the hayedes cluster.",
    "Find me the stars in the orion cluster.",
    "Find me the stars in the andromeda cluster.",
    "Find me the stars in the milky way.",
    "Find me the stars in the virgo cluster.",
    "Find me the stars in the bootes cluster.",
    "Find me the stars in the perseus cluster.",
    "Find me the stars in the hydra cluster.",
    "Find me the stars in the Pleiades cluster.",
    "Retrieve stars within 50 light-years of Earth.",
    "List all stars in the Orion Nebula.",
    "Get data on stars with a radial velocity greater than 100 km/s.",
    "Fetch all known white dwarfs in the Sirius star system.",
    "Identify stars in the globular cluster Omega Centauri.",
    "Display stars from the Andromeda Galaxy that are visible from Earth.",
    "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",
        "Authorization": f"Bearer {openai_key}"
    }

    prompt_text = f'''
    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 fields like source_id, positional data (RA, DEC), motion data (pmra, pmdec), light parameters, and placeholders for additional parameters. 

    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. Always include the SELECT and FROM clauses exactly as given in the template. 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 " '#' || 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)
    ```

    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 and astrsonomy to create their request.>",
        "where_clause": "<a properly formatted ADQL query that will return the stars the user is asking for>"
    }}

    The users prompt is "{user_prompt}"
    '''

    data = {
        "model": "gpt-4-turbo",
        "response_format": { "type": "json_object" },
        "messages": [
            {"role": "system", "content": prompt_text},
        ],
        "temperature": .6
    }

    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:
        where_clause = output_json['where_clause']
        reasoning = output_json['reasoning']
        print(f"dialog: {reasoning}")
        print(f"query: {where_clause}")
        # Update theme on the server
        return reasoning, where_clause
    except KeyError as e:
        print(f"Key error: {e}")
        return "Failed to generate query."
    except json.JSONDecodeError:
        print("JSON decoding failed")
        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
    '#' || source_id as source_id,
    -- Measured Position
    ra,
    dec,
    (1 / parallax)*1000 as dist,
    -- Measured Motion
    pmra,
    pmdec,
    radial_velocity as rv,
    --Key source light params for HR diagram
    phot_g_mean_mag,
    bp_rp as bp_rp_mag,
    --Placeholders to be filled in by StarGateVR
    0 as l,
    0 as b,
    0 as X,
    0 as Y,
    0 as Z,
    0 as U,
    0 as V,
    0 as W,
    0 as abs_mag,
    0 as zero,
    1 as one,
    --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,
    g_rp as g_rp_mag,
    bp_g as bp_g_mag,
    radial_velocity_error as rv_error,
    parallax_error,

    -- Note: No comma after this last SELECT item
    parallax_over_error
    -- Use DR3
    FROM gaiadr3.gaia_source    
    {partial_query}
    '''

    print ("completed query")
    print (query_template)
    print ("  ")
    
    return query_template

def download_url_from_query(query, user_prompt):
    # Create the TAP URL
    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):
    # Create the TAP URL
    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

def download_from_tap(url, output_path):
    try:
        response = requests.get(url)
        response.raise_for_status()  # Raises an HTTPError for bad responses (4xx, 5xx)
        with open(output_path, 'wb') as f:
            f.write(response.content)
        print(f"Data successfully downloaded to {output_path}.")
    except requests.exceptions.HTTPError as err:
        print(f"HTTP error occurred: {err}")  # Handle specific HTTP errors
    except Exception as err:
        print(f"An error occurred: {err}")  # Handle other possible errors

def download_data(tap_url):
    return f"[Download Data]({tap_url})"

# Main function to process all queries
def process_queries():
    results = {}
    for prompt in user_prompt_examples:
        print (f"Processing prompt: {prompt}")
        reasoning, where_clause = talk_to_llm(prompt)
        query = complete_query(where_clause)
        download_url_from_query(query, prompt)
    return results

with gr.Blocks() as demo:
    with gr.Row():
        user_prompt = gr.Textbox(label="Enter your query for the LLM", value="Find me the stars in the hayedes cluster.")
        submit_btn = gr.Button("Ask LLM")
    reasoning_output = gr.Textbox(label="Reasoning")
    where_clause_output = gr.Textbox(label="Where Clause")

    
    make_query_complete_bt  = gr.Button("Complete Query")
    complete_query_output = gr.Textbox(label="Complete Query")
    make_query_complete_bt.click(fn=complete_query, inputs=where_clause_output, outputs=[complete_query_output])
    submit_btn.click(fn=talk_to_llm, inputs=user_prompt, outputs=[reasoning_output, where_clause_output])

    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=complete_query_output, outputs=[download_data_output])

#demo.launch(server_name="0.0.0.0", server_port=7861, share=True, debug=True)
demo.launch()


#process_queries()