wizardchen
397689d2f3
- Added a new `.env.lite.example` file for the Lite version, providing a minimal configuration template. - Updated `.env.example` to remove deprecated variables and include new Docreader settings. - Enhanced Docker configurations to support the Lite version, including a new Dockerfile for the Docreader service. - Introduced a Makefile target for building and running the Lite version, along with packaging capabilities. - Created GitHub workflows for building and releasing Lite binaries, including Homebrew formula support. - Implemented a new service file for managing the Lite version as a system service. This update enables a streamlined, single-binary deployment of WeKnora, reducing external dependencies and simplifying setup.
259 lines
8.9 KiB
Python
259 lines
8.9 KiB
Python
import logging
|
|
import os
|
|
import re
|
|
import sys
|
|
import traceback
|
|
import uuid
|
|
from concurrent import futures
|
|
from typing import Optional
|
|
|
|
import grpc
|
|
from grpc_health.v1 import health_pb2_grpc
|
|
from grpc_health.v1.health import HealthServicer
|
|
|
|
from docreader import config
|
|
from docreader.config import CONFIG
|
|
from docreader.parser import Parser
|
|
from docreader.proto import docreader_pb2_grpc
|
|
from docreader.parser.registry import registry
|
|
from docreader.proto.docreader_pb2 import (
|
|
ReadRequest,
|
|
ReadResponse,
|
|
ImageRef,
|
|
ListEnginesResponse,
|
|
ParserEngineInfo,
|
|
)
|
|
from docreader.utils.request import init_logging_request_id, request_id_context
|
|
|
|
_SURROGATE_RE = re.compile(r"[\ud800-\udfff]")
|
|
|
|
|
|
def to_valid_utf8_text(s: Optional[str]) -> str:
|
|
if not s:
|
|
return ""
|
|
s = _SURROGATE_RE.sub("\ufffd", s)
|
|
return s.encode("utf-8", errors="replace").decode("utf-8")
|
|
|
|
|
|
for handler in logging.root.handlers[:]:
|
|
logging.root.removeHandler(handler)
|
|
|
|
handler = logging.StreamHandler(sys.stdout)
|
|
logging.root.addHandler(handler)
|
|
|
|
_level_name = (os.environ.get("LOG_LEVEL") or "INFO").upper()
|
|
_level = getattr(logging, _level_name, logging.INFO)
|
|
logging.root.setLevel(_level)
|
|
|
|
logger = logging.getLogger(__name__)
|
|
logger.info("Initializing server logging, level=%s", _level_name)
|
|
|
|
init_logging_request_id()
|
|
|
|
|
|
def _resolve_images(images: dict, request_id: str, storage_map: dict | None = None) -> tuple[str, list]:
|
|
"""Resolve document images with priority: shared storage > temp dir > inline bytes.
|
|
|
|
``images`` is a dict of {relative_path: raw_data} where raw_data is
|
|
base64-encoded string or raw bytes.
|
|
|
|
``storage_map`` is a dict from the request's config.image_storage map.
|
|
When it contains a valid "provider" key, images are uploaded to shared object
|
|
storage and storage_key is set (no inline bytes, saving gRPC message size).
|
|
|
|
Returns (image_dir_path, list[ImageRef]).
|
|
"""
|
|
import base64
|
|
|
|
if not images:
|
|
return "", []
|
|
|
|
mime_map = {
|
|
".png": "image/png", ".jpg": "image/jpeg", ".jpeg": "image/jpeg",
|
|
".gif": "image/gif", ".webp": "image/webp", ".bmp": "image/bmp",
|
|
}
|
|
|
|
storage_client = None
|
|
use_storage = False
|
|
provider = (storage_map or {}).get("provider", "")
|
|
if provider:
|
|
try:
|
|
from docreader.parser.storage import create_storage
|
|
storage_client = create_storage(storage_map)
|
|
use_storage = True
|
|
logger.info("Using shared storage (%s) for image upload", provider)
|
|
except Exception as e:
|
|
logger.warning("Failed to init shared storage, falling back to inline bytes: %s", e)
|
|
|
|
base_dir = CONFIG.image_output_dir
|
|
output_dir = os.path.join(base_dir, request_id, "images")
|
|
wrote_to_dir = False
|
|
if not use_storage:
|
|
try:
|
|
os.makedirs(output_dir, exist_ok=True)
|
|
wrote_to_dir = True
|
|
except OSError:
|
|
logger.warning("Cannot write to image output dir %s", output_dir)
|
|
|
|
refs = []
|
|
for ref_path, b64data in images.items():
|
|
try:
|
|
img_bytes = base64.b64decode(b64data)
|
|
except Exception:
|
|
img_bytes = b64data.encode("utf-8") if isinstance(b64data, str) else b64data
|
|
|
|
fname = os.path.basename(ref_path) or f"{uuid.uuid4().hex}.png"
|
|
ext = os.path.splitext(fname)[1].lower()
|
|
mime = mime_map.get(ext, "application/octet-stream")
|
|
|
|
if use_storage and storage_client:
|
|
storage_key = storage_client.upload_bytes(img_bytes, ext)
|
|
if storage_key:
|
|
refs.append(ImageRef(
|
|
filename=fname,
|
|
original_ref=ref_path,
|
|
mime_type=mime,
|
|
storage_key=storage_key,
|
|
))
|
|
continue
|
|
logger.warning("Storage upload failed for %s, falling back to inline bytes", fname)
|
|
|
|
if wrote_to_dir:
|
|
dest = os.path.join(output_dir, fname)
|
|
try:
|
|
with open(dest, "wb") as f:
|
|
f.write(img_bytes)
|
|
except OSError as e:
|
|
logger.warning("Failed to write image %s: %s", dest, e)
|
|
|
|
refs.append(ImageRef(
|
|
filename=fname,
|
|
original_ref=ref_path,
|
|
mime_type=mime,
|
|
image_data=img_bytes,
|
|
))
|
|
|
|
image_dir = os.path.join(base_dir, request_id) if wrote_to_dir else ""
|
|
mode = "storage" if use_storage else ("dir+inline" if wrote_to_dir else "inline")
|
|
logger.info("Resolved %d images (mode=%s)", len(refs), mode)
|
|
return image_dir, refs
|
|
|
|
|
|
class DocReaderServicer(docreader_pb2_grpc.DocReaderServicer):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.parser = Parser()
|
|
|
|
def Read(self, request: ReadRequest, context):
|
|
"""Unified read: file mode (file_content set) or URL mode (url set)."""
|
|
request_id = request.request_id or str(uuid.uuid4())
|
|
is_url = bool(request.url)
|
|
|
|
with request_id_context(request_id):
|
|
try:
|
|
cfg = request.config
|
|
parser_engine = cfg.parser_engine if cfg else ""
|
|
engine_overrides = dict(cfg.parser_engine_overrides) if cfg else {}
|
|
storage_map = dict(cfg.image_storage) if cfg and cfg.image_storage else None
|
|
|
|
if is_url:
|
|
logger.info("Read(URL): url=%s", request.url)
|
|
result = self.parser.parse_url(
|
|
request.url,
|
|
request.title,
|
|
parser_engine=parser_engine,
|
|
engine_overrides=engine_overrides,
|
|
)
|
|
source_desc = request.url
|
|
else:
|
|
file_type = (
|
|
request.file_type or os.path.splitext(request.file_name)[1][1:]
|
|
)
|
|
logger.info(
|
|
"Read(File): file=%s, type=%s, size=%d bytes",
|
|
request.file_name, file_type, len(request.file_content),
|
|
)
|
|
result = self.parser.parse_file(
|
|
request.file_name,
|
|
file_type,
|
|
request.file_content,
|
|
parser_engine=parser_engine,
|
|
engine_overrides=engine_overrides,
|
|
)
|
|
source_desc = request.file_name
|
|
|
|
if not result or not result.content:
|
|
error_msg = f"Failed to parse: {source_desc}"
|
|
logger.error(error_msg)
|
|
return ReadResponse(error=error_msg)
|
|
|
|
_c = to_valid_utf8_text
|
|
image_dir, image_refs = _resolve_images(
|
|
result.images, request_id, storage_map=storage_map
|
|
)
|
|
|
|
response = ReadResponse(
|
|
markdown_content=_c(result.content),
|
|
image_refs=image_refs,
|
|
image_dir_path=image_dir,
|
|
)
|
|
logger.info(
|
|
"Read response: content_len=%d, images=%d",
|
|
len(result.content), len(image_refs),
|
|
)
|
|
return response
|
|
|
|
except Exception as e:
|
|
error_msg = f"Error reading document: {e}"
|
|
logger.error(error_msg)
|
|
logger.info("Traceback: %s", traceback.format_exc())
|
|
return ReadResponse(error=str(e))
|
|
|
|
def ListEngines(self, request, context):
|
|
overrides = dict(getattr(request, "config_overrides", None) or {})
|
|
engines_data = registry.list_engines(overrides=overrides or None)
|
|
engines = [
|
|
ParserEngineInfo(
|
|
name=e["name"],
|
|
description=e["description"],
|
|
file_types=e["file_types"],
|
|
available=e.get("available", True),
|
|
unavailable_reason=e.get("unavailable_reason", ""),
|
|
)
|
|
for e in engines_data
|
|
]
|
|
return ListEnginesResponse(engines=engines)
|
|
|
|
|
|
def main():
|
|
config.print_config()
|
|
|
|
server = grpc.server(
|
|
futures.ThreadPoolExecutor(max_workers=CONFIG.grpc_max_workers),
|
|
options=[
|
|
("grpc.max_send_message_length", CONFIG.grpc_max_file_size_mb),
|
|
("grpc.max_receive_message_length", CONFIG.grpc_max_file_size_mb),
|
|
],
|
|
)
|
|
|
|
docreader_pb2_grpc.add_DocReaderServicer_to_server(DocReaderServicer(), server)
|
|
|
|
health_servicer = HealthServicer()
|
|
health_pb2_grpc.add_HealthServicer_to_server(health_servicer, server)
|
|
|
|
server.add_insecure_port(f"[::]:{CONFIG.grpc_port}")
|
|
server.start()
|
|
|
|
logger.info("Server started on port %d", CONFIG.grpc_port)
|
|
logger.info("Server is ready to accept connections")
|
|
|
|
try:
|
|
server.wait_for_termination()
|
|
except KeyboardInterrupt:
|
|
logger.info("Received termination signal, shutting down server")
|
|
server.stop(0)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|