Карго-Трекер без нервов: Парсинг китайских накладных (1688 / Taobao / WeChat) с фото через Document AI + Gemini 2.5 Flash

import os
import json
import logging
from typing import Dict, Any, Optional
from google.api_core.client_options import ClientOptions
from google.cloud import documentai
from google import genai
from google.genai import types

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Настройки Document AI (OCR)
PROJECT_ID = os.environ.get("GOOGLE_CLOUD_PROJECT")
LOCATION = os.environ.get("DOCAI_LOCATION", "eu")
PROCESSOR_ID = os.environ.get("DOCAI_PROCESSOR_ID")

# Инициализация клиентов
docai_client = documentai.DocumentProcessorServiceClient(
    client_options=ClientOptions(api_endpoint=f"{LOCATION}-documentai.googleapis.com")
)
ai_client = genai.Client(api_key=os.environ.get("GEMINI_API_KEY"))

SYSTEM_PROMPT = """Ты — экспертный логистический ассистент карго-компании (Казахстан-Китай).
На вход поступает сырой распознанный текст (OCR) с китайской накладной (快递单) от курьеров ZTO, YTO, SF Express, STO или скриншот из 1688/Taobao/WeChat.
Твоя задача — извлечь следующие данные, игнорируя визуальный мусор и нерелевантный текст:
1. Идентификатор/Код клиента (часто содержит латиницу и цифры, например "ALM-882", "KZ-01-M", "K89").
2. Трек-номер посылки (快递单号 / 运单号) — обычно длинный цифровой или буквенно-цифровой код (12-15 символов).
3. Курьерская служба (ZTO/中通, YTO/圆通, SF/顺丰 и т.д.).
4. Вес в кг (重量), если указан.
5. Описание товара (переведи с китайского на русский).

Верни строгий JSON:
{
  "client_code": "string или null",
  "tracking_number": "string",
  "courier": "string",
  "weight_kg": number или null,
  "description_ru": "string"
}"""

async def process_cargo_waybill(image_bytes: bytes, mime_type: str = "image/jpeg") -> Dict[str, Any]:
    """
    Двухэтапный парсинг:
    1. Document AI для высокоточного OCR мелкого/мятого китайского текста.
    2. Gemini 2.5 Flash для семантического извлечения сущностей из текста OCR.
    """
    
    # ЭТАП 1: OCR через Document AI (Pretrained OCR Processor)
    name = docai_client.processor_path(PROJECT_ID, LOCATION, PROCESSOR_ID)
    raw_document = documentai.RawDocument(content=image_bytes, mime_type=mime_type)
    request = documentai.ProcessRequest(name=name, raw_document=raw_document)
    
    logger.info("Отправка в Document AI...")
    result = docai_client.process_document(request=request)
    document = result.document
    raw_text = document.text
    logger.info(f"Document AI извлек {len(raw_text)} символов текста.")

    # Если OCR не нашел текст, пробуем отдать картинку напрямую в Gemini (Fallback)
    if len(raw_text.strip()) < 10:
        logger.warning("Document AI не нашел текст. Используем Gemini Multimodal Fallback.")
        prompt_content = [
            types.Part.from_bytes(data=image_bytes, mime_type=mime_type),
            "Проанализируй накладную и верни JSON."
        ]
    else:
        # ЭТАП 2: Gemini 2.5 Flash для структурирования
        logger.info("Передача сырого текста в Gemini 2.5 Flash...")
        prompt_content = [
            f"Вот сырой текст с накладной, полученный через OCR:\n\n{raw_text}\n\n",
            "Извлеки из него данные и верни JSON согласно системной инструкции."
        ]

    response = ai_client.models.generate_content(
        model="gemini-2.5-flash",
        contents=prompt_content,
        config=types.GenerateContentConfig(
            system_instruction=SYSTEM_PROMPT,
            response_mime_type="application/json",
            temperature=0.0
        )
    )
    
    try:
        parsed_data = json.loads(response.text)
        return parsed_data
    except json.JSONDecodeError:
        logger.error("Ошибка парсинга JSON от Gemini.")
        return {"error": "Failed to parse"}

from fastapi import FastAPI, UploadFile, File, HTTPException
import os
from google.cloud import firestore
from cargo_parser_service import process_cargo_waybill

app = FastAPI(title="OZAT Cargo Tracker API")
db = firestore.Client(database=os.environ.get("FIRESTORE_DATABASE", "(default)"))

@app.post("/api/v1/parse-waybill")
async def parse_waybill_endpoint(file: UploadFile = File(...)):
    if not file.content_type.startswith("image/"):
        raise HTTPException(status_code=400, detail="Only images are supported")
        
    image_bytes = await file.read()
    
    # Вызов конвейера Document AI + Gemini
    try:
        parsed_data = await process_cargo_waybill(image_bytes, file.content_type)
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Processing failed: {str(e)}")

    if "error" in parsed_data:
        raise HTTPException(status_code=422, detail="AI could not extract structured data")

    tracking_number = parsed_data.get("tracking_number")
    if not tracking_number:
        raise HTTPException(status_code=400, detail="Could not find tracking number on image")

    # Сохраняем результат в Firestore
    doc_ref = db.collection("cargo_waybills").document(tracking_number)
    doc_ref.set({
        "tracking_number": tracking_number,
        "client_code": parsed_data.get("client_code"),
        "courier": parsed_data.get("courier"),
        "weight_kg": parsed_data.get("weight_kg"),
        "description": parsed_data.get("description_ru"),
        "status": "ARRIVED_AT_WAREHOUSE",
        "created_at": firestore.SERVER_TIMESTAMP
    })

    return {
        "status": "success",
        "data": parsed_data,
        "message": f"Waybill {tracking_number} saved to database."
    }

terraform {
  required_providers {
    google = {
      source  = "hashicorp/google"
      version = "~> 5.0"
    }
  }
}

variable "project_id" { type = string }
variable "region" { default = "europe-west1" }

provider "google" {
  project = var.project_id
  region  = var.region
}

# Активация API Document AI
resource "google_project_service" "documentai" {
  service = "documentai.googleapis.com"
  disable_on_destroy = false
}

# Создание Document AI Processor (OCR)
resource "google_document_ai_processor" "ocr_processor" {
  display_name = "cargo-ocr-processor"
  location     = "eu"
  type         = "OCR_PROCESSOR"
  depends_on   = [google_project_service.documentai]
}

# Секрет Gemini API
resource "google_secret_manager_secret" "gemini_key" {
  secret_id = "gemini-api-key"
  replication { auto {} }
}

# Cloud Run Service
resource "google_cloud_run_v2_service" "cargo_parser" {
  name     = "cargo-waybill-parser"
  location = var.region
  ingress  = "INGRESS_TRAFFIC_ALL"

  template {
    scaling { max_instance_count = 10 }
    containers {
      image = "${var.region}-docker.pkg.dev/${var.project_id}/apps/cargo-parser:latest"
      
      env {
        name  = "DOCAI_PROCESSOR_ID"
        value = google_document_ai_processor.ocr_processor.id
      }
      env {
        name  = "DOCAI_LOCATION"
        value = "eu"
      }
      env {
        name = "GEMINI_API_KEY"
        value_source { secret_key_ref { secret = google_secret_manager_secret.gemini_key.secret_id; version = "latest" } }
      }
    }
  }
}

# 1. Авторизация и настройка
gcloud auth login
gcloud config set project ozat-cloud-kz

# 2. Создание процессора Document AI (если через CLI)
# Альтернативно используется Terraform (см. листинг выше)

# 3. Сборка образа
gcloud builds submit --tag europe-west1-docker.pkg.dev/ozat-cloud-kz/apps/cargo-parser:latest

# 4. Деплой Cloud Run
gcloud run deploy cargo-waybill-parser \
  --image europe-west1-docker.pkg.dev/ozat-cloud-kz/apps/cargo-parser:latest \
  --region europe-west1 \
  --platform managed \
  --allow-unauthenticated \
  --set-secrets="GEMINI_API_KEY=gemini-api-key:latest" \
  --set-env-vars="DOCAI_LOCATION=eu,FIRESTORE_DATABASE=(default)" \
  --memory=512Mi \
  --cpu=1