feat:新增sci_doc_parsing 测试脚本

README.md

@@ -1,2 +1,17 @@
 # TriModalParser
 
+## 添加依赖
+```
+pip imstall uv
+```
+```
+uv pip install .
+```
+## 测试doc_parsing
+1. 如果src没有message文件夹,那么请使用如下命令进行proto编译
+```
+python -m grpc_tools.protoc -Iproto --python_out=src/message --grpc_python_out=src/message proto/common/v1/s3.proto
+python -m grpc_tools.protoc -Iproto --python_out=src/message --grpc_python_out=src/message proto/parse/v1/pdf_parser.proto
+```
+2. 运行tests/utils.py 上传需要解析的文件到s3，这一步需要保留key值
+3. 运行tests/test_doc_parsing.py, 修改main函数中send_task函数的参数，将任务推入队列中，等待任务完成结果输出

compose.yaml

@@ -0,0 +1,119 @@
+name: sci-assistant
+
+services:
+  minio:
+    image: minio/minio:${MINIO_IMAGE_TAG:-RELEASE.2025-04-08T15-41-24Z}
+    restart: always
+    ports:
+      - "9000:9000"
+      - "9002:9001"
+    healthcheck:
+      test: mc ready local
+      interval: 10s
+      timeout: 10s
+      retries: 3
+    command: server --console-address ":9001"
+    environment:
+      MINIO_ROOT_USER: ${MINIO_ACCESS_KEY:-minio}
+      MINIO_ROOT_PASSWORD: ${MINIO_SECRET_KEY:-minio123}
+      MINIO_VOLUMES: /data
+      MINIO_PROMETHEUS_AUTH_TYPE: public
+    volumes:
+      - ./minio/data:/data
+
+  minio-init-bucket:
+    image: minio/mc:${MC_IMAGE_TAG:-RELEASE.2025-04-08T15-39-49Z}
+    command: mb -p minio/${MINIO_BUCKET:-sci-assistant}
+    environment:
+      MC_HOST_minio: http://${MINIO_ACCESS_KEY:-minio}:${MINIO_SECRET_KEY:-minio123}@minio:9000
+    depends_on:
+      minio:
+        condition: service_healthy
+
+  rabbitmq:
+    image: rabbitmq:${RABBITMQ_IMAGE_TAG:-4.1-management}
+    restart: always
+    healthcheck:
+      test: rabbitmq-diagnostics check_port_connectivity
+      interval: 10s
+      timeout: 3s
+      retries: 3
+    ports:
+      - "5671:5672"
+      - "15671:15672"
+    volumes:
+      - ./rabbitmq/definitions.json:/etc/rabbitmq/definitions.json:ro
+      - ./rabbitmq/conf.d:/etc/rabbitmq/conf.d:ro
+      - ./rabbitmq/data:/var/lib/rabbitmq
+
+
+  mysql:
+    image: mysql:${MYSQL_IMAGE_TAG:-8.4}
+    restart: always
+    healthcheck:
+      test: mysqladmin ping -h localhost -u $$MYSQL_USER --password=$$MYSQL_PASSWORD
+      interval: 10s
+      timeout: 5s
+      retries: 3
+    command:
+      --character-set-server=utf8mb4
+      --collation-server=utf8mb4_general_ci
+    environment:
+      MYSQL_ROOT_PASSWORD: ${MYSQL_ROOT_PASSWORD:-sci-1234}
+      MYSQL_USER: ${MYSQL_USER:-sci-assistant}
+      MYSQL_PASSWORD: ${MYSQL_PASSWORD:-sci-1234}
+      MYSQL_DATABASE: ${MYSQL_DATABASE:-sci-assistant}
+      TZ: Asia/Shanghai
+    volumes:
+      - ./mysql/initdb.d:/docker-entrypoint-initdb.d:ro
+      - ./mysql/data:/var/lib/mysql
+
+  redis:
+    image: redis:${REDIS_IMAGE_TAG:-7.4}
+    restart: always
+    healthcheck:
+      test: |
+        if [ -z "$$REDIS_PASSWORD" ]; then
+          redis-cli ping | grep PONG
+        else
+          redis-cli -a $$REDIS_PASSWORD ping | grep PONG
+        fi
+      interval: 10s
+      timeout: 1s
+      retries: 3
+    command:
+      --requirepass "${REDIS_PASSWORD:-}"
+    environment:
+      REDIS_PASSWORD: ${REDIS_PASSWORD:-}
+
+  sci-assistant-docparsing:
+    image: git.kclab.cloud/sci-agent/sci-assistant-docparsing:${DOCPARSING_IMAGE_TAG:-0.0.5}
+    devices: 
+      - "nvidia.com/gpu=7"
+    environment:
+      # MQ Configuration
+      AMQP_URI: amqp://sci-assistant:sci-1234@rabbitmq:5672/%2F?heartbeat=60
+      # Queue Names
+      PARSING_INPUT_QUEUE: PARSINGINPUT
+      PARSING_OUTPUT_QUEUE: PARSINGOUTPUT
+      # Concurrency Control
+      MAX_PARSING_CONCURRENCY: ${MAX_PARSING_CONCURRENCY:-1}
+      # S3 Configuration
+      S3_ENDPOINT: http://minio:9000
+      S3_ACCESS_KEY: ${MINIO_ACCESS_KEY:-minio}
+      S3_SECRET_KEY: ${MINIO_SECRET_KEY:-minio123}
+      S3_BUCKET_NAME: ${MINIO_BUCKET:-sci-assistant}
+      S3_REGION_NAME: us-east-1
+      S3_ADDRESSING_STYLE: auto
+      # allocation size (MiB)
+      PARSING_ALLOCATION_SIZE: ${PARSING_ALLOCATION_SIZE:-10000}
+      PARSING_ALLOCATION_SCALE: ${PARSING_ALLOCATION_SCALE:-0.85}
+    depends_on:
+      rabbitmq:
+        condition: service_healthy
+      minio:
+        condition: service_healthy
+      minio-init-bucket:
+        condition: service_completed_successfully
+    profiles:
+      - disable_until_fixed

proto/common/v1/s3.proto

@@ -0,0 +1,7 @@
+syntax = "proto3";
+
+package common.v1;
+
+message S3File {
+  string key = 1 ;
+}

proto/parse/v1/pdf_parser.proto

@@ -0,0 +1,44 @@
+syntax = "proto3";
+
+package parse.v1;
+
+import "common/v1/s3.proto";
+
+
+message PDFCover {
+  uint32 page_num = 1 ;
+  repeated double bbox = 2 ;
+}
+
+message PDFSegment {
+  enum SegmentType {
+    TYPE_UNSPECIFIED = 0;
+    TYPE_TEXT = 1;
+    TYPE_IMAGE = 2;
+    TYPE_FORMULA = 3; // 特指单行公式
+    TYPE_TABLE = 4;
+  }
+  string chapter_name = 1 ; // 图片不存在 chapter name
+  string id = 2 ;
+  SegmentType type = 3 ;
+  string content = 4 ; // 图片的这个部分为空，表格为解析版
+  common.v1.S3File image = 5 ; // 图片格式数据
+  repeated PDFCover covers = 6 ;
+  repeated string related_segment_ids = 7 ;
+  string reference_id = 8 ; // 图片和表格具有该字段
+  string caption = 9 ; // 图片和表格具有该字段
+}
+
+message PDFMetadata {
+  string title = 1 ;
+}
+
+message PDFParserInput {
+  common.v1.S3File file = 1 ;
+}
+
+message PDFParserOutput {
+  common.v1.S3File file = 1 ;
+  PDFMetadata metadata = 2 ;
+  repeated PDFSegment segments = 3 ;
+}

pyproject.toml

@@ -4,7 +4,11 @@ version = "0.1.0"
 description = "Add your description here"
 readme = "README.md"
 requires-python = ">=3.12"
-dependencies = []
+dependencies = [
+    "aio-pika>=9.5.5",
+    "grpcio-tools>=1.73.0",
+    "protobuf>=6.31.1",
+]
 
 [dependency-groups]
 bandit = [

src/message/common/v1/s3_pb2_grpc.py

@@ -0,0 +1,24 @@
+# Generated by the gRPC Python protocol compiler plugin. DO NOT EDIT!
+"""Client and server classes corresponding to protobuf-defined services."""
+import grpc
+import warnings
+
+
+GRPC_GENERATED_VERSION = '1.73.0'
+GRPC_VERSION = grpc.__version__
+_version_not_supported = False
+
+try:
+    from grpc._utilities import first_version_is_lower
+    _version_not_supported = first_version_is_lower(GRPC_VERSION, GRPC_GENERATED_VERSION)
+except ImportError:
+    _version_not_supported = True
+
+if _version_not_supported:
+    raise RuntimeError(
+        f'The grpc package installed is at version {GRPC_VERSION},'
+        + f' but the generated code in common/v1/s3_pb2_grpc.py depends on'
+        + f' grpcio>={GRPC_GENERATED_VERSION}.'
+        + f' Please upgrade your grpc module to grpcio>={GRPC_GENERATED_VERSION}'
+        + f' or downgrade your generated code using grpcio-tools<={GRPC_VERSION}.'
+    )

src/message/parse/v1/pdf_parser_pb2_grpc.py

@@ -0,0 +1,24 @@
+# Generated by the gRPC Python protocol compiler plugin. DO NOT EDIT!
+"""Client and server classes corresponding to protobuf-defined services."""
+import grpc
+import warnings
+
+
+GRPC_GENERATED_VERSION = '1.73.0'
+GRPC_VERSION = grpc.__version__
+_version_not_supported = False
+
+try:
+    from grpc._utilities import first_version_is_lower
+    _version_not_supported = first_version_is_lower(GRPC_VERSION, GRPC_GENERATED_VERSION)
+except ImportError:
+    _version_not_supported = True
+
+if _version_not_supported:
+    raise RuntimeError(
+        f'The grpc package installed is at version {GRPC_VERSION},'
+        + f' but the generated code in parse/v1/pdf_parser_pb2_grpc.py depends on'
+        + f' grpcio>={GRPC_GENERATED_VERSION}.'
+        + f' Please upgrade your grpc module to grpcio>={GRPC_GENERATED_VERSION}'
+        + f' or downgrade your generated code using grpcio-tools<={GRPC_VERSION}.'
+    )

tests/test_doc_parsing.py

@@ -0,0 +1,58 @@
+import asyncio
+import aio_pika 
+from aio_pika import Message, IncomingMessage
+from src.message.parse.v1.pdf_parser_pb2 import PDFParserInput, PDFParserOutput
+from src.message.common.v1.s3_pb2 import S3File
+import logging
+logging.basicConfig(level=logging.INFO)
+
+AMQP_URI = "amqp://sci-assistant:sci-1234@localhost:5671/%2F?heartbeat=60"
+INPUT_QUEUE = "PARSINGINPUT"
+OUTPUT_QUEUE = "PARSINGOUTPUT"
+
+
+async def send_task(channel, s3_key: str):
+    message = PDFParserInput()
+    message.file.CopyFrom(S3File(key=s3_key))
+    body = message.SerializeToString()
+
+    await channel.default_exchange.publish(
+        Message(body=body, content_type="application/x-protobuf"),
+        routing_key=INPUT_QUEUE
+    )
+    print(f"[SEND] Sent PDFParserInput with key: {s3_key}")
+
+
+async def on_result(message: IncomingMessage):
+    async with message.process():
+        output = PDFParserOutput()
+        output.ParseFromString(message.body)
+
+        print(f"\n[RECEIVED] Output file key: {output.file.key}")
+        print(f"Title: {output.metadata.title}")
+        print(f"Segments ({len(output.segments)}):")
+        for seg in output.segments[:5]:  # 打印前5个
+            print(f"  ID: {seg.id}, Type: {seg.type}, Content: {seg.content[:100]}, Images: {seg.image}")
+
+
+async def main():
+    connection = await aio_pika.connect_robust(AMQP_URI)
+    channel = await connection.channel()
+
+    # 开始监听输出队列
+    await channel.set_qos(prefetch_count=10)
+    await channel.declare_queue(OUTPUT_QUEUE, durable=True)
+    await channel.declare_queue(INPUT_QUEUE, durable=True)
+
+    queue = await channel.declare_queue(OUTPUT_QUEUE, durable=True)
+    await queue.consume(on_result)
+
+    # 发送任务（替换成你需要解析的 PDF 文件 key）
+    await send_task(channel, "my-uploaded-file.pdf")
+
+    print(" [*] 等待解析结果。按 Ctrl+C 退出")
+    await asyncio.Future()
+
+
+if __name__ == "__main__":
+    asyncio.run(main())