基于Llama 3和LangChain，使用自然语言进行SQL查询精华

小虎哦哦

发布于 2024-7-16 08:48

浏览

0收藏

在数据泛滥的今天，如何高效地从海量信息中提取有价值的洞察，已成为开发者面临的共同挑战。本文介绍一个创新项目——基于Streamlit的应用程序，它能够理解自然语言并直接与SQL数据库进行交互，从而简化数据分析流程。

1 环境设置

在本项目中，我们选择PostgreSQL作为后端数据库。用户需自行安装并配置PostgreSQL，同时需要搭建Ollama环境以及OpenWebUI界面。具体的安装和配置步骤，建议参考相应的官方文档或社区指南进行操作。

2 创建用于生成SQL的自定义模型

接着是创建自定义LLM，根据用户输入生成SQL查询。这里会使用OpenWebUI完成这项任务。以下是设置方法：

创建新模型文件：

命名并描述您的自定义模型。
定义模型内容，基于Llama 3，温度设置为零。
在系统消息中，指示模型为PostgreSQL生成SQL，整合数据库架构和关系。

模型内容定义：

提供数据库架构。
包括有关查询结构和表连接的指南。
专注于SQL查询生成，忽略响应中的格式。

保存模型：

添加提示建议和相关类别。
保存模型文件，使其在Ollama生态系统中可用。

3 构建LangChain集成

LangChain是开源的框架，用于构建基于大型语言模型的应用程序。它通过将语言模型与内部数据源相结合，自动化地串联起一系列操作，生成智能响应。

利用LangChain，我们能够实现数据库与自然语言之间的无缝交互。这一解决方案的实现，是通过在Python的app.py文件中编写代码来完成的。

首先，确保环境已安装以下包：

langchain
langchain-community
streamlit

4 建立数据库连接

将数据库凭据保存在本地变量中，并创建一个使用LangChain的SQLDatabase包装器建立连接的函数。

pwd = os.environ['PGPASS']
uid = os.environ['PGUID']
server = "localhost"
db = "adventureworks"
port = 5432

def init_database() -> SQLDatabase:
  db_uri = f"postgresql://{uid}:{pwd}@{server}:{port}/{db}"
  return SQLDatabase.from_uri(db_uri, schema="sales")

5 使用自定义LLM生成SQL查询

创建自定义函数，根据用户输入生成SQL。这个函数使用Ollama访问自定义模型。

def llm_query(question):
  # 第一个sql llm
  llm = ChatOllama(model="llama-sql")
  prompt = ChatPromptTemplate.from_template(" {topic}")
  # 链
  chain = prompt | llm | StrOutputParser()
  # 链调用
  sql = chain.invoke({"topic": f"{question}"})
  sql = re.sub(r'(?:(?<=_) | (?=_))','',sql)
  # 返回sql查询
  return sql

6 创建完整链

开发一个功能强大的函数，它能够实现与数据库的自然语言交互。该函数能够接收用户的查询、数据库的连接信息以及之前的聊天记录，从而为用户提供一个连贯且上下文相关的交互体验。

def get_response(user_query: str, db: SQLDatabase, chat_history: list):
  sql_query = llm_query(user_query)

  template = """
  您是公司中经验丰富的数据分析师。您正在与询问公司数据库问题的用户进行交互。
  根据表架构、问题、SQL查询和SQL响应，用自然语言编写回答。

  对话历史：{chat_history}
  用户问题：{question}
  SQL响应：{response}
  """

  prompt = ChatPromptTemplate.from_template(template)

  # llm
  llm = ChatOllama(model="llama3",  temperature=0)

  chain = (
    RunnablePassthrough.assign(
        response=lambda vars: db.run(sql_query),
    )
    | prompt
    | llm
    | StrOutputParser()
  )

  return chain.invoke({
    "question": user_query,
    "chat_history": chat_history,
  })

7 构建Streamlit应用程序

创建Streamlit组件以存储聊天历史和管理用户交互。

初始化数据库：

# 初始化数据库
db = init_database()
st.session_state.db = db
st.success("已连接到数据库！")

处理用户输入：

user_query = st.chat_input("输入一个问题...")
if user_query is not None and user_query.strip() != "":
    st.session_state.chat_history.append(HumanMessage(content=user_query))
    
    with st.chat_message("用户"):
        st.markdown(user_query)
        
    with st.chat_message("AI"):
        response = get_response(user_query, st.session_state.db, st.session_state.chat_history)
        st.markdown(response)