@marimo-team/frontend 0.14.18-dev9 → 0.15.1-dev10

package/src/core/ai/ids/ids.ts

@@ -0,0 +1,71 @@
+/* Copyright 2024 Marimo. All rights reserved. */
+
+import type { TypedString } from "@/utils/typed";
+
+export const PROVIDERS = [
+  "openai",
+  "anthropic",
+  "google",
+  "ollama",
+  "bedrock",
+  "deepseek",
+  "azure",
+  "github",
+  "marimo",
+] as const;
+export type ProviderId = (typeof PROVIDERS)[number];
+
+export type ShortModelId = TypedString<"ShortModelId">;
+
+/**
+ * Qualified model id
+ * `provider_id/short_model_id`;
+ */
+export type QualifiedModelId = `${ProviderId}/${ShortModelId}`;
+
+export class AiModelId {
+  readonly providerId: ProviderId;
+  readonly shortModelId: ShortModelId;
+
+  constructor(providerId: ProviderId, shortModelId: ShortModelId) {
+    this.providerId = providerId;
+    this.shortModelId = shortModelId;
+  }
+
+  get id(): QualifiedModelId {
+    return `${this.providerId}/${this.shortModelId}`;
+  }
+
+  static parse(id: string): AiModelId {
+    if (!id.includes("/")) {
+      const providerId = guessProviderId(id);
+      return new AiModelId(providerId, id as ShortModelId);
+    }
+
+    const [providerId, ...shortModelId] = id.split("/");
+    return new AiModelId(
+      providerId as ProviderId,
+      shortModelId.join("/") as ShortModelId,
+    );
+  }
+}
+
+function guessProviderId(id: string): ProviderId {
+  if (id.startsWith("gpt") || id.startsWith("o3") || id.startsWith("o1")) {
+    return "openai";
+  }
+  if (id.startsWith("claude")) {
+    return "anthropic";
+  }
+  if (id.startsWith("gemini") || id.startsWith("google")) {
+    return "google";
+  }
+  if (id.startsWith("deepseek")) {
+    return "deepseek";
+  }
+  return "ollama";
+}
+
+export function isKnownAIProvider(providerId: ProviderId): boolean {
+  return PROVIDERS.includes(providerId);
+}

package/src/core/ai/model-registry.ts

@@ -0,0 +1,158 @@
+/* Copyright 2024 Marimo. All rights reserved. */
+
+import type {
+  AiModel as AiModelType,
+  AiProvider,
+  Role,
+} from "@marimo-team/llm-info";
+import { models } from "@marimo-team/llm-info/models.json";
+import { providers } from "@marimo-team/llm-info/providers.json";
+import { MultiMap } from "@/utils/multi-map";
+import { once } from "@/utils/once";
+import type { ProviderId } from "./ids/ids";
+import { AiModelId, type QualifiedModelId, type ShortModelId } from "./ids/ids";
+
+export interface AiModel extends AiModelType {
+  roles: Role[];
+  providers: ProviderId[];
+  /** Whether this is a custom model. */
+  custom: boolean;
+}
+
+const getKnownModelMap = once((): ReadonlyMap<QualifiedModelId, AiModel> => {
+  const modelMap = new Map<QualifiedModelId, AiModel>();
+  for (const model of models) {
+    const modelId = model.model as ShortModelId;
+    const modelInfo: AiModel = {
+      ...model,
+      roles: model.roles.map((role) => role as Role),
+      providers: model.providers as ProviderId[],
+      custom: false,
+    };
+
+    for (const provider of modelInfo.providers) {
+      const qualifiedModelId: QualifiedModelId = `${provider}/${modelId}`;
+      modelMap.set(qualifiedModelId, modelInfo);
+    }
+  }
+  return modelMap;
+});
+
+const getProviderMap = once((): ReadonlyMap<ProviderId, AiProvider> => {
+  const providerMap = new Map<ProviderId, AiProvider>();
+  for (const provider of providers) {
+    providerMap.set(provider.id as ProviderId, provider);
+  }
+  return providerMap;
+});
+
+export class AiModelRegistry {
+  private modelsByProviderMap = new MultiMap<ProviderId, AiModel>();
+  private modelsMap: ReadonlyMap<QualifiedModelId, AiModel>;
+  private customModels: ReadonlySet<QualifiedModelId>;
+  private displayedModels: ReadonlySet<QualifiedModelId>;
+
+  private constructor(
+    customModels: QualifiedModelId[],
+    displayedModels: QualifiedModelId[],
+  ) {
+    this.customModels = new Set(customModels);
+    this.displayedModels = new Set(displayedModels);
+    this.modelsMap = new Map<QualifiedModelId, AiModel>();
+    this.buildMaps();
+  }
+
+  static getProviderInfo(providerId: ProviderId) {
+    return getProviderMap().get(providerId);
+  }
+
+  /**
+   * @param customModels - A list of custom models to use that are not from the default list.
+   * @param displayedModels - A list of models to display in the UI. If empty, all models will be displayed.
+   *
+   * Models should be in the format of `provider_id/short_model_id`.
+   */
+  static create(opts: { customModels?: string[]; displayedModels?: string[] }) {
+    const { customModels = [], displayedModels = [] } = opts;
+    return new AiModelRegistry(
+      customModels.map((model) => AiModelId.parse(model).id),
+      displayedModels.map((model) => AiModelId.parse(model).id),
+    );
+  }
+
+  /**
+   * Builds the maps of models by provider and custom models.
+   */
+  private buildMaps() {
+    const displayedModels = this.displayedModels;
+    const hasDisplayedModels = displayedModels.size > 0;
+    const knownModelMap = getKnownModelMap();
+    const customModelsMap = new Map<QualifiedModelId, AiModel>();
+
+    let modelsMap = new Map<QualifiedModelId, AiModel>();
+
+    for (const model of this.customModels) {
+      if (hasDisplayedModels && !displayedModels.has(model)) {
+        continue;
+      }
+      const modelId = AiModelId.parse(model);
+      const modelInfo: AiModel = {
+        name: modelId.shortModelId,
+        model: modelId.shortModelId,
+        description: "Custom model",
+        providers: [modelId.providerId],
+        roles: [],
+        thinking: false,
+        custom: true,
+      };
+      customModelsMap.set(model, modelInfo);
+    }
+
+    // Add known models
+    if (hasDisplayedModels) {
+      for (const model of displayedModels) {
+        const modelInfo = knownModelMap.get(model);
+        if (modelInfo) {
+          modelsMap.set(model, modelInfo);
+        }
+      }
+    } else {
+      modelsMap = new Map(knownModelMap);
+    }
+
+    // Set custom models first, then known models
+    modelsMap = new Map([...customModelsMap, ...modelsMap]);
+
+    // Group by provider
+    for (const [qualifiedModelId, model] of modelsMap.entries()) {
+      const modelId = AiModelId.parse(qualifiedModelId);
+      this.modelsByProviderMap.add(modelId.providerId, model);
+    }
+
+    this.modelsMap = modelsMap;
+  }
+
+  getDisplayedModels() {
+    return this.displayedModels;
+  }
+
+  getCustomModels() {
+    return this.customModels;
+  }
+
+  getModelsByProvider(provider: ProviderId) {
+    return this.modelsByProviderMap.get(provider) || [];
+  }
+
+  getGroupedModelsByProvider() {
+    return this.modelsByProviderMap;
+  }
+
+  getModelsMap() {
+    return this.modelsMap;
+  }
+
+  getModel(qualifiedModelId: QualifiedModelId) {
+    return this.modelsMap.get(qualifiedModelId);
+  }
+}

package/src/core/ai/state.ts

@@ -0,0 +1,58 @@
+/* Copyright 2024 Marimo. All rights reserved. */
+
+import type { Message as AIMessage } from "@ai-sdk/react";
+import { atom } from "jotai";
+import { atomWithStorage } from "jotai/utils";
+import type { CellId } from "../cells/ids";
+
+const KEY = "marimo:ai:chatState:v3";
+
+export const aiCompletionCellAtom = atom<{
+  cellId: CellId;
+  initialPrompt?: string;
+} | null>(null);
+
+const INCLUDE_OTHER_CELLS_KEY = "marimo:ai:includeOtherCells";
+export const includeOtherCellsAtom = atomWithStorage<boolean>(
+  INCLUDE_OTHER_CELLS_KEY,
+  true,
+);
+
+export interface Message {
+  id: string;
+  role: "user" | "assistant" | "data" | "system";
+  content: string;
+  timestamp: number;
+  parts?: AIMessage["parts"];
+}
+
+export interface Chat {
+  id: string;
+  title: string;
+  messages: Message[];
+  createdAt: number;
+  updatedAt: number;
+}
+
+export interface ChatState {
+  chats: Chat[];
+  activeChatId: string | null;
+}
+
+export const chatStateAtom = atomWithStorage<ChatState>(KEY, {
+  chats: [],
+  activeChatId: null,
+});
+
+export const activeChatAtom = atom(
+  (get) => {
+    const state = get(chatStateAtom);
+    return state.chats.find((chat) => chat.id === state.activeChatId);
+  },
+  (get, set, chatId: string | null) => {
+    set(chatStateAtom, (prev) => ({
+      ...prev,
+      activeChatId: chatId,
+    }));
+  },
+);

package/src/core/alerts/state.ts

@@ -0,0 +1,75 @@
+/* Copyright 2024 Marimo. All rights reserved. */
+
+import { useAtomValue } from "jotai";
+import type { PackageInstallationStatus } from "@/core/kernel/messages";
+import { createReducerAndAtoms } from "@/utils/createReducer";
+import { generateUUID } from "@/utils/uuid";
+
+type Identified<T> = { id: string } & T;
+
+export interface MissingPackageAlert {
+  kind: "missing";
+  packages: string[];
+  isolated: boolean;
+}
+
+export interface InstallingPackageAlert {
+  kind: "installing";
+  packages: PackageInstallationStatus;
+}
+
+export function isMissingPackageAlert(
+  alert: MissingPackageAlert | InstallingPackageAlert,
+): alert is MissingPackageAlert {
+  return alert.kind === "missing";
+}
+
+export function isInstallingPackageAlert(
+  alert: MissingPackageAlert | InstallingPackageAlert,
+): alert is InstallingPackageAlert {
+  return alert.kind === "installing";
+}
+
+/** Prominent alerts.
+ *
+ * Right now we only have one type of alert.
+ */
+interface AlertState {
+  packageAlert:
+    | Identified<MissingPackageAlert>
+    | Identified<InstallingPackageAlert>
+    | null;
+}
+
+const { valueAtom: alertAtom, useActions } = createReducerAndAtoms(
+  () => ({ packageAlert: null }) as AlertState,
+  {
+    addPackageAlert: (
+      state,
+      alert: MissingPackageAlert | InstallingPackageAlert,
+    ) => {
+      return {
+        ...state,
+        packageAlert: { id: generateUUID(), ...alert },
+      };
+    },
+
+    clearPackageAlert: (state, id: string) => {
+      return state.packageAlert !== null && state.packageAlert.id === id
+        ? { ...state, packageAlert: null }
+        : state;
+    },
+  },
+);
+
+/**
+ * React hook to get the Alert state.
+ */
+export const useAlerts = () => useAtomValue(alertAtom);
+
+/**
+ * React hook to get the Alerts actions.
+ */
+export function useAlertActions() {
+  return useActions();
+}