1
+ import os
2
+ import asyncio
3
+ import streamlit as st
4
+ from langchain_community .utilities import SQLDatabase
5
+ from langchain_google_genai import ChatGoogleGenerativeAI
6
+ from langchain_community .agent_toolkits import create_sql_agent
7
+
8
+ from dotenv import load_dotenv
9
+ load_dotenv ()
10
+
11
+ os .environ ["GOOGLE_API_KEY" ]= os .getenv ("GOOGLE_API_KEY" )
12
+ groq_api_key = os .getenv ('GROQ_API_KEY' )
13
+
14
+ async def get_answer_by_sql_query (db ,input ):
15
+ llm = ChatGoogleGenerativeAI (model = "gemini-1.5-pro-latest" ,temperature = 0 )
16
+ agent_executor = create_sql_agent (llm , db = db , verbose = True )
17
+
18
+ response = agent_executor .invoke ({"input" : input })
19
+ return response
20
+
21
+
22
+ def main ():
23
+ st .header ('Chat with your Postgres Sql Database...' )
24
+
25
+ if "conversation" not in st .session_state :
26
+ st .session_state .conversation = None
27
+
28
+ if "activate_chat" not in st .session_state :
29
+ st .session_state .activate_chat = False
30
+
31
+ if "messages" not in st .session_state :
32
+ st .session_state .messages = []
33
+ st .session_state .chat_history = []
34
+
35
+ for message in st .session_state .messages :
36
+ with st .chat_message (message ["role" ], avatar = message ['avatar' ]):
37
+ st .markdown (message ["content" ])
38
+
39
+ with st .sidebar :
40
+ st .subheader ('Please Enter your credentials' )
41
+ username = st .text_input ("enter your username" )
42
+ password = st .text_input ("enter your password" )
43
+ db_name = st .text_input ("enter your database name" )
44
+ if st .button ('Process' ):
45
+ if username is not None and password is not None :
46
+ postgres_uri = f"postgresql+psycopg2://{ username } :{ password } @localhost:5432/{ db_name } "
47
+ st .write (f'Processed database: { postgres_uri } ' )
48
+ st .write (f'Connection to database successfully..' )
49
+
50
+ db = SQLDatabase .from_uri (postgres_uri )
51
+ if "db" not in st .session_state :
52
+ st .session_state .db = db
53
+ st .session_state .activate_chat = True
54
+
55
+ if st .session_state .activate_chat == True :
56
+ if prompt := st .chat_input ("Ask your question from the Database" ):
57
+ with st .chat_message ("user" , avatar = '👨🏻' ):
58
+ st .markdown (prompt )
59
+ st .session_state .messages .append ({"role" : "user" , "avatar" :'👨🏻' , "content" : prompt })
60
+ db = st .session_state .db
61
+ response = asyncio .run (get_answer_by_sql_query (db ,prompt ))
62
+ cleaned_response = response ["output" ]
63
+ with st .chat_message ("assistant" , avatar = '🤖' ):
64
+ st .markdown (cleaned_response )
65
+ st .session_state .messages .append ({"role" : "assistant" , "avatar" :'🤖' , "content" : cleaned_response })
66
+ else :
67
+ st .markdown ('Please Enter your credentials to chat' )
68
+
69
+
70
+ if __name__ == '__main__' :
71
+ main ()
0 commit comments