Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,40 +1,68 @@
-import altair as alt
-import numpy as np
-import pandas as pd
 import streamlit as st
+from langchain_community.utilities import SQLDatabase
+from langchain_google_genai import ChatGoogleGenerativeAI
+from langchain_community.agent_toolkits import SQLDatabaseToolkit
+from langgraph.prebuilt import create_react_agent
+import os
 
-"""
-# Welcome to Streamlit!
-
-Edit `/streamlit_app.py` to customize this app to your heart's desire :heart:.
-If you have any questions, checkout our [documentation](https://docs.streamlit.io) and [community
-forums](https://discuss.streamlit.io).
-
-In the meantime, below is an example of what you can do with just a few lines of code:
-"""
-
-num_points = st.slider("Number of points in spiral", 1, 10000, 1100)
-num_turns = st.slider("Number of turns in spiral", 1, 300, 31)
-
-indices = np.linspace(0, 1, num_points)
-theta = 2 * np.pi * num_turns * indices
-radius = indices
-
-x = radius * np.cos(theta)
-y = radius * np.sin(theta)
-
-df = pd.DataFrame({
-    "x": x,
-    "y": y,
-    "idx": indices,
-    "rand": np.random.randn(num_points),
-})
-
-st.altair_chart(alt.Chart(df, height=700, width=700)
-    .mark_point(filled=True)
-    .encode(
-        x=alt.X("x", axis=None),
-        y=alt.Y("y", axis=None),
-        color=alt.Color("idx", legend=None, scale=alt.Scale()),
-        size=alt.Size("rand", legend=None, scale=alt.Scale(range=[1, 150])),
-    ))
+
+db = SQLDatabase.from_uri("sqlite:///SQLite.db")
+print(db.dialect)
+print(db.get_usable_table_names())
+db.run("SELECT * FROM Artist LIMIT 10;")
+
+
+api_key = os.getenv("GOOGLE_API_KEY")
+# Ensure GOOGLE_API_KEY is set in your environment
+
+# Initialize the Chat Model
+llm = ChatGoogleGenerativeAI(model="gemini-2.5-flash", temperature=0, api_key = api_key)
+messages = [
+    ("system", "You are a helpful assistant that translates English to French."),
+    ("human", "I love programming."),
+]
+llm.invoke(messages)
+
+
+toolkit = SQLDatabaseToolkit(db=db, llm=llm)
+
+tools = toolkit.get_tools()
+
+system_message = """
+You are an agent designed to interact with a SQL database.
+Given an input question, create a syntactically correct {dialect} query to run,
+then look at the results of the query and return the answer. Unless the user
+specifies a specific number of examples they wish to obtain, always limit your
+query to at most {top_k} results.
+
+You can order the results by a relevant column to return the most interesting
+examples in the database. Never query for all the columns from a specific table,
+only ask for the relevant columns given the question.
+
+You MUST double check your query before executing it. If you get an error while
+executing a query, rewrite the query and try again.
+
+DO NOT make any DML statements (INSERT, UPDATE, DELETE, DROP etc.) to the
+database.
+
+To start you should ALWAYS look at the tables in the database to see what you
+can query. Do NOT skip this step.
+
+Then you should query the schema of the most relevant tables.
+""".format(
+    dialect="SQLite",
+    top_k=5,
+)
+
+agent_executor = create_react_agent(llm, tools, prompt=system_message)
+st.title("SQL AI Agent")
+question = st.text_input("Enter your question about the database:")
+if question:
+    for step in agent_executor.stream(
+        {"messages": [{"role": "user", "content": question}]},
+        stream_mode='values'
+    ):
+        # Display each step's message content in Streamlit
+        print(step["messages"][-1].pretty_print())
+
+    st.write(step["messages"][-1].content)