1 files changed, 79 insertions, 42 deletions

diff --git a/src/gemini.ts b/src/gemini.ts
index 3e636c2..5c7267b 100644
--- a/src/gemini.ts
+++ b/src/gemini.ts
@@ -1,21 +1,19 @@
+import mime from "mime-types";
 import {
   Chat,
+  createPartFromBase64,
+  createPartFromUri,
+  createUserContent,
   GoogleGenAI,
   type Content,
   type GeneratedImage,
   type GeneratedVideo,
+  type Part,
 } from "@google/genai";
-import { RESPONSE_LENGTH } from "./logic/constants";
-import type {
-  AIModelAPI,
-  ChatMessage,
-  OChoice,
-  OChunk,
-  OMessage,
-} from "./types";
-import type { AsyncRes } from "sortug";
+import type { AIModelAPI, InputToken } from "./types";
+import type { AsyncRes, Result } from "sortug";
 
-export default class GeminiAPI {
+export default class GeminiAPI implements AIModelAPI {
   tokenizer: (text: string) => number;
   maxTokens: number;
   private model: string;
@@ -23,70 +21,109 @@ export default class GeminiAPI {
   chats: Map<string, Chat> = new Map<string, Chat>();
 
   constructor(
+    model?: string,
     maxTokens = 200_000,
     tokenizer: (text: string) => number = (text) => text.length / 3,
-    model?: string,
   ) {
     this.maxTokens = maxTokens;
     this.tokenizer = tokenizer;
 
     const gem = new GoogleGenAI({ apiKey: Bun.env["GEMINI_API_KEY"]! });
     this.api = gem;
-    this.model = model || "gemini-2.5-pro-preview-05-06	";
+    this.model = model || "gemini-2.5-pro";
   }
 
-  createChat({ name, history }: { name?: string; history?: Content[] }) {
-    const chat = this.api.chats.create({ model: this.model, history });
-    this.chats.set(name ? name : Date.now().toString(), chat);
+  // input data in  gemini gets pretty involved
+  //
+  // data
+  // Union type
+  // data can be only one of the following:
+  // text
+  // string
+  // Inline text.
+
+  // inlineData
+  // object (Blob)
+  // Inline media bytes.
+
+  // functionCall
+  // object (FunctionCall)
+  // A predicted FunctionCall returned from the model that contains a string representing the FunctionDeclaration.name with the arguments and their values.
+
+  // functionResponse
+  // object (FunctionResponse)
+  // The result output of a FunctionCall that contains a string representing the FunctionDeclaration.name and a structured JSON object containing any output from the function is used as context to the model.
+
+  // fileData
+  // object (FileData)
+  // URI based data.
+
+  // executableCode
+  // object (ExecutableCode)
+  // Code generated by the model that is meant to be executed.
+
+  // codeExecutionResult
+  // object (CodeExecutionResult)
+  // Result of executing the ExecutableCode.
+
+  // metadata
+  // Union type
+  public setModel(model: string) {
+    this.model = model;
   }
-  async followChat(name: string, message: string): AsyncRes<string> {
-    const chat = this.chats.get(name);
-    if (!chat) return { error: "no chat with that name" };
-    else {
-      const response = await chat.sendMessage({ message });
-      const text = response.text;
-      return { ok: text || "" };
-    }
+  private contentFromImage(imageString: string): Result<Part> {
+    // TODO
+    const mimeType = mime.lookup(imageString);
+    if (!mimeType) return { error: "no mimetype" };
+    const url = URL.parse(imageString);
+    if (url) {
+      const part = createPartFromUri(imageString, mimeType);
+      return { ok: part };
+    } else return { ok: createPartFromBase64(imageString, mimeType) };
   }
-  async followChatStream(
-    name: string,
-    message: string,
-    handler: (data: string) => void,
-  ) {
-    const chat = this.chats.get(name);
-    if (!chat) throw new Error("no chat!");
-    else {
-      const response = await chat.sendMessageStream({ message });
-      for await (const chunk of response) {
-        const text = chunk.text;
-        handler(text || "");
-      }
+  public buildInput(tokens: InputToken[]): Result<Content> {
+    try {
+      const input = createUserContent(
+        tokens.map((t) => {
+          if ("text" in t) return t.text;
+          if ("img" in t) {
+            const imagePart = this.contentFromImage(t.img);
+            if ("error" in imagePart) throw new Error("image failed");
+            else return imagePart.ok;
+          }
+          return "oy vey";
+        }),
+      );
+      return { ok: input };
+    } catch (e) {
+      return { error: `${e}` };
     }
   }
 
-  async send(message: string, systemPrompt?: string): AsyncRes<string> {
+  async send(input: string | Content, systemPrompt?: string): AsyncRes<string> {
     try {
       const opts = {
         model: this.model,
-        contents: message,
+        contents: input,
       };
       const fopts = systemPrompt
         ? { ...opts, config: { systemInstruction: systemPrompt } }
         : opts;
       const response = await this.api.models.generateContent(fopts);
-      return { ok: response.text || "" };
+      if (!response.text) return { error: "no text in response" };
+      return { ok: response.text };
     } catch (e) {
       return { error: `${e}` };
     }
   }
-  async sendStream(
+  async stream(
+    input: string | Content,
     handler: (s: string) => void,
-    message: string,
     systemPrompt?: string,
   ) {
     const opts = {
       model: this.model,
-      contents: message,
+      contents: input,
     };
     const fopts = systemPrompt
       ? { ...opts, config: { systemInstruction: systemPrompt } }