Browse Source

tool meta

Dax Raad 9 months ago
parent
commit
9b564f0b73

+ 0 - 3
js/bun.lock

@@ -12,7 +12,6 @@
         "clipanion": "^4.0.0-rc.4",
         "hono": "^4.7.10",
         "hono-openapi": "^0.4.8",
-        "shell-quote": "^1.8.2",
         "zod": "^3.25.3",
         "zod-openapi": "^4.2.4",
       },
@@ -168,8 +167,6 @@
 
     "scheduler": ["[email protected]", "", { "dependencies": { "loose-envify": "^1.1.0" } }, "sha512-UOShsPwz7NrMUqhR6t0hWjFduvOzbtv7toDH1/hIrfRNIDBnnBWd0CwJTGvTpngVlmwGCdP9/Zl/tVrDqcuYzQ=="],
 
-    "shell-quote": ["[email protected]", "", {}, "sha512-AzqKpGKjrj7EM6rKVQEPpB288oCfnrEIuyoT9cyF4nmGa7V8Zk6f7RRqYisX8X9m+Q7bd632aZW4ky7EhbQztA=="],
-
     "signal-exit": ["[email protected]", "", {}, "sha512-wnD2ZE+l+SPC/uoS0vXeE9L1+0wuaMqKlfz9AMUo38JsyLSBWSFcHR1Rri62LZc12vLr1gb3jl7iwQhgwpAbGQ=="],
 
     "slice-ansi": ["[email protected]", "", { "dependencies": { "ansi-styles": "^6.2.1", "is-fullwidth-code-point": "^5.0.0" } }, "sha512-bSiSngZ/jWeX93BqeIAbImyTbEihizcwNjFoRUIY/T1wWQsfsm2Vw1agPKylXvQTU7iASGdHhyqRlqQzfz+Htg=="],

+ 1 - 1
js/src/app/index.ts

@@ -12,7 +12,6 @@ export namespace App {
   const ctx = Context.create<Info>("app");
 
   export async function create(input: { directory: string }) {
-    Log.file(input.directory);
     log.info("creating");
 
     const config = await Config.load(input.directory);
@@ -20,6 +19,7 @@ export namespace App {
     const dataDir = AppPath.data(input.directory);
     await fs.mkdir(dataDir, { recursive: true });
     log.info("created", { path: dataDir });
+    Log.file(input.directory);
 
     const services = new Map<any, any>();
 

+ 1 - 1
js/src/bus/index.ts

@@ -33,7 +33,7 @@ export namespace Bus {
   export function specs() {
     const children = {} as any;
     for (const [type, def] of registry.entries()) {
-      children[def.type] = def.properties;
+      children["event." + def.type] = def.properties;
     }
     const result = z.toJSONSchema(z.object(children)) as any;
     result.definitions = result.properties;

+ 30 - 1
js/src/server/server.ts

@@ -6,6 +6,7 @@ import { streamSSE } from "hono/streaming";
 import { Session } from "../session/session";
 import { resolver, validator as zValidator } from "hono-openapi/zod";
 import { z } from "zod";
+import "zod-openapi/extend";
 
 export namespace Server {
   const log = Log.create({ service: "server" });
@@ -58,7 +59,35 @@ export namespace Server {
               description: "Successfully created session",
               content: {
                 "application/json": {
-                  schema: resolver(Session.Info),
+                  schema: resolver(
+                    Session.Info.openapi({
+                      ref: "Session.Info",
+                    }),
+                  ),
+                },
+              },
+            },
+          },
+        }),
+        async (c) => {
+          const session = await Session.create();
+          return c.json(session);
+        },
+      )
+      .post(
+        "/session_messages",
+        describeRoute({
+          description: "Get messages for a session",
+          responses: {
+            200: {
+              description: "Successfully created session",
+              content: {
+                "application/json": {
+                  schema: resolver(
+                    Session.Info.openapi({
+                      ref: "Session.Info",
+                    }),
+                  ),
                 },
               },
             },

+ 27 - 5
js/src/session/session.ts

@@ -7,7 +7,6 @@ import { Log } from "../util/log";
 import {
   convertToModelMessages,
   streamText,
-  tool,
   type TextUIPart,
   type ToolInvocationUIPart,
   type UIDataTypes,
@@ -17,9 +16,11 @@ import {
 import { z } from "zod";
 import { BashTool } from "../tool/bash";
 import { EditTool } from "../tool/edit";
-import ANTHROPIC_PROMPT from "./prompt/anthropic.txt";
 import { ViewTool } from "../tool/view";
 
+import ANTHROPIC_PROMPT from "./prompt/anthropic.txt";
+import type { Tool } from "../tool/tool";
+
 export namespace Session {
   const log = Log.create({ service: "session" });
 
@@ -34,7 +35,13 @@ export namespace Session {
   });
   export type Info = z.output<typeof Info>;
 
-  export type Message = UIMessage<{ sessionID: string }>;
+  export type Message = UIMessage<{
+    time: {
+      created: number;
+    };
+    sessionID: string;
+    tool: Record<string, Tool.Metadata>;
+  }>;
 
   const state = App.state("session", () => {
     const sessions = new Map<string, Info>();
@@ -125,7 +132,7 @@ export namespace Session {
       );
     }
     if (msgs.length === 0) {
-      const system: UIMessage<{ sessionID: string }> = {
+      const system: Message = {
         id: Identifier.ascending("message"),
         role: "system",
         parts: [
@@ -136,6 +143,10 @@ export namespace Session {
         ],
         metadata: {
           sessionID,
+          time: {
+            created: Date.now(),
+          },
+          tool: {},
         },
       };
       msgs.push(system);
@@ -147,7 +158,11 @@ export namespace Session {
       id: Identifier.ascending("message"),
       parts,
       metadata: {
+        time: {
+          created: Date.now(),
+        },
         sessionID,
+        tool: {},
       },
     };
     msgs.push(msg);
@@ -170,9 +185,14 @@ export namespace Session {
       role: "assistant",
       parts: [],
       metadata: {
+        time: {
+          created: Date.now(),
+        },
         sessionID,
+        tool: {},
       },
     };
+    const metadata = next.metadata!;
     msgs.push(next);
     let text: TextUIPart | undefined;
     const reader = result.toUIMessageStream().getReader();
@@ -217,10 +237,12 @@ export namespace Session {
               p.toolInvocation.toolCallId === value.toolCallId,
           ) as ToolInvocationUIPart | undefined;
           if (match) {
+            const { output, metadata } = value.result as any;
+            next.metadata!.tool[value.toolCallId] = metadata;
             match.toolInvocation = {
               ...match.toolInvocation,
               state: "result",
-              result: value.result,
+              result: output,
             };
           }
           break;

+ 5 - 3
js/src/tool/bash.ts

@@ -1,5 +1,5 @@
 import { z } from "zod";
-import { tool } from "./tool";
+import { Tool, tool } from "./tool";
 
 const MAX_OUTPUT_LENGTH = 30000;
 const BANNED_COMMANDS = [
@@ -170,7 +170,7 @@ Important:
 - Return an empty response - the user will see the gh output directly
 - Never update git config`;
 
-export const BashTool = tool({
+export const BashTool = Tool.define({
   name: "bash",
   description: DESCRIPTION,
   parameters: z.object({
@@ -193,7 +193,9 @@ export const BashTool = tool({
       timeout: timeout,
     });
     return {
-      content: process.stdout.toString("utf-8"),
+      output: {
+        content: process.stdout.toString("utf-8"),
+      },
     };
   },
 });

+ 13 - 7
js/src/tool/edit.ts

@@ -1,11 +1,11 @@
 import { z } from "zod";
-import { tool } from "./tool";
 import * as fs from "fs";
 import * as path from "path";
 import { Log } from "../util/log";
-import { App } from "../app";
+import { Tool } from "./tool";
+import { FileTimes } from "./util/file-times";
 
-const log = Log.create({ service: "edit-tool" });
+const log = Log.create({ service: "tool.edit" });
 
 // Simple diff generation
 function generateDiff(
@@ -116,7 +116,7 @@ When making edits:
 
 Remember: when making multiple file edits in a row to the same file, you should prefer to send all edits in a single message with multiple calls to this tool, rather than multiple messages with a single call each.`;
 
-export const EditTool = tool({
+export const EditTool = Tool.define({
   name: "edit",
   description: DESCRIPTION,
   parameters: z.object({
@@ -136,14 +136,20 @@ export const EditTool = tool({
 
     // Handle different operations based on parameters
     if (params.old_string === "") {
-      return createNewFile(filePath, params.new_string);
+      return {
+        output: createNewFile(filePath, params.new_string),
+      };
     }
 
     if (params.new_string === "") {
-      return deleteContent(filePath, params.old_string);
+      return {
+        output: deleteContent(filePath, params.old_string),
+      };
     }
 
-    return replaceContent(filePath, params.old_string, params.new_string);
+    return {
+      output: replaceContent(filePath, params.old_string, params.new_string),
+    };
   },
 });
 

+ 44 - 23
js/src/tool/tool.ts

@@ -1,28 +1,49 @@
-import { type Tool, tool as AITool } from "ai";
+import { tool, type Tool as AITool } from "ai";
 import { Log } from "../util/log";
 
 const log = Log.create({ service: "tool" });
 
-export function tool<Params, Result>(
-  tool: Tool<Params, Result> & {
-    name: string;
-  },
-) {
-  return {
-    [tool.name]: AITool({
-      ...tool,
-      execute: async (params, opts) => {
-        log.info("invoking", {
-          id: opts.toolCallId,
-          name: tool.name,
-          ...params,
-        });
-        try {
-          return tool.execute!(params, opts);
-        } catch (e: any) {
-          return "An error occurred: " + e.toString();
-        }
-      },
-    }),
-  };
+export namespace Tool {
+  export interface Metadata {
+    properties: Record<string, any>;
+    time: {
+      start: number;
+      end: number;
+    };
+  }
+  export function define<Params, Output>(
+    input: AITool<Params, { metadata?: any; output: Output }> & {
+      name: string;
+    },
+  ) {
+    return {
+      [input.name]: tool({
+        ...input,
+        execute: async (params, opts) => {
+          log.info("invoking", {
+            id: opts.toolCallId,
+            name: input.name,
+            ...params,
+          });
+          try {
+            const start = Date.now();
+            const result = await input.execute!(params, opts);
+            const metadata: Metadata = {
+              properties: result.metadata,
+              time: {
+                start,
+                end: Date.now(),
+              },
+            };
+            return {
+              metadata,
+              output: result.output,
+            };
+          } catch (e: any) {
+            return "An error occurred: " + e.toString();
+          }
+        },
+      }),
+    };
+  }
 }

+ 5 - 49
js/src/tool/view.ts

@@ -1,7 +1,7 @@
 import { z } from "zod";
-import { tool } from "./tool";
 import * as fs from "fs";
 import * as path from "path";
+import { Tool } from "./tool";
 
 const MAX_READ_SIZE = 250 * 1024;
 const DEFAULT_READ_LIMIT = 2000;
@@ -38,7 +38,7 @@ TIPS:
 - For code exploration, first use Grep to find relevant files, then View to examine them
 - When viewing large files, use the offset parameter to read specific sections`;
 
-export const ViewTool = tool({
+export const ViewTool = Tool.define({
   name: "view",
   description: DESCRIPTION,
   parameters: z.object({
@@ -117,56 +117,12 @@ export const ViewTool = tool({
     }
     output += "\n</file>";
 
-    return output;
+    return {
+      output: output,
+    };
   },
 });
 
-function addLineNumbers(content: string, startLine: number): string {
-  if (!content) {
-    return "";
-  }
-
-  const lines = content.split("\n");
-
-  return lines
-    .map((line, i) => {
-      const lineNum = i + startLine;
-      const numStr = lineNum.toString();
-
-      if (numStr.length >= 6) {
-        return `${numStr}|${line}`;
-      } else {
-        const paddedNum = numStr.padStart(6, " ");
-        return `${paddedNum}|${line}`;
-      }
-    })
-    .join("\n");
-}
-
-function readTextFile(
-  filePath: string,
-  offset: number,
-  limit: number,
-): { content: string; lineCount: number } {
-  const fileContent = fs.readFileSync(filePath, "utf8");
-  const allLines = fileContent.split("\n");
-  let lineCount = allLines.length;
-
-  // Get the lines we want based on offset and limit
-  const selectedLines = allLines
-    .slice(offset, offset + limit)
-    .map((line) =>
-      line.length > MAX_LINE_LENGTH
-        ? line.substring(0, MAX_LINE_LENGTH) + "..."
-        : line,
-    );
-
-  return {
-    content: selectedLines.join("\n"),
-    lineCount,
-  };
-}
-
 function isImageFile(filePath: string): string | false {
   const ext = path.extname(filePath).toLowerCase();
   switch (ext) {

+ 1 - 1
pkg/client/event.go

@@ -10,7 +10,7 @@ import (
 )
 
 var EventMap = map[string]any{
-	"storage.write": StorageWrite{},
+	"storage.write": EventStorageWrite{},
 }
 
 type EventMessage struct {