基于Cognee实现PDF图文并茂解析与检索的实践方案

在各类文档处理场景中，文本与图表结合的内容十分常见，单纯的文本解析无法满足完整的信息提取需求。基于OCR/PDF解析接口和Cognee的API规范，可搭建一套通用的PDF图文解析与检索方案，以下是详细的实践流程。

一、核心流程梳理

这套方案的核心是将PDF中的文字和图表分离后重新建立关联，最终实现精准的图文检索，核心步骤如下：

调用OCR/PDF解析接口，提取PDF中的文字内容和各类图表图片资源；
解压解析结果包，得到包含图片相对路径的Markdown文档和独立的图片文件夹；
将提取的图片上传至图片服务器，获取线上访问地址，同时替换Markdown文档中的本地图片路径；
调用Cognee的系列接口，完成「创建数据集→上传MD文件→构建知识图谱→检索文档信息」的全链路操作；
将检索结果渲染为HTML页面，实现文字数据与对应图表的同步展示。

二、环境准备

开始编码前，需安装必备的Python依赖库，用于处理HTTP请求、文件解压和Markdown转HTML等操作：

Bash 复制代码

pip install requests python-dotenv markdown

三、具体实现步骤

3.1 调用OCR/PDF解析接口提取文档内容

按照curl命令的参数格式编写代码，确保请求体和命令行调用完全一致，以此精准解析含图表的PDF文件：

Python 复制代码

import requests
import os
from zipfile import ZipFile
import time

# 全局配置（根据自身环境修改）
OCR_SERVER_URL = "https://your_ocr_server/ocr/file_parse"
PDF_FILE_PATH = "sample_with_images.pdf"  # 通用含图片PDF文件
OUTPUT_ZIP_PATH = "parsed_results.zip"

# Cognee配置（从登录接口获取Token）
COGNEE_BASE_URL = "https://your-cognee-server/api/v1"
COGNEE_TOKEN = "your-bearer-token"
COGNEE_HEADERS = {
    "Authorization": f"Bearer {COGNEE_TOKEN}",
    "Content-Type": "application/json"
}

def call_ocr_pdf_parse_api():
    """调用OCR/PDF解析接口，提取PDF的文字和图片"""
    if not os.path.exists(PDF_FILE_PATH):
        print(f"错误：PDF文件 {PDF_FILE_PATH} 不存在")
        return False

    # 严格对齐curl的-F参数构造form_data
    form_data = {
        "return_middle_json": (None, "false"),
        "return_model_output": (None, "false"),
        "return_md": (None, "true"),
        "return_images": (None, "true"),
        "end_page_id": (None, "99999"),
        "parse_method": (None, "auto"),
        "start_page_id": (None, "0"),
        "lang_list": (None, "ch"),
        "output_dir": (None, "./output"),
        "server_url": (None, "string"),
        "return_content_list": (None, "false"),
        "backend": (None, "pipeline"),
        "table_enable": (None, "true"),
        "response_format_zip": (None, "true"),
        "formula_enable": (None, "true"),
        "files": (
            os.path.basename(PDF_FILE_PATH),
            open(PDF_FILE_PATH, "rb"),
            "application/pdf"
        )
    }

    try:
        response = requests.post(
            url=OCR_SERVER_URL,
            headers={"accept": "application/json"},
            files=form_data,
            timeout=600
        )
        response.raise_for_status()

        with open(OUTPUT_ZIP_PATH, "wb") as f:
            f.write(response.content)
        print(f"PDF解析完成，结果压缩包已保存至：{OUTPUT_ZIP_PATH}")
        return True
    except requests.exceptions.HTTPError as e:
        print(f"OCR接口调用失败：{e.response.status_code} - {e.response.text}")
        return False
    except Exception as e:
        print(f"OCR接口调用失败：{str(e)}")
        return False
    finally:
        if "files" in form_data:
            form_data["files"][1].close()

3.2 解压结果包，提取MD文件与图片

解析后的压缩包包含文档文字内容和各类图表图片，通过以下代码可快速定位并提取关键文件：

Python 复制代码

def unzip_results(zip_path, target_dir="./parsed_result"):
    """解压结果压缩包，返回MD文件路径和图片目录路径"""
    if not os.path.exists(zip_path):
        print(f"错误：压缩包 {zip_path} 不存在")
        return None, None

    if not os.path.exists(target_dir):
        os.makedirs(target_dir)

    try:
        with ZipFile(zip_path, "r") as zip_ref:
            zip_ref.extractall(target_dir)
        print(f"压缩包已解压至：{target_dir}")

        md_file_path = None
        images_dir = None
        for root, dirs, files in os.walk(target_dir):
            for file in files:
                if file.endswith(".md") and os.path.basename(PDF_FILE_PATH).split(".")[0] in file:
                    md_file_path = os.path.join(root, file)
            if "images" in dirs:
                images_dir = os.path.join(root, "images")
            if md_file_path and images_dir:
                break

        return md_file_path, images_dir
    except Exception as e:
        print(f"解压失败：{str(e)}")
        return None, None

3.3 上传图表图片并替换MD文档路径

为实现Markdown文档中图片的在线展示，需将解析出的各类图表上传至图片服务器，同时替换文档中的本地路径：

Python 复制代码

def upload_images_to_server(images_dir, image_server_url="https://your_images_server/upload"):
    """上传文档图片到服务器，返回{本地文件名: 线上地址}映射"""
    image_url_map = {}
    if not os.path.exists(images_dir):
        return image_url_map

    for img_file in os.listdir(images_dir):
        if not img_file.endswith(".jpg"):
            continue
        img_path = os.path.join(images_dir, img_file)
        try:
            with open(img_path, "rb") as f:
                files = {"file": (img_file, f, "image/jpeg")}
                response = requests.post(image_server_url, files=files, timeout=30)
            response.raise_for_status()
            img_url = response.json().get("url")
            if img_url:
                image_url_map[img_file] = img_url
                print(f"图片 {img_file} 上传成功，线上地址：{img_url}")
        except Exception as e:
            print(f"图片 {img_file} 上传失败：{str(e)}")
    return image_url_map

def replace_md_image_path(md_file_path, image_url_map):
    """替换MD文档中的图片相对路径为线上地址"""
    if not md_file_path or not image_url_map:
        return None, None

    with open(md_file_path, "r", encoding="utf-8") as f:
        md_content = f.read()

    for img_file, img_url in image_url_map.items():
        old_path = f"./images/{img_file}"
        md_content = md_content.replace(old_path, img_url)

    new_md_path = md_file_path.replace(".md", "_online.md")
    with open(new_md_path, "w", encoding="utf-8") as f:
        f.write(md_content)
    print(f"替换路径后的MD文档已保存至：{new_md_path}")
    return new_md_path, md_content

3.4 基于Cognee构建知识图谱并实现检索

遵循Cognee API规范，通过一系列接口操作，可构建包含文档文本、图片及其关联关系的知识图谱，最终实现精准检索。

3.4.1 创建文档专属数据集

在Cognee中，数据集是组织文档的基本单元，为解析后的文档创建独立数据集，可方便后续管理和检索：

Python 复制代码

def create_cognee_dataset(dataset_name="pdf-documents-dataset"):
    """创建文档专属数据集"""
    url = f"{COGNEE_BASE_URL}/datasets"
    payload = {
        "dataset_data": {
            "name": dataset_name
        }
    }
    try:
        response = requests.post(url, headers=COGNEE_HEADERS, json=payload, timeout=30)
        response.raise_for_status()
        dataset_info = response.json()
        dataset_id = dataset_info.get("id")
        print(f"数据集创建成功，ID：{dataset_id}")
        return dataset_id
    except Exception as e:
        print(f"创建数据集失败：{str(e)}")
        return None

3.4.2 上传MD文档到Cognee数据集

使用Cognee的/api/v1/add接口，将替换好图片路径的MD文档上传到指定数据集：

Python 复制代码

def upload_md_to_cognee(dataset_id, md_file_path):
    """上传MD文档到指定Cognee数据集"""
    if not dataset_id or not os.path.exists(md_file_path):
        return None

    url = f"{COGNEE_BASE_URL}/add"
    with open(md_file_path, "rb") as f:
        files = {
            "data": (os.path.basename(md_file_path), f, "text/markdown"),
            "datasetId": (None, dataset_id)
        }
        try:
            response = requests.post(
                url,
                headers={"Authorization": COGNEE_HEADERS["Authorization"]},
                files=files,
                timeout=30
            )
            response.raise_for_status()
            print("MD文档上传成功")
            return response.json()
        except Exception as e:
            print(f"上传MD文件失败：{str(e)}")
            return None

3.4.3 调用cognify接口构建通用知识图谱

这是生成包含文档文本、图片及其关联关系的知识图谱和向量库的核心步骤，使用通用配置即可适配各类文档，建议设置异步执行避免超时：

Python 复制代码

def call_cognee_cognify(dataset_id):
    """调用cognify接口，构建文档文本与图片的关联知识图谱"""
    if not dataset_id:
        return False

    url = f"{COGNEE_BASE_URL}/cognify"
    payload = {
        "dataset_ids": [dataset_id],
        "run_in_background": True  # 异步执行适配大文件处理
        # 通用场景下无需自定义prompt，默认配置即可完成文本与图片的关联提取
    }
    try:
        response = requests.post(url, headers=COGNEE_HEADERS, json=payload, timeout=30)
        response.raise_for_status()
        task_info = response.json()
        print(f"cognify任务启动成功，任务ID：{task_info.get('task_id')}")
        return True
    except Exception as e:
        print(f"调用cognify失败：{str(e)}")
        return False

3.4.4 查询数据集处理状态

异步执行cognify后，可通过以下接口轮询查询文档的处理进度：

Python 复制代码

def check_cognee_dataset_status(dataset_id):
    """查询数据集的cognify处理状态"""
    url = f"{COGNEE_BASE_URL}/datasets/status"
    params = {"dataset": dataset_id}
    try:
        response = requests.get(url, headers=COGNEE_HEADERS, params=params, timeout=30)
        response.raise_for_status()
        status_info = response.json()
        current_status = status_info.get("statuses")[0].get("status")
        print(f"当前数据集状态：{current_status}")
        return current_status
    except Exception as e:
        print(f"查询数据集状态失败：{str(e)}")
        return None

3.4.5 检索文档信息（以基金相关问题为例）

基于构建好的知识图谱，可检索各类文档相关问题，以下以基金净值与规模的查询为例，展示图文并茂的检索效果：

Python 复制代码

def search_cognee_doc_info(dataset_id, query="某基金的净值与规模是怎样的？", top_k=5):
    """调用search接口检索文档信息，返回含图片的结果"""
    url = f"{COGNEE_BASE_URL}/search"
    payload = {
        "search_type": "GRAPH_COMPLETION",  # 基于知识图谱实现图文关联检索
        "dataset_ids": [dataset_id],
        "query": query,
        "top_k": top_k
    }
    try:
        response = requests.post(url, headers=COGNEE_HEADERS, json=payload, timeout=30)
        response.raise_for_status()
        return response.json()
    except Exception as e:
        print(f"检索文档信息失败：{str(e)}")
        return None

3.5 完整流程串联与图文结果展示

将所有步骤串联，可实现从PDF解析到图文并茂检索结果展示的全流程，同时将结果转为HTML方便直接查看：

Python 复制代码

import markdown

def render_result_to_html(search_result, output_html="doc_result.html"):
    """将检索结果渲染为HTML，实现图片直接展示"""
    if not search_result or not search_result.get("results"):
        print("无检索结果可渲染")
        return False

    # 拼接检索结果内容
    html_content = "<html><head><meta charset='utf-8'><title>文档检索结果</title></head><body>"
    html_content += f"<h1>查询问题：{search_result.get('query')}</h1>"
    for idx, res in enumerate(search_result.get("results", [])):
        html_content += f"<h2>结果 {idx+1}</h2>"
        # 将Markdown内容转为HTML，图片会自动加载线上地址
        html_content += markdown.markdown(res.get("content", "无相关内容"))
    html_content += "</body></html>"

    # 保存HTML文件
    with open(output_html, "w", encoding="utf-8") as f:
        f.write(html_content)
    print(f"图文结果已保存为HTML文件：{output_html}")
    print("直接打开该文件即可查看包含图表的文档检索信息")
    return True

if __name__ == "__main__":
    # 1. OCR解析PDF
    if not call_ocr_pdf_parse_api():
        exit(1)

    # 2. 解压文件
    md_file, images_dir = unzip_results(OUTPUT_ZIP_PATH)
    if not md_file or not images_dir:
        exit(1)

    # 3. 上传图片并替换MD路径
    image_url_map = upload_images_to_server(images_dir)
    new_md_file, _ = replace_md_image_path(md_file, image_url_map)
    if not new_md_file:
        exit(1)

    # 4. Cognee全流程
    dataset_id = create_cognee_dataset()
    if not dataset_id:
        exit(1)

    if upload_md_to_cognee(dataset_id, new_md_file):
        if call_cognee_cognify(dataset_id):
            # 轮询等待文档处理完成
            while True:
                current_status = check_cognee_dataset_status(dataset_id)
                if current_status == "DATASET_PROCESSING_COMPLETED":
                    break
                elif current_status is None:
                    print("查询状态失败，退出轮询")
                    exit(1)
                time.sleep(10)  # 每10秒查询一次
            
            # 执行检索（以基金相关问题为例，可替换为任意文档相关查询）
            doc_search_result = search_cognee_doc_info(dataset_id)
            if doc_search_result:
                # 渲染为HTML，直接查看图文结果
                render_result_to_html(doc_search_result)

四、关键说明

接口合规性
- 全程使用Cognee标准接口，核心依赖cognify接口构建文档知识图谱和向量库；
- search接口选择GRAPH_COMPLETION类型，可精准关联文档中的文字数据和对应的图表图片。
通用场景适配
- cognify接口采用默认配置即可适配各类含图文的文档，无需单独设置场景化提示词；
- 检索时仅需替换query参数，即可实现不同类型文档的图文关联查询，基金相关问题仅作为示例。
鉴权与异常处理
- 所有Cognee接口均携带Bearer Token鉴权，Token过期需重新登录获取；
- 增加轮询状态查询、文件存在性校验等逻辑，保证文档处理流程的稳定性。

五、总结

这套基于OCR解析接口和Cognee API规范的方案，是一套通用的PDF图文解析与检索解决方案。该方案可实现文字与图表的精准关联，输出的HTML文件可直接查看图文并茂的检索结果，适用于技术手册、行业报告、金融文档等各类含图文的文档处理场景，为文档的智能化处理提供了可落地的实践路径。