Merge branch 'main' of https://github.com/devforth/adminforth-agent

Author: yaroslav8765

agent/middleware/apiBasedTools.ts

@@ -11,6 +11,8 @@ import {
 } from "../toolCallEvents.js";
 import { ALWAYS_AVAILABLE_API_TOOL_NAMES } from "../tools/index.js";
 import { createApiTool } from "../tools/apiTool.js";
+import type { AgentEventEmitter } from "../../agentEvents.js";
+import type { SequenceDebugCollector } from "./sequenceDebug.js";
 
 function getEnabledApiToolNames(messages: unknown[]) {
   const enabledToolNames = new Set<string>();
@@ -80,11 +82,19 @@ export function createApiBasedToolsMiddleware(
     async wrapToolCall(request, handler) {
       const startedAt = Date.now();
       const toolInput = JSON.stringify(request.toolCall.args ?? {});
-      const { adminUser, emitToolCallEvent, userTimeZone } = request.runtime.context as {
+      const { adminUser, emit, sequenceDebugSink, userTimeZone } = request.runtime.context as {
         adminUser: AdminUser;
-        emitToolCallEvent: ToolCallEventSink;
+        emit?: AgentEventEmitter;
+        sequenceDebugSink: SequenceDebugCollector;
         userTimeZone: string;
       };
+      const emitToolCall: ToolCallEventSink = (event) => {
+        sequenceDebugSink.handleToolCallEvent(event);
+        void emit?.({
+          type: "tool-call",
+          data: event,
+        });
+      };
       const toolArgs = (request.toolCall.args ?? {}) as Record<string, unknown>;
       let toolInfo: string | undefined;
 
@@ -102,7 +112,7 @@ export function createApiBasedToolsMiddleware(
         });
       }
       const toolCallTracker = createToolCallTracker({
-        emit: emitToolCallEvent,
+        emit: emitToolCall,
         toolCallId: request.toolCall.id,
         toolName: request.toolCall.name,
         toolInfo,

agent/middleware/sequenceDebug.ts

@@ -26,6 +26,9 @@ export type SequenceDebug = {
   reasoningTokens: number;
   text: string;
   textTokens: number;
+  uncachedInputTokens: number;
+  cachedInputTokens: number;
+  outputTokens: number;
   cachedTokens: number;
   responseId: string | null;
   toolCalls: SequenceDebugToolCall[];
@@ -45,13 +48,17 @@ type SequenceDebugModelCall = {
   reasoningTokens: number;
   text: string;
   textTokens: number;
+  uncachedInputTokens: number;
+  cachedInputTokens: number;
+  outputTokens: number;
   cachedTokens: number;
   responseId: string | null;
   resultType: SequenceDebugResultType;
 };
 
 type OpenAiUsageMetadata = {
   input_tokens?: number;
+  output_tokens?: number;
   input_token_details?: {
     cache_read?: number;
   };
@@ -82,6 +89,9 @@ function createPendingSequenceDebug(sequenceId: number): PendingSequenceDebug {
     reasoningTokens: 0,
     text: "",
     textTokens: 0,
+    uncachedInputTokens: 0,
+    cachedInputTokens: 0,
+    outputTokens: 0,
     cachedTokens: 0,
     responseId: null,
     toolCalls: [],
@@ -112,6 +122,9 @@ function finalizeSequenceDebug(sequence: PendingSequenceDebug): SequenceDebug {
     reasoningTokens: sequence.reasoningTokens,
     text: sequence.text,
     textTokens: sequence.textTokens,
+    uncachedInputTokens: sequence.uncachedInputTokens,
+    cachedInputTokens: sequence.cachedInputTokens,
+    outputTokens: sequence.outputTokens,
     cachedTokens: sequence.cachedTokens,
     responseId: sequence.responseId,
     toolCalls: sequence.toolCalls.map(({ completed: _completed, ...toolCall }) => toolCall),
@@ -133,6 +146,21 @@ function getDebugModelName(model: SequenceDebugPromptModel) {
   return typeof model.getName === "function" ? model.getName() : undefined;
 }
 
+function getDebugToolName(tool: unknown) {
+  if (!tool || typeof tool !== "object") {
+    return null;
+  }
+
+  const name = (tool as { name?: unknown }).name;
+  return typeof name === "string" ? name : null;
+}
+
+function formatToolsForDebug(tools: unknown[]) {
+  return tools.map((tool) => ({
+    name: getDebugToolName(tool),
+  }));
+}
+
 function stringifyPromptForDebug(params: {
   model: SequenceDebugPromptModel;
   systemMessage: { text: string };
@@ -160,7 +188,7 @@ function stringifyPromptForDebug(params: {
     },
     systemMessage,
     messages,
-    ...(tools.length > 0 ? { tools } : {}),
+    ...(tools.length > 0 ? { tools: formatToolsForDebug(tools) } : {}),
     ...(toolChoice !== undefined ? { toolChoice } : {}),
     ...(modelSettings ? { modelSettings } : {}),
     ...(invocationParams ? { invocationParams } : {}),
@@ -219,6 +247,9 @@ async function countTokens(model: unknown, content: string) {
 }
 
 function extractSequenceResponseDebug(message: AIMessage): SequenceDebugModelCall {
+  const usageMetadata = message.usage_metadata as OpenAiUsageMetadata | undefined;
+  const promptTokens = usageMetadata?.input_tokens ?? 0;
+  const cachedInputTokens = usageMetadata?.input_token_details?.cache_read ?? 0;
   const blocks = getMessageBlocks(message);
   const reasoning = blocks
     .filter((block: any) => block?.type === "reasoning")
@@ -230,16 +261,15 @@ function extractSequenceResponseDebug(message: AIMessage): SequenceDebugModelCal
     .join("");
 
   return {
-    promptTokens:
-      (message.usage_metadata as OpenAiUsageMetadata | undefined)?.input_tokens ??
-      0,
+    promptTokens,
     reasoning,
     reasoningTokens: 0,
     text: textFromBlocks || (typeof message.content === "string" ? message.content : ""),
     textTokens: 0,
-    cachedTokens:
-      (message.usage_metadata as OpenAiUsageMetadata | undefined)
-        ?.input_token_details?.cache_read ?? 0,
+    uncachedInputTokens: Math.max(promptTokens - cachedInputTokens, 0),
+    cachedInputTokens,
+    outputTokens: usageMetadata?.output_tokens ?? 0,
+    cachedTokens: cachedInputTokens,
     responseId:
       (message.response_metadata as OpenAiResponseMetadata | undefined)?.id ??
       null,
@@ -290,6 +320,9 @@ export function createSequenceDebugCollector(): SequenceDebugCollector {
       sequenceDebug.reasoningTokens = params.reasoningTokens;
       sequenceDebug.text = params.text;
       sequenceDebug.textTokens = params.textTokens;
+      sequenceDebug.uncachedInputTokens = params.uncachedInputTokens;
+      sequenceDebug.cachedInputTokens = params.cachedInputTokens;
+      sequenceDebug.outputTokens = params.outputTokens;
       sequenceDebug.cachedTokens = params.cachedTokens;
       sequenceDebug.responseId = params.responseId;
       sequenceDebug.resultType = params.resultType;
@@ -387,6 +420,8 @@ export function createSequenceDebugMiddleware(
         promptTokens,
         reasoningTokens,
         textTokens,
+        uncachedInputTokens: debug.promptTokens ? debug.uncachedInputTokens : promptTokens,
+        outputTokens: debug.outputTokens || reasoningTokens + textTokens,
       });
       return response;
     },

agent/models/AgentModeResolver.ts

@@ -0,0 +1,9 @@
+import type { PluginOptions } from "../../types.js";
+
+export class AgentModeResolver {
+  constructor(private readonly options: PluginOptions) {}
+
+  resolve(modeName?: string | null) {
+    return this.options.modes.find((mode) => mode.name === modeName) ?? this.options.modes[0];
+  }
+}

agent/models/AgentModelFactory.ts

@@ -0,0 +1,28 @@
+import type { CompletionAdapter } from "adminforth";
+import { createAgentChatModel } from "../simpleAgent.js";
+import type { AgentTurnModels } from "../turn/turnTypes.js";
+
+export class AgentModelFactory {
+  constructor(private readonly maxTokens: number) {}
+
+  async create(completionAdapter: CompletionAdapter): Promise<AgentTurnModels> {
+    const [primaryModelSpec, summaryModelSpec] = await Promise.all([
+      createAgentChatModel({
+        adapter: completionAdapter,
+        maxTokens: this.maxTokens,
+        purpose: "primary",
+      }),
+      createAgentChatModel({
+        adapter: completionAdapter,
+        maxTokens: this.maxTokens,
+        purpose: "summary",
+      }),
+    ]);
+
+    return {
+      model: primaryModelSpec.model,
+      summaryModel: summaryModelSpec.model,
+      modelMiddleware: primaryModelSpec.middleware,
+    };
+  }
+}

agent/runtime/AgentContext.ts

@@ -0,0 +1,30 @@
+import type { AdminUser } from "adminforth";
+import { z } from "zod";
+import type { AgentEventEmitter } from "../../agentEvents.js";
+import type { SequenceDebugCollector } from "../middleware/sequenceDebug.js";
+import type { CurrentPageContext } from "../tools/getUserLocation.js";
+import type { AgentTurnContext } from "../turn/turnTypes.js";
+
+export const contextSchema = z.object({
+  adminUser: z.custom<AdminUser>(),
+  userTimeZone: z.string(),
+  sessionId: z.string(),
+  turnId: z.string(),
+  abortSignal: z.custom<AbortSignal>().optional(),
+  currentPage: z.custom<CurrentPageContext>().optional(),
+  chatSurface: z.string().optional(),
+  adminBaseUrl: z.string().optional(),
+  adminPublicOrigin: z.string().optional(),
+  emit: z.custom<AgentEventEmitter>().optional(),
+  sequenceDebugSink: z.custom<SequenceDebugCollector>(),
+});
+
+export function toLangchainAgentContext(
+  context: AgentTurnContext & {
+    adminBaseUrl: string;
+    emit?: AgentEventEmitter;
+    sequenceDebugSink: SequenceDebugCollector;
+  },
+) {
+  return context;
+}

agent/runtime/AgentRuntime.ts

@@ -0,0 +1,68 @@
+import type { IAdminForth } from "adminforth";
+import { createAgent, summarizationMiddleware } from "langchain";
+import type { BaseCheckpointSaver } from "@langchain/langgraph";
+import { createApiBasedToolsMiddleware } from "../middleware/apiBasedTools.js";
+import { createSequenceDebugMiddleware } from "../middleware/sequenceDebug.js";
+import { createAgentLlmMetricsLogger } from "../simpleAgent.js";
+import type { AgentToolProvider } from "../tools/AgentToolProvider.js";
+import type { AgentRuntimeRunInput } from "../turn/turnTypes.js";
+import { contextSchema, toLangchainAgentContext } from "./AgentContext.js";
+
+export type AgentRuntimeOptions = {
+  name: string;
+  getAdminforth: () => IAdminForth;
+  getCheckpointer: () => BaseCheckpointSaver;
+  toolProvider: AgentToolProvider;
+};
+
+export class AgentRuntime {
+  constructor(private readonly options: AgentRuntimeOptions) {}
+
+  async stream(input: AgentRuntimeRunInput) {
+    const apiBasedTools = this.options.toolProvider.getApiBasedTools();
+    const tools = await this.options.toolProvider.getTools(apiBasedTools);
+    const adminforth = this.options.getAdminforth();
+    const apiBasedToolsMiddleware = createApiBasedToolsMiddleware(
+      apiBasedTools,
+      adminforth,
+    );
+    const sequenceDebugMiddleware = createSequenceDebugMiddleware(
+      input.observability.sequenceDebugSink,
+    );
+    const middleware = [
+      apiBasedToolsMiddleware,
+      ...(input.models.modelMiddleware ?? []),
+      sequenceDebugMiddleware,
+      summarizationMiddleware({
+        model: input.models.summaryModel,
+        trigger: { tokens: 1024 * 64 },
+        keep: { messages: 10 },
+      }),
+    ] as const;
+
+    const agent = createAgent({
+      name: this.options.name,
+      model: input.models.model,
+      checkpointer: this.options.getCheckpointer(),
+      tools,
+      contextSchema,
+      middleware,
+    });
+
+    return agent.stream({ messages: input.messages } as any, {
+      streamMode: "messages",
+      recursionLimit: 100,
+      callbacks: [createAgentLlmMetricsLogger()],
+      signal: input.context.abortSignal,
+      configurable: {
+        thread_id: input.context.sessionId,
+      },
+      context: toLangchainAgentContext({
+        ...input.context,
+        adminBaseUrl: adminforth.config.baseUrlSlashed,
+        emit: input.observability.emit,
+        sequenceDebugSink: input.observability.sequenceDebugSink,
+      }),
+    });
+  }
+}