1 files changed, 34 insertions, 14 deletions

diff --git a/src/claude.ts b/src/claude.ts
index 629f9c2..8269378 100644
--- a/src/claude.ts
+++ b/src/claude.ts
@@ -1,9 +1,13 @@
 import Claude from "@anthropic-ai/sdk";
 import { RESPONSE_LENGTH } from "./logic/constants";
-import type { AIModelAPI, ChatMessage } from "./types";
+import type { AIModelAPI, ChatMessage, InputToken } from "./types";
 import { BOOKWORM_SYS } from "./prompts";
 import type { AsyncRes } from "sortug";
-import type { MessageCreateParamsStreaming } from "@anthropic-ai/sdk/resources";
+import type {
+  ImageBlockParam,
+  MessageCreateParamsStreaming,
+  TextBlockParam,
+} from "@anthropic-ai/sdk/resources";
 
 type Message = Claude.Messages.MessageParam;
 
@@ -30,14 +34,31 @@ export default class ClaudeAPI implements AIModelAPI {
       return { role, content: m.text };
     });
   }
+  private buildInput(tokens: InputToken[]): Message[] {
+    // can do base64 for images too
+    const content = tokens.map((t) => {
+      const content =
+        "text" in t
+          ? ({ type: "text", text: t.text } as TextBlockParam)
+          : "img" in t
+            ? ({
+                type: "image",
+                source: { type: "url", url: t.img },
+              } as ImageBlockParam)
+            : ({ type: "text", text: "oy vey" } as TextBlockParam);
+      return content;
+    });
+
+    return [{ role: "user", content }];
+  }
 
-  public async send(input: string | ChatMessage[], sys?: string) {
-    const msgs: ChatMessage[] =
+  // https://docs.anthropic.com/en/api/messages-examples#vision
+  public async send(input: string | InputToken[], sys?: string) {
+    const msgs: Message[] =
       typeof input === "string"
-        ? [{ author: "user", text: input, sent: 0 }]
-        : input;
-    const messages = this.mapMessages(msgs);
-    const truncated = this.truncateHistory(messages);
+        ? [{ role: "user", content: input }]
+        : this.buildInput(input);
+    const truncated = this.truncateHistory(msgs);
     const res = await this.apiCall(truncated, sys);
     return res;
   }
@@ -62,16 +83,15 @@ export default class ClaudeAPI implements AIModelAPI {
   }
 
   public async stream(
-    input: string | ChatMessage[],
+    input: string | InputToken[],
     handle: (c: any) => void,
     sys?: string,
   ) {
-    const msgs: ChatMessage[] =
+    const msgs: Message[] =
       typeof input === "string"
-        ? [{ author: "user", text: input, sent: 0 }]
-        : input;
-    const messages = this.mapMessages(msgs);
-    const truncated = this.truncateHistory(messages);
+        ? [{ role: "user", content: input }]
+        : this.buildInput(input);
+    const truncated = this.truncateHistory(msgs);
     await this.apiCallStream(truncated, handle, sys);
   }