解决LangChain+OpenAI操作PostgreSQL大数据库的上下文长度超限错误

问题描述 #

使用LangChain结合OpenAI与PostgreSQL数据库交互，小型数据库运行正常，但数据库规模超过10000行时，触发以下错误：
This model's maximum context length is 4097 tokens. However, your messages resulted in 5100 tokens. Please reduce the length of the messages

可复现该错误的公开数据库为MindsDB房产数据库，连接字符串：

postgresql+psycopg2://demo_user:demo_password@REDACTED:5432/demo

问题代码如下：

import sys
from langchain import OpenAI
from langchain import SQLDatabase
from langchain.chat_models import ChatOpenAI
from langchain_experimental.sql import SQLDatabaseChain

import environ
env = environ.Env()
environ.Env.read_env()

API_KEY = env('OPENAI_API_KEY')

if API_KEY == "":
    print("Missing OpenAPI key")
    exit()

if len(sys.argv) < 2:
    print("Missing db connection string.  Example 'postgresql+psycopg2://postgres:1234@localhost:6667/mydb'")
    exit()

dbstring = sys.argv[1]

print("Using OpenAPI with key ["+API_KEY+"] and Database ["+dbstring+"]")

# Setup database
db = SQLDatabase.from_uri(
    dbstring,
)

# setup llm
llm = ChatOpenAI(model_name="gpt-3.5-turbo",
    temperature=0,
    max_tokens=1000,
    openai_api_key=API_KEY)

# Create db chain
QUERY = """
Given an input question, first create a syntactically correct postgresql query to run, then look at the results of the query and return the answer.
Use the following format:

Question: Question here
SQLQuery: SQL Query to run
SQLResult: Result of the SQLQuery
Answer: Final answer here

{question}
"""

# Setup the database chain
db_chain = SQLDatabaseChain(llm=llm, database=db, verbose=True)

def get_prompt():
    print("Type 'exit' to quit")

    while True:
        prompt = input("Enter a prompt: ")

        if prompt.lower() == 'exit':
            print('Exiting...')
            break
        else:
            try:
                question = QUERY.format(question=prompt)
                print(db_chain.run(question))
            except Exception as e:
                print(e)

get_prompt()

核心原因 #

SQLDatabaseChain默认会将全量表结构+完整查询结果集打包传入LLM上下文，当数据库表字段多、行数多，或查询返回大量数据时，总token数会超出gpt-3.5-turbo的4097上限。

解决方案 #

1. 限制加载的数据库表 #

初始化SQLDatabase时，仅加载业务相关的表，避免无关表结构占用token：

db = SQLDatabase.from_uri(
    dbstring,
    include_tables=['target_table_name']  # 替换为实际需要的表名
)

2. 限制查询结果行数 #

通过配置限制单次查询返回的行数，减少结果集的token占用：

db_chain = SQLDatabaseChain.from_llm(
    llm=llm,
    database=db,
    verbose=True,
    sql_database_kwargs={"max_rows": 50}  # 单次查询最多返回50行，可按需调整
)

3. 切换大上下文模型 #

改用支持更大token容量的模型，直接提升上下文处理能力：

llm = ChatOpenAI(model_name="gpt-3.5-turbo-16k",  # 支持16384 token
    temperature=0,
    max_tokens=1000,
    openai_api_key=API_KEY)

4. 精简Prompt模板 #

删除冗余描述，减少固定模板的token开销：

QUERY = """
根据问题生成合法PostgreSQL查询，执行后返回答案：
问题：{question}
SQL查询：
执行结果：
答案：
"""

5. 使用多步链式查询 #

用SQLDatabaseSequentialChain将复杂任务拆分为多步执行，每步仅传递必要上下文：

from langchain.chains import SQLDatabaseSequentialChain

db_chain = SQLDatabaseSequentialChain.from_llm(
    llm=llm,
    database=db,
    verbose=True,
    top_k=50  # 限制每步查询结果行数
)

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