Streamlit调用OCI生成式AI Cohere模型多轮提示触发Token超限错误

使用Streamlit开发基于OCI Generative AI的Cohere模型聊天机器人，首次输入提示可正常运行，但第二次及后续输入会触发token超限错误。即使移除create_history_aware_retriever改为无记忆模式，错误依然存在。

错误信息 #

oci.exceptions.ServiceError: {'target_service': 'generative_ai_inference', 'status': 400, 'code': '400', 'opc-request-id': '62F4609E6A4349A3B312ED1D73380147/94883F1D68C1BD57118373ED41B882E3/A7E92505DFB5E7811684526204E71E51', 'message': '{"message":"too many tokens: total number of tokens in the prompt cannot exceed 4081 - received 20861. Try using a shorter prompt, or enabling prompt truncating. See https://docs.cohere.com/reference/generate for more details."}', 'operation_name': 'generate_text', 'timestamp': '2024-05-25T16:12:53.250556+00:00', 'client_version': 'Oracle-PythonSDK/2.126.4', 'request_endpoint': 'POST https://inference.generativeai.us-chicago-1.oci.oraclecloud.com/20231130/actions/generateText', 'logging_tips': 'To get more info on the failing request, refer to https://docs.oracle.com/en-us/iaas/tools/python/latest/logging.html for ways to log the request/response details.', 'troubleshooting_tips': "See https://docs.oracle.com/iaas/Content/API/References/apierrors.htm#apierrors_400__400_400 for more information about resolving this error. Also see https://docs.oracle.com/iaas/api/#/en/generative-ai-inference/20231130/GenerateTextResult/GenerateText for details on this operation's requirements. If you are unable to resolve this generative_ai_inference issue, please contact Oracle support and provide them this full error message."}

相关代码 #

import oci
import streamlit as st
from dotenv import load_dotenv
from langchain.chains import (
    create_retrieval_chain,
)
from langchain.chains.combine_documents import create_stuff_documents_chain
from langchain_community.embeddings import OCIGenAIEmbeddings

# """
# import fix worked:
# Before: from langchain_community.llms import OCIGenAI
# After: from langchain_community.llms.oci_generative_ai import OCIGenAI
# """
from langchain_community.llms.oci_generative_ai import OCIGenAI
from langchain_community.vectorstores import Chroma
from langchain_core.prompts import ChatPromptTemplate

import chromadb

if "history" not in st.session_state:
    st.session_state.history = []

load_dotenv()

llm = OCIGenAI(
    model_id="cohere.command-light",
    service_endpoint="https://inference.generativeai.us-chicago-1.oci.oraclecloud.com",
    compartment_id="ocid1.compartment.oc1..aaaaaaaacz66k7qusk5kg5wc4keajwvi2meiauw6wmyztmrb2tm6gt7tzqsa",
    model_kwargs={"max_tokens": 1000}
)

embeddings = OCIGenAIEmbeddings(
    model_id="cohere.embed-english-v3.0",
    service_endpoint="https://inference.generativeai.us-chicago-1.oci.oraclecloud.com",
    compartment_id="ocid1.compartment.oc1..aaaaaaaacz66k7qusk5kg5wc4keajwvi2meiauw6wmyztmrb2tm6gt7tzqsa",
)

client = chromadb.HttpClient(host="localhost", port=8000)
db = Chroma(client=client, embedding_function=embeddings, collection_name="ncert-eng-chromadb")
retv = db.as_retriever(search_type="similarity", search_kwargs={"k": 3})

from langchain.chains import create_history_aware_retriever
from langchain_core.prompts import MessagesPlaceholder
from langchain_core.messages import HumanMessage, AIMessage

prompt_template = ChatPromptTemplate.from_messages([
    ("system", "Answer the user's questions based on the below context:\n\n{context}"),
    MessagesPlaceholder(variable_name="chat_history"),
    ("user", "{input}"),
])

retriever_chain = create_history_aware_retriever(llm, retv, prompt_template)

chat_history = []

document_chain = create_stuff_documents_chain(llm, prompt_template)
chain = create_retrieval_chain(retriever_chain, document_chain)

st.title('🤖 Welcome to the ChatBot')
for msg in st.session_state.history:
    with st.chat_message(msg['role']):
        st.markdown(msg['content'])

prompt = st.chat_input("Say something")
if prompt:
    st.session_state.history.append({
        'role': 'user',
        'content': prompt
    })

    with st.chat_message("user"):
        st.markdown(prompt)

    with st.spinner('💡Thinking'):
        response = chain.invoke({
            "chat_history": chat_history,
            "input": prompt
        })

        chat_history = [HumanMessage(content=prompt), AIMessage(content=response["answer"])]

        st.session_state.history.append({
            'role': 'Assistant',
            'content': response["answer"]
        })

        with st.chat_message("Assistant"):
            st.markdown(response["answer"])

更新发现 #

问题疑似与本地Docker部署的ChromaDB中存储的NCERT CBSE 12年级英文课本章节数据有关：当询问与AI最新回复相关的内容时，检索到的文档包含大量换行符，导致提示token数暴增，需优化嵌入规则及RAG代码。

1. 清理ChromaDB中的冗余文档内容 #

对原始课本数据做预处理，移除冗余换行和空白字符，再重新生成嵌入写入ChromaDB：

import re
def clean_text(text):
    # 替换连续换行、制表符为单个换行
    text = re.sub(r'\n\s*\n', '\n', text)
    # 移除多余空格
    text = re.sub(r'\s+', ' ', text).strip()
    return text

2. 优化RAG链的token控制 #

• 减少检索文档数量：将retv = db.as_retriever(search_type="similarity", search_kwargs={"k": 3})中的k值调整为2或1，降低上下文总token数。
• 替换文档合并链：用create_map_reduce_documents_chain或create_refine_documents_chain替代create_stuff_documents_chain，避免一次性将所有检索文档塞入提示。
• 开启自动截断：在LLM初始化时添加truncate参数，让Cohere自动截断过长提示：

  llm = OCIGenAI(
      model_id="cohere.command-light",
      service_endpoint="https://inference.generativeai.us-chicago-1.oci.oraclecloud.com",
      compartment_id="ocid1.compartment.oc1..aaaaaaaacz66k7qusk5kg5wc4keajwvi2meiauw6wmyztmrb2tm6gt7tzqsa",
      model_kwargs={"max_tokens": 1000, "truncate": "END"}
  )
  

3. 修复会话历史管理 #

将chat_history存入Streamlit的session_state，同时限制历史消息长度，避免累积过多token：

# 初始化会话历史
if "chat_history" not in st.session_state:
    st.session_state.chat_history = []

# 调用链时使用session_state中的历史
response = chain.invoke({
    "chat_history": st.session_state.chat_history,
    "input": prompt
})

# 更新历史，仅保留最近3轮对话
st.session_state.chat_history.extend([HumanMessage(content=prompt), AIMessage(content=response["answer"])])
if len(st.session_state.chat_history) > 6:  # 3轮对话对应6条消息（用户+助理各3次）
    st.session_state.chat_history = st.session_state.chat_history[-6:]

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