File size: 854 Bytes
c92ed68
fde297f
 
 
 
 
6ce6f34
fde297f
 
 
db949df
fde297f
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
import os
from langchain import OpenAI, SQLDatabase
from langchain.agents.agent_toolkits import SQLDatabaseToolkit
from langchain.agents import create_sql_agent
from langchain.agents import AgentExecutor
from langchain.agents.agent_types import AgentType
# from env import OPENAI_API_KEY

dburi = "sqlite:///output.db"
db = SQLDatabase.from_uri(dburi)
llm = OpenAI(temperature=0, openai_api_key=os.environ['OPENAI_API_KEY'])

toolkit = SQLDatabaseToolkit(db=db, llm=llm)
agent_executor = create_sql_agent(
    llm=llm,
    toolkit=toolkit,
    verbose=True,
    agent_type=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
)


def driver1(industry_type):
    response = agent_executor.run(
        "What are the customers in"
        + industry_type
        + " industry that chose salesforce? Give the answer in form of bullet points ->"
    )
    return response