pensiev 0.25.5__py3-none-any.whl

memos/read_metadata.py

@@ -0,0 +1,32 @@
+import json
+import argparse
+from .utils import get_image_metadata
+
+
+def read_metadata(image_path):
+    try:
+        metadata = get_image_metadata(image_path)
+
+        if metadata is None:
+            print("No metadata found or unsupported file format.")
+            return None
+
+        return metadata if metadata else None
+
+    except Exception as e:
+        print(f"An error occurred: {str(e)}")
+        return None
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Read metadata from a screenshot")
+    parser.add_argument("image_path", type=str, help="Path to the screenshot image")
+    args = parser.parse_args()
+
+    metadata = read_metadata(args.image_path)
+    if metadata is not None:
+        print(json.dumps(metadata, indent=4))
+
+
+if __name__ == "__main__":
+    main()

memos/record.py

@@ -0,0 +1,358 @@
+import json
+import os
+import time
+import logging
+import platform
+import subprocess
+import argparse
+from PIL import Image
+import imagehash
+from memos.utils import write_image_metadata
+import ctypes
+from mss import mss
+from pathlib import Path
+from memos.config import settings
+
+# Import platform-specific modules
+if platform.system() == "Windows":
+    import win32gui
+    import win32process
+    import psutil
+elif platform.system() == "Darwin":
+    from AppKit import NSWorkspace
+    from Quartz import (
+        CGWindowListCopyWindowInfo,
+        kCGWindowListOptionOnScreenOnly,
+        kCGNullWindowID,
+        CGSessionCopyCurrentDictionary,
+    )
+
+logging.basicConfig(
+    level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s"
+)
+
+
+# Functions moved from common.py
+def load_screen_sequences(base_dir, date):
+    try:
+        with open(os.path.join(base_dir, date, ".screen_sequences"), "r") as f:
+            return json.load(f)
+    except FileNotFoundError:
+        return {}
+
+
+def save_screen_sequences(base_dir, screen_sequences, date):
+    with open(os.path.join(base_dir, date, ".screen_sequences"), "w") as f:
+        json.dump(screen_sequences, f)
+        f.flush()
+        os.fsync(f.fileno())
+
+
+def load_previous_hashes(base_dir):
+    date = time.strftime("%Y%m%d")
+    hash_file = os.path.join(base_dir, date, ".previous_hashes")
+    try:
+        with open(hash_file, "r") as f:
+            return json.load(f)
+    except FileNotFoundError:
+        return {}
+
+
+def save_previous_hashes(base_dir, previous_hashes):
+    date = time.strftime("%Y%m%d")
+    hash_file = os.path.join(base_dir, date, ".previous_hashes")
+    os.makedirs(os.path.dirname(hash_file), exist_ok=True)
+    with open(hash_file, "w") as f:
+        json.dump(previous_hashes, f)
+
+
+def get_active_window_info_darwin():
+    active_app = NSWorkspace.sharedWorkspace().activeApplication()
+    app_name = active_app["NSApplicationName"]
+    app_pid = active_app["NSApplicationProcessIdentifier"]
+
+    windows = CGWindowListCopyWindowInfo(
+        kCGWindowListOptionOnScreenOnly, kCGNullWindowID
+    )
+    for window in windows:
+        if window["kCGWindowOwnerPID"] == app_pid:
+            window_title = window.get("kCGWindowName", "")
+            if window_title:
+                return app_name, window_title
+
+    return app_name, ""  # 如果没有找到窗口标题，则返回空字符串作为标题
+
+
+def get_active_window_info_windows():
+    try:
+        window = win32gui.GetForegroundWindow()
+        _, pid = win32process.GetWindowThreadProcessId(window)
+        app_name = psutil.Process(pid).name()
+        window_title = win32gui.GetWindowText(window)
+        return app_name, window_title
+    except:
+        return "", ""
+
+
+def get_active_window_info():
+    if platform.system() == "Darwin":
+        return get_active_window_info_darwin()
+    elif platform.system() == "Windows":
+        return get_active_window_info_windows()
+
+
+def take_screenshot_macos(
+    base_dir,
+    previous_hashes,
+    threshold,
+    screen_sequences,
+    date,
+    timestamp,
+    app_name,
+    window_title,
+):
+    screenshots = []
+    result = subprocess.check_output(["system_profiler", "SPDisplaysDataType", "-json"])
+    displays_info = json.loads(result)["SPDisplaysDataType"][0]["spdisplays_ndrvs"]
+    screen_names = {}
+
+    for display_index, display_info in enumerate(displays_info):
+        base_screen_name = display_info["_name"].replace(" ", "_").lower()
+        if base_screen_name in screen_names:
+            screen_names[base_screen_name] += 1
+            screen_name = f"{base_screen_name}_{screen_names[base_screen_name]}"
+        else:
+            screen_names[base_screen_name] = 1
+            screen_name = base_screen_name
+
+        temp_filename = os.path.join(
+            base_dir, date, f"temp_screenshot-{timestamp}-of-{screen_name}.png"
+        )
+        subprocess.run(
+            ["screencapture", "-C", "-x", "-D", str(display_index + 1), temp_filename]
+        )
+
+        with Image.open(temp_filename) as img:
+            img = img.convert("RGB")
+            current_hash = str(imagehash.phash(img))
+
+            if (
+                screen_name in previous_hashes
+                and imagehash.hex_to_hash(current_hash)
+                - imagehash.hex_to_hash(previous_hashes[screen_name])
+                < threshold
+            ):
+                logging.info(
+                    f"Screenshot for {screen_name} is similar to the previous one. Skipping."
+                )
+                os.remove(temp_filename)
+                yield screen_name, None, "Skipped (similar to previous)"
+                continue
+
+            previous_hashes[screen_name] = current_hash
+            screen_sequences[screen_name] = screen_sequences.get(screen_name, 0) + 1
+
+            metadata = {
+                "timestamp": timestamp,
+                "active_app": app_name,
+                "active_window": window_title,
+                "screen_name": screen_name,
+                "sequence": screen_sequences[screen_name],
+            }
+
+            # Save as WebP with metadata included
+            webp_filename = os.path.join(
+                base_dir, date, f"screenshot-{timestamp}-of-{screen_name}.webp"
+            )
+            img.save(webp_filename, format="WebP", quality=85)
+            write_image_metadata(webp_filename, metadata)
+
+            save_screen_sequences(base_dir, screen_sequences, date)
+
+        os.remove(temp_filename)
+        screenshots.append(webp_filename)
+        yield screen_name, webp_filename, "Saved"
+
+
+def take_screenshot_windows(
+    base_dir,
+    previous_hashes,
+    threshold,
+    screen_sequences,
+    date,
+    timestamp,
+    app_name,
+    window_title,
+):
+    with mss() as sct:
+        for i, monitor in enumerate(
+            sct.monitors[1:], 1
+        ):  # Skip the first monitor (entire screen)
+            safe_monitor_name = f"monitor_{i}"
+            logging.info(f"Processing monitor: {safe_monitor_name}")
+
+            webp_filename = os.path.join(
+                base_dir, date, f"screenshot-{timestamp}-of-{safe_monitor_name}.webp"
+            )
+
+            img = sct.grab(monitor)
+            img = Image.frombytes("RGB", img.size, img.bgra, "raw", "BGRX")
+            current_hash = str(imagehash.phash(img))
+
+            if (
+                safe_monitor_name in previous_hashes
+                and imagehash.hex_to_hash(current_hash)
+                - imagehash.hex_to_hash(previous_hashes[safe_monitor_name])
+                < threshold
+            ):
+                logging.info(
+                    f"Screenshot for {safe_monitor_name} is similar to the previous one. Skipping."
+                )
+                yield safe_monitor_name, None, "Skipped (similar to previous)"
+                continue
+
+            previous_hashes[safe_monitor_name] = current_hash
+            screen_sequences[safe_monitor_name] = (
+                screen_sequences.get(safe_monitor_name, 0) + 1
+            )
+
+            metadata = {
+                "timestamp": timestamp,
+                "active_app": app_name,
+                "active_window": window_title,
+                "screen_name": safe_monitor_name,
+                "sequence": screen_sequences[safe_monitor_name],
+            }
+
+            img.save(webp_filename, format="WebP", quality=85)
+            write_image_metadata(webp_filename, metadata)
+            save_screen_sequences(base_dir, screen_sequences, date)
+
+            yield safe_monitor_name, webp_filename, "Saved"
+
+
+def take_screenshot(
+    base_dir, previous_hashes, threshold, screen_sequences, date, timestamp
+):
+    app_name, window_title = get_active_window_info()
+    os.makedirs(os.path.join(base_dir, date), exist_ok=True)
+    worklog_path = os.path.join(base_dir, date, "worklog")
+
+    with open(worklog_path, "a") as worklog:
+        if platform.system() == "Darwin":
+            screenshot_generator = take_screenshot_macos(
+                base_dir,
+                previous_hashes,
+                threshold,
+                screen_sequences,
+                date,
+                timestamp,
+                app_name,
+                window_title,
+            )
+        elif platform.system() == "Windows":
+            screenshot_generator = take_screenshot_windows(
+                base_dir,
+                previous_hashes,
+                threshold,
+                screen_sequences,
+                date,
+                timestamp,
+                app_name,
+                window_title,
+            )
+        else:
+            raise NotImplementedError(
+                f"Unsupported operating system: {platform.system()}"
+            )
+
+        screenshots = []
+        for screen_name, screenshot_file, status in screenshot_generator:
+            worklog.write(f"{timestamp} - {screen_name} - {status}\n")
+            if screenshot_file:
+                screenshots.append(screenshot_file)
+
+    return screenshots
+
+
+def is_screen_locked():
+    if platform.system() == "Darwin":
+        session_dict = CGSessionCopyCurrentDictionary()
+        if session_dict:
+            screen_locked = session_dict.get("CGSSessionScreenIsLocked", 0)
+            return bool(screen_locked)
+        return False
+    elif platform.system() == "Windows":
+        user32 = ctypes.windll.User32
+        return user32.GetForegroundWindow() == 0
+
+
+def run_screen_recorder_once(threshold, base_dir, previous_hashes):
+    if not is_screen_locked():
+        date = time.strftime("%Y%m%d")
+        timestamp = time.strftime("%Y%m%d-%H%M%S")
+        screen_sequences = load_screen_sequences(base_dir, date)
+        screenshot_files = take_screenshot(
+            base_dir, previous_hashes, threshold, screen_sequences, date, timestamp
+        )
+        for screenshot_file in screenshot_files:
+            logging.info(f"Screenshot saved: {screenshot_file}")
+        save_previous_hashes(base_dir, previous_hashes)
+    else:
+        logging.info("Screen is locked. Skipping screenshot.")
+
+
+def run_screen_recorder(threshold, base_dir, previous_hashes):
+    while True:
+        try:
+            if not is_screen_locked():
+                date = time.strftime("%Y%m%d")
+                timestamp = time.strftime("%Y%m%d-%H%M%S")
+                screen_sequences = load_screen_sequences(base_dir, date)
+                screenshot_files = take_screenshot(
+                    base_dir,
+                    previous_hashes,
+                    threshold,
+                    screen_sequences,
+                    date,
+                    timestamp,
+                )
+                for screenshot_file in screenshot_files:
+                    logging.info(f"Screenshot saved: {screenshot_file}")
+            else:
+                logging.info("Screen is locked. Skipping screenshot.")
+        except Exception as e:
+            logging.error(f"An error occurred: {str(e)}. Skipping this iteration.")
+
+        time.sleep(settings.record_interval)
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Screen Recorder")
+    parser.add_argument(
+        "--threshold", type=int, default=4, help="Threshold for image similarity"
+    )
+    parser.add_argument("--base-dir", type=str, help="Base directory for screenshots")
+    parser.add_argument("--once", action="store_true", help="Run once and exit")
+    args = parser.parse_args()
+
+    base_dir = (
+        os.path.expanduser(args.base_dir) if args.base_dir else settings.resolved_screenshots_dir
+    )
+    previous_hashes = load_previous_hashes(base_dir)
+
+    if args.once:
+        run_screen_recorder_once(args, base_dir, previous_hashes)
+    else:
+        while True:
+            try:
+                run_screen_recorder(args, base_dir, previous_hashes)
+            except Exception as e:
+                logging.error(
+                    f"Critical error occurred, program will restart in 10 seconds: {str(e)}"
+                )
+                time.sleep(10)
+
+
+if __name__ == "__main__":
+    main()

memos/schemas.py

@@ -0,0 +1,289 @@
+from pydantic import (
+    BaseModel,
+    ConfigDict,
+    DirectoryPath,
+    HttpUrl,
+    Field,
+    model_validator,
+)
+from typing import List, Optional, Any, Dict
+from datetime import datetime
+from enum import Enum
+
+
+class FolderType(Enum):
+    DEFAULT = "DEFAULT"
+    DUMMY = "DUMMY"
+
+
+class MetadataSource(Enum):
+    USER_GENERATED = "user_generated"
+    SYSTEM_GENERATED = "system_generated"
+    PLUGIN_GENERATED = "plugin_generated"
+
+
+class MetadataType(Enum):
+    JSON_DATA = "json"
+    TEXT_DATA = "text"
+    NUMBER_DATA = "number"
+
+
+class NewFolderParam(BaseModel):
+    path: DirectoryPath
+    last_modified_at: datetime
+    type: str = FolderType.DEFAULT
+
+
+class NewLibraryParam(BaseModel):
+    name: str
+    folders: List[NewFolderParam] = []
+
+
+class NewFoldersParam(BaseModel):
+    folders: List[NewFolderParam] = []
+
+
+class EntityMetadataParam(BaseModel):
+    key: str
+    value: str
+    source: str
+    data_type: MetadataType
+
+
+class NewEntityParam(BaseModel):
+    filename: str
+    filepath: str
+    size: int
+    file_created_at: datetime
+    file_last_modified_at: datetime
+    file_type: str
+    file_type_group: str
+    folder_id: int
+    tags: List[str] | None = None
+    metadata_entries: List[EntityMetadataParam] | None = None
+
+
+class UpdateEntityParam(BaseModel):
+    size: int | None = None
+    file_created_at: datetime | None = None
+    file_last_modified_at: datetime | None = None
+    file_type: str | None = None
+    file_type_group: str | None = None
+    tags: List[str] | None = None
+    metadata_entries: List[EntityMetadataParam] | None = None
+
+
+class UpdateTagParam(BaseModel):
+    description: str | None
+    color: str | None
+
+
+class UpdateEntityTagsParam(BaseModel):
+    tags: List[str] = []
+
+
+class UpdateEntityMetadataParam(BaseModel):
+    metadata_entries: List[EntityMetadataParam]
+
+
+class NewPluginParam(BaseModel):
+    name: str
+    description: str | None
+    webhook_url: HttpUrl
+
+
+class NewLibraryPluginParam(BaseModel):
+    plugin_id: Optional[int] = None
+    plugin_name: Optional[str] = None
+
+    @model_validator(mode="after")
+    def check_either_id_or_name(self):
+        plugin_id = self.plugin_id
+        plugin_name = self.plugin_name
+        if not (plugin_id or plugin_name):
+            raise ValueError("Either plugin_id or plugin_name must be provided")
+        if plugin_id is not None and plugin_name is not None:
+            raise ValueError("Only one of plugin_id or plugin_name should be provided")
+        return self
+
+
+class Folder(BaseModel):
+    id: int
+    path: str
+    last_modified_at: datetime
+    type: FolderType
+
+    model_config = ConfigDict(from_attributes=True)
+
+
+class Plugin(BaseModel):
+    id: int
+    name: str
+    description: str | None
+    webhook_url: str
+
+    model_config = ConfigDict(from_attributes=True)
+
+
+class Library(BaseModel):
+    id: int
+    name: str
+    folders: List[Folder] = []
+    plugins: List[Plugin] = []
+
+    model_config = ConfigDict(from_attributes=True)
+
+
+class Tag(BaseModel):
+    id: int
+    name: str
+    description: str | None
+    color: str | None
+    created_at: datetime
+    # source: str
+
+    model_config = ConfigDict(from_attributes=True)
+
+
+class EntityMetadata(BaseModel):
+    id: int
+    entity_id: int
+    key: str
+    value: str
+    source: str
+    data_type: MetadataType
+
+    model_config = ConfigDict(from_attributes=True)
+
+
+class EntityPluginStatus(BaseModel):
+    plugin_id: int
+    processed_at: datetime
+
+
+class Entity(BaseModel):
+    id: int
+    filepath: str
+    filename: str
+    size: int
+    file_created_at: datetime
+    file_last_modified_at: datetime
+    file_type: str
+    file_type_group: str
+    last_scan_at: datetime | None
+    folder_id: int
+    library_id: int
+    tags: List[Tag] = []
+    metadata_entries: List[EntityMetadata] = []
+    plugin_status: List[EntityPluginStatus] = []
+
+    model_config = ConfigDict(from_attributes=True)
+
+    def get_metadata_by_key(self, key: str) -> Optional[EntityMetadata]:
+        """
+        Get EntityMetadata by key.
+
+        Args:
+            key (str): The key to search for in metadata entries.
+
+        Returns:
+            Optional[EntityMetadata]: The EntityMetadata if found, None otherwise.
+        """
+        for metadata in self.metadata_entries:
+            if metadata.key == key:
+                return metadata
+        return None
+
+
+class MetadataIndexItem(BaseModel):
+    key: str
+    value: Any
+    source: str
+
+
+class EntitySearchResult(BaseModel):
+    id: str
+    filepath: str
+    filename: str
+    size: int
+    file_created_at: int = Field(..., description="Unix timestamp")
+    created_date: Optional[str] = None
+    created_month: Optional[str] = None
+    created_year: Optional[str] = None
+    file_last_modified_at: int = Field(..., description="Unix timestamp")
+    file_type: str
+    file_type_group: str
+    last_scan_at: Optional[int] = Field(None, description="Unix timestamp")
+    library_id: int
+    folder_id: int
+    tags: List[str]
+    metadata_entries: List[MetadataIndexItem]
+    facets: Optional[Dict[str, Any]] = None
+
+
+class FacetCount(BaseModel):
+    count: int
+    highlighted: str
+    value: str
+
+
+class FacetStats(BaseModel):
+    total_values: int
+
+
+class Facet(BaseModel):
+    counts: List[FacetCount]
+    field_name: str
+    sampled: bool
+    stats: FacetStats
+
+
+class TextMatchInfo(BaseModel):
+    best_field_score: str
+    best_field_weight: int
+    fields_matched: int
+    num_tokens_dropped: int
+    score: str
+    tokens_matched: int
+    typo_prefix_score: int
+
+
+class HybridSearchInfo(BaseModel):
+    rank_fusion_score: float
+
+
+class SearchHit(BaseModel):
+    document: EntitySearchResult
+    highlight: Dict[str, Any] = {}
+    highlights: List[Any] = []
+    hybrid_search_info: Optional[HybridSearchInfo] = None
+    text_match: Optional[int] = None
+    text_match_info: Optional[TextMatchInfo] = None
+
+
+class RequestParams(BaseModel):
+    collection_name: str
+    first_q: str
+    per_page: int
+    q: str
+    app_names: List[str] | None = None
+
+
+class SearchResult(BaseModel):
+    facet_counts: List[Facet]
+    found: int
+    hits: List[SearchHit]
+    out_of: int
+    page: int
+    request_params: RequestParams
+    search_cutoff: bool
+    search_time_ms: int
+
+
+class EntityContext(BaseModel):
+    prev: List[Entity]
+    next: List[Entity]
+
+
+class BatchIndexRequest(BaseModel):
+    entity_ids: List[int]