Update utils.py

utils.py

@@ -0,0 +1,104 @@
+from langchain_community.utilities import SQLDatabase
+from langchain_core.callbacks import BaseCallbackHandler
+from typing import TYPE_CHECKING, Any, Optional, TypeVar, Union
+from uuid import UUID
+from langchain_community.agent_toolkits import create_sql_agent
+from langchain_openai import ChatOpenAI
+from langchain_community.vectorstores import Chroma
+from langchain_core.example_selectors import SemanticSimilarityExampleSelector
+from langchain_openai import OpenAIEmbeddings
+from langchain.agents.agent_toolkits import create_retriever_tool
+from langchain_core.output_parsers import JsonOutputParser
+import os
+from langchain_core.prompts import (
+    ChatPromptTemplate,
+    FewShotPromptTemplate,
+    MessagesPlaceholder,
+    PromptTemplate,
+    SystemMessagePromptTemplate,
+)
+import ast
+import re
+
+
+def query_as_list(db, query):
+    res = db.run(query)
+    res = [el for sub in ast.literal_eval(res) for el in sub if el]
+    res = [re.sub(r"\b\d+\b", "", string).strip() for string in res]
+    return list(set(res))
+
+
+def get_answer(user_query):
+
+    global retriever_tool, example_selector, db, llm
+
+
+    system_prefix = """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 the question does not seem related to the database, just return "I don't know" as the answer.
+
+    Here are some examples of user inputs and their corresponding SQL queries:"""
+
+    few_shot_prompt = FewShotPromptTemplate(
+        example_selector=example_selector,
+        example_prompt=PromptTemplate.from_template(
+            "User input: {input}\nSQL query: {query}"
+        ),
+        input_variables=["input", "dialect", "top_k"],
+        prefix=system_prefix,
+        suffix="",
+    )
+
+    employee = query_as_list(db, "SELECT FullName FROM Employee")
+    system_unique_name_prompt = """
+    If you need to filter on a proper noun, you must ALWAYS first look up the filter value using the "search_proper_nouns" tool!
+
+    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.
+
+    """
+
+
+    prompt_val = few_shot_prompt.invoke(
+        {
+            "input": user_query,
+            "top_k": 5,
+            "dialect": "SQLite",
+
+            "agent_scratchpad": [],
+        }
+    )
+
+    final_prompt = prompt_val.to_string() + '\n' + system_unique_name_prompt
+    full_prompt = ChatPromptTemplate.from_messages(
+        [
+            ("system",final_prompt),
+            ("human", "{input}"),
+            MessagesPlaceholder("agent_scratchpad"),
+        ]
+    )
+
+
+    agent = create_sql_agent(
+        llm=llm,
+        db=db,
+        max_iterations = 40,
+        extra_tools=[retriever_tool],
+        prompt=full_prompt,
+        agent_type="openai-tools",
+        verbose=True,
+    )
+
+    result = agent.invoke({'input': user_query})
+
+    return result['output']