1 files changed, 173 insertions, 0 deletions
diff --git a/packages/ai/src/claude.ts b/packages/ai/src/claude.ts
new file mode 100644
index 0000000..a411030
--- /dev/null
+++ b/packages/ai/src/claude.ts
@@ -0,0 +1,173 @@
+import Claude from "@anthropic-ai/sdk";
+import { RESPONSE_LENGTH } from "./logic/constants";
+import type { AIModelAPI, ChatMessage, InputToken } from "./types";
+import { BOOKWORM_SYS } from "./prompts";
+import type { AsyncRes } from "@sortug/lib";
+import type {
+  ImageBlockParam,
+  MessageCreateParamsStreaming,
+  TextBlockParam,
+} from "@anthropic-ai/sdk/resources";
+
+type Message = Claude.Messages.MessageParam;
+
+export default class ClaudeAPI implements AIModelAPI {
+  private model: string = "claude-opus-4-20250514";
+  tokenizer: (text: string) => number;
+  maxTokens: number;
+  // model: string = "claude-3-5-sonnet-20241022";
+  constructor(
+    model: string,
+    maxTokens = 200_000,
+    tokenizer: (text: string) => number = (text) => text.length / 3,
+  ) {
+    this.maxTokens = maxTokens;
+    this.tokenizer = tokenizer;
+    if (model) this.model = model;
+  }
+  public setModel(model: string) {
+    this.model = model;
+  }
+  private mapMessages(input: ChatMessage[]): Message[] {
+    return input.map((m) => {
+      const role = m.author === "claude" ? "assistant" : "user";
+      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 }];
+  }
+
+  // https://docs.anthropic.com/en/api/messages-examples#vision
+  public async send(input: string | InputToken[], sys?: string) {
+    const msgs: Message[] =
+      typeof input === "string"
+        ? [{ role: "user", content: input }]
+        : this.buildInput(input);
+    const truncated = this.truncateHistory(msgs);
+    const res = await this.apiCall(truncated, sys);
+    return res;
+  }
+
+  public async sendDoc(data: string) {
+    const sys = BOOKWORM_SYS;
+    const msg: Message = {
+      role: "user",
+      content: [
+        {
+          type: "document",
+          source: { type: "base64", data, media_type: "application/pdf" },
+        },
+        {
+          type: "text",
+          text: "Please analyze this according to your system prompt. Be thorough.",
+        },
+      ],
+    };
+    const res = await this.apiCall([msg], sys);
+    return res;
+  }
+
+  public async stream(
+    input: string | InputToken[],
+    handle: (c: any) => void,
+    sys?: string,
+  ) {
+    const msgs: Message[] =
+      typeof input === "string"
+        ? [{ role: "user", content: input }]
+        : this.buildInput(input);
+    const truncated = this.truncateHistory(msgs);
+    await this.apiCallStream(truncated, handle, sys);
+  }
+
+  private truncateHistory(messages: Message[]): Message[] {
+    const totalTokens = messages.reduce((total, message) => {
+      return total + this.tokenizer(message.content as string);
+    }, 0);
+    while (totalTokens > this.maxTokens && messages.length > 1) {
+      messages.splice(0, 1);
+    }
+    return messages;
+  }
+
+  // TODO
+  // https://docs.anthropic.com/en/api/messages-examples#putting-words-in-claudes-mouth
+  private async apiCall(
+    messages: Message[],
+    system?: string,
+  ): Promise<AsyncRes<string>> {
+    try {
+      const claud = new Claude();
+      const params = {
+        model: this.model,
+        max_tokens: RESPONSE_LENGTH,
+        messages,
+      };
+      const res = await claud.messages.create(
+        system ? { ...params, system } : params,
+      );
+      const resm: string = res.content.reduce((acc: string, item) => {
+        if (item.type === "tool_use") return acc;
+        else if (item.type === "text") return `${acc}\n${item.text}`;
+        else return acc;
+      }, "");
+      // const resm = res.content.reduce((acc: string[], item) => {
+      //     if (item.type === "tool_use") return acc;
+      //     else if (item.type === "text") return [...acc, item.text]
+      //     else return acc;
+      //   }, []);
+      return { ok: resm };
+    } catch (e) {
+      console.log(e, "error in claude api");
+      return { error: `${e}` };
+    }
+  }
+
+  private async apiCallStream(
+    messages: Message[],
+    handle: (c: any) => void,
+    system?: string,
+  ): Promise<void> {
+    try {
+      const claud = new Claude();
+      const params = {
+        model: this.model,
+        max_tokens: RESPONSE_LENGTH,
+        messages,
+        stream: true as true,
+      };
+      const fparams: MessageCreateParamsStreaming = system
+        ? { ...params, system }
+        : params;
+      const stream = await claud.messages.create(fparams);
+
+      for await (const part of stream) {
+        if (part.type === "message_start") continue;
+        if (part.type === "content_block_start") continue;
+        if (part.type === "content_block_delta") {
+          console.log("delta", part.delta);
+          const delta: any = part.delta;
+          handle(delta.text);
+        }
+      }
+    } catch (e) {
+      console.log(e, "error in claude api");
+      handle(`Error streaming Claude, ${e}`);
+    }
+  }
+}