pensiev 0.25.5__py3-none-any.whl

memos/migrations.py

@@ -0,0 +1,19 @@
+from pathlib import Path
+from alembic.config import Config
+from alembic import command
+from .config import settings
+
+
+def run_migrations():
+    """Run all pending database migrations."""
+    # Get the directory containing the migrations
+    migrations_dir = Path(__file__).parent / "migrations"
+    alembic_dir = migrations_dir / "alembic"
+
+    # Create an Alembic configuration
+    alembic_cfg = Config(str(migrations_dir / "alembic.ini"))
+    alembic_cfg.set_main_option("script_location", str(alembic_dir))
+    alembic_cfg.set_main_option("sqlalchemy.url", settings.database_url)
+
+    # Run the migration
+    command.upgrade(alembic_cfg, "head")

memos/models.py

@@ -0,0 +1,199 @@
+from sqlalchemy import (
+    Integer,
+    String,
+    Text,
+    DateTime,
+    Enum,
+    ForeignKey,
+    func,
+    Index,
+)
+from datetime import datetime
+from sqlalchemy.orm import relationship, DeclarativeBase, Mapped, mapped_column, Session
+from typing import List
+from .schemas import MetadataSource, MetadataType, FolderType
+
+
+class RawBase(DeclarativeBase):
+    pass
+
+
+class Base(RawBase):
+    __abstract__ = True
+    id: Mapped[int] = mapped_column(Integer, primary_key=True)
+    created_at: Mapped[datetime] = mapped_column(
+        DateTime, server_default=func.now(), nullable=False
+    )
+    updated_at: Mapped[datetime] = mapped_column(
+        DateTime, server_default=func.now(), onupdate=func.now(), nullable=False
+    )
+
+
+class LibraryModel(Base):
+    __tablename__ = "libraries"
+    name: Mapped[str] = mapped_column(String, nullable=False, unique=True)
+    folders: Mapped[List["FolderModel"]] = relationship(
+        "FolderModel",
+        back_populates="library",
+        lazy="joined",
+        primaryjoin="and_(LibraryModel.id==FolderModel.library_id, FolderModel.type=='DEFAULT')",
+    )
+    plugins: Mapped[List["PluginModel"]] = relationship(
+        "PluginModel", secondary="library_plugins", lazy="joined"
+    )
+
+
+class FolderModel(Base):
+    __tablename__ = "folders"
+    path: Mapped[str] = mapped_column(String, nullable=False)
+    library_id: Mapped[int] = mapped_column(
+        Integer, ForeignKey("libraries.id"), nullable=False
+    )
+    library: Mapped["LibraryModel"] = relationship(
+        "LibraryModel", back_populates="folders"
+    )
+    entities: Mapped[List["EntityModel"]] = relationship(
+        "EntityModel", back_populates="folder"
+    )
+    type: Mapped[FolderType] = mapped_column(Enum(FolderType), nullable=False)
+    last_modified_at: Mapped[datetime | None] = mapped_column(DateTime, nullable=False)
+
+
+class EntityPluginStatusModel(RawBase):
+    __tablename__ = "entity_plugin_status"
+
+    entity_id: Mapped[int] = mapped_column(
+        Integer, ForeignKey("entities.id", ondelete="CASCADE"), primary_key=True
+    )
+    plugin_id: Mapped[int] = mapped_column(
+        Integer, ForeignKey("plugins.id", ondelete="CASCADE"), primary_key=True
+    )
+    processed_at: Mapped[datetime] = mapped_column(
+        DateTime, server_default=func.now(), nullable=False
+    )
+
+    __table_args__ = (
+        Index("idx_entity_plugin_entity_id", "entity_id"),
+        Index("idx_entity_plugin_plugin_id", "plugin_id"),
+    )
+
+
+class EntityModel(Base):
+    __tablename__ = "entities"
+    filepath: Mapped[str] = mapped_column(String, nullable=False)
+    filename: Mapped[str] = mapped_column(String, nullable=False)
+    size: Mapped[int] = mapped_column(Integer, nullable=False)
+    file_created_at: Mapped[datetime] = mapped_column(DateTime, nullable=False)
+    file_last_modified_at: Mapped[datetime] = mapped_column(DateTime, nullable=False)
+    file_type: Mapped[str] = mapped_column(String, nullable=False)
+    file_type_group: Mapped[str] = mapped_column(String, nullable=False)
+    last_scan_at: Mapped[datetime | None] = mapped_column(
+        DateTime, server_default=func.now(), onupdate=func.now(), nullable=True
+    )
+    library_id: Mapped[int] = mapped_column(
+        Integer, ForeignKey("libraries.id"), nullable=False
+    )
+    folder_id: Mapped[int] = mapped_column(
+        Integer, ForeignKey("folders.id"), nullable=False
+    )
+    folder: Mapped["FolderModel"] = relationship(
+        "FolderModel", 
+        back_populates="entities",
+        lazy="select"
+    )
+    metadata_entries: Mapped[List["EntityMetadataModel"]] = relationship(
+        "EntityMetadataModel", 
+        lazy="select",
+        cascade="all, delete-orphan"
+    )
+    tags: Mapped[List["TagModel"]] = relationship(
+        "TagModel",
+        secondary="entity_tags",
+        lazy="select",
+        cascade="all, delete",
+        overlaps="entities",
+    )
+    plugin_status: Mapped[List["EntityPluginStatusModel"]] = relationship(
+        "EntityPluginStatusModel", 
+        cascade="all, delete-orphan",
+        lazy="select"
+    )
+
+    # 添加索引
+    __table_args__ = (
+        Index("idx_filepath", "filepath"),
+        Index("idx_filename", "filename"),
+        Index("idx_file_type", "file_type"),
+        Index("idx_library_id", "library_id"),
+        Index("idx_folder_id", "folder_id"),
+        Index("idx_file_type_group", "file_type_group"),
+        Index("idx_file_created_at", "file_created_at"),
+    )
+
+    @classmethod
+    def update_last_scan_at(cls, session: Session, entity: "EntityModel"):
+        entity.last_scan_at = func.now()
+        session.add(entity)
+
+    @property
+    def tag_names(self) -> List[str]:
+        return [tag.name for tag in self.tags]
+
+
+class TagModel(Base):
+    __tablename__ = "tags"
+    name: Mapped[str] = mapped_column(String, nullable=False)
+    description: Mapped[str | None] = mapped_column(Text, nullable=True)
+    color: Mapped[str | None] = mapped_column(String, nullable=True)
+    # source: Mapped[str | None] = mapped_column(String, nullable=True)
+
+
+class EntityTagModel(Base):
+    __tablename__ = "entity_tags"
+    entity_id: Mapped[int] = mapped_column(
+        Integer, ForeignKey("entities.id", ondelete="CASCADE"), nullable=False
+    )
+    tag_id: Mapped[int] = mapped_column(Integer, ForeignKey("tags.id"), nullable=False)
+    source: Mapped[MetadataSource] = mapped_column(Enum(MetadataSource), nullable=False)
+
+    __table_args__ = (
+        Index("idx_entity_tag_entity_id", "entity_id"),
+        Index("idx_entity_tag_tag_id", "tag_id"),
+    )
+
+
+class EntityMetadataModel(Base):
+    __tablename__ = "metadata_entries"
+    entity_id: Mapped[int] = mapped_column(
+        Integer, ForeignKey("entities.id"), nullable=False
+    )
+    key: Mapped[str] = mapped_column(String, nullable=False)
+    value: Mapped[str] = mapped_column(Text, nullable=False)
+    source_type: Mapped[MetadataSource] = mapped_column(
+        Enum(MetadataSource), nullable=False
+    )
+    source: Mapped[str | None] = mapped_column(String, nullable=True)
+    data_type: Mapped[MetadataType] = mapped_column(Enum(MetadataType), nullable=False)
+    entity = relationship("EntityModel", back_populates="metadata_entries")
+
+    __table_args__ = (
+        Index("idx_metadata_entity_id", "entity_id"),
+        Index("idx_metadata_key", "key"),
+    )
+
+
+class PluginModel(Base):
+    __tablename__ = "plugins"
+    name: Mapped[str] = mapped_column(String, nullable=False)
+    description: Mapped[str | None] = mapped_column(Text, nullable=True)
+    webhook_url: Mapped[str] = mapped_column(String, nullable=False)
+
+
+class LibraryPluginModel(Base):
+    __tablename__ = "library_plugins"
+    library_id: Mapped[int] = mapped_column(
+        Integer, ForeignKey("libraries.id"), nullable=False
+    )
+    plugin_id: Mapped[int] = mapped_column(
+        Integer, ForeignKey("plugins.id"), nullable=False
+    )

memos/plugins/ocr/main.py

@@ -0,0 +1,251 @@
+import asyncio
+import logging
+import os
+from typing import Optional
+import httpx
+import json
+import base64
+from PIL import Image
+import numpy as np
+from concurrent.futures import ThreadPoolExecutor
+from functools import partial
+import yaml
+import io
+import platform
+import cpuinfo
+
+MAX_THUMBNAIL_SIZE = (1920, 1920)
+
+from fastapi import APIRouter, Request, HTTPException
+from memos.schemas import Entity, MetadataType
+
+METADATA_FIELD_NAME = "ocr_result"
+PLUGIN_NAME = "ocr"
+
+router = APIRouter(tags=[PLUGIN_NAME], responses={404: {"description": "Not found"}})
+endpoint = None
+token = None
+concurrency = None
+semaphore = None
+use_local = False
+ocr = None
+thread_pool = None
+
+# Configure logger
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+
+def get_metadata_name() -> str:
+    """Return the metadata field name used by this plugin."""
+    return METADATA_FIELD_NAME
+
+
+def image2base64(img_path):
+    try:
+        with Image.open(img_path) as img:
+            img = img.convert("RGB")
+            img.thumbnail(MAX_THUMBNAIL_SIZE)
+            buffered = io.BytesIO()
+            img.save(buffered, format="JPEG")
+            encoded_string = base64.b64encode(buffered.getvalue()).decode("utf-8")
+        return encoded_string
+    except Exception as e:
+        logger.error(f"Error processing image {img_path}: {str(e)}")
+        return None
+
+
+async def fetch(endpoint: str, client, image_base64, headers: Optional[dict] = None):
+    async with semaphore:  # 使用信号量控制并发
+        response = await client.post(
+            f"{endpoint}",
+            json={"image_base64": image_base64},
+            timeout=60,
+            headers=headers,
+        )
+        if response.status_code != 200:
+            return None
+        return response.json()
+
+
+def convert_ocr_results(results):
+    if results is None:
+        return []
+    
+    converted = []
+    for result in results:
+        item = {"dt_boxes": result[0], "rec_txt": result[1], "score": result[2]}
+        converted.append(item)
+    return converted
+
+
+def convert_ocr_data(ocr_data):
+    converted_data = []
+    for text, score, bbox in ocr_data:
+        x_min, y_min, x_max, y_max = bbox
+        dt_boxes = [
+            [x_min, y_min],
+            [x_max, y_min],
+            [x_max, y_max],
+            [x_min, y_max]
+        ]
+        entry = {
+            'dt_boxes': dt_boxes,
+            'rec_txt': text,
+            'score': float(score)
+        }
+        converted_data.append(entry)
+    return converted_data
+
+
+def predict_local(img_path):
+    try:
+        if platform.system() == 'Darwin':  # Check if the OS is macOS
+            from ocrmac import ocrmac
+            result = ocrmac.OCR(img_path, language_preference=['zh-Hans']).recognize(px=True)
+            return convert_ocr_data(result)
+        else:
+            with Image.open(img_path) as img:
+                img = img.convert("RGB")
+                img.thumbnail(MAX_THUMBNAIL_SIZE)
+                img_array = np.array(img)
+            results, _ = ocr(img_array)
+            return convert_ocr_results(results)
+    except Exception as e:
+        logger.error(f"Error processing image {img_path}: {str(e)}")
+        return None
+
+
+async def async_predict_local(img_path):
+    loop = asyncio.get_running_loop()
+    results = await loop.run_in_executor(thread_pool, partial(predict_local, img_path))
+    return results
+
+
+# Modify the predict function to use semaphore
+async def predict(img_path):
+    if use_local:
+        return await async_predict_local(img_path)
+    
+    image_base64 = image2base64(img_path)
+    if not image_base64:
+        return None
+
+    async with httpx.AsyncClient() as client:
+        headers = {"Authorization": f"Bearer {token.get_secret_value()}"} if token else {}
+        return await fetch(endpoint, client, image_base64, headers)
+
+
+@router.get("/")
+async def read_root():
+    return {"healthy": True}
+
+
+@router.post("", include_in_schema=False)
+@router.post("/")
+async def ocr(entity: Entity, request: Request):
+    metadata_field_name = get_metadata_name()
+    if not entity.file_type_group == "image":
+        return {metadata_field_name: "{}"}
+
+    # Check if the metadata field already exists and has a non-empty value
+    existing_metadata = entity.get_metadata_by_key(metadata_field_name)
+    if existing_metadata and existing_metadata.value and existing_metadata.value.strip():
+        logger.info(f"Skipping OCR processing for file: {entity.filepath} due to existing metadata")
+        return {metadata_field_name: existing_metadata.value}
+
+    # Check if the entity contains the tag "low_info"
+    if any(tag.name == "low_info" for tag in entity.tags):
+        logger.info(f"Skipping OCR processing for file: {entity.filepath} due to 'low_info' tag")
+        return {metadata_field_name: "{}"}
+
+    location_url = request.headers.get("Location")
+    if not location_url:
+        raise HTTPException(status_code=400, detail="Location header is missing")
+
+    patch_url = f"{location_url}/metadata"
+
+    ocr_result = await predict(entity.filepath)
+    if ocr_result:
+        filtered_results = [r for r in ocr_result if r['score'] > 0.5][:10]
+        texts = [f"{r['rec_txt']}({r['score']:.2f})" for r in filtered_results]
+        total = len(ocr_result)
+        logger.info(f"First {len(texts)}/{total} OCR results: {texts}")
+    else:
+        logger.info(f"No OCR result found for file: {entity.filepath}")
+        return {metadata_field_name: "{}"}
+
+    # Call the URL to patch the entity's metadata
+    async with httpx.AsyncClient() as client:
+        response = await client.patch(
+            patch_url,
+            json={
+                "metadata_entries": [
+                    {
+                        "key": metadata_field_name,
+                        "value": json.dumps(
+                            ocr_result,
+                            default=lambda o: o.item() if hasattr(o, "item") else o,
+                        ),
+                        "source": PLUGIN_NAME,
+                        "data_type": MetadataType.JSON_DATA.value,
+                    }
+                ]
+            },
+            timeout=30,
+        )
+
+    # Check if the patch request was successful
+    if response.status_code != 200:
+        raise HTTPException(
+            status_code=response.status_code, detail="Failed to patch entity metadata"
+        )
+
+    return {
+        metadata_field_name: json.dumps(
+            ocr_result,
+            default=lambda o: o.item() if hasattr(o, "item") else o,
+        )
+    }
+
+
+def init_plugin(config):
+    global endpoint, token, concurrency, semaphore, use_local, ocr, thread_pool
+    endpoint = config.endpoint
+    token = config.token
+    concurrency = config.concurrency
+    use_local = config.use_local
+    semaphore = asyncio.Semaphore(concurrency)
+    
+    if use_local:
+        config_path = os.path.join(os.path.dirname(__file__), "ppocr.yaml")
+        
+        # Load and update the config file with absolute model paths
+        with open(config_path, 'r') as f:
+            ocr_config = yaml.safe_load(f)
+        
+        model_dir = os.path.join(os.path.dirname(__file__), "models")
+        ocr_config['Det']['model_path'] = os.path.join(model_dir, os.path.basename(ocr_config['Det']['model_path']))
+        ocr_config['Cls']['model_path'] = os.path.join(model_dir, os.path.basename(ocr_config['Cls']['model_path']))
+        ocr_config['Rec']['model_path'] = os.path.join(model_dir, os.path.basename(ocr_config['Rec']['model_path']))
+        
+        # Save the updated config to a temporary file with strings wrapped in double quotes
+        temp_config_path = os.path.join(os.path.dirname(__file__), "temp_ppocr.yaml")
+        with open(temp_config_path, 'w') as f:
+            yaml.safe_dump(ocr_config, f)
+        
+        if platform.system() == 'Windows' and 'Intel' in cpuinfo.get_cpu_info()['brand_raw']:
+            from rapidocr_openvino import RapidOCR
+            ocr = RapidOCR(config_path=temp_config_path)
+        else:
+            from rapidocr_onnxruntime import RapidOCR
+            ocr = RapidOCR(config_path=temp_config_path)
+        thread_pool = ThreadPoolExecutor(max_workers=concurrency)
+
+    logger.info("OCR plugin initialized")
+    logger.info(f"Endpoint: {endpoint}")
+    logger.info(f"Token: {token}")
+    logger.info(f"Concurrency: {concurrency}")
+    logger.info(f"Use local: {use_local}")
+    if use_local:
+        logger.info(f"OCR library: {'rapidocr_openvino' if platform.system() == 'Windows' and 'Intel' in cpuinfo.get_cpu_info()['brand_raw'] else 'rapidocr_onnxruntime'}")

memos/plugins/ocr/ppocr-gpu.yaml

@@ -0,0 +1,43 @@
+Global:
+    text_score: 0.5
+    use_det: true
+    use_cls: true
+    use_rec: true
+    print_verbose: false
+    min_height: 30
+    width_height_ratio: 40
+    max_side_len: 2000
+    min_side_len: 30
+
+Det:
+    use_cuda: true
+
+    model_path: models/ch_PP-OCRv4_det_infer.onnx
+
+    limit_side_len: 1500
+    limit_type: min
+
+    thresh: 0.3
+    box_thresh: 0.3
+    max_candidates: 1000
+    unclip_ratio: 1.6
+    use_dilation: true
+    score_mode: fast
+
+Cls:
+    use_cuda: true
+
+    model_path: models/ch_ppocr_mobile_v2.0_cls_train.onnx
+
+    cls_image_shape: [3, 48, 192]
+    cls_batch_num: 6
+    cls_thresh: 0.9
+    label_list: ['0', '180']
+
+Rec:
+    use_cuda: true
+
+    model_path: models/ch_PP-OCRv4_rec_infer.onnx
+
+    rec_img_shape: [3, 48, 320]
+    rec_batch_num: 6

memos/plugins/ocr/ppocr.yaml

@@ -0,0 +1,44 @@
+Global:
+    text_score: 0.5
+    use_det: true
+    use_cls: true
+    use_rec: true
+    print_verbose: false
+    min_height: 30
+    width_height_ratio: 40
+    use_space_char: true
+    max_side_len: 2000
+    min_side_len: 30
+
+Det:
+    use_cuda: false
+
+    model_path: models/ch_PP-OCRv4_det_infer.onnx
+
+    limit_side_len: 1500
+    limit_type: min
+
+    thresh: 0.3
+    box_thresh: 0.3
+    max_candidates: 1000
+    unclip_ratio: 1.6
+    use_dilation: true
+    score_mode: fast
+
+Cls:
+    use_cuda: false
+
+    model_path: models/ch_ppocr_mobile_v2.0_cls_train.onnx
+
+    cls_image_shape: [3, 48, 192]
+    cls_batch_num: 6
+    cls_thresh: 0.9
+    label_list: ['0', '180']
+
+Rec:
+    use_cuda: false
+
+    model_path: models/ch_PP-OCRv4_rec_infer.onnx
+
+    rec_img_shape: [3, 48, 320]
+    rec_batch_num: 6