Memory

import { stepCountIs, streamText, tool, type ModelMessage } from "ai"
import { z } from "zod"

import type { MemoryInput, MemoryRecord } from "@/components/blocks/memory/memory"

type MemoryStore = {
  list: (userId: string) => Promise<MemoryRecord[]>
  add: (userId: string, input: MemoryInput & { source: { label: string; href: string } }) => Promise<MemoryRecord>
  markUsed: (userId: string, ids: string[]) => Promise<void>
  settings: (userId: string) => Promise<{ enabled: boolean; paused: boolean }>
}

function memoryInstructions(memories: MemoryRecord[]) {
  if (memories.length === 0) return ""
  return [
    "What you remember about the user. Use it when it helps, and never mention it unprompted.",
    ...memories.map((memory) => `- [${memory.id}] ${memory.text}`),
  ].join("\n")
}

export async function respond({
  model,
  messages,
  store,
  userId,
  chat,
}: {
  model: Parameters<typeof streamText>[0]["model"]
  messages: ModelMessage[]
  store: MemoryStore
  userId: string
  chat: { id: string; title: string }
}) {
  const { enabled, paused } = await store.settings(userId)
  const active = enabled && !paused
  const memories = active ? await store.list(userId) : []
  if (memories.length > 0) await store.markUsed(userId, memories.map((memory) => memory.id))

  return streamText({
    model,
    messages,
    system: memoryInstructions(memories),
    stopWhen: stepCountIs(5),
    tools: active
      ? {
          save_memory: tool({
            description:
              "Save one durable fact about the user that will help in future chats, such as a preference, their role or an ongoing project. Write it as a short sentence about the user. Don't save secrets, health details or anything they asked you to forget.",
            inputSchema: z.object({
              text: z.string().max(500),
              category: z.enum(["about", "preferences", "work", "projects"]),
            }),
            execute: async ({ text, category }) => {
              const memory = await store.add(userId, {
                text,
                category,
                scope: null,
                source: { label: chat.title, href: `/chat/${chat.id}` },
              })
              return { id: memory.id, text: memory.text }
            },
          }),
        }
      : undefined,
  })
}

"use client"

import * as React from "react"
import { useChat } from "@ai-sdk/react"
import { getToolName, isToolUIPart, type UIMessage } from "ai"

import { Dialog, DialogBody, DialogContent, DialogHeader, DialogTitle } from "@/components/ui/dialog"

import { ChatAssistantMessage, ChatThread, ChatUserMessage } from "../chat-thread/chat-thread"
import { Markdown } from "../markdown/markdown"
import { PromptInput, PromptInputBody, PromptInputSubmit, PromptInputTextarea } from "../prompt-input/prompt-input"
import { Memory, MemoryUpdatedChip, type MemoryChanges, type MemoryRecord } from "@/components/blocks/memory/memory"

function textOf(message: UIMessage) {
  return message.parts.flatMap((part) => (part.type === "text" ? [part.text] : [])).join("")
}

async function request<T>(input: string, init?: RequestInit) {
  const response = await fetch(input, { ...init, headers: { "content-type": "application/json" } })
  if (!response.ok) {
    const body = (await response.json().catch(() => ({}))) as { error?: string }
    throw new Error(body.error ?? "Something went wrong.")
  }
  return (response.status === 204 ? null : await response.json()) as T
}

export function ChatWithMemory() {
  const { messages, status, sendMessage } = useChat()
  const [open, setOpen] = React.useState(false)
  const [highlight, setHighlight] = React.useState<string | null>(null)
  const [memories, setMemories] = React.useState<MemoryRecord[]>([])
  const [settings, setSettings] = React.useState({ enabled: true, history: false, paused: false })

  const manage = async (id?: string) => {
    const data = await request<{ memories: MemoryRecord[]; settings: typeof settings }>("/api/memories")
    setMemories(data.memories)
    setSettings(data.settings)
    setHighlight(id ?? null)
    setOpen(true)
  }

  return (
    <>
      <ChatThread
        busy={status === "submitted" || status === "streaming"}
        composer={
          <PromptInput status={status} onSubmit={({ text }) => sendMessage({ text })}>
            <PromptInputBody>
              <PromptInputTextarea />
              <PromptInputSubmit />
            </PromptInputBody>
          </PromptInput>
        }
      >
        {messages.map((message) =>
          message.role === "user" ? (
            <ChatUserMessage key={message.id} id={message.id} text={textOf(message)} />
          ) : (
            <ChatAssistantMessage key={message.id} id={message.id} text={textOf(message)}>
              <Markdown text={textOf(message)} />
              {message.parts.filter(isToolUIPart).flatMap((part) => {
                if (getToolName(part) !== "save_memory") return []
                const input = part.input as { text?: string } | undefined
                const output = part.state === "output-available" ? (part.output as { id: string; text: string }) : null
                return [
                  <MemoryUpdatedChip
                    key={part.toolCallId}
                    memory={output?.text ?? input?.text ?? ""}
                    action="saved"
                    status={part.state === "output-error" ? "error" : output ? "done" : "saving"}
                    onManage={() => void manage(output?.id)}
                  />,
                ]
              })}
            </ChatAssistantMessage>
          )
        )}
      </ChatThread>
      <Dialog open={open} onOpenChange={setOpen}>
        <DialogContent size="lg">
          <DialogHeader>
            <DialogTitle>Memory</DialogTitle>
          </DialogHeader>
          <DialogBody>
            <div className="pb-6">
            <Memory
              memories={memories}
              enabled={settings.enabled}
              historyEnabled={settings.history}
              paused={settings.paused}
              highlight={highlight}
              onToggle={async (setting, value) => {
                await request("/api/memories/settings", { method: "PATCH", body: JSON.stringify({ [setting]: value }) })
                setSettings((current) => ({ ...current, [setting]: value }))
              }}
              onAdd={async (input) => {
                const memory = await request<MemoryRecord>("/api/memories", { method: "POST", body: JSON.stringify(input) })
                setMemories((current) => [memory, ...current])
                return memory
              }}
              onEdit={async (memory, changes: MemoryChanges) => {
                await request(`/api/memories/${memory.id}`, { method: "PATCH", body: JSON.stringify(changes) })
                setMemories((current) => current.map((entry) => (entry.id === memory.id ? { ...entry, ...changes } : entry)))
              }}
              onDelete={async (memory) => {
                await request(`/api/memories/${memory.id}`, { method: "DELETE" })
                setMemories((current) => current.filter((entry) => entry.id !== memory.id))
              }}
              onDeleteMany={async (targets) => {
                const ids = targets.map((entry) => entry.id)
                await request("/api/memories", { method: "DELETE", body: JSON.stringify({ ids }) })
                setMemories((current) => current.filter((entry) => !ids.includes(entry.id)))
              }}
              onForgetAll={async () => {
                await request("/api/memories?all=true", { method: "DELETE" })
                setMemories([])
              }}
            />
            </div>
          </DialogBody>
        </DialogContent>
      </Dialog>
    </>
  )
}

"use client"

import * as React from "react"
import { IconBookmark } from "@tabler/icons-react"

import { SettingsSection, SettingsShell, type SettingsSectionItem } from "../settings/settings"
import { Memory, MemorySkeleton, type MemoryRecord } from "@/components/blocks/memory/memory"

const sections: SettingsSectionItem[] = [
  {
    id: "personalization",
    label: "Personalization",
    description: "What Hexta remembers about you across chats.",
    icon: <IconBookmark />,
  },
]

type Settings = { enabled: boolean; history: boolean; paused: boolean }

export function PersonalizationSettings({
  load,
  save,
}: {
  load: () => Promise<{ memories: MemoryRecord[]; settings: Settings; used: number; limit: number }>
  save: {
    toggle: (settings: Partial<Settings>) => Promise<void>
    edit: (id: string, text: string) => Promise<void>
    remove: (id: string) => Promise<void>
  }
}) {
  const [data, setData] = React.useState<Awaited<ReturnType<typeof load>> | null>(null)
  const [failed, setFailed] = React.useState(false)

  const reload = React.useCallback(() => {
    setFailed(false)
    load()
      .then(setData)
      .catch(() => setFailed(true))
  }, [load])

  React.useEffect(reload, [reload])

  return (
    <SettingsShell sections={sections}>
      <SettingsSection
        id="personalization"
        status={failed ? "error" : data ? "ready" : "loading"}
        skeleton={<MemorySkeleton rows={3} />}
        onRetry={reload}
      >
        {data ? (
          <Memory
            variant="compact"
            memories={data.memories}
            enabled={data.settings.enabled}
            limit={data.limit}
            used={data.used}
            onToggle={(setting, value) => save.toggle({ [setting]: value })}
            onEdit={async (memory, changes) => {
              if (changes.text) await save.edit(memory.id, changes.text)
            }}
            onDelete={async (memory) => {
              await save.remove(memory.id)
              setData((current) =>
                current ? { ...current, memories: current.memories.filter((entry) => entry.id !== memory.id) } : current
              )
            }}
          />
        ) : null}
      </SettingsSection>
    </SettingsShell>
  )
}