feat: enhance memory API responses with limits and update UI components for memory limit handling

This commit is contained in:
Anish Sarkar 2026-05-20 03:17:05 +05:30
parent fa6d7c60bf
commit 73043a0756
9 changed files with 132 additions and 51 deletions

View file

@ -8,7 +8,9 @@ from sqlalchemy.ext.asyncio import AsyncSession
from app.db import User, get_async_session
from app.services.memory import (
MemoryRead,
MemoryScope,
memory_limits,
read_memory,
reset_memory,
save_memory,
@ -18,10 +20,6 @@ from app.users import current_active_user
router = APIRouter()
class MemoryRead(BaseModel):
memory_md: str
class MemoryUpdate(BaseModel):
memory_md: str
@ -36,7 +34,7 @@ async def get_user_memory(
target_id=user.id,
session=session,
)
return MemoryRead(memory_md=memory_md)
return MemoryRead(memory_md=memory_md, limits=memory_limits())
@router.put("/users/me/memory", response_model=MemoryRead)
@ -53,7 +51,7 @@ async def update_user_memory(
)
if result.status == "error":
raise HTTPException(status_code=400, detail=result.message)
return MemoryRead(memory_md=result.memory_md)
return MemoryRead(memory_md=result.memory_md, limits=memory_limits())
@router.post("/users/me/memory/reset", response_model=MemoryRead)
@ -68,4 +66,4 @@ async def reset_user_memory(
)
if result.status == "error":
raise HTTPException(status_code=400, detail=result.message)
return MemoryRead(memory_md=result.memory_md)
return MemoryRead(memory_md=result.memory_md, limits=memory_limits())

View file

@ -8,7 +8,9 @@ from sqlalchemy.ext.asyncio import AsyncSession
from app.db import User, get_async_session
from app.services.memory import (
MemoryRead,
MemoryScope,
memory_limits,
read_memory,
reset_memory,
save_memory,
@ -19,15 +21,11 @@ from app.utils.rbac import check_search_space_access
router = APIRouter()
class TeamMemoryRead(BaseModel):
memory_md: str
class TeamMemoryUpdate(BaseModel):
memory_md: str
@router.get("/searchspaces/{search_space_id}/memory", response_model=TeamMemoryRead)
@router.get("/searchspaces/{search_space_id}/memory", response_model=MemoryRead)
async def get_team_memory(
search_space_id: int,
session: AsyncSession = Depends(get_async_session),
@ -39,10 +37,10 @@ async def get_team_memory(
target_id=search_space_id,
session=session,
)
return TeamMemoryRead(memory_md=memory_md)
return MemoryRead(memory_md=memory_md, limits=memory_limits())
@router.put("/searchspaces/{search_space_id}/memory", response_model=TeamMemoryRead)
@router.put("/searchspaces/{search_space_id}/memory", response_model=MemoryRead)
async def update_team_memory(
search_space_id: int,
body: TeamMemoryUpdate,
@ -58,10 +56,10 @@ async def update_team_memory(
)
if result.status == "error":
raise HTTPException(status_code=400, detail=result.message)
return TeamMemoryRead(memory_md=result.memory_md)
return MemoryRead(memory_md=result.memory_md, limits=memory_limits())
@router.post("/searchspaces/{search_space_id}/memory/reset", response_model=TeamMemoryRead)
@router.post("/searchspaces/{search_space_id}/memory/reset", response_model=MemoryRead)
async def reset_team_memory(
search_space_id: int,
session: AsyncSession = Depends(get_async_session),
@ -75,4 +73,4 @@ async def reset_team_memory(
)
if result.status == "error":
raise HTTPException(status_code=400, detail=result.message)
return TeamMemoryRead(memory_md=result.memory_md)
return MemoryRead(memory_md=result.memory_md, limits=memory_limits())

View file

@ -1,9 +1,11 @@
"""First-class memory service for user and team markdown memory."""
from .schemas import MemoryLimits, MemoryRead
from .service import (
MemoryScope,
SaveResult,
extract_and_save,
memory_limits,
read_memory,
reset_memory,
save_memory,
@ -18,9 +20,12 @@ from .validation import (
__all__ = [
"MEMORY_HARD_LIMIT",
"MEMORY_SOFT_LIMIT",
"MemoryLimits",
"MemoryRead",
"MemoryScope",
"SaveResult",
"extract_and_save",
"memory_limits",
"read_memory",
"reset_memory",
"save_memory",

View file

@ -1,4 +1,4 @@
"""Structured output schemas for memory extraction."""
"""Schemas for memory API responses and structured extraction."""
from __future__ import annotations
@ -7,6 +7,20 @@ from typing import Literal
from pydantic import BaseModel, Field
class MemoryLimits(BaseModel):
"""Canonical memory size limits exposed to clients."""
soft: int
hard: int
class MemoryRead(BaseModel):
"""Memory document payload returned by user and team memory APIs."""
memory_md: str
limits: MemoryLimits
class MemoryExtractionDecision(BaseModel):
"""Structured extraction result; avoids string sentinel parsing."""

View file

@ -9,7 +9,6 @@ from typing import Any, Literal
from uuid import UUID
from langchain_core.messages import HumanMessage
from pydantic import BaseModel
from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession
@ -19,9 +18,10 @@ from app.services.memory.prompts import (
USER_MEMORY_EXTRACT_PROMPT,
)
from app.services.memory.rewrite import forced_rewrite
from app.services.memory.schemas import MemoryExtractionDecision
from app.services.memory.schemas import MemoryExtractionDecision, MemoryLimits
from app.services.memory.validation import (
MEMORY_HARD_LIMIT,
MEMORY_SOFT_LIMIT,
soft_limit_warning,
strip_preamble_to_first_heading,
validate_bullet_format,
@ -68,8 +68,8 @@ class SaveResult:
return data
class MemoryRead(BaseModel):
memory_md: str
def memory_limits() -> MemoryLimits:
return MemoryLimits(soft=MEMORY_SOFT_LIMIT, hard=MEMORY_HARD_LIMIT)
def _normalize_scope(scope: MemoryScope | str) -> MemoryScope:

View file

@ -14,11 +14,11 @@ import {
DropdownMenuTrigger,
} from "@/components/ui/dropdown-menu";
import { Spinner } from "@/components/ui/spinner";
import { MEMORY_HARD_LIMIT, useUserMemory } from "@/hooks/use-memory";
import { getMemoryLimitState, useUserMemory } from "@/hooks/use-memory";
export function MemoryContent() {
const activeSearchSpaceId = useAtomValue(activeSearchSpaceIdAtom);
const { memory, displayMemory, loading, saving, reset } = useUserMemory(
const { memory, displayMemory, limits, loading, saving, reset } = useUserMemory(
Number(activeSearchSpaceId)
);
@ -59,11 +59,11 @@ export function MemoryContent() {
};
const charCount = memory.length;
const limitState = getMemoryLimitState(charCount, limits);
const getCounterColor = () => {
if (charCount > MEMORY_HARD_LIMIT) return "text-red-500";
if (charCount > 15_000) return "text-orange-500";
if (charCount > 10_000) return "text-yellow-500";
if (limitState.level === "error") return "text-red-500";
if (limitState.level === "warning") return "text-orange-500";
return "text-muted-foreground";
};
@ -112,13 +112,7 @@ export function MemoryContent() {
</div>
<div className="flex items-center justify-between gap-2">
<span className={`text-xs shrink-0 ${getCounterColor()}`}>
{charCount.toLocaleString()} / {MEMORY_HARD_LIMIT.toLocaleString()}
<span className="hidden sm:inline"> characters</span>
<span className="sm:hidden"> chars</span>
{charCount > 15_000 && charCount <= MEMORY_HARD_LIMIT && " - Approaching limit"}
{charCount > MEMORY_HARD_LIMIT && " - Exceeds limit"}
</span>
<span className={`text-xs shrink-0 ${getCounterColor()}`}>{limitState.label}</span>
<div className="flex items-center gap-1.5 sm:gap-2">
<Button
type="button"

View file

@ -21,8 +21,10 @@ import { MarkdownViewer } from "@/components/markdown-viewer";
import { Alert, AlertDescription } from "@/components/ui/alert";
import { Button } from "@/components/ui/button";
import { Drawer, DrawerContent, DrawerHandle, DrawerTitle } from "@/components/ui/drawer";
import { Separator } from "@/components/ui/separator";
import { Spinner } from "@/components/ui/spinner";
import { useMediaQuery } from "@/hooks/use-media-query";
import { getMemoryLimitState, type MemoryLimits } from "@/hooks/use-memory";
import { useElectronAPI } from "@/hooks/use-platform";
import { authenticatedFetch, getBearerToken, redirectToLogin } from "@/lib/auth-utils";
import { inferMonacoLanguageFromPath } from "@/lib/editor-language";
@ -127,6 +129,7 @@ export function EditorPanelContent({
const [saving, setSaving] = useState(false);
const [downloading, setDownloading] = useState(false);
const [isEditing, setIsEditing] = useState(false);
const [memoryLimits, setMemoryLimits] = useState<MemoryLimits | null>(null);
const [editedMarkdown, setEditedMarkdown] = useState<string | null>(null);
const [localFileContent, setLocalFileContent] = useState("");
@ -168,6 +171,7 @@ export function EditorPanelContent({
setLocalFileContent("");
setHasCopied(false);
setIsEditing(false);
setMemoryLimits(null);
initialLoadDone.current = false;
changeCountRef.current = 0;
@ -221,7 +225,11 @@ export function EditorPanelContent({
.catch(() => ({ detail: "Failed to fetch memory" }));
throw new Error(errorData.detail || "Failed to fetch memory");
}
const data = (await response.json()) as { memory_md?: string };
const data = (await response.json()) as {
memory_md?: string;
limits?: MemoryLimits;
};
setMemoryLimits(data.limits ?? null);
const content: EditorContent = {
document_id: memoryScope === "team" ? -1002 : -1001,
title: title || (memoryScope === "team" ? "Team Memory" : "Personal Memory"),
@ -376,9 +384,13 @@ export function EditorPanelContent({
.catch(() => ({ detail: "Failed to save memory" }));
throw new Error(errorData.detail || "Failed to save memory");
}
const data = (await response.json()) as { memory_md?: string };
const data = (await response.json()) as {
memory_md?: string;
limits?: MemoryLimits;
};
const savedContent = data.memory_md ?? markdownRef.current;
markdownRef.current = savedContent;
setMemoryLimits(data.limits ?? memoryLimits);
setEditorDoc((prev) => (prev ? { ...prev, source_markdown: savedContent } : prev));
setEditedMarkdown(null);
if (!options?.silent) {
@ -434,6 +446,7 @@ export function EditorPanelContent({
isLocalFileMode,
isMemoryMode,
localFilePath,
memoryLimits,
memoryScope,
resolveLocalVirtualPath,
searchSpaceId,
@ -455,6 +468,17 @@ export function EditorPanelContent({
const showDesktopHeader = !!onClose;
const showEditingActions = isEditableType && isEditing;
const localFileLanguage = inferMonacoLanguageFromPath(localFilePath);
const activeMarkdown = editedMarkdown ?? editorDoc?.source_markdown ?? "";
const memoryLimitState = isMemoryMode
? getMemoryLimitState(activeMarkdown.length, memoryLimits)
: null;
const memoryCounterClassName =
memoryLimitState?.level === "error"
? "text-red-500"
: memoryLimitState?.level === "warning"
? "text-orange-500"
: "text-muted-foreground";
const saveDisabled = saving || !hasUnsavedChanges || (memoryLimitState?.isOverLimit ?? false);
const handleCancelEditing = useCallback(() => {
const savedContent = editorDoc?.source_markdown ?? "";
@ -540,6 +564,17 @@ export function EditorPanelContent({
<div className="grid h-10 grid-cols-[minmax(0,1fr)_auto] items-center gap-3 border-b px-4">
<div className="min-w-0 flex flex-1 items-center gap-2">
<p className="truncate text-sm text-muted-foreground">{displayTitle}</p>
{memoryLimitState && (
<>
<Separator
orientation="vertical"
className="mx-1 bg-border data-[orientation=vertical]:h-4 data-[orientation=vertical]:w-px dark:bg-white/10"
/>
<span className={`shrink-0 text-xs ${memoryCounterClassName}`}>
{memoryLimitState.label}
</span>
</>
)}
</div>
<div className="flex items-center gap-1 shrink-0">
{showEditingActions ? (
@ -561,7 +596,7 @@ export function EditorPanelContent({
const saveSucceeded = await handleSave({ silent: true });
if (saveSucceeded) setIsEditing(false);
}}
disabled={saving || !hasUnsavedChanges}
disabled={saveDisabled}
>
<span className={saving ? "opacity-0" : ""}>Save</span>
{saving && <Spinner size="xs" className="absolute" />}
@ -613,6 +648,17 @@ export function EditorPanelContent({
<div className="flex h-14 items-center justify-between border-b px-4 shrink-0">
<div className="flex flex-1 min-w-0 items-center gap-2">
<h2 className="text-sm font-semibold truncate">{displayTitle}</h2>
{memoryLimitState && (
<>
<Separator
orientation="vertical"
className="mx-1 bg-border data-[orientation=vertical]:h-4 data-[orientation=vertical]:w-px dark:bg-white/10"
/>
<span className={`shrink-0 text-xs ${memoryCounterClassName}`}>
{memoryLimitState.label}
</span>
</>
)}
</div>
<div className="flex items-center gap-1 shrink-0">
{showEditingActions ? (
@ -634,7 +680,7 @@ export function EditorPanelContent({
const saveSucceeded = await handleSave({ silent: true });
if (saveSucceeded) setIsEditing(false);
}}
disabled={saving || !hasUnsavedChanges}
disabled={saveDisabled}
>
<span className={saving ? "opacity-0" : ""}>Save</span>
{saving && <Spinner size="xs" className="absolute" />}

View file

@ -12,14 +12,14 @@ import {
DropdownMenuTrigger,
} from "@/components/ui/dropdown-menu";
import { Spinner } from "@/components/ui/spinner";
import { MEMORY_HARD_LIMIT, useTeamMemory } from "@/hooks/use-memory";
import { getMemoryLimitState, useTeamMemory } from "@/hooks/use-memory";
interface TeamMemoryManagerProps {
searchSpaceId: number;
}
export function TeamMemoryManager({ searchSpaceId }: TeamMemoryManagerProps) {
const { memory, displayMemory, loading, saving, reset } = useTeamMemory(searchSpaceId);
const { memory, displayMemory, limits, loading, saving, reset } = useTeamMemory(searchSpaceId);
const handleClear = async () => {
try {
@ -58,11 +58,11 @@ export function TeamMemoryManager({ searchSpaceId }: TeamMemoryManagerProps) {
};
const charCount = memory.length;
const limitState = getMemoryLimitState(charCount, limits);
const getCounterColor = () => {
if (charCount > MEMORY_HARD_LIMIT) return "text-red-500";
if (charCount > 15_000) return "text-orange-500";
if (charCount > 10_000) return "text-yellow-500";
if (limitState.level === "error") return "text-red-500";
if (limitState.level === "warning") return "text-orange-500";
return "text-muted-foreground";
};
@ -113,13 +113,7 @@ export function TeamMemoryManager({ searchSpaceId }: TeamMemoryManagerProps) {
</div>
<div className="flex items-center justify-between gap-2">
<span className={`text-xs shrink-0 ${getCounterColor()}`}>
{charCount.toLocaleString()} / {MEMORY_HARD_LIMIT.toLocaleString()}
<span className="hidden sm:inline"> characters</span>
<span className="sm:hidden"> chars</span>
{charCount > 15_000 && charCount <= MEMORY_HARD_LIMIT && " - Approaching limit"}
{charCount > MEMORY_HARD_LIMIT && " - Exceeds limit"}
</span>
<span className={`text-xs shrink-0 ${getCounterColor()}`}>{limitState.label}</span>
<div className="flex items-center gap-1.5 sm:gap-2">
<Button
type="button"

View file

@ -4,13 +4,19 @@ import { useCallback, useEffect, useState } from "react";
import { z } from "zod";
import { baseApiService } from "@/lib/apis/base-api.service";
export const MEMORY_HARD_LIMIT = 25_000;
const MemoryLimitsSchema = z.object({
soft: z.number(),
hard: z.number(),
});
const MemoryReadSchema = z.object({
memory_md: z.string(),
limits: MemoryLimitsSchema,
});
type MemoryScope = "user" | "team";
export type MemoryLimits = z.infer<typeof MemoryLimitsSchema>;
export type MemoryLimitLevel = "ok" | "warning" | "error";
interface UseMemoryOptions {
scope: MemoryScope;
@ -31,8 +37,30 @@ export function stripMemoryDisplayPrefixes(memory: string) {
);
}
export function getMemoryLimitState(length: number, limits?: MemoryLimits | null) {
if (!limits) {
return {
level: "ok" as MemoryLimitLevel,
label: `${length.toLocaleString()} chars`,
isOverLimit: false,
};
}
const isOverLimit = length > limits.hard;
const isNearLimit = length > limits.soft;
const level: MemoryLimitLevel = isOverLimit ? "error" : isNearLimit ? "warning" : "ok";
const suffix = isOverLimit ? " - Exceeds limit" : isNearLimit ? " - Approaching limit" : "";
return {
level,
label: `${length.toLocaleString()}/${limits.hard.toLocaleString()} chars${suffix}`,
isOverLimit,
};
}
export function useMemory({ scope, searchSpaceId, autoLoad = true }: UseMemoryOptions) {
const [memory, setMemory] = useState("");
const [limits, setLimits] = useState<MemoryLimits | null>(null);
const [loading, setLoading] = useState(autoLoad);
const [saving, setSaving] = useState(false);
@ -41,6 +69,7 @@ export function useMemory({ scope, searchSpaceId, autoLoad = true }: UseMemoryOp
try {
const data = await baseApiService.get(getMemoryPath(scope, searchSpaceId), MemoryReadSchema);
setMemory(data.memory_md);
setLimits(data.limits);
return data.memory_md;
} finally {
setLoading(false);
@ -66,6 +95,7 @@ export function useMemory({ scope, searchSpaceId, autoLoad = true }: UseMemoryOp
}
);
setMemory(data.memory_md);
setLimits(data.limits);
return data.memory_md;
} finally {
setSaving(false);
@ -82,6 +112,7 @@ export function useMemory({ scope, searchSpaceId, autoLoad = true }: UseMemoryOp
MemoryReadSchema
);
setMemory(data.memory_md);
setLimits(data.limits);
return data.memory_md;
} finally {
setSaving(false);
@ -91,6 +122,7 @@ export function useMemory({ scope, searchSpaceId, autoLoad = true }: UseMemoryOp
return {
memory,
setMemory,
limits,
displayMemory: stripMemoryDisplayPrefixes(memory),
loading,
saving,