inspiration-collector/tools/memos_client.py

"""Memos API client - fetch memos using Connect RPC protocol."""

import logging
from collections import Counter
from datetime import datetime, timezone, timedelta
from re import findall

import requests

TZ_BEIJING = timezone(timedelta(hours=8))

logger = logging.getLogger(__name__)


class MemosClient:
    """Client for self-hosted Memos API (Connect RPC protocol)."""

    def __init__(self, base_url, access_token):
        self.base_url = base_url.rstrip("/")
        self.session = requests.Session()
        self.session.headers.update({
            "Authorization": f"Bearer {access_token}",
            "Content-Type": "application/json",
        })
        self._user_id = None

    def get_user(self):
        """Get current user info via Connect RPC."""
        resp = self.session.post(
            f"{self.base_url}/memos.api.v1.AuthService/GetCurrentUser",
            json={},
            timeout=10
        )
        resp.raise_for_status()
        data = resp.json()
        user = data.get("user", data)
        self._user_id = user.get("name", "")
        logger.info("Memos user: %s (%s)", user.get("username"), self._user_id)
        return user

    def get_user_id(self):
        """Get current user's resource name (e.g. 'users/FXY')."""
        if not self._user_id:
            self.get_user()
        return self._user_id

    def list_memos(self, days=1, page_size=100):
        """Fetch memos from the last N days (Beijing time reference) via Connect RPC."""
        user = self.get_user_id()
        now_bj = datetime.now(TZ_BEIJING)
        since_bj = now_bj - timedelta(days=days)
        since_utc = since_bj.astimezone(timezone.utc)
        since = since_utc.isoformat()

        payload = {
            "pageSize": page_size,
            "filter": f"creator == '{user}'",
        }

        resp = self.session.post(
            f"{self.base_url}/memos.api.v1.MemoService/ListMemos",
            json=payload,
            timeout=15
        )
        resp.raise_for_status()
        data = resp.json()

        memos = data.get("memos", [])
        logger.info("Memos API OK | fetched %d memos", len(memos))

        results = []
        for m in memos:
            created = m.get("createTime", "")
            try:
                created_dt = datetime.fromisoformat(created.replace("Z", "+00:00"))
            except (ValueError, AttributeError):
                created_dt = datetime.now(timezone.utc)

            if created_dt < since_utc:
                continue

            results.append({
                "id": m.get("name", "").split("/")[-1],
                "content": m.get("content", ""),
                "created_at": created,
                "tags": self._extract_tags(m.get("content", "")),
                "visibility": m.get("visibility", ""),
            })

        logger.info("Filtered to %d memos since %s", len(results), since[:10])
        return results

    def list_all_memos_from_range(self, start_date, end_date, page_size=200):
        """Fetch memos within a date range via Connect RPC."""
        user = self.get_user_id()

        payload = {"pageSize": page_size}

        resp = self.session.post(
            f"{self.base_url}/memos.api.v1.MemoService/ListMemos",
            json=payload,
            timeout=15
        )
        resp.raise_for_status()
        data = resp.json()

        memos = data.get("memos", [])
        logger.info("Memos API OK | fetched %d memos total", len(memos))

        results = []
        for m in memos:
            created = m.get("createTime", "")
            try:
                created_dt = datetime.fromisoformat(created.replace("Z", "+00:00"))
            except (ValueError, AttributeError):
                continue

            if created_dt < start_date or created_dt > end_date:
                continue

            results.append({
                "id": m.get("name", "").split("/")[-1],
                "content": m.get("content", ""),
                "created_at": created,
                "tags": self._extract_tags(m.get("content", "")),
                "visibility": m.get("visibility", ""),
            })

        logger.info(
            "Filtered to %d memos in range %s - %s",
            len(results), start_date.strftime("%m/%d"), end_date.strftime("%m/%d")
        )
        return results

    def list_memos_by_tag(self, tag, page_size=200):
        """Fetch all memos containing a specific tag.

        Useful for retrieving memos across multiple days related to a specific
        topic (e.g. a training course that spans a full week).
        """
        payload = {"pageSize": page_size}

        resp = self.session.post(
            f"{self.base_url}/memos.api.v1.MemoService/ListMemos",
            json=payload,
            timeout=15
        )
        resp.raise_for_status()
        data = resp.json()
        memos = data.get("memos", [])

        results = []
        for m in memos:
            content_text = m.get("content", "")
            tags = self._extract_tags(content_text)
            if tag not in tags:
                continue

            results.append({
                "id": m.get("name", "").split("/")[-1],
                "content": content_text,
                "created_at": m.get("createTime", ""),
                "tags": tags,
                "visibility": m.get("visibility", ""),
            })

        logger.info("Found %d memos with tag #%s", len(results), tag)
        return results

    @staticmethod
    def aggregate_tags(memos):
        """Count tag frequency across a list of memos.

        Returns dict of {tag: count}, sorted by frequency descending.
        """
        counter = Counter()
        for m in memos:
            for tag in m.get("tags", []):
                counter[tag] += 1
        return dict(counter.most_common())

    @staticmethod
    def group_memos_by_tag(memos):
        """Group memos by their tags.

        Returns dict of {tag: [memo_list]}, with "无标签" key for memos
        without any tags. A memo can appear under multiple tags.
        """
        groups = {}
        untagged = []
        for m in memos:
            tags = m.get("tags", [])
            if not tags:
                untagged.append(m)
                continue
            for tag in tags:
                if tag not in groups:
                    groups[tag] = []
                groups[tag].append(m)
        if untagged:
            groups["无标签"] = untagged
        return groups

    @staticmethod
    def _extract_tags(content):
        """Extract #tags from memo content."""
        return findall(r"#(\w[\w\-]*)", content)