memlord 0.2.3__py3-none-any.whl

memlord/auth.py

@@ -0,0 +1,42 @@
+from contextlib import asynccontextmanager
+
+import bcrypt
+from fastmcp.server.dependencies import get_access_token
+from sqlalchemy import select
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from fastmcp.dependencies import Depends as MCPDepends
+
+from memlord.config import settings
+from memlord.db import MCPSessionDep
+from memlord.models.oauth_client import OAuthClient
+
+
+def hash_password(password: str) -> str:
+    return bcrypt.hashpw(password.encode(), bcrypt.gensalt()).decode()
+
+
+def verify_password(plain: str, hashed: str) -> bool:
+    return bcrypt.checkpw(plain.encode(), hashed.encode())
+
+
+async def _current_user_gen(
+    s: AsyncSession = MCPSessionDep,  # type: ignore[assignment]
+):
+    access_token = get_access_token()
+    if access_token is None:
+        if settings.stdio_user_id is None:
+            raise PermissionError("Authentication required")
+        yield settings.stdio_user_id
+        return
+    user_id = await s.scalar(
+        select(OAuthClient.user_id).where(
+            OAuthClient.client_id == access_token.client_id
+        )
+    )
+    if user_id is None:
+        raise PermissionError("Unauthenticated")
+    yield user_id
+
+
+MCPUserDep = MCPDepends(asynccontextmanager(_current_user_gen))

memlord/config.py

@@ -0,0 +1,34 @@
+from pathlib import Path
+
+from pydantic import EmailStr, Field
+from pydantic_settings import BaseSettings, SettingsConfigDict
+
+
+class Settings(BaseSettings):
+    model_config = SettingsConfigDict(
+        env_prefix="MEMLORD_", env_file=".env", extra="ignore"
+    )
+
+    db_url: str = "postgresql+asyncpg://postgres:postgres@localhost/memlord"
+    db_echo: bool = False
+
+    model_dir: Path = Path("src/memlord/onnx")
+    host: str = "0.0.0.0"
+    port: int = 8000
+    base_url: str = "http://localhost:8000"
+    rrf_k: int = 60
+    default_limit: int = 10
+    sim_threshold: float = Field(0.25, ge=0.0, le=1.0)
+    dedup_threshold: float = Field(0.85, ge=0.0, le=1.0)
+    oauth_jwt_secret: str = "memlord-dev-secret-please-change"
+    stdio_user_id: int | None = Field(None, description="use for stdio mode")
+
+    smtp_host: str | None = None
+    smtp_port: int = 587
+    smtp_user: str | None = None
+    smtp_password: str | None = None
+    smtp_from: EmailStr | None = None
+    smtp_tls: bool = True
+
+
+settings = Settings()

memlord/dao/__init__.py

@@ -0,0 +1,4 @@
+from .email_token import EmailTokenDao
+from .memory import MemoryDao
+from .user import UserDao
+from .workspace import WorkspaceDao

memlord/dao/email_token.py

@@ -0,0 +1,65 @@
+import hashlib
+import secrets
+from datetime import timedelta
+
+from sqlalchemy import delete, insert, select
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from memlord.models.email_token import EmailToken, TokenPurpose
+from memlord.utils.dt import utcnow
+
+_TTL: dict[TokenPurpose, timedelta] = {
+    TokenPurpose.verify: timedelta(hours=24),
+    TokenPurpose.reset: timedelta(hours=1),
+}
+
+
+def _hash(token: str) -> str:
+    return hashlib.sha256(token.encode()).hexdigest()
+
+
+class EmailTokenDao:
+    def __init__(self, s: AsyncSession) -> None:
+        self._s = s
+
+    async def create(self, user_id: int, purpose: TokenPurpose) -> str:
+        """Create a new token, delete any existing one for same user+purpose, return raw token."""
+        await self._s.execute(
+            delete(EmailToken).where(
+                EmailToken.user_id == user_id,
+                EmailToken.purpose == purpose,
+            )
+        )
+        raw = secrets.token_urlsafe(32)
+        await self._s.execute(
+            insert(EmailToken).values(
+                token_hash=_hash(raw),
+                user_id=user_id,
+                purpose=purpose,
+                expires_at=utcnow() + _TTL[purpose],
+            )
+        )
+        return raw
+
+    async def consume(self, raw: str, purpose: TokenPurpose) -> int | None:
+        """Validate token; on success delete it and return user_id, else return None."""
+        row = (
+            (
+                await self._s.execute(
+                    select(EmailToken.user_id, EmailToken.expires_at).where(
+                        EmailToken.token_hash == _hash(raw),
+                        EmailToken.purpose == purpose,
+                    )
+                )
+            )
+            .mappings()
+            .one_or_none()
+        )
+        if row is None:
+            return None
+        if utcnow() > row["expires_at"]:
+            return None
+        await self._s.execute(
+            delete(EmailToken).where(EmailToken.token_hash == _hash(raw))
+        )
+        return row["user_id"]

memlord/dao/memory.py

@@ -0,0 +1,289 @@
+from pgvector.sqlalchemy import Vector
+from sqlalchemy import bindparam, delete, Float, insert, select, update
+from sqlalchemy.dialects.postgresql import insert as pg_insert
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from memlord.config import settings
+from memlord.dao.workspace import WorkspaceDao
+from memlord.embeddings import embed
+from memlord.models import Memory, MemoryTag, Tag
+from memlord.schemas import MemoryListItem, MemoryType
+
+_UNSET = object()
+
+
+def _embed_text(content: str, tags: set[str]) -> str:
+    return f"{content} {' '.join(sorted(tags))}" if tags else content
+
+
+class MemoryDao:
+    def __init__(self, s: AsyncSession, uid: int) -> None:
+        self._s = s
+        self._uid = uid
+
+    async def _upsert_tags(self, memory_id: int, tags: set[str]) -> None:
+        for tag_name in tags:
+            normalized = tag_name.lower().strip()
+            if not normalized:
+                continue
+            await self._s.execute(
+                pg_insert(Tag).values(name=normalized).on_conflict_do_nothing()
+            )
+            tag_id = await self._s.scalar(select(Tag.id).where(Tag.name == normalized))
+            await self._s.execute(
+                pg_insert(MemoryTag)
+                .values(memory_id=memory_id, tag_id=tag_id)
+                .on_conflict_do_nothing()
+            )
+
+    async def _fetch_tag_names(self, memory_id: int) -> list[str]:
+        rows = await self._s.execute(
+            select(Tag.name)
+            .join(MemoryTag, MemoryTag.tag_id == Tag.id)
+            .where(MemoryTag.memory_id == memory_id)
+        )
+        return [row[0] for row in rows.fetchall()]
+
+    async def _cleanup_orphan_tags(self) -> None:
+        await self._s.execute(delete(Tag).where(~Tag.id.in_(select(MemoryTag.tag_id))))
+
+    async def _replace_tags(self, memory_id: int, tags: set[str]) -> None:
+        await self._s.execute(delete(MemoryTag).where(MemoryTag.memory_id == memory_id))
+        await self._upsert_tags(memory_id, tags)
+        await self._cleanup_orphan_tags()
+
+    async def _check_near_duplicate(
+        self, vector: list[float], workspace_id: int
+    ) -> None:
+        """Raise ValueError if a near-duplicate exists in the workspace."""
+        vec_param = bindparam("vec", type_=Vector(384))
+        distance_expr = Memory.embedding.op("<=>", return_type=Float)(vec_param)
+        dup_row = (
+            (
+                await self._s.execute(
+                    select(Memory.id, distance_expr.label("distance"))
+                    .where(
+                        Memory.embedding.isnot(None),
+                        Memory.workspace_id == workspace_id,
+                    )
+                    .order_by(distance_expr)
+                    .limit(1),
+                    {"vec": vector},
+                )
+            )
+            .mappings()
+            .one_or_none()
+        )
+        if dup_row is None:
+            return
+        similarity = 1.0 - dup_row["distance"]
+        if similarity >= settings.dedup_threshold:
+            raise ValueError(
+                f"Near-duplicate found (id={dup_row['id']}, similarity={round(similarity, 4):.4f}). "
+                f"Review with get_memory({dup_row['id']}). Pass force=True to store anyway."
+            )
+
+    async def _personal_workspace_id(self) -> int:
+        ws = await WorkspaceDao(self._s).get_personal(self._uid)
+        return ws.id
+
+    async def _accessible_workspace_ids(self) -> list[int]:
+        return await WorkspaceDao(self._s).get_accessible_workspace_ids(self._uid)
+
+    async def create(
+        self,
+        content: str,
+        memory_type: MemoryType,
+        metadata: dict,
+        tags: set[str],
+        workspace_id: int | None = None,
+        force: bool = False,
+    ) -> tuple[int, bool]:
+        if workspace_id is None:
+            workspace_id = await self._personal_workspace_id()
+
+        memory_id = await self._s.scalar(
+            select(Memory.id).where(
+                Memory.content == content,
+                Memory.workspace_id == workspace_id,
+            )
+        )
+        if memory_id is not None:
+            return memory_id, False
+
+        vector = await embed(_embed_text(content, tags or []))
+
+        if not force:
+            await self._check_near_duplicate(vector, workspace_id)
+
+        memory_id = await self._s.scalar(
+            insert(Memory)
+            .values(
+                content=str(content),
+                memory_type=MemoryType(memory_type),
+                extra_data=metadata or {},
+                embedding=vector,
+                created_by=self._uid,
+                workspace_id=workspace_id,
+            )
+            .returning(Memory.id)
+        )
+        assert memory_id is not None
+
+        await self._upsert_tags(memory_id, tags or set())
+        return memory_id, True
+
+    async def update(
+        self,
+        id: int,
+        workspace_ids: list[int] | None = None,
+        content: str = _UNSET,  # type: ignore[assignment]
+        memory_type: MemoryType = _UNSET,  # type: ignore[assignment]
+        metadata: dict = _UNSET,  # type: ignore[assignment]
+        tags: set[str] = _UNSET,  # type: ignore[assignment]
+    ) -> int:
+        """Update memory fields. Pass _UNSET to leave a field unchanged; None sets it to NULL."""
+        if workspace_ids is None:
+            workspace_ids = await self._accessible_workspace_ids()
+        access_check = Memory.workspace_id.in_(workspace_ids)
+
+        memory_id = await self._s.scalar(
+            select(Memory.id).where(Memory.id == id, access_check)
+        )
+        if memory_id is None:
+            raise ValueError(f"Memory with id={id} not found")
+
+        values: dict = {}
+        if memory_type is not _UNSET:
+            values["memory_type"] = MemoryType(memory_type)
+        if metadata is not _UNSET:
+            values["extra_data"] = metadata or {}
+
+        if content is not _UNSET or tags is not _UNSET:
+            new_content = (
+                content
+                if content is not _UNSET
+                else (
+                    await self._s.scalar(
+                        select(Memory.content).where(Memory.id == memory_id)
+                    )
+                    or ""
+                )
+            )
+            new_tags = (
+                list(tags)
+                if tags is not _UNSET
+                else await self._fetch_tag_names(memory_id)
+            )
+            if content is not _UNSET:
+                values["content"] = content
+            values["embedding"] = await embed(_embed_text(new_content, new_tags))
+
+        if values:
+            await self._s.execute(
+                update(Memory).where(Memory.id == memory_id).values(**values)
+            )
+
+        if tags is not _UNSET:
+            await self._replace_tags(memory_id, tags)  # type: ignore[arg-type]
+
+        return memory_id
+
+    async def delete(self, id: int, workspace_ids: list[int] | None = None) -> None:
+        if workspace_ids is None:
+            workspace_ids = await self._accessible_workspace_ids()
+        access_check = Memory.workspace_id.in_(workspace_ids)
+
+        result = await self._s.scalar(
+            delete(Memory).where(Memory.id == id, access_check).returning(Memory.id)
+        )
+        if result is None:
+            raise ValueError(f"Memory with id={id} not found")
+        await self._cleanup_orphan_tags()
+
+    async def get(
+        self, id: int, workspace_ids: list[int] | None = None
+    ) -> MemoryListItem | None:
+        if workspace_ids is None:
+            workspace_ids = await self._accessible_workspace_ids()
+        access_check = Memory.workspace_id.in_(workspace_ids)
+
+        row = (
+            (
+                await self._s.execute(
+                    select(
+                        Memory.id,
+                        Memory.content,
+                        Memory.memory_type,
+                        Memory.extra_data.label("metadata"),
+                        Memory.created_at,
+                        Memory.workspace_id,
+                    ).where(Memory.id == id, access_check)
+                )
+            )
+            .mappings()
+            .one_or_none()
+        )
+        if row is None:
+            return None
+        tags = (await self.fetch_tags([id])).get(id, [])
+        return MemoryListItem(
+            **row,
+            tags=tags,
+        )
+
+    async def move(self, id: int, target_workspace_id: int) -> None:
+        """Move memory to a different workspace. Raises ValueError if not found or duplicate."""
+        workspace_ids = await self._accessible_workspace_ids()
+        row = (
+            (
+                await self._s.execute(
+                    select(Memory.id, Memory.content).where(
+                        Memory.id == id, Memory.workspace_id.in_(workspace_ids)
+                    )
+                )
+            )
+            .mappings()
+            .one_or_none()
+        )
+        if row is None:
+            raise ValueError(f"Memory with id={id} not found")
+        if target_workspace_id not in workspace_ids:
+            raise PermissionError(
+                f"No access to target workspace {target_workspace_id}"
+            )
+        duplicate = await self._s.scalar(
+            select(Memory.id).where(
+                Memory.content == row["content"],
+                Memory.workspace_id == target_workspace_id,
+                Memory.id != id,
+            )
+        )
+        if duplicate is not None:
+            raise ValueError(
+                "A memory with the same content already exists in the target workspace"
+            )
+        await self._s.execute(
+            update(Memory)
+            .where(Memory.id == id)
+            .values(workspace_id=target_workspace_id)
+        )
+
+    async def fetch_tags(self, memory_ids: list[int]) -> dict[int, list[str]]:
+        rows = await self._s.execute(
+            select(MemoryTag.memory_id, Tag.name)
+            .join(Tag, MemoryTag.tag_id == Tag.id)
+            .where(MemoryTag.memory_id.in_(memory_ids))
+        )
+        result: dict[int, list[str]] = {i: [] for i in memory_ids}
+        for mid, name in rows.fetchall():
+            result[mid].append(name)
+        return result
+
+    async def fetch_metadata(self, memory_ids: list[int]) -> dict[int, tuple]:
+        rows = await self._s.execute(
+            select(Memory.id, Memory.extra_data, Memory.created_at).where(
+                Memory.id.in_(memory_ids)
+            )
+        )
+        return {row.id: (row.extra_data, row.created_at) for row in rows.fetchall()}

memlord/dao/user.py

@@ -0,0 +1,105 @@
+from sqlalchemy import insert, select, update
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from memlord.auth import verify_password
+from memlord.dao.workspace import WorkspaceDao
+from memlord.models.user import User
+from memlord.schemas.user import UserInfo
+
+
+class UserDao:
+    def __init__(self, s: AsyncSession) -> None:
+        self._s = s
+
+    async def authenticate(self, email: str, password: str) -> UserInfo | None:
+        row = (
+            (
+                await self._s.execute(
+                    select(
+                        User.id,
+                        User.display_name,
+                        User.email,
+                        User.email_verified,
+                        User.hashed_password,
+                    ).where(User.email == email.strip().lower())
+                )
+            )
+            .mappings()
+            .one_or_none()
+        )
+        if row is None or not verify_password(password, row["hashed_password"]):
+            return None
+        return UserInfo(
+            id=row["id"],
+            display_name=row["display_name"],
+            email=row["email"],
+            email_verified=row["email_verified"],
+        )
+
+    async def exists_by_email(self, email: str) -> bool:
+        result = await self._s.scalar(
+            select(User.id).where(User.email == email.strip().lower())
+        )
+        return result is not None
+
+    async def get_by_id(self, id: int) -> UserInfo | None:
+        row = (
+            (
+                await self._s.execute(
+                    select(
+                        User.id, User.display_name, User.email, User.email_verified
+                    ).where(User.id == id)
+                )
+            )
+            .mappings()
+            .one_or_none()
+        )
+        if row is None:
+            return None
+        return UserInfo(
+            id=row["id"],
+            display_name=row["display_name"],
+            email=row["email"],
+            email_verified=row["email_verified"],
+        )
+
+    async def get_email_by_id(self, id: int) -> str | None:
+        return await self._s.scalar(select(User.email).where(User.id == id))
+
+    async def get_id_by_email(self, email: str) -> int | None:
+        return await self._s.scalar(
+            select(User.id).where(User.email == email.strip().lower())
+        )
+
+    async def set_email_verified(self, user_id: int) -> None:
+        await self._s.execute(
+            update(User).where(User.id == user_id).values(email_verified=True)
+        )
+
+    async def set_password(self, user_id: int, hashed_password: str) -> None:
+        await self._s.execute(
+            update(User)
+            .where(User.id == user_id)
+            .values(hashed_password=hashed_password)
+        )
+
+    async def create(
+        self, email: str, display_name: str, hashed_password: str
+    ) -> UserInfo:
+        user_id = await self._s.scalar(
+            insert(User)
+            .values(
+                email=email.strip().lower(),
+                display_name=display_name.strip(),
+                hashed_password=hashed_password,
+            )
+            .returning(User.id)
+        )
+        assert user_id is not None
+        await WorkspaceDao(self._s).create_personal(user_id)
+        return UserInfo(
+            id=user_id,
+            display_name=display_name.strip(),
+            email=email.strip().lower(),
+            email_verified=False,
+        )