sqlbot / temp.py
edithram23's picture
Update temp.py
1b14c2f verified
from langchain_community.utilities import SerpAPIWrapper
from langgraph.prebuilt import create_react_agent
import getpass
import os
import sqlite3
from langchain_openai import ChatOpenAI
from langchain_community.tools.tavily_search import TavilySearchResults
from langchain_core.messages import HumanMessage,AIMessage
from langchain.agents import initialize_agent
from langchain.agents.agent_types import AgentType
from langchain_community.agent_toolkits.sql.toolkit import SQLDatabaseToolkit
from langchain_community.utilities.sql_database import SQLDatabase
from langchain_openai import ChatOpenAI
from langchain_groq import ChatGroq
from dotenv import load_dotenv
load_dotenv('.env')
class Script:
def __init__(self):
self.model = ChatOpenAI(model="gpt-4o-mini")
self.groq = ChatGroq(model='llama3-70b-8192')
self.db = SQLDatabase.from_uri("sqlite:///sample_database.db")
self.toolkit = SQLDatabaseToolkit(db=self.db, llm=self.model)
search = TavilySearchResults(name='Search-Tool',max_results=1,max_retries=3, description='A Search Tool that can be used to get zip codes, city, state, etc.')
tools = self.toolkit.get_tools()+[search]
self.agent = create_react_agent(self.model, tools)
def history(self,hist):
message = []
for i in hist[-5:]:
if(i['role']=='user'):
message+=[HumanMessage(content=i['content'])]
else:
message+=[AIMessage(content=i['content'])]
return message
def request(self,message):
message = self.history(message)
try :
response = self.agent.invoke({"messages": message})
return response["messages"][-1].content
except:
return 'Sorry unable to process the request'