RAGFlow部分文件下载功能实现

RAGFlow 目前已经有完整文件下载功能，通过 api/apps/document_app.py 中的 get 方法实现,这个方法允许用户通过文档 ID 获取整个文档的内容。但目前没有直接支持获取文件部分内容的 API。所以需要我们用户自己手动实现, 在此记录一下

实现方案

可以扩展现有的 API，添加一个新的端点，允许用户指定要获取的文件范围。这个功能对于大型文件特别有用，用户可能只需要文件的特定部分。

方案设计

• 创建一个新的 API 端点，接受文档 ID 和范围参数
• 根据文件类型和范围参数，返回文件的指定部分
• 支持不同类型的范围指定（页码范围、字节范围、块范围等）

代码实现示例

添加新端点的逻辑：

@manager.route('/get_partial/<doc_id>', methods=['GET'])  
def get_partial(doc_id):  
    try:  
        # 获取范围参数  
        start_pos = request.args.get('start', type=int, default=0)  
        end_pos = request.args.get('end', type=int)  # 如果不指定，则到文件末尾  
        range_type = request.args.get('range_type', default='bytes')  # 范围类型：bytes, pages, chunks  
          
        # 获取文档信息  
        e, doc = DocumentService.get_by_id(doc_id)  
        if not e:  
            return get_data_error_result(message="Document not found!")  
  
        # 获取文件存储地址  
        b, n = File2DocumentService.get_storage_address(doc_id=doc_id)  
          
        # 根据不同的范围类型处理  
        if range_type == 'bytes':  
            # 字节范围处理  
            file_content = STORAGE_IMPL.get_partial(b, n, start_pos, end_pos)  
        elif range_type == 'pages' and doc.type == FileType.PDF.value:  
            # PDF页面范围处理  
            file_content = get_pdf_pages(b, n, start_pos, end_pos)  
        elif range_type == 'chunks':  
            # 获取特定的块范围  
            chunks = get_document_chunks(doc_id, start_pos, end_pos)  
            return get_json_result(data=chunks)  
        else:  
            return get_data_error_result(message=f"Unsupported range type: {range_type} for file type: {doc.type}")  
          
        # 创建响应  
        response = flask.make_response(file_content)  
          
        # 设置适当的Content-Type  
        ext = re.search(r".([^.]+)$", doc.name)  
        if ext:  
            if doc.type == FileType.VISUAL.value:  
                response.headers.set('Content-Type', 'image/%s' % ext.group(1))  
            else:  
                response.headers.set('Content-Type', 'application/%s' % ext.group(1))  
                  
        # 设置Content-Range头  
        if range_type == 'bytes' and end_pos:  
            response.headers.set('Content-Range', f'bytes {start_pos}-{end_pos}')  
              
        return response  
    except Exception as e:  
        return server_error_response(e)  
  
# 辅助函数：获取PDF的特定页面  
def get_pdf_pages(bucket, name, start_page, end_page):  
    from io import BytesIO  
    import PyPDF2  
      
    # 获取完整PDF  
    pdf_data = STORAGE_IMPL.get(bucket, name)  
    pdf_stream = BytesIO(pdf_data)  
      
    # 创建PDF读取器和写入器  
    reader = PyPDF2.PdfReader(pdf_stream)  
    writer = PyPDF2.PdfWriter()  
      
    # 验证页码范围  
    max_page = len(reader.pages)  
    start_page = max(0, start_page)  
    end_page = min(max_page, end_page) if end_page is not None else max_page  
      
    # 添加指定范围的页面  
    for page_num in range(start_page, end_page):  
        writer.add_page(reader.pages[page_num])  
      
    # 写入新的PDF  
    output = BytesIO()  
    writer.write(output)  
    output.seek(0)  
      
    return output.getvalue()  
  
# 辅助函数：获取文档的特定块  
def get_document_chunks(doc_id, start_chunk, end_chunk):  
    # 获取文档的所有块  
    tenant_id = DocumentService.get_tenant_id(doc_id)  
    if not tenant_id:  
        return []  
      
    # 从文档存储中获取块  
    chunks = settings.docStoreConn.search(  
        {"doc_id": doc_id},   
        search.index_name(tenant_id),   
        None,  # kb_id  
        0,     # offset  
        1000,  # limit，可以根据需要调整  
        None,  # sort_by  
        None   # sort_order  
    )  
      
    # 验证块范围  
    total_chunks = len(chunks)  
    start_chunk = max(0, start_chunk)  
    end_chunk = min(total_chunks, end_chunk) if end_chunk is not None else total_chunks  
      
    # 返回指定范围的块  
    return chunks[start_chunk:end_chunk]

存储实现扩展

# 在存储实现类中添加获取部分内容的方法  
def get_partial(self, bucket, name, start_pos=0, end_pos=None):  
    """  
    获取对象的部分内容  
      
    Args:  
        bucket: 存储桶名称  
        name: 对象名称  
        start_pos: 起始位置（字节）  
        end_pos: 结束位置（字节），如果为None则到文件末尾  
      
    Returns:  
        部分文件内容的字节数据  
    """  
    try:  
        # 构建范围字符串  
        range_header = f'bytes={start_pos}-'  
        if end_pos is not None:  
            range_header = f'bytes={start_pos}-{end_pos}'  
              
        # 获取部分内容  
        response = self.client.get_object(  
            Bucket=bucket,  
            Key=name,  
            Range=range_header  
        )  
          
        # 读取内容  
        return response['Body'].read()  
    except Exception as e:  
        logging.error(f"Error getting partial object {bucket}/{name}: {e}")  
        raise

使用示例

HTTP API 调用示例

# 获取文档的前1000字节  
GET /api/document/get_partial/doc_123?start=0&end=999&range_type=bytes  
  
# 获取PDF文档的第2-5页  
GET /api/document/get_partial/doc_123?start=1&end=5&range_type=pages  
  
# 获取文档的第10-20个块  
GET /api/document/get_partial/doc_123?start=10&end=20&range_type=chunks  

优化方案

为了避免将整个文件加载到内存中，可以采用以下几种策略，根据不同的文件类型和范围类型进行优化：

1. 字节范围请求的优化

对于字节范围请求，我们可以直接利用底层存储系统的范围请求功能，这样就不需要先将整个文件加载到内存中：

def get_partial(self, bucket, name, start_pos=0, end_pos=None):  
    """  
    获取对象的部分内容，直接使用存储系统的范围请求功能  
    """  
    try:  
        # 构建范围字符串  
        range_header = f'bytes={start_pos}-'  
        if end_pos is not None:  
            range_header = f'bytes={start_pos}-{end_pos}'  
              
        # 使用MinIO的范围请求功能直接获取部分内容  
        response = self.client.get_object(  
            Bucket=bucket,  
            Key=name,  
            Range=range_header  
        )  
          
        # 返回响应流，而不是读取全部内容到内存  
        return response['Body']  
    except Exception as e:  
        logging.error(f"Error getting partial object {bucket}/{name}: {e}")  
        raise

2. PDF页面范围请求的优化

对于PDF文件，我们可以使用支持流式处理的库，如PyPDF2或pdfrw，并结合范围请求来只加载需要的页面：

def get_pdf_pages_streaming(bucket, name, start_page, end_page):  
    """  
    流式获取PDF的特定页面，避免将整个PDF加载到内存  
    """  
    from io import BytesIO  
    import PyPDF2  
      
    # 创建一个新的PDF写入器  
    writer = PyPDF2.PdfWriter()  
      
    # 使用RAGFlow现有的PDF解析器，它支持页面级别的处理  
    pdf_parser = RAGFlowPdfParser()  
      
    # 验证页码范围  
    total_pages = pdf_parser.total_page_number(name, STORAGE_IMPL.get_head(bucket, name, 1024))  
    start_page = max(0, start_page)  
    end_page = min(total_pages, end_page) if end_page is not None else total_pages  
      
    # 逐页处理，每次只加载一页到内存  
    for page_num in range(start_page, end_page):  
        # 计算当前页面在文件中的大致位置（这需要文件的索引信息）  
        page_data = pdf_parser.get_page(bucket, name, page_num)  
          
        # 从页面数据创建PDF页面对象  
        reader = PyPDF2.PdfReader(BytesIO(page_data))  
        if len(reader.pages) > 0:  
            writer.add_page(reader.pages[0])  
      
    # 写入新的PDF  
    output = BytesIO()  
    writer.write(output)  
    output.seek(0)  
      
    return output

3. 块范围请求的优化

对于文档块，我们可以使用分页查询来避免一次性加载所有块：

def get_document_chunks_paged(doc_id, start_chunk, end_chunk, page_size=100):  
    """  
    分页获取文档的特定块，避免一次性加载所有块  
    """  
    # 获取文档的租户ID  
    tenant_id = DocumentService.get_tenant_id(doc_id)  
    if not tenant_id:  
        return []  
      
    # 计算需要获取的块数量  
    chunk_count = end_chunk - start_chunk if end_chunk is not None else None  
      
    # 如果没有指定结束块，先获取总块数  
    if chunk_count is None:  
        total_count = settings.docStoreConn.count(  
            {"doc_id": doc_id},   
            search.index_name(tenant_id)  
        )  
        chunk_count = total_count - start_chunk  
      
    # 分页获取块  
    chunks = []  
    current_offset = start_chunk  
    remaining = chunk_count  
      
    while remaining > 0:  
        # 计算当前页大小  
        current_page_size = min(page_size, remaining)  
          
        # 获取当前页的块  
        page_chunks = settings.docStoreConn.search(  
            {"doc_id": doc_id},   
            search.index_name(tenant_id),   
            None,  # kb_id  
            current_offset,  # offset  
            current_page_size,  # limit  
            None,  # sort_by  
            None   # sort_order  
        )  
          
        # 添加到结果中  
        chunks.extend(page_chunks)  
          
        # 更新偏移量和剩余数量  
        current_offset += len(page_chunks)  
        remaining -= len(page_chunks)  
          
        # 如果返回的块数少于请求的数量，说明已经到达末尾  
        if len(page_chunks) < current_page_size:  
            break  
      
    return chunks

API端点实现

结合上述优化策略，我们可以修改API端点实现：

@manager.route('/get_partial/<doc_id>', methods=['GET'])  
def get_partial(doc_id):  
    try:  
        # 获取范围参数  
        start_pos = request.args.get('start', type=int, default=0)  
        end_pos = request.args.get('end', type=int)  # 如果不指定，则到文件末尾  
        range_type = request.args.get('range_type', default='bytes')  # 范围类型：bytes, pages, chunks  
          
        # 获取文档信息  
        e, doc = DocumentService.get_by_id(doc_id)  
        if not e:  
            return get_data_error_result(message="Document not found!")  
  
        # 获取文件存储地址  
        b, n = File2DocumentService.get_storage_address(doc_id=doc_id)  
          
        # 根据不同的范围类型处理  
        if range_type == 'bytes':  
            # 字节范围处理 - 流式返回  
            file_stream = STORAGE_IMPL.get_partial(b, n, start_pos, end_pos)  
              
            # 创建流式响应  
            response = flask.Response(  
                stream_with_context(iter_file(file_stream)),  
                mimetype='application/octet-stream'  
            )  
              
            # 设置适当的Content-Type  
            ext = re.search(r".([^.]+)$", doc.name)  
            if ext:  
                if doc.type == FileType.VISUAL.value:  
                    response.headers.set('Content-Type', f'image/{ext.group(1)}')  
                else:  
                    response.headers.set('Content-Type', f'application/{ext.group(1)}')  
                      
            # 设置Content-Range头  
            if end_pos:  
                response.headers.set('Content-Range', f'bytes {start_pos}-{end_pos}')  
                  
            return response  
              
        elif range_type == 'pages' and doc.type == FileType.PDF.value:  
            # PDF页面范围处理 - 流式处理  
            pdf_stream = get_pdf_pages_streaming(b, n, start_pos, end_pos)  
              
            # 创建流式响应  
            response = flask.Response(  
                stream_with_context(iter_file(pdf_stream)),  
                mimetype='application/pdf'  
            )  
              
            # 设置下载文件名  
            filename = f"{doc.name.split('.')[0]}_pages_{start_pos}-{end_pos}.pdf"  
            response.headers.set('Content-Disposition', f'attachment; filename="{filename}"')  
              
            return response  
              
        elif range_type == 'chunks':  
            # 获取特定的块范围 - 分页处理  
            chunks = get_document_chunks_paged(doc_id, start_pos, end_pos)  
            return get_json_result(data=chunks)  
              
        else:  
            return get_data_error_result(message=f"Unsupported range type: {range_type} for file type: {doc.type}")  
              
    except Exception as e:  
        return server_error_response(e)  
  
# 辅助函数：迭代文件流  
def iter_file(file_stream, chunk_size=8192):  
    """  
    迭代文件流，每次返回指定大小的块  
    """  
    while True:  
        chunk = file_stream.read(chunk_size)  
        if not chunk:  
            break  
        yield chunk

存储实现扩展

为了支持上述优化，我们需要在存储实现中添加一些辅助方法：

# 在MinIO存储实现中添加获取文件头部的方法  
def get_head(self, bucket, name, size=1024):  
    """  
    获取对象的头部数据  
      
    Args:  
        bucket: 存储桶名称  
        name: 对象名称  
        size: 要获取的字节数  
      
    Returns:  
        对象头部的字节数据  
    """  
    try:  
        response = self.client.get_object(  
            Bucket=bucket,  
            Key=name,  
            Range=f'bytes=0-{size-1}'  
        )  
        return response['Body'].read()  
    except Exception as e:  
        logging.error(f"Error getting object head {bucket}/{name}: {e}")  
        raise

针对PDF文件的特殊优化

对于PDF文件，我们可以利用RAGFlow现有的PDF解析器进行更高效的处理。从代码中可以看到，RAGFlowPdfParser类 pdf_parser.py:48-52 已经实现了PDF文件的解析功能。

我们可以扩展这个类，添加一个方法来获取特定页面的内容：

# 在RAGFlowPdfParser类中添加获取特定页面的方法  
def get_page(self, bucket, name, page_num):  
    """  
    获取PDF文件的特定页面  
      
    Args:  
        bucket: 存储桶名称  
        name: 对象名称  
        page_num: 页码（从0开始）  
      
    Returns:  
        页面的PDF数据  
    """  
    # 使用pdfplumber打开文件  
    with sys.modules[LOCK_KEY_pdfplumber]:  
        file_data = STORAGE_IMPL.get(bucket, name)  
        with pdfplumber.open(BytesIO(file_data)) as pdf:  
            if page_num < 0 or page_num >= len(pdf.pages):  
                raise ValueError(f"Page number {page_num} out of range (0-{len(pdf.pages)-1})")  
              
            # 提取单页  
            page = pdf.pages[page_num]  
              
            # 创建一个新的PDF，只包含这一页  
            from pikepdf import Pdf, Page  
            new_pdf = Pdf.new()  
            src_pdf = Pdf.open(BytesIO(file_data))  
            new_pdf.pages.append(Page(src_pdf.pages[page_num]))  
              
            # 保存到内存  
            output = BytesIO()  
            new_pdf.save(output)  
            output.seek(0)  
              
            return output.getvalue()

与任务执行系统的集成

从代码中可以看到，RAGFlow使用任务执行系统来处理文档 task_service.py:292-307 。我们可以扩展这个系统，使其支持部分文件处理：

# 在TaskService中添加创建部分处理任务的方法  
def queue_partial_tasks(doc, bucket, name, start, end, range_type, priority=0):  
    """  
    创建并排队处理文档部分内容的任务  
      
    Args:  
        doc: 文档字典  
        bucket: 存储桶名称  
        name: 文件名  
        start: 起始位置  
        end: 结束位置  
        range_type: 范围类型（bytes, pages, chunks）  
        priority: 任务优先级  
    """  
    task = {  
        "id": get_uuid(),  
        "doc_id": doc["id"],  
        "progress": 0.0,  
        "range_type": range_type,  
        "start": start,  
        "end": end  
    }  
      
    # 根据范围类型设置特定参数  
    if range_type == 'pages':  
        task["from_page"] = start  
        task["to_page"] = end  
      
    # 计算任务摘要  
    chunking_config = DocumentService.get_chunking_config(doc["id"])  
    hasher = xxhash.xxh64()  
    for field in sorted(chunking_config.keys()):  
        if field == "parser_config":  
            for k in ["raptor", "graphrag"]:  
                if k in chunking_config[field]:  
                    del chunking_config[field][k]  
        hasher.update(str(chunking_config[field]).encode("utf-8"))  
    for field in ["doc_id", "range_type", "start", "end"]:  
        hasher.update(str(task.get(field, "")).encode("utf-8"))  
    task["digest"] = hasher.hexdigest()  
    task["priority"] = priority  
      
    # 插入任务到数据库  
    bulk_insert_into_db(Task, [task], True)  
      
    # 将任务加入队列  
    REDIS_CONN.queue_product(  
        get_svr_queue_name(priority), message=task  
    )  
      
    return task

完整的API实现

结合上述所有优化，以下是完整的API实现：

@manager.route('/get_partial/<doc_id>', methods=['GET'])  
def get_partial(doc_id):  
    """  
    获取文档的部分内容  
      
    参数:  
        doc_id: 文档ID  
        start: 起始位置（默认为0）  
        end: 结束位置（如果不指定，则到文件末尾）  
        range_type