基于LangChain调用自定义Llama2 API实现Text-to-SQL查询PostgreSQL

依赖安装 #

先安装所需的Python包：

pip install langchain langchain-ollama psycopg2-binary langchain-experimental

完整修改代码 #

# 导入依赖库
from langchain_ollama import OllamaLLM
from langchain.sql_database import SQLDatabase
from langchain.agents import create_sql_agent, AgentType
from langchain.agents.agent_toolkits import SQLDatabaseToolkit

# 初始化PostgreSQL数据库连接
db = SQLDatabase.from_uri("postgresql://postgres:password@localhost:5432/postgres")

# 初始化Ollama部署的Llama2模型
llm = OllamaLLM(
    model="llama2",
    base_url="http://localhost:11434",  # 对应你的Ollama API服务地址
    temperature=0  # 固定输出，减少随机性
)

# 创建SQL工具包与智能代理
toolkit = SQLDatabaseToolkit(db=db, llm=llm)
agent_executor = create_sql_agent(
    llm=llm,
    toolkit=toolkit,
    verbose=True,
    agent_type=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
)

# 执行对话式查询
agent_executor.run("Describe the transaction table")

关键说明 #

• 确保Ollama服务已启动，且已通过ollama pull llama2命令拉取对应模型
• 替换数据库URI中的用户名、密码、数据库名、端口为你的实际配置
• 若Llama2输出不稳定，可尝试将agent_type改为AgentType.CHAT_ZERO_SHOT_REACT_DESCRIPTION，该类型更适配支持对话格式的大模型

内容的提问来源于stack exchange，提问作者A_K