feat: add and update model aliases for Agent and Memory 🚀

microsoft · Dec 13, 2024 · 8270c9a · 8270c9a
1 parent 2e9ade5
commit 8270c9a
Show file tree

Hide file tree

Showing 4 changed files with 13 additions and 3 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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),