From ac4078d4af77e1ca08d716cbad2712e3016015d3 Mon Sep 17 00:00:00 2001 From: pgmsoul Date: Mon, 22 Jun 2026 11:06:01 +0800 Subject: [PATCH] add openai api --- .env.example | 3 + README.md | 76 ++++++++++++++++++ src/openocr_markdown/api.py | 2 + src/openocr_markdown/config.py | 2 + src/openocr_markdown/engine.py | 13 +++ src/openocr_markdown/openai_routes.py | 88 +++++++++++++++++++++ src/openocr_markdown/openai_schemas.py | 65 +++++++++++++++ src/openocr_markdown/openai_utils.py | 105 +++++++++++++++++++++++++ tests/test_openai_utils.py | 43 ++++++++++ 9 files changed, 397 insertions(+) create mode 100644 src/openocr_markdown/openai_routes.py create mode 100644 src/openocr_markdown/openai_schemas.py create mode 100644 src/openocr_markdown/openai_utils.py create mode 100644 tests/test_openai_utils.py diff --git a/.env.example b/.env.example index 5b52462..a756584 100644 --- a/.env.example +++ b/.env.example @@ -7,4 +7,7 @@ OPENOCR_AUTO_DOWNLOAD=true OPENOCR_LAYOUT_THRESHOLD=0.4 OPENOCR_MAX_LENGTH=2048 OPENOCR_CORS_ORIGINS=* +# Optional. When set, /v1/* endpoints require: Authorization: Bearer +OPENOCR_API_KEY= +OPENOCR_OPENAI_MODEL=opendoc-0.1b OPENOCR_LOG_LEVEL=info diff --git a/README.md b/README.md index 82d1794..f974aa7 100644 --- a/README.md +++ b/README.md @@ -11,6 +11,8 @@ - `GET /output/...`：访问生成的 Markdown、JSON 和原始文件 - `GET /health`：健康检查 - `GET /docs`：Swagger API 文档 +- `GET /v1/models`：OpenAI 兼容模型列表 +- `POST /v1/chat/completions`：OpenAI 兼容文档转 Markdown 接口 --- @@ -104,6 +106,78 @@ curl -F "file=@/path/to/page.png" -F "force=true" http://127.0.0.1:8000/api/reco Markdown 访问地址：`http://服务器IP:8000/output/<文档名>/<文档名>.md` +### OpenAI 兼容接口 + +适用于已有 OpenAI SDK / LangChain / 各类 Agent 框架的调用方式。 +若配置了 `OPENOCR_ROOT_PATH=/openocr`，则路径前缀为 `/openocr/v1/...`。 + +**列出模型** + +```bash +curl https://jsuse.com/openocr/v1/models +``` + +**文档转 Markdown（chat/completions）** + +```bash +curl https://jsuse.com/openocr/v1/chat/completions \ + -H "Content-Type: application/json" \ + -d '{ + "model": "opendoc-0.1b", + "messages": [ + { + "role": "user", + "content": [ + {"type": "text", "text": "Convert this document to markdown"}, + { + "type": "image_url", + "image_url": {"url": "https://example.com/page.png"} + } + ] + } + ] + }' +``` + +也支持 `data:image/png;base64,...` 形式的图片 URL。 + +响应中的 `choices[0].message.content` 即为 Markdown 文本。 + +若设置了 `OPENOCR_API_KEY`，需附加请求头： + +```bash +-H "Authorization: Bearer your-api-key" +``` + +Python 示例（OpenAI SDK）： + +```python +from openai import OpenAI + +client = OpenAI( + base_url="https://jsuse.com/openocr/v1", + api_key="your-api-key-or-empty", +) + +response = client.chat.completions.create( + model="opendoc-0.1b", + messages=[ + { + "role": "user", + "content": [ + {"type": "text", "text": "Convert to markdown"}, + { + "type": "image_url", + "image_url": {"url": "https://example.com/page.png"}, + }, + ], + } + ], +) + +print(response.choices[0].message.content) +``` + --- ## 日常开发与更新 @@ -203,6 +277,8 @@ docker compose up -d --build # 重新构建并启动 | `OPENOCR_LAYOUT_THRESHOLD` | `0.4` | 版面分析阈值 | | `OPENOCR_MAX_LENGTH` | `2048` | 识别最大长度 | | `OPENOCR_CORS_ORIGINS` | `*` | 跨域来源，多个用逗号分隔 | +| `OPENOCR_API_KEY` | *(空)* | OpenAI 兼容接口 Bearer Token，留空则不校验 | +| `OPENOCR_OPENAI_MODEL` | `opendoc-0.1b` | OpenAI 兼容接口使用的模型 ID | | `OPENOCR_LOG_LEVEL` | `info` | 日志级别 | --- diff --git a/src/openocr_markdown/api.py b/src/openocr_markdown/api.py index fe4583b..a14dad7 100644 --- a/src/openocr_markdown/api.py +++ b/src/openocr_markdown/api.py @@ -12,6 +12,7 @@ from fastapi.staticfiles import StaticFiles from .config import settings from .engine import get_doc_parser, parse_to_outputs +from .openai_routes import router as openai_router from .schemas import HistoryItem, HistoryResponse, RecognizeResponse @@ -44,6 +45,7 @@ def create_app() -> FastAPI: allow_methods=["*"], allow_headers=["*"], ) + app.include_router(openai_router) @app.on_event("startup") async def warm_up_model() -> None: diff --git a/src/openocr_markdown/config.py b/src/openocr_markdown/config.py index a23a57c..800f0c1 100644 --- a/src/openocr_markdown/config.py +++ b/src/openocr_markdown/config.py @@ -19,6 +19,8 @@ class Settings(BaseSettings): layout_threshold: float = Field(default=0.4, alias="OPENOCR_LAYOUT_THRESHOLD") max_length: int = Field(default=2048, alias="OPENOCR_MAX_LENGTH") cors_origins: str = Field(default="*", alias="OPENOCR_CORS_ORIGINS") + api_key: str = Field(default="", alias="OPENOCR_API_KEY") + openai_model_id: str = Field(default="opendoc-0.1b", alias="OPENOCR_OPENAI_MODEL") log_level: Literal["debug", "info", "warning", "error"] = Field( default="info", alias="OPENOCR_LOG_LEVEL", diff --git a/src/openocr_markdown/engine.py b/src/openocr_markdown/engine.py index 4264a60..d0cfc7a 100644 --- a/src/openocr_markdown/engine.py +++ b/src/openocr_markdown/engine.py @@ -29,3 +29,16 @@ def parse_to_outputs(input_path: Path, output_dir: Path) -> dict[str, Any]: parser.save_to_markdown(result, str(output_dir)) parser.save_to_json(result, str(output_dir)) return result + + +def parse_to_markdown_text(input_path: Path, output_dir: Path) -> tuple[str, str]: + """Parse a document and return `(doc_id, markdown_text)`.""" + parse_to_outputs(input_path, output_dir) + doc_id = input_path.stem + markdown_path = output_dir / doc_id / f"{doc_id}.md" + if not markdown_path.exists(): + candidates = sorted((output_dir / doc_id).glob("*.md")) + if not candidates: + raise FileNotFoundError(f"Markdown output not found for document: {doc_id}") + markdown_path = candidates[0] + return doc_id, markdown_path.read_text(encoding="utf-8") diff --git a/src/openocr_markdown/openai_routes.py b/src/openocr_markdown/openai_routes.py new file mode 100644 index 0000000..8d208de --- /dev/null +++ b/src/openocr_markdown/openai_routes.py @@ -0,0 +1,88 @@ +import time +from uuid import uuid4 + +from fastapi import APIRouter, Depends, Header, HTTPException + +from .config import settings +from .engine import parse_to_markdown_text +from .openai_schemas import ( + ChatCompletionChoice, + ChatCompletionMessage, + ChatCompletionRequest, + ChatCompletionResponse, + OpenAIModel, + OpenAIModelList, +) +from .openai_utils import cleanup_temp_file, extract_image_urls, materialize_image_url + +router = APIRouter(prefix="/v1", tags=["OpenAI Compatible"]) + + +async def verify_api_key(authorization: str | None = Header(default=None)) -> None: + if not settings.api_key: + return + if not authorization or not authorization.startswith("Bearer "): + raise HTTPException(status_code=401, detail="Missing or invalid Authorization header.") + token = authorization.removeprefix("Bearer ").strip() + if token != settings.api_key: + raise HTTPException(status_code=401, detail="Invalid API key.") + + +@router.get("/models", response_model=OpenAIModelList) +async def list_models(_: None = Depends(verify_api_key)) -> OpenAIModelList: + created = int(time.time()) + return OpenAIModelList( + data=[ + OpenAIModel( + id=settings.openai_model_id, + created=created, + owned_by="openocr", + ) + ] + ) + + +@router.post("/chat/completions", response_model=ChatCompletionResponse) +async def chat_completions( + request: ChatCompletionRequest, + _: None = Depends(verify_api_key), +) -> ChatCompletionResponse: + if request.stream: + raise HTTPException(status_code=400, detail="Streaming is not supported.") + + if request.model != settings.openai_model_id: + raise HTTPException( + status_code=404, + detail=f"Model '{request.model}' not found. Use '{settings.openai_model_id}'.", + ) + + image_urls = extract_image_urls(request.messages) + if not image_urls: + raise HTTPException( + status_code=400, + detail="At least one user message with image_url content is required.", + ) + if len(image_urls) > 1: + raise HTTPException(status_code=400, detail="Only one image per request is supported.") + + upload_path = materialize_image_url(image_urls[0], settings.upload_dir) + try: + _, markdown_text = parse_to_markdown_text(upload_path, settings.output_dir) + except Exception as exc: + raise HTTPException(status_code=500, detail=str(exc)) from exc + finally: + cleanup_temp_file(upload_path) + + created = int(time.time()) + return ChatCompletionResponse( + id=f"chatcmpl-{uuid4().hex}", + created=created, + model=settings.openai_model_id, + choices=[ + ChatCompletionChoice( + index=0, + message=ChatCompletionMessage(content=markdown_text), + finish_reason="stop", + ) + ], + ) diff --git a/src/openocr_markdown/openai_schemas.py b/src/openocr_markdown/openai_schemas.py new file mode 100644 index 0000000..b686e88 --- /dev/null +++ b/src/openocr_markdown/openai_schemas.py @@ -0,0 +1,65 @@ +from typing import Any, Literal + +from pydantic import BaseModel, Field + + +class OpenAIModel(BaseModel): + id: str + object: Literal["model"] = "model" + created: int + owned_by: str = "openocr" + + +class OpenAIModelList(BaseModel): + object: Literal["list"] = "list" + data: list[OpenAIModel] + + +class ChatMessageContentImageUrl(BaseModel): + url: str + detail: str | None = None + + +class ChatMessageContentPart(BaseModel): + type: Literal["text", "image_url"] + text: str | None = None + image_url: ChatMessageContentImageUrl | None = None + + +class ChatMessage(BaseModel): + role: Literal["system", "user", "assistant", "tool"] + content: str | list[ChatMessageContentPart | dict[str, Any]] + + +class ChatCompletionRequest(BaseModel): + model: str + messages: list[ChatMessage] + stream: bool = False + temperature: float | None = None + max_tokens: int | None = None + + +class ChatCompletionMessage(BaseModel): + role: Literal["assistant"] = "assistant" + content: str + + +class ChatCompletionChoice(BaseModel): + index: int + message: ChatCompletionMessage + finish_reason: Literal["stop", "length"] = "stop" + + +class ChatCompletionUsage(BaseModel): + prompt_tokens: int = 0 + completion_tokens: int = 0 + total_tokens: int = 0 + + +class ChatCompletionResponse(BaseModel): + id: str + object: Literal["chat.completion"] = "chat.completion" + created: int + model: str + choices: list[ChatCompletionChoice] + usage: ChatCompletionUsage = Field(default_factory=ChatCompletionUsage) diff --git a/src/openocr_markdown/openai_utils.py b/src/openocr_markdown/openai_utils.py new file mode 100644 index 0000000..f08c479 --- /dev/null +++ b/src/openocr_markdown/openai_utils.py @@ -0,0 +1,105 @@ +import base64 +import re +import shutil +from pathlib import Path +from urllib.parse import urlparse +from urllib.request import urlopen +from uuid import uuid4 + +from fastapi import HTTPException + +from .openai_schemas import ChatMessage, ChatMessageContentPart + +SUPPORTED_SUFFIXES = {".jpg", ".jpeg", ".png", ".bmp", ".webp", ".pdf"} +DATA_URL_PATTERN = re.compile( + r"^data:(?P[\w/+.-]+);base64,(?P.+)$", + re.DOTALL, +) +MIME_SUFFIX_MAP = { + "image/jpeg": ".jpg", + "image/jpg": ".jpg", + "image/png": ".png", + "image/bmp": ".bmp", + "image/webp": ".webp", + "application/pdf": ".pdf", +} + + +def _suffix_from_url(url: str) -> str: + path = urlparse(url).path + suffix = Path(path).suffix.lower() + return suffix if suffix in SUPPORTED_SUFFIXES else ".png" + + +def _save_bytes(content: bytes, upload_dir: Path, suffix: str) -> Path: + upload_dir.mkdir(parents=True, exist_ok=True) + upload_path = upload_dir / f"openai_{uuid4().hex}{suffix}" + upload_path.write_bytes(content) + return upload_path + + +def _decode_data_url(data_url: str, upload_dir: Path) -> Path: + match = DATA_URL_PATTERN.match(data_url.strip()) + if not match: + raise HTTPException(status_code=400, detail="Invalid image data URL.") + + mime = match.group("mime").lower() + suffix = MIME_SUFFIX_MAP.get(mime) + if not suffix: + raise HTTPException(status_code=400, detail=f"Unsupported data URL mime type: {mime}") + + try: + raw = base64.b64decode(match.group("data"), validate=True) + except Exception as exc: + raise HTTPException(status_code=400, detail="Invalid base64 image data.") from exc + + return _save_bytes(raw, upload_dir, suffix) + + +def _download_http_url(url: str, upload_dir: Path) -> Path: + parsed = urlparse(url) + if parsed.scheme not in {"http", "https"}: + raise HTTPException(status_code=400, detail=f"Unsupported image URL scheme: {parsed.scheme}") + + suffix = _suffix_from_url(url) + try: + with urlopen(url, timeout=60) as response: + content = response.read() + except Exception as exc: + raise HTTPException(status_code=400, detail=f"Failed to download image URL: {exc}") from exc + + return _save_bytes(content, upload_dir, suffix) + + +def materialize_image_url(image_url: str, upload_dir: Path) -> Path: + if image_url.startswith("data:"): + return _decode_data_url(image_url, upload_dir) + if image_url.startswith(("http://", "https://")): + return _download_http_url(image_url, upload_dir) + raise HTTPException( + status_code=400, + detail="image_url must be an http(s) URL or a data: URL.", + ) + + +def extract_image_urls(messages: list[ChatMessage]) -> list[str]: + urls: list[str] = [] + for message in messages: + if message.role != "user" or isinstance(message.content, str): + continue + + for part in message.content: + parsed = ( + part + if isinstance(part, ChatMessageContentPart) + else ChatMessageContentPart.model_validate(part) + ) + if parsed.type == "image_url" and parsed.image_url and parsed.image_url.url: + urls.append(parsed.image_url.url) + + return urls + + +def cleanup_temp_file(path: Path) -> None: + if path.exists(): + path.unlink() diff --git a/tests/test_openai_utils.py b/tests/test_openai_utils.py new file mode 100644 index 0000000..02a12b2 --- /dev/null +++ b/tests/test_openai_utils.py @@ -0,0 +1,43 @@ +import base64 + +import pytest +from fastapi import HTTPException + +from openocr_markdown.openai_schemas import ChatMessage, ChatMessageContentImageUrl, ChatMessageContentPart +from openocr_markdown.openai_utils import DATA_URL_PATTERN, extract_image_urls, materialize_image_url + + +def test_extract_image_urls_from_multimodal_message() -> None: + messages = [ + ChatMessage( + role="user", + content=[ + ChatMessageContentPart(type="text", text="Convert to markdown"), + ChatMessageContentPart( + type="image_url", + image_url=ChatMessageContentImageUrl(url="https://example.com/page.png"), + ), + ], + ) + ] + assert extract_image_urls(messages) == ["https://example.com/page.png"] + + +def test_materialize_data_url(tmp_path) -> None: + png_bytes = base64.b64decode( + "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8z8BQDwAEhQGAhKmMIQAAAABJRU5ErkJggg==" + ) + data_url = "data:image/png;base64," + base64.b64encode(png_bytes).decode() + path = materialize_image_url(data_url, tmp_path) + assert path.exists() + assert path.suffix == ".png" + + +def test_data_url_pattern_rejects_invalid_input() -> None: + assert DATA_URL_PATTERN.match("not-a-data-url") is None + + +def test_materialize_image_url_rejects_unsupported_scheme(tmp_path) -> None: + with pytest.raises(HTTPException) as exc: + materialize_image_url("ftp://example.com/a.png", tmp_path) + assert exc.value.status_code == 400