Update app.py

app.py

@@ -44,6 +44,13 @@ Instructions:
 
 Begin extraction.
 """
+def get_llm_response(entity, query):
+    # Format the prompt with the entity and query
+    formatted_prompt = PROMPT_TEMPLATE.format(entity=entity, query=query)
+    
+    # Pass the formatted prompt as a SystemMessage and get the response
+    response = llm([SystemMessage(content=formatted_prompt)])
+    return response[0].content
 
 # Set up the page
 st.set_page_config(page_title="DataScribe", page_icon=":notebook_with_decorative_cover:", layout="wide")
@@ -134,6 +141,9 @@ elif selected == "Upload Data":
 elif selected == "Define Query":
     st.header("Define Your Custom Query")
 
+    entity = st.text_input("Enter the entity name")
+    query = st.text_input("Enter the query (e.g., 'contact information for {entity}')")
+
     if st.session_state["data"] is not None:
         column_selection = st.selectbox("Select the primary column for entities", options=st.session_state["data"].columns)
         query_template = st.text_input("Define your query template", "Get me the email for {company}")
@@ -148,6 +158,10 @@ elif selected == "Define Query":
             st.code(example_query)
     else:
         st.warning("Please upload data first.")
+    if st.button("Extract Information"):
+        if entity and query:
+            response_text = get_llm_response(entity, query)
+            st.write(reponse_text)
 
 # Extract Information Section with Progress Bar
 elif selected == "Extract Information":