// memory-tool.ts
import { createAgent } from "langchain";
import { MemorySaver } from "@langchain/langgraph";
import { MemoryVectorStore } from "langchain/vectorstores/memory";
import { OpenAIEmbeddings } from "@langchain/openai";
import { Document } from "@langchain/core/documents";
import { tool } from "@langchain/core/tools";
import { z } from "zod";
 
const embeddings = new OpenAIEmbeddings({ model: "text-embedding-3-small" });
const vectorStore = new MemoryVectorStore(embeddings);
 
// 工具 1：检索历史
const recallMemory = tool(
  async ({ query }) => {
    const docs = await vectorStore.similaritySearch(query, 3);
    if (docs.length === 0) return "（没有找到相关的历史信息）";
    return docs.map((d, i) => `[${i + 1}] ${d.pageContent}`).join("\n");
  },
  {
    name: "recall_memory",
    description:
      "在用户的长期记忆里检索与当前问题语义相关的片段。当用户提到自己之前说过的内容、或问题涉及用户个人偏好/背景时调用。",
    schema: z.object({
      query: z.string().describe("要检索的关键词或问题，越具体越好"),
    }),
  }
);
 
// 工具 2：把值得长期记住的事实写入记忆
const saveMemory = tool(
  async ({ fact }) => {
    await vectorStore.addDocuments([
      new Document({
        pageContent: fact,
        metadata: { timestamp: Date.now() },
      }),
    ]);
    return `已记忆：${fact}`;
  },
  {
    name: "save_memory",
    description:
      "把用户主动透露的、值得长期记住的事实写入记忆。例如姓名、职业、偏好、关键决定。寒暄和临时信息不要写入。",
    schema: z.object({
      fact: z.string().describe("一句话事实，不要带语气词"),
    }),
  }
);
 
const agent = createAgent({
  model: "anthropic:claude-sonnet-4-6",
  tools: [recallMemory, saveMemory],
  systemPrompt: `你是一个有长期记忆的助手。
- 用户透露身份、偏好等信息时，主动调 save_memory 记下来
- 用户的问题涉及他过去说过的内容时，先调 recall_memory 检索
- 不要凭印象回答，靠工具拿事实`,
  checkpointer: new MemorySaver(),
});
 
const config = { configurable: { thread_id: "demo" } };
 
await agent.invoke(
  { messages: [{ role: "user", content: "我叫李华，在上海做全栈，主要用 React + Node" }] },
  config
);
 
await agent.invoke(
  { messages: [{ role: "user", content: "我周末喜欢跑步和看科幻小说" }] },
  config
);
 
// 100 轮后……
const r = await agent.invoke(
  { messages: [{ role: "user", content: "推荐一本适合我的技术书" }] },
  config
);
console.log(r.messages.at(-1)?.text);
// Agent 会先调 recall_memory("用户的技术背景") → 拿到"React + Node + 全栈"
// 然后给出针对性建议

import { OpenAIEmbeddings } from "@langchain/openai";
 
const embeddings = new OpenAIEmbeddings({
  model: "text-embedding-3-small",
  dimensions: 1024, // text-embedding-3 支持降维
});

import { PGVectorStore } from "@langchain/community/vectorstores/pgvector";
import { OpenAIEmbeddings } from "@langchain/openai";
 
const embeddings = new OpenAIEmbeddings({ model: "text-embedding-3-small" });
 
const vectorStore = await PGVectorStore.initialize(embeddings, {
  postgresConnectionOptions: {
    connectionString: process.env.DATABASE_URL!,
  },
  tableName: "user_memories",
  columns: {
    idColumnName: "id",
    vectorColumnName: "embedding",
    contentColumnName: "content",
    metadataColumnName: "metadata",
  },
});

import { Document } from "@langchain/core/documents";
 
await vectorStore.addDocuments([
  new Document({
    pageContent: `用户：${userMsg}\n助手：${aiMsg}`,
    metadata: { userId, timestamp: Date.now() },
  }),
]);

import { createMiddleware } from "langchain";
import { ChatAnthropic } from "@langchain/anthropic";
 
const extractor = new ChatAnthropic({ model: "claude-haiku-4-5", temperature: 0 });
 
export const autoExtractMw = createMiddleware({
  name: "auto-extract",
  afterModel: async (state) => {
    const lastUser = [...state.messages].reverse().find((m) => m.getType() === "human");
    if (!lastUser) return;
 
    const extract = await extractor.invoke([
      {
        role: "system",
        content:
          "提取用户消息中值得长期记住的事实（身份/偏好/决定）。没有则输出 NONE。一句话。",
      },
      { role: "user", content: lastUser.text ?? "" },
    ]);
    const fact = (typeof extract.text === "string" ? extract.text : "").trim();
    if (fact && fact !== "NONE") {
      await vectorStore.addDocuments([
        new Document({ pageContent: fact, metadata: { auto: true, ts: Date.now() } }),
      ]);
    }
  },
});

const retriever = vectorStore.asRetriever({
  k: 5,
  searchType: "mmr",
  searchKwargs: { fetchK: 20, lambda: 0.7 },
});

const docs = await vectorStore.similaritySearch("加班", 3, {
  userId: "user-001", // 只看这个用户的记忆
});

const results = await vectorStore.similaritySearchWithScore("加班可以吗", 5);
const filtered = results.filter(([_, score]) => score < 0.4); // OpenAI cosine distance，越小越像

const agent = createAgent({
  model: "anthropic:claude-sonnet-4-6",
  tools: [recallMemory, saveMemory],
  systemPrompt: "...",
  middleware: [autoExtractMw], // 自动沉淀
  checkpointer: new PostgresSaver(...), // 短期
});

// 写入时记 ttl
await vectorStore.addDocuments([
  new Document({
    pageContent: fact,
    metadata: { ts: Date.now(), ttlDays: 90 },
  }),
]);
 
// 定期任务（pgvector 直接 SQL）
await pool.query(`DELETE FROM user_memories WHERE (metadata->>'ts')::bigint < $1`, [
  Date.now() - 90 * 24 * 3600 * 1000,
]);

const importance = await scorer.invoke(`评估事实重要性 1-10：${fact}`);
if (parseInt(importance.text ?? "0") >= 6) {
  await vectorStore.addDocuments([new Document({ pageContent: fact, metadata: { importance } })]);
}