如何优化OpenAI-ChatGPT Python脚本的对话记忆存储功能？

原代码核心问题 #

• 记忆读取逻辑混乱：遍历目录所有文件（含非文本文件），重复打开文件可能引发IO错误，还会读取无关内容干扰模型
• 对话上下文不完整：只存储模型回复，不记录用户提问，模型无法理解完整对话流程
• Prompt超限导致输出中断：text-davinci-003最大支持4097token，多次交互后记忆内容持续累加，很快超过上限触发模型停止响应
• 全局目录切换风险：os.chdir修改全局工作目录，可能影响后续文件操作
• 缺乏错误处理：文件读取、API调用失败时无容错机制，直接崩溃
• 模型选型过时：text-davinci-003已被GPT-3.5-turbo替代，后者更高效、成本更低，原生支持对话格式

优化方案 #

1. 规范记忆存储格式 #

采用单文件存储完整对话，每条对话明确区分用户和助手角色，格式示例：

用户: Welcome Duncan. I am so happy to meet you.
助手: 很高兴认识你！
用户: Hi Duncan! I am testing your talking function. What memories can you read?
助手: 我能读取我们之前所有的对话记录...

2. 优化记忆读取逻辑 #

仅读取指定的对话存储文件，跳过目录中其他无关文件，避免重复打开文件操作。

3. 实现Token超限处理 #

使用tiktoken库计算Prompt的token数，当接近模型上限时，自动裁剪最早的对话内容，保证Prompt在有效范围内。

4. 替换为GPT-3.5-turbo模型 #

使用更适合对话场景的ChatCompletion接口，原生支持多轮对话上下文传递，效率和成本更优。

5. 移除全局目录切换 #

通过绝对路径操作文件，避免修改全局工作目录。

6. 添加错误处理 #

捕获文件IO异常、API调用异常，输出错误信息避免程序崩溃。

优化后完整代码 #

import openai
import os
import tiktoken

# 配置API密钥
openAIKey = '你的API密钥'
openai.api_key = openAIKey

# 模型配置
MODEL = "gpt-3.5-turbo"
MAX_TOKENS = 4096  # 模型最大token限制
RESERVE_TOKENS = 1000  # 预留token给新提问和模型回复

def count_tokens(text, model=MODEL):
    """计算文本的token数量"""
    encoding = tiktoken.encoding_for_model(model)
    return len(encoding.encode(text))

def load_memory(memory_file):
    """加载对话记忆，返回完整对话文本"""
    if not os.path.exists(memory_file):
        return ""
    try:
        with open(memory_file, 'r', encoding='utf-8') as f:
            return f.read()
    except Exception as e:
        print(f"读取记忆文件失败: {e}")
        return ""

def trim_memory(memory_text, model=MODEL):
    """当记忆内容token数超过阈值时，裁剪最早的对话"""
    total_tokens = count_tokens(memory_text)
    if total_tokens <= MAX_TOKENS - RESERVE_TOKENS:
        return memory_text
    
    # 按对话分割，从最早的开始删除
    dialogues = memory_text.split("\n\n")
    trimmed_dialogues = []
    current_tokens = 0
    
    # 从最新对话开始累加，直到接近阈值
    for dialogue in reversed(dialogues):
        dialogue_tokens = count_tokens(dialogue)
        if current_tokens + dialogue_tokens <= MAX_TOKENS - RESERVE_TOKENS:
            trimmed_dialogues.insert(0, dialogue)
            current_tokens += dialogue_tokens
        else:
            break
    
    return "\n\n".join(trimmed_dialogues)

def save_memory(memory_file, user_input, assistant_response):
    """将新的对话对写入记忆文件"""
    dialogue = f"用户: {user_input}\n助手: {assistant_response}"
    try:
        with open(memory_file, 'a', encoding='utf-8') as f:
            f.write(dialogue + "\n\n")
    except Exception as e:
        print(f"保存记忆失败: {e}")

def talk(prompt, n=1, temperature=0.5, memory=True, record_memory=True, 
         memory_address=r'G:\My Drive\0- Personal\07- Duncan\memory', verbose=True):
    """核心对话函数"""
    memory_file = os.path.join(memory_address, "dialogue_memory.txt")
    memory_txt = ""
    
    if memory:
        memory_txt = load_memory(memory_file)
        # 裁剪过长的记忆
        memory_txt = trim_memory(memory_txt)
    
    # 构建ChatCompletion的消息列表
    messages = []
    if memory_txt:
        # 将历史对话拆分为消息对
        dialogues = memory_txt.strip().split("\n\n")
        for d in dialogues:
            parts = d.split("\n")
            if len(parts) >=2:
                user_msg = parts[0].replace("用户: ", "")
                assistant_msg = parts[1].replace("助手: ", "")
                messages.append({"role": "user", "content": user_msg})
                messages.append({"role": "assistant", "content": assistant_msg})
    
    # 添加当前用户提问
    messages.append({"role": "user", "content": prompt})
    
    response_array = []
    try:
        for _ in range(n):
            completion = openai.ChatCompletion.create(
                model=MODEL,
                messages=messages,
                temperature=temperature,
                max_tokens=RESERVE_TOKENS // 2  # 限制单次回复长度
            )
            response = completion.choices[0].message['content'].strip()
            response_array.append(response)
            
            if verbose:
                print(f"\n助手: {response}\n" + "-"*80)
            
            if record_memory:
                save_memory(memory_file, prompt, response)
    
    except Exception as e:
        print(f"API调用失败: {e}")
        return []
    
    return response_array

if __name__ == '__main__':
    memory = True
    record_memory = True
    temperature = 0.5
    n = 1
    verbose = True
    memory_address = r'G:\My Drive\0- Personal\07- Duncan\memory'
    prompt = 'Welcome Duncan. I am so happy to meet you.\nHi Duncan! I am testing your talking function. What memories can you read?'
    
    ra = talk(prompt=prompt, n=n, temperature=temperature, memory=memory, 
              memory_address=memory_address, verbose=verbose)

额外说明 #

• 需先安装依赖库：pip install openai tiktoken
• 建议将API密钥通过环境变量传递（如os.getenv("OPENAI_API_KEY")），避免硬编码泄露
• 可根据需求调整RESERVE_TOKENS参数，平衡历史对话长度和单次回复空间

内容的提问来源于stack exchange，提问作者Mario del Pino