向RAGFlow中上传文档到对应的知识库

# -*- coding: utf-8 -*-
"""
Created on Sun Oct 12 21:09:04 2025

@author: 18309
"""

import requests
import json
import os
import logging

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s'
)

# 全局配置 - 请根据实际情况修改
RAGFLOW_SERVER = "http://IP:端口"  # Ragflow服务器地址
API_KEY = "ragflow-U5N************Mm"             # API密钥
HEADERS = {
    "Authorization": f"Bearer {API_KEY}",
    "Content-Type": "application/json"
}

import os
## 获取一个文件夹下的子文件夹和文件
def get_all_filenames(folder_path):
    """
    获取指定文件夹下的所有文件名称
    
    参数:
        folder_path: 目标文件夹路径
        
    返回:
        字典，包含文件列表和子文件夹列表
    """
    # 检查文件夹是否存在
    if not os.path.exists(folder_path):
        print(f"错误：文件夹 '{folder_path}' 不存在")
        return None
    
    # 检查是否为有效文件夹
    if not os.path.isdir(folder_path):
        print(f"错误：'{folder_path}' 不是一个文件夹")
        return None
    
    # 初始化存储列表
    file_names = []
    subfolder_names = []
    
    # 遍历文件夹内容
    for entry in os.listdir(folder_path):
        entry_full_path = os.path.join(folder_path, entry)
        if os.path.isfile(entry_full_path):
            file_names.append(entry)
        elif os.path.isdir(entry_full_path):
            subfolder_names.append(entry)
    
    return {
        "files": file_names,
        "subfolders": subfolder_names
    }

## 获取知识库
def get_knowledge_bases():
    """获取所有知识库列表"""
    url = f"{RAGFLOW_SERVER}/api/v1/datasets"
    
    try:
        response = requests.get(url, headers=HEADERS)
        response.raise_for_status()
        
        result = response.json()
        # 处理不同可能的响应结构
        if isinstance(result, list):
            return result
        elif isinstance(result, dict) and "data" in result:
            return result["data"]
        else:
            logging.error("知识库列表响应结构不符合预期")
            return []
            
    except requests.exceptions.RequestException as e:
        logging.error(f"获取知识库失败: {str(e)}")
        return []
## 创建知识库
def create_knowledge_base(name, description=""):
    """创建新的知识库，返回知识库ID或None"""
    url = f"{RAGFLOW_SERVER}/api/v1/datasets"
    
    payload = {
        "name": name,
        "description": description
    }
    
    try:
        response = requests.post(
            url,
            headers=HEADERS,
            data=json.dumps(payload)
        )
        
        if response.status_code == 200:
            result = response.json()
            # 检查可能的ID位置
            if "id" in result:
                return result["id"]
            elif "data" in result and "id" in result["data"]:
                return result["data"]["id"]
            else:
                logging.error("创建知识库响应结构异常")
                return None
        elif response.status_code == 409:
            logging.error(f"知识库 '{name}' 已存在")
            return None
        else:
            logging.error(f"创建失败，状态码: {response.status_code}, 响应: {response.text}")
            return None
            
    except requests.exceptions.RequestException as e:
        logging.error(f"创建知识库请求异常: {str(e)}")
        return None

## 上传文件
def upload_and_parse_file(kb_id, file_path, language="zh"):
    """上传文件到指定知识库并解析，返回文档ID或None"""
    # 检查文件是否存在
    if not os.path.exists(file_path):
        logging.error(f"文件不存在: {file_path}")
        return None
        
    if not os.path.isfile(file_path):
        logging.error(f"{file_path} 不是有效的文件")
        return None
    
    url = f"{RAGFLOW_SERVER}/api/v1/datasets/{kb_id}/documents"
    
    # 上传文件使用单独的请求头（不需要Content-Type）
    upload_headers = {
        "Authorization": f"Bearer {API_KEY}"
    }
    
    # 准备文件数据
    try:
        with open(file_path, 'rb') as f:
            files = {
                "file": (os.path.basename(file_path), f)
            }
            
            data = {
                "language": language
            }
            
            response = requests.post(
                url,
                headers=upload_headers,
                files=files,
                data=data
            )
            
            response.raise_for_status()
            result = response.json()
            
            # 处理不同的响应结构
            if "id" in result:
                return result["id"]
            elif "data" in result and isinstance(result["data"], list) and len(result["data"]) > 0:
                return result["data"][0]["id"]
            else:
                logging.error("文件上传响应结构不符合预期")
                return None
                
    except requests.exceptions.RequestException as e:
        logging.error(f"文件上传失败: {str(e)}")
        return None
    except Exception as e:
        logging.error(f"处理文件时出错: {str(e)}")
        return None

# 主执行逻辑
if __name__ == "__main__":
    ## 0 准备待上传的文件
    target_folder  = "D:/WORK/目标文件"      # ← 本地文件
    # 获取文件和子文件夹名称
    result = get_all_filenames(target_folder)
    ### 文件名列表
    file_lst = result['files']
    # 1. 获取现有知识库
    logging.info("获取所有知识库...")
    knowledge_bases = get_knowledge_bases()
    
    if knowledge_bases:
        logging.info(f"发现 {len(knowledge_bases)} 个知识库:")
        for kb in knowledge_bases:
            logging.info(f" - 名称: {kb.get('name')}, ID: {kb.get('id')}")
    else:
        logging.info("没有找到现有知识库")
    
    # 2. 准备目标知识库
    target_kb_name = "【成果库】-test"
    target_kb_id = None
    
    # 检查是否已存在
    if knowledge_bases:
        for kb in knowledge_bases:
            if kb.get('name') == target_kb_name:
                target_kb_id = kb.get('id')
                break
    
    # 如果不存在则创建
    if not target_kb_id:
        logging.info(f"创建新知识库: {target_kb_name}")
        target_kb_id = create_knowledge_base(
            name=target_kb_name,
            description="存储科技项目历年成果"
        )
        
        if not target_kb_id:
            logging.error("无法创建目标知识库，程序退出")
            exit(1)
        else:
            logging.info(f"知识库创建成功，ID: {target_kb_id}")
    
    # 3. 上传文件并解析
    file_full_path = []
    for file in file_lst:
        file_path = target_folder + '/' + file
        file_full_path.append(file_path)
    files_to_upload = file_full_path
    
    for file_path in files_to_upload:
        if os.path.exists(file_path):
            logging.info(f"正在上传并解析: {file_path}")
            doc_id = upload_and_parse_file(target_kb_id, file_path)
            
            if doc_id:
                logging.info(f"文件处理成功，文档ID: {doc_id}")
            else:
                logging.warning(f"文件 {file_path} 处理失败")
        else:
            logging.warning(f"文件 {file_path} 不存在，跳过")
    
    logging.info("所有操作完成")