add openai api

2026-06-22 11:06:01 +08:00
parent 30ead8fa43
commit ac4078d4af
9 changed files with 397 additions and 0 deletions
@@ -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 <key>
+OPENOCR_API_KEY=
+OPENOCR_OPENAI_MODEL=opendoc-0.1b
 OPENOCR_LOG_LEVEL=info
@@ -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` | 日志级别 |

 ---
@@ -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:
@@ -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",
@@ -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")
@@ -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",
+            )
+        ],
+    )
@@ -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)
@@ -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<mime>[\w/+.-]+);base64,(?P<data>.+)$",
+    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()
@@ -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