From 8270c9a4eb00cf9031095355d10cf1c75afbd293 Mon Sep 17 00:00:00 2001
From: Peli de Halleux <pelikhan@users.noreply.github.com>
Date: Fri, 13 Dec 2024 15:31:32 +0000
Subject: [PATCH] =?UTF-8?q?feat:=20add=20and=20update=20model=20aliases=20?=
 =?UTF-8?q?for=20Agent=20and=20Memory=20=F0=9F=9A=80?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 docs/src/content/docs/reference/scripts/model-aliases.mdx | 5 +++++
 packages/cli/src/nodehost.ts                              | 4 ++++
 packages/core/src/agent.ts                                | 4 ++--
 packages/core/src/runpromptcontext.ts                     | 3 ++-
 4 files changed, 13 insertions(+), 3 deletions(-)

diff --git a/docs/src/content/docs/reference/scripts/model-aliases.mdx b/docs/src/content/docs/reference/scripts/model-aliases.mdx
index ef7b01792..d3ccae406 100644
--- a/docs/src/content/docs/reference/scripts/model-aliases.mdx
+++ b/docs/src/content/docs/reference/scripts/model-aliases.mdx
@@ -81,3 +81,8 @@ By default, GenAIScript supports the following model aliases:
 - `large`: `gpt-4o like` model
 - `small`: `gpt-4o-mini` model or similar. A smaller, cheaper faster model
 - `vision`: `gpt-4o-mini`. A model that can analyze images
+
+The following aliases are also set so that you can override LLMs used by GenAIScript itself.
+
+- `agent`: `large`. Model used by the Agent LLM.
+- `memory`: `small`. Moel used by the agent short term memory.
diff --git a/packages/cli/src/nodehost.ts b/packages/cli/src/nodehost.ts
index 24a4d7d2e..2d3531dfc 100644
--- a/packages/cli/src/nodehost.ts
+++ b/packages/cli/src/nodehost.ts
@@ -31,6 +31,8 @@ import {
     MODEL_PROVIDER_AZURE_SERVERLESS_OPENAI,
     DOT_ENV_FILENAME,
     DEFAULT_VISION_MODEL,
+    LARGE_MODEL_ID,
+    SMALL_MODEL_ID,
 } from "../../core/src/constants"
 import { tryReadText } from "../../core/src/fs"
 import {
@@ -154,6 +156,8 @@ export class NodeHost implements RuntimeHost {
             small: { model: DEFAULT_SMALL_MODEL, source: "default" },
             vision: { model: DEFAULT_VISION_MODEL, source: "default" },
             embeddings: { model: DEFAULT_EMBEDDINGS_MODEL, source: "default" },
+            agent: { model: LARGE_MODEL_ID, source: "default" },
+            memory: { model: SMALL_MODEL_ID, source: "default" },
         },
         cli: {},
         env: {},
diff --git a/packages/core/src/agent.ts b/packages/core/src/agent.ts
index 52dcb1db2..e57ee0787 100644
--- a/packages/core/src/agent.ts
+++ b/packages/core/src/agent.ts
@@ -30,8 +30,8 @@ export async function agentQueryMemory(
             await defMemory(_)
         },
         {
-            model: "small",
-            system: ["system"],
+            model: "memory",
+            system: [],
             flexTokens: 20000,
             label: "agent memory query",
         }
diff --git a/packages/core/src/runpromptcontext.ts b/packages/core/src/runpromptcontext.ts
index f0a76f220..345fb95aa 100644
--- a/packages/core/src/runpromptcontext.ts
+++ b/packages/core/src/runpromptcontext.ts
@@ -169,7 +169,7 @@ export function createChatTurnGenerationContext(
                 cacheControl: (cc) => {
                     current.ephemeral = cc === "ephemeral"
                     return res
-                }
+                },
             } satisfies PromptTemplateString)
             return res
         },
@@ -496,6 +496,7 @@ export function createChatGenerationContext(
                             })
                     },
                     {
+                        model: "agent",
                         label: agentLabel,
                         system: agentSystem,
                         tools: agentTools.map(({ id }) => id),