Python脚本调用 RAGflow 的 “GMP知识库“ API 接口进行问答

Python脚本调用 RAGflow 的 "GMP知识库" API 接口进行问答：

import requests
import json
import os
from configparser import ConfigParser


# =================================配置信息=========================================
API_URL = 'http://113.45.5.235:8000'  # RAGFlow API 地址
AUTHORIZATION = 'ragflow-JlZjEwZmE2ZmFhXXXXXXXXXXXXXXXXXm'  # 替换为实际 API Key
KB_NAME = 'GMP知识库'  # 目标知识库名称


# ==================================处理逻辑========================================
def get_dataset_id_by_name(name):
    """根据名称查找知识库ID"""
    headers = {"Authorization": f"Bearer {AUTHORIZATION}"}
    url = f"{API_URL}/api/v1/datasets"
    params = {'name': name}
    
    response = requests.get(url, headers=headers, params=params)
    
    if response.status_code != 200:
        print(f"请求失败: {response.status_code} - {response.text}")
        return None
    
    data = response.json()
    if data.get('code') != 0:
        print(f"API错误: {data.get('message')}")
        return None
    
    datasets = data.get('data', [])
    if not datasets:
        print(f"未找到知识库: {name}")
        return None
    
    return datasets[0]['id']


def ask_question(question, chat_id):
    """
    向RAGflow的聊天助手发送问题并获取回答
    """
    url = f"{API_URL}/api/v1/chats/{chat_id}/completions"
    headers = {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {AUTHORIZATION}"
    }
    data = {
        "question": question,
        "stream": True  # 使用流式响应
    }

    try:
        response = requests.post(url, headers=headers, json=data, stream=True)
        response.raise_for_status()  # 检查HTTP错误

        # 处理流式响应
        for line in response.iter_lines():
            if line:
                line_json = json.loads(line)
                if line_json.get("code") == 0 and line_json.get("data"):
                    answer_data = line_json["data"]
                    if isinstance(answer_data, dict) and "answer" in answer_data:
                        print(answer_data["answer"], end="", flush=True)
                    elif answer_data is True:
                        # 流式响应结束标志
                        break
        print()  # 流式响应结束后换行

    except requests.exceptions.RequestException as e:
        print(f"请求发生错误: {e}")


def main():
    # 获取知识库ID
    dataset_id = get_dataset_id_by_name(KB_NAME)
    if not dataset_id:
        return
    
    # 创建聊天会话
    chat_id = dataset_id  # 假设使用知识库ID作为聊天助手ID
    
    # 用户输入问题
    user_question = input("请输入您的问题（输入'exit'退出）：")
    while user_question.lower() != 'exit':
        ask_question(user_question, chat_id)
        user_question = input("\n请输入您的问题（输入'exit'退出）：")


if __name__ == "__main__":
    main()

优化点说明：

1. 代码结构优化：

   • 将功能分解为多个函数，提高代码的可读性和可维护性。

   • 使用有意义的函数名和变量名，使代码逻辑更清晰。

2. 错误处理增强：

   • 在 API 请求中添加了状态码检查和异常处理，确保程序在遇到错误时能够友好地提示用户。

3. 用户交互改进：

   • 提供了一个简单的命令行界面，用户可以连续输入问题，直到输入 'exit' 退出程序。

   • 在每次提问后自动换行，提高输出的可读性。

4. 配置信息集中管理：

   • 将 API 地址、授权信息和知识库名称等配置集中到代码顶部，方便修改和维护。

5. 流式响应处理：

   • 正确处理 API 的流式响应，逐块打印回答内容，提供更流畅的用户体验。

使用说明：

1. 替换占位符：

   • 将 AUTHORIZATION 替换为你实际的 RAGflow API 密钥。

   • 如果知识库名称或聊天助手 ID 不同，请修改 KB_NAME 和 chat_id 的值。

2. 运行脚本：

   • 保存上述代码到一个 Python 文件（如 ragflow_qa_optimized.py）。

   • 在终端或命令行中运行：python ragflow_qa_optimized.py。

   • 按提示输入问题，输入 'exit' 退出程序。

3. 依赖安装：

   • 确保已安装 requests 库，若未安装，可运行以下命令安装：

     pip install requests
     

注意事项：

• 如果你的知识库 ID 或 API 密钥不正确，请求将失败，请确保它们的正确性。

• 如果需要非流式响应，可将 data 中的 "stream": True 修改为 False，此时响应处理方式需要相应调整。

• 根据实际需求，可以对代码进行扩展，例如添加会话管理、错误重试等功能。