@m6d/cortex-server 1.1.1 → 1.2.0

package/src/ai/index.ts

@@ -1,10 +1,10 @@
 import {
-  type ToolSet,
-  convertToModelMessages,
-  generateId,
-  generateText,
-  safeValidateUIMessages,
-  streamText,
+    type ToolSet,
+    convertToModelMessages,
+    generateId,
+    generateText,
+    safeValidateUIMessages,
+    streamText,
 } from "ai";
 import { HTTPException } from "hono/http-exception";
 import type { ResolvedCortexAgentConfig } from "../config.ts";
@@ -15,131 +15,123 @@ import { createQueryGraphTool } from "./tools/query-graph.tool.ts";
 import { createCallEndpointTool } from "./tools/call-endpoint.tool.ts";
 import { createExecuteCodeTool } from "./tools/execute-code.tool.ts";
 import { captureFilesTool } from "./tools/capture-files.tool.ts";
-import { createCapturedFileInterceptor } from "./interceptors/resolve-captured-files.ts";
+import { createRequestInterceptor } from "./interceptors/request-interceptor.ts";
 import { createNeo4jClient } from "../graph/neo4j.ts";
 import { resolveFromGraph } from "../graph/resolver.ts";
 import { notify } from "../ws/index.ts";
 
 export async function stream(
-  messages: unknown[],
-  thread: Thread,
-  userId: string,
-  token: string,
-  config: ResolvedCortexAgentConfig,
+    messages: unknown[],
+    thread: Thread,
+    userId: string,
+    token: string,
+    config: ResolvedCortexAgentConfig,
 ) {
-  const validationResult = await safeValidateUIMessages({ messages });
-  if (!validationResult.success) {
-    throw new HTTPException(423, { message: "Invalid messages format" });
-  }
-
-  const validatedMessages = validationResult.data;
-  await config.db.messages.upsert(thread.id, validatedMessages);
-
-  const originalMessages = await config.db.messages
-    .list(userId, thread.id, { limit: 20 })
-    .then((x) => x.map((y) => y.content));
-
-  const recentMessages = await convertToModelMessages(originalMessages);
-
-  const prompt =
-    originalMessages
-      .filter((x) => x.role === "user")
-      .at(-1)
-      ?.parts.find((x) => x.type === "text")?.text ?? "";
-
-  const model = createModel(config.model);
-  const embeddingModel = createEmbeddingModel(config.embedding);
-  const neo4j = createNeo4jClient(config.neo4j, embeddingModel);
-
-  // Pre-resolve graph context
-  const resolved = await resolveFromGraph(prompt, {
-    neo4j,
-    embeddingModel,
-    reranker: config.reranker,
-  });
-
-  // Build tools
-  const builtInTools: Record<string, unknown> = {
-    captureFiles: captureFilesTool,
-    queryGraph: createQueryGraphTool(neo4j),
-  };
-
-  if (config.backendFetch) {
-    const backendFetchWithInterceptor = {
-      ...config.backendFetch,
-      transformRequestBody:
-        config.backendFetch.transformRequestBody ??
-        createCapturedFileInterceptor(config.db, config.storage),
+    const validationResult = await safeValidateUIMessages({ messages });
+    if (!validationResult.success) {
+        throw new HTTPException(423, { message: "Invalid messages format" });
+    }
+
+    const validatedMessages = validationResult.data;
+    await config.db.messages.upsert(thread.id, validatedMessages);
+
+    const originalMessages = await config.db.messages
+        .list(userId, thread.id, { limit: 20 })
+        .then((x) => x.map((y) => y.content));
+
+    const recentMessages = await convertToModelMessages(originalMessages);
+
+    const prompt =
+        originalMessages
+            .filter((x) => x.role === "user")
+            .at(-1)
+            ?.parts.find((x) => x.type === "text")?.text ?? "";
+
+    const model = createModel(config.model);
+    const embeddingModel = createEmbeddingModel(config.embedding);
+    const neo4j = createNeo4jClient(config.neo4j, embeddingModel);
+
+    // Pre-resolve graph context
+    const resolved = await resolveFromGraph(prompt, {
+        neo4j,
+        embeddingModel,
+        reranker: config.reranker,
+    });
+
+    // Build tools
+    const builtInTools: Record<string, unknown> = {
+        captureFiles: captureFilesTool,
+        queryGraph: createQueryGraphTool(neo4j),
     };
 
-    builtInTools["callEndpoint"] = createCallEndpointTool(
-      backendFetchWithInterceptor,
-      token,
-    );
-    builtInTools["executeCode"] = createExecuteCodeTool(
-      backendFetchWithInterceptor,
-      token,
-    );
-  }
-
-  const tools = {
-    ...builtInTools,
-    ...config.tools,
-  } as ToolSet;
-
-  const systemPrompt = await buildSystemPrompt(
-    config,
-    prompt,
-    thread,
-    token,
-    resolved,
-  );
-
-  const result = streamText({
-    model,
-    system: systemPrompt,
-    tools,
-    messages: recentMessages,
-    onAbort: () => {
-      console.log("Stream aborted");
-    },
-  });
-
-  return result.toUIMessageStreamResponse({
-    originalMessages,
-    generateMessageId: generateId,
-    onFinish: async ({ messages: finishedMessages, isAborted }) => {
-      if (isAborted) {
-        console.log("Stream was aborted");
-      }
-      await config.db.messages.upsert(thread.id, finishedMessages);
-      config.onStreamFinish?.({ messages: finishedMessages, isAborted });
-    },
-  });
+    if (config.backendFetch) {
+        const backendFetchWithInterceptor = {
+            ...config.backendFetch,
+            transformRequestBody:
+                config.backendFetch.transformRequestBody ??
+                createRequestInterceptor(
+                    config.db,
+                    config.storage,
+                    config.backendFetch.interceptor,
+                ),
+        };
+
+        builtInTools["callEndpoint"] = createCallEndpointTool(backendFetchWithInterceptor, token);
+        builtInTools["executeCode"] = createExecuteCodeTool(backendFetchWithInterceptor, token);
+    }
+
+    const tools = {
+        ...builtInTools,
+        ...config.tools,
+    } as ToolSet;
+
+    const systemPrompt = await buildSystemPrompt(config, prompt, thread, token, resolved);
+
+    const result = streamText({
+        model,
+        system: systemPrompt,
+        tools,
+        messages: recentMessages,
+        onAbort: () => {
+            console.log("Stream aborted");
+        },
+    });
+
+    return result.toUIMessageStreamResponse({
+        originalMessages,
+        generateMessageId: generateId,
+        onFinish: async ({ messages: finishedMessages, isAborted }) => {
+            if (isAborted) {
+                console.log("Stream was aborted");
+            }
+            await config.db.messages.upsert(thread.id, finishedMessages);
+            config.onStreamFinish?.({ messages: finishedMessages, isAborted });
+        },
+    });
 }
 
 export async function generateTitle(
-  threadId: string,
-  prompt: string,
-  userId: string,
-  config: ResolvedCortexAgentConfig,
+    threadId: string,
+    prompt: string,
+    userId: string,
+    config: ResolvedCortexAgentConfig,
 ) {
-  const model = createModel(config.model);
+    const model = createModel(config.model);
 
-  const { output } = await generateText({
-    model,
-    system: `You are an expert in generating titles for threads of chats
+    const { output } = await generateText({
+        model,
+        system: `You are an expert in generating titles for threads of chats
 given the first message in that thread.
 
 When asked, only respond with the title without saying you're
 going to do so or any other speech. Spit out only the title.`,
-    prompt: `Generate a title for this prompt: ${prompt}`,
-  });
+        prompt: `Generate a title for this prompt: ${prompt}`,
+    });
 
-  await config.db.threads.updateTitle(threadId, output ?? "");
+    await config.db.threads.updateTitle(threadId, output ?? "");
 
-  notify(userId, {
-    type: "thread:title-updated",
-    payload: { threadId, title: output ?? "" },
-  });
+    notify(userId, {
+        type: "thread:title-updated",
+        payload: { threadId, title: output ?? "" },
+    });
 }

package/src/ai/interceptors/request-interceptor.ts

@@ -0,0 +1,61 @@
+import type { DatabaseAdapter } from "../../adapters/database.ts";
+import type { StorageAdapter } from "../../adapters/storage.ts";
+import { streamToBase64, tokenToUserId } from "../helpers.ts";
+
+export type ResolvedFile = {
+    name: string;
+    bytes: string;
+};
+
+export type RequestInterceptorOptions = {
+    transformFile?: (file: ResolvedFile) => unknown;
+};
+
+export function createRequestInterceptor(
+    db: DatabaseAdapter,
+    storage: StorageAdapter,
+    options?: RequestInterceptorOptions,
+) {
+    const transformFile = options?.transformFile ?? ((file: ResolvedFile) => file);
+
+    return async function resolveRequestBody(
+        body: Record<string, unknown>,
+        context: { token: string },
+    ) {
+        const userId = tokenToUserId(context.token);
+
+        async function traverse(obj: Record<string, unknown>) {
+            await Promise.all(
+                Object.keys(obj).map(async (key) => {
+                    const value = obj[key];
+                    if (typeof value === "string" && value.startsWith("capturedFile")) {
+                        const [_, uploadId] = value.split("#");
+                        const file = await db.capturedFiles.getById(uploadId!, userId);
+                        if (file) {
+                            const fileStream = await storage.stream(`captured_files/${uploadId}`);
+                            const bytes = await streamToBase64(fileStream);
+                            obj[key] = transformFile({ name: file.name, bytes });
+                        }
+                    }
+
+                    if (Array.isArray(value)) {
+                        await Promise.all(
+                            value.map(async (child) => {
+                                if (typeof child === "object" && child !== null) {
+                                    await traverse(child as Record<string, unknown>);
+                                }
+                            }),
+                        );
+                    }
+
+                    if (typeof value === "object" && value !== null && !Array.isArray(value)) {
+                        await traverse(value as Record<string, unknown>);
+                    }
+                }),
+            );
+        }
+
+        await traverse(body);
+        return body;
+    };
+}

package/src/ai/prompt.ts

@@ -3,118 +3,115 @@ import type { ResolvedCortexAgentConfig } from "../config.ts";
 import type { Thread } from "../types.ts";
 
 export async function buildSystemPrompt(
-  config: ResolvedCortexAgentConfig,
-  prompt: string,
-  thread: Thread,
-  token: string,
-  resolved: ResolvedContext | null,
+    config: ResolvedCortexAgentConfig,
+    prompt: string,
+    thread: Thread,
+    token: string,
+    resolved: ResolvedContext | null,
 ) {
-  // Resolve the consumer's base system prompt
-  let basePrompt: string;
-  if (typeof config.systemPrompt === "function") {
-    // Resolve session data with caching
-    let session: Record<string, unknown> | null = thread.session as Record<
-      string,
-      unknown
-    > | null;
+    // Resolve the consumer's base system prompt
+    let basePrompt: string;
+    if (typeof config.systemPrompt === "function") {
+        // Resolve session data with caching
+        let session: Record<string, unknown> | null = thread.session as Record<
+            string,
+            unknown
+        > | null;
 
-    if (!session && config.loadSessionData) {
-      session = await config.loadSessionData(token);
-      // Persist to DB for future cache hits
-      await config.db.threads.updateSession(thread.id, session);
-      thread.session = session;
-    }
+        if (!session && config.loadSessionData) {
+            session = await config.loadSessionData(token);
+            // Persist to DB for future cache hits
+            await config.db.threads.updateSession(thread.id, session);
+            thread.session = session;
+        }
 
-    basePrompt = await config.systemPrompt(session);
-  } else {
-    basePrompt = config.systemPrompt;
-  }
+        basePrompt = await config.systemPrompt(session);
+    } else {
+        basePrompt = config.systemPrompt;
+    }
 
-  const sections: string[] = [basePrompt];
+    const sections: string[] = [basePrompt];
 
-  // Pre-resolved endpoints from knowledge graph
-  if (resolved) {
-    sections.push(buildResolvedSection(resolved));
-  }
+    // Pre-resolved endpoints from knowledge graph
+    if (resolved) {
+        sections.push(buildResolvedSection(resolved));
+    }
 
-  return sections.join("\n");
+    return sections.join("\n");
 }
 
 function buildResolvedSection(resolved: ResolvedContext) {
-  const parts: string[] = [
-    `
+    const parts: string[] = [
+        `
 ## Pre-resolved API Endpoints
 The following endpoints were automatically matched to the user's message.`,
-  ];
+    ];
 
-  for (const ep of resolved.readEndpoints) {
-    const rules =
-      ep.rules.length > 0 ? `\n  Rules: ${ep.rules.join("; ")}` : "";
-    const deps =
-      ep.dependencies.length > 0
-        ? "\n  Dependencies:\n" +
-          ep.dependencies
-            .map(
-              (d) =>
-                `    - Call ${d.depMethod} ${d.depPath} first → use its "${d.fromField}" as "${d.paramName}"`,
-            )
-            .join("\n")
-        : "";
-    parts.push(
-      `
+    for (const ep of resolved.readEndpoints) {
+        const rules = ep.rules.length > 0 ? `\n  Rules: ${ep.rules.join("; ")}` : "";
+        const deps =
+            ep.dependencies.length > 0
+                ? "\n  Dependencies:\n" +
+                  ep.dependencies
+                      .map(
+                          (d) =>
+                              `    - Call ${d.depMethod} ${d.depPath} first → use its "${d.fromField}" as "${d.paramName}"`,
+                      )
+                      .join("\n")
+                : "";
+        parts.push(
+            `
 ### ${ep.concept} (read)
 - Endpoint: ${ep.name}
 - ${ep.method} ${ep.path}
 - Params: ${ep.params}
 - Body: ${ep.body}
 - Response: ${ep.response}${rules}${deps}`,
-    );
-  }
+        );
+    }
 
-  for (const ep of resolved.writeEndpoints) {
-    const rules =
-      ep.rules.length > 0 ? `\n  Rules: ${ep.rules.join("; ")}` : "";
-    const deps =
-      ep.dependencies.length > 0
-        ? "\n  Dependencies:\n" +
-          ep.dependencies
-            .map(
-              (d) =>
-                `    - Call ${d.depMethod} ${d.depPath} first → use its "${d.fromField}" as "${d.paramName}"`,
-            )
-            .join("\n")
-        : "";
-    parts.push(
-      `
+    for (const ep of resolved.writeEndpoints) {
+        const rules = ep.rules.length > 0 ? `\n  Rules: ${ep.rules.join("; ")}` : "";
+        const deps =
+            ep.dependencies.length > 0
+                ? "\n  Dependencies:\n" +
+                  ep.dependencies
+                      .map(
+                          (d) =>
+                              `    - Call ${d.depMethod} ${d.depPath} first → use its "${d.fromField}" as "${d.paramName}"`,
+                      )
+                      .join("\n")
+                : "";
+        parts.push(
+            `
 ### ${ep.concept} (write)
 - Endpoint: ${ep.name}
 - ${ep.method} ${ep.path}
 - Params: ${ep.params}
 - Body: ${ep.body}
 - Response: ${ep.response}${rules}${deps}`,
-    );
-  }
+        );
+    }
 
-  for (const svc of resolved.services) {
-    const rules =
-      svc.rules.length > 0 ? `\n  Rules: ${svc.rules.join("; ")}` : "";
-    parts.push(
-      `
+    for (const svc of resolved.services) {
+        const rules = svc.rules.length > 0 ? `\n  Rules: ${svc.rules.join("; ")}` : "";
+        parts.push(
+            `
 ### ${svc.concept} via ${svc.serviceName} (service)
 - Built-in ID: ${svc.builtInId}
 - Description: ${svc.description || "N/A"}${rules}`,
-    );
-  }
+        );
+    }
 
-  if (
-    resolved.readEndpoints.length === 0 &&
-    resolved.writeEndpoints.length === 0 &&
-    resolved.services.length === 0
-  ) {
-    parts.push(
-      "\nNo matching endpoints found. Use queryGraph to search the knowledge graph manually.",
-    );
-  }
+    if (
+        resolved.readEndpoints.length === 0 &&
+        resolved.writeEndpoints.length === 0 &&
+        resolved.services.length === 0
+    ) {
+        parts.push(
+            "\nNo matching endpoints found. Use queryGraph to search the knowledge graph manually.",
+        );
+    }
 
-  return parts.join("");
+    return parts.join("");
 }