app.py

import json
import os
import chainlit as cl
import pandas as pd
from langchain.schema.runnable.config import RunnableConfig
from langchain.text_splitter import CharacterTextSplitter
from langchain.tools.retriever import create_retriever_tool
from langchain_community.agent_toolkits import SQLDatabaseToolkit
from langchain_community.agent_toolkits.sql.base import create_sql_agent
from langchain_community.utilities import SQLDatabase
from langchain_community.vectorstores import FAISS
from langchain_core.prompts import (
    ChatPromptTemplate,
    MessagesPlaceholder,
)
from langchain_core.runnables import RunnableLambda
from langchain_openai.chat_models import ChatOpenAI
from langchain_openai.embeddings import OpenAIEmbeddings
from sqlalchemy import create_engine

system = """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 have access to tools for interacting with the database.
Only use the given tools. Only use the information returned by the tools to construct your final answer.
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.

If you are prompted for the specific TEXT of a review, use the search_reviews tool to find the relevant review(s).

You have access to the following tables: {table_names}

If the question does not seem related to the database, just return "I don't know" as the answer."""


def setup():
    lf = pd.read_csv("Main.csv")
    engine = create_engine("sqlite:///TrustWhole1.db")
    lf.to_sql("MainTable", engine, index=True, if_exists="replace")
    db = SQLDatabase(engine=engine)

    llm = ChatOpenAI(
        model="gpt-4-turbo-preview", temperature=0, api_key=os.getenv("OPENAI_API_KEY")
    )
    sql_toolkit = SQLDatabaseToolkit(db=db, llm=llm)

    texts_to_split = ""
    with open("combined_extracted.json", "r") as file:
        data = json.load(file)

    for item in data:
        texts_to_split += (
            item["companyName"]
            + "\n"
            + item["reviewTitle"]
            + "\n"
            + item["reviewDescription"]
            + "\n\n"
        )

    text_splitter = CharacterTextSplitter(
        separator="\n",
        chunk_size=2500,
        chunk_overlap=100,
    )
    texts = text_splitter.split_text(texts_to_split)
    embeddings = OpenAIEmbeddings(model="text-embedding-ada-002")
    vectorstore = FAISS.from_texts(texts, embeddings)
    retriever = vectorstore.as_retriever()

    tool_retiv = create_retriever_tool(
        retriever,
        name="search_reviews",
        description="Use this tool to search for and retrieve relevant reviews and their titles for the companies mentioned in the user's query. This tool should be used when the question cannot be directly answered using SQL, particularly when the query involves non-numeric information or requires contextual understanding from the reviews. However, if the question involves aggregation operations or primarily focuses on numerical data, the SQL tool should be used instead.",
    )

    prompt = ChatPromptTemplate.from_messages(
        [
            ("system", system),
            ("human", "{input}"),
            MessagesPlaceholder("agent_scratchpad"),
        ]
    )

    sql_agent = create_sql_agent(
        llm,
        toolkit=sql_toolkit,
        prompt=prompt,
        verbose=True,
        agent_type="openai-tools",
        extra_tools=[tool_retiv],
    )

    return sql_agent | RunnableLambda(lambda x: x["output"])


agent_executor = setup()


@cl.on_chat_start
async def on_chat_start():
    cl.user_session.set("agent_executor", agent_executor)


@cl.on_message
async def main(message):
    agent_executor = cl.user_session.get("agent_executor")

    msg = cl.Message(content="")
    stream = agent_executor.astream(
        {"input": message.content},
        config=RunnableConfig(callbacks=[cl.LangchainCallbackHandler()]),
    )
    async for chunk in stream:
        await msg.stream_token(chunk)

    await msg.send()