@tablestore/tablestore-agent-cli 0.3.1 → 0.4.0

package/dist/dashboard-runtime/dashboard/src/app/api/v1/llm/chat/server-file-memory.ts

@@ -0,0 +1,113 @@
+import type { MemoryClient } from "@tablestore/tablestore-agent-sdk";
+
+import {
+  FileMemoryCache,
+  fileMemoryCacheKey,
+  type FileMemoryTransport,
+  normalizeFileMemoryScope,
+  normalizeFileMemoryStoreName,
+} from "@/lib/file-memory-core";
+import type { ChatRequestBody } from "@/types/llm-chat";
+import { throwIfMemorySecurity } from "./memory-security";
+
+const cache = new FileMemoryCache();
+const clientIds = new WeakMap<MemoryClient, number>();
+let nextClientId = 1;
+
+function clientNamespace(client: MemoryClient): string {
+  const existing = clientIds.get(client);
+  if (existing !== undefined) return `memory-client:${existing}`;
+  const assigned = nextClientId;
+  nextClientId += 1;
+  clientIds.set(client, assigned);
+  return `memory-client:${assigned}`;
+}
+
+function unwrap<T>(
+  response: {
+    code: string;
+    data?: unknown;
+    message?: string;
+    details?: unknown;
+    upstreamStatus?: unknown;
+    httpStatus?: unknown;
+    status?: unknown;
+  },
+  operation: string,
+): T {
+  if (response.code !== "SUCCESS") {
+    throwIfMemorySecurity(response);
+    throw new Error(
+      response.message || `${operation} failed (${response.code})`,
+    );
+  }
+  if (!response.data || typeof response.data !== "object") {
+    throw new Error(`${operation} omitted data`);
+  }
+  return response.data as T;
+}
+
+export async function loadServerFileMemory(input: {
+  client: MemoryClient;
+  binding: NonNullable<ChatRequestBody["memory"]>;
+  query: string;
+  signal: AbortSignal;
+}) {
+  const scope = normalizeFileMemoryScope(input.binding.scope);
+  const memoryStoreName = normalizeFileMemoryStoreName(
+    input.binding.memoryStoreName,
+  );
+  const key = fileMemoryCacheKey(
+    clientNamespace(input.client),
+    memoryStoreName,
+    scope,
+  );
+  const transport: FileMemoryTransport = {
+    async listItems(request, signal) {
+      try {
+        const response = await input.client.listItems(request, signal);
+        return unwrap<{ items: unknown[]; nextToken?: string }>(
+          response,
+          "ListItems",
+        );
+      } catch (error) {
+        throwIfMemorySecurity(error);
+        throw error;
+      }
+    },
+    async getItem(request, signal) {
+      try {
+        const response = await input.client.getItem(request, signal);
+        return unwrap<{
+          content?: unknown;
+          itemId?: unknown;
+          scope?: unknown;
+          path?: unknown;
+          contentSha256?: unknown;
+          contentSizeBytes?: unknown;
+        }>(response, "GetItem");
+      } catch (error) {
+        throwIfMemorySecurity(error);
+        throw error;
+      }
+    },
+  };
+  const loaded = await cache.load({
+    key,
+    memoryStoreName,
+    scope,
+    transport,
+    signal: input.signal,
+  });
+  input.signal.throwIfAborted();
+  const search = cache.search({
+    key,
+    query: input.query,
+    topK: input.binding.topK ?? 5,
+  });
+  return { ...loaded, hits: search.results };
+}
+
+export function clearServerFileMemoryForTests(): void {
+  cache.clear();
+}