From 705a1a66a1f1dea2dc9ae6f78b8ebc1f8fd6505e Mon Sep 17 00:00:00 2001
From: Henry Fontanier <h.fontanier@gmail.com>
Date: Fri, 22 Dec 2023 10:26:12 +0100
Subject: [PATCH] feat: release gemini-pro (custom assistants + global agent
 (#3006)

Co-authored-by: Henry Fontanier <henry@dust.tt>
---
 .../assistant_builder/AssistantBuilder.tsx    |  2 +
 front/lib/api/assistant/global_agents.ts      | 40 ++++++++++++++++++-
 front/lib/assistant.ts                        |  2 +
 3 files changed, 43 insertions(+), 1 deletion(-)

diff --git a/front/components/assistant_builder/AssistantBuilder.tsx b/front/components/assistant_builder/AssistantBuilder.tsx
index 1772aa441add..c600f7f32b57 100644
--- a/front/components/assistant_builder/AssistantBuilder.tsx
+++ b/front/components/assistant_builder/AssistantBuilder.tsx
@@ -19,6 +19,7 @@ import {
   AgentConfigurationScope,
   ConnectorProvider,
   DataSourceType,
+  GEMINI_PRO_DEFAULT_MODEL_CONFIG,
 } from "@dust-tt/types";
 import { UserType, WorkspaceType } from "@dust-tt/types";
 import {
@@ -79,6 +80,7 @@ const usedModelConfigs = [
   CLAUDE_INSTANT_DEFAULT_MODEL_CONFIG,
   MISTRAL_MEDIUM_MODEL_CONFIG,
   MISTRAL_SMALL_MODEL_CONFIG,
+  GEMINI_PRO_DEFAULT_MODEL_CONFIG,
 ];
 
 // Avatar URLs
diff --git a/front/lib/api/assistant/global_agents.ts b/front/lib/api/assistant/global_agents.ts
index 678e45afdbdf..6128cd3125a4 100644
--- a/front/lib/api/assistant/global_agents.ts
+++ b/front/lib/api/assistant/global_agents.ts
@@ -4,7 +4,11 @@ import { promisify } from "util";
 
 const readFileAsync = promisify(fs.readFile);
 
-import { ConnectorProvider, DataSourceType } from "@dust-tt/types";
+import {
+  ConnectorProvider,
+  DataSourceType,
+  GEMINI_PRO_DEFAULT_MODEL_CONFIG,
+} from "@dust-tt/types";
 import {
   CLAUDE_DEFAULT_MODEL_CONFIG,
   CLAUDE_INSTANT_DEFAULT_MODEL_CONFIG,
@@ -307,6 +311,37 @@ async function _getMistralSmallGlobalAgent({
   };
 }
 
+async function _getGeminiProGlobalAgent({
+  settings,
+}: {
+  settings: GlobalAgentSettings | null;
+}): Promise<AgentConfigurationType> {
+  const status = settings ? settings.status : "disabled_by_admin";
+  return {
+    id: -1,
+    sId: GLOBAL_AGENTS_SID.GEMINI_PRO,
+    version: 0,
+    versionAuthorId: null,
+    name: "gemini-pro",
+    description:
+      "Google's our best model for scaling across a wide range of tasks (8k context).",
+    pictureUrl: "https://dust.tt/static/systemavatar/gemini_avatar_full.png",
+    status,
+    scope: "global",
+    userListStatus: status === "active" ? "in-list" : "not-in-list",
+    generation: {
+      id: -1,
+      prompt: `Never start your messages with "[assistant:"`,
+      model: {
+        providerId: GEMINI_PRO_DEFAULT_MODEL_CONFIG.providerId,
+        modelId: GEMINI_PRO_DEFAULT_MODEL_CONFIG.modelId,
+      },
+      temperature: 0.7,
+    },
+    action: null,
+  };
+}
+
 async function _getManagedDataSourceAgent(
   auth: Authenticator,
   {
@@ -688,6 +723,9 @@ export async function getGlobalAgent(
     case GLOBAL_AGENTS_SID.MISTRAL_SMALL:
       agentConfiguration = await _getMistralSmallGlobalAgent({ settings });
       break;
+    case GLOBAL_AGENTS_SID.GEMINI_PRO:
+      agentConfiguration = await _getGeminiProGlobalAgent({ settings });
+      break;
     case GLOBAL_AGENTS_SID.SLACK:
       agentConfiguration = await _getSlackGlobalAgent(auth, {
         settings,
diff --git a/front/lib/assistant.ts b/front/lib/assistant.ts
index efdf82b0bf6a..6016a6542f8a 100644
--- a/front/lib/assistant.ts
+++ b/front/lib/assistant.ts
@@ -44,6 +44,7 @@ export enum GLOBAL_AGENTS_SID {
   // it interferes with the retrieval of ongoing conversations involving this agent.
   // Needed to preserve ongoing chat integrity due to 'sId=mistral' references in legacy messages.
   MISTRAL_SMALL = "mistral",
+  GEMINI_PRO = "gemini-pro",
 }
 
 const CUSTOM_ORDER: string[] = [
@@ -58,6 +59,7 @@ const CUSTOM_ORDER: string[] = [
   GLOBAL_AGENTS_SID.CLAUDE_INSTANT,
   GLOBAL_AGENTS_SID.MISTRAL_MEDIUM,
   GLOBAL_AGENTS_SID.MISTRAL_SMALL,
+  GLOBAL_AGENTS_SID.GEMINI_PRO,
   GLOBAL_AGENTS_SID.HELPER,
 ];