cleanup

Author: timothycarambat

server/__tests__/utils/agents/defaults.test.js

@@ -39,14 +39,13 @@ describe("WORKSPACE_AGENT.getDefinition", () => {
     };
     const user = { id: 1 };
     const provider = "openai";
-
+    const expectedPrompt = await Provider.systemPrompt({ provider, workspace, user });
     const definition = await WORKSPACE_AGENT.getDefinition(
       provider,
       workspace,
       user
     );
-
-    expect(definition.role).toBe(Provider.systemPrompt(provider));
+    expect(definition.role).toBe(expectedPrompt);
     expect(SystemPromptVariables.expandSystemPromptVariables).not.toHaveBeenCalled();
   });
 
@@ -84,7 +83,6 @@ describe("WORKSPACE_AGENT.getDefinition", () => {
     };
     const user = null;
     const provider = "lmstudio";
-
     const expandedPrompt = "You are a helpful assistant. Today is January 1, 2024.";
     SystemPromptVariables.expandSystemPromptVariables.mockResolvedValue(expandedPrompt);
 
@@ -118,15 +116,15 @@ describe("WORKSPACE_AGENT.getDefinition", () => {
 
   it("should use LMStudio specific prompt when workspace has no openAiPrompt", async () => {
     const workspace = { id: 1, openAiPrompt: null };
+    const user = null;
     const provider = "lmstudio";
-
     const definition = await WORKSPACE_AGENT.getDefinition(
       provider,
       workspace,
       null
     );
 
-    expect(definition.role).toBe(Provider.systemPrompt(provider));
+    expect(definition.role).toBe(await Provider.systemPrompt({ provider, workspace, user }));
     expect(definition.role).toContain("helpful ai assistant");
   });
 });

server/utils/agents/aibitat/providers/ai-provider.js

@@ -19,6 +19,9 @@ const { toValidNumber, safeJsonParse } = require("../../../http");
 const { getLLMProviderClass } = require("../../../helpers");
 const { parseLMStudioBasePath } = require("../../../AiProviders/lmStudio");
 const { parseFoundryBasePath } = require("../../../AiProviders/foundry");
+const {
+  SystemPromptVariables,
+} = require("../../../../models/systemPromptVariables");
 
 const DEFAULT_WORKSPACE_PROMPT =
   "You are a helpful ai assistant who can assist the user and use tools available to help answer the users prompts and questions.";
@@ -288,10 +291,7 @@ class Provider {
     return llm.promptWindowLimit(modelName);
   }
 
-  // For some providers we may want to override the system prompt to be more verbose.
-  // Currently we only do this for lmstudio, but we probably will want to expand this even more
-  // to any Untooled LLM.
-  static systemPrompt(provider = null) {
+  static defaultSystemPromptForProvider(provider = null) {
     switch (provider) {
       case "lmstudio":
         return "You are a helpful ai assistant who can assist the user and use tools available to help answer the users prompts and questions. Tools will be handled by another assistant and you will simply receive their responses to help answer the user prompt - always try to answer the user's prompt the best you can with the context available to you and your general knowledge.";
@@ -300,6 +300,27 @@ class Provider {
     }
   }
 
+  /**
+   * Get the system prompt for a provider.
+   * @param {string} provider
+   * @param {import("@prisma/client").workspaces | null} workspace
+   * @param {import("@prisma/client").users | null} user
+   * @returns {Promise<string>}
+   */
+  static async systemPrompt({
+    provider = null,
+    workspace = null,
+    user = null,
+  }) {
+    if (!workspace?.openAiPrompt)
+      return Provider.defaultSystemPromptForProvider(provider);
+    return await SystemPromptVariables.expandSystemPromptVariables(
+      workspace.openAiPrompt,
+      user?.id || null,
+      workspace.id
+    );
+  }
+
   /**
    * Whether the provider supports agent streaming.
    * Disabled by default and needs to be explicitly enabled in the provider

server/utils/agents/defaults.js

@@ -16,7 +16,7 @@ const DEFAULT_SKILLS = [
 
 const USER_AGENT = {
   name: "USER",
-  getDefinition: async () => {
+  getDefinition: () => {
     return {
       interrupt: "ALWAYS",
       role: "I am the human monitor and oversee this chat. Any questions on action or decision making should be directed to me.",
@@ -26,20 +26,16 @@ const USER_AGENT = {
 
 const WORKSPACE_AGENT = {
   name: "@agent",
+  /**
+   * Get the definition for the workspace agent with its role (prompt) and functions in Aibitat format
+   * @param {string} provider
+   * @param {import("@prisma/client").workspaces | null} workspace
+   * @param {import("@prisma/client").users | null} user
+   * @returns {Promise<{ role: string, functions: object[] }>}
+   */
   getDefinition: async (provider = null, workspace = null, user = null) => {
-    // If workspace has a system prompt, use it with variable expansion
-    // Otherwise fall back to provider default
-    let role = Provider.systemPrompt(provider);
-    if (workspace?.openAiPrompt) {
-      role = await SystemPromptVariables.expandSystemPromptVariables(
-        workspace.openAiPrompt,
-        user?.id || null,
-        workspace.id
-      );
-    }
-
     return {
-      role,
+      role: await Provider.systemPrompt({ provider, workspace, user }),
       functions: [
         ...(await agentSkillsFromSystemSettings()),
         ...ImportedPlugin.activeImportedPlugins(),

server/utils/agents/ephemeral.js

@@ -4,6 +4,7 @@ const ImportedPlugin = require("./imported");
 const MCPCompatibilityLayer = require("../MCP");
 const { AgentFlows } = require("../agentFlows");
 const { httpSocket } = require("./aibitat/plugins/http-socket.js");
+const { User } = require("../../models/user");
 const { WorkspaceChats } = require("../../models/workspaceChats");
 const { safeJsonParse } = require("../http");
 const {
@@ -26,7 +27,7 @@ class EphemeralAgentHandler extends AgentHandler {
   #invocationUUID = null;
   /** @type {import("@prisma/client").workspaces|null} the workspace to use for the agent */
   #workspace = null;
-  /** @type {import("@prisma/client").users|null} the user id to use for the agent */
+  /** @type {import("@prisma/client").users["id"]|null} the user id to use for the agent */
   #userId = null;
   /** @type {import("@prisma/client").workspace_threads|null} the workspace thread id to use for the agent */
   #threadId = null;
@@ -69,6 +70,9 @@ class EphemeralAgentHandler extends AgentHandler {
     this.#workspace = workspace;
     this.#prompt = prompt;
 
+    // Note: userId for ephemeral agent is only available
+    // via the workspace-thread chat endpoints for the API
+    // since workspaces can belong to multiple users.
     this.#userId = userId;
     this.#threadId = threadId;
     this.#sessionId = sessionId;
@@ -319,10 +323,10 @@ class EphemeralAgentHandler extends AgentHandler {
   async #loadAgents() {
     // Default User agent and workspace agent
     this.log(`Attaching user and default agent to Agent cluster.`);
-    this.aibitat.agent(USER_AGENT.name, await USER_AGENT.getDefinition());
-
-    // Get user object from invocation for variable expansion
-    const user = this.#userId ? { id: this.#userId } : null;
+    this.aibitat.agent(USER_AGENT.name, USER_AGENT.getDefinition());
+    const user = this.#userId
+      ? await User.get({ id: Number(this.#userId) })
+      : null;
 
     this.aibitat.agent(
       WORKSPACE_AGENT.name,

server/utils/agents/index.js

@@ -3,6 +3,7 @@ const AgentPlugins = require("./aibitat/plugins");
 const {
   WorkspaceAgentInvocation,
 } = require("../../models/workspaceAgentInvocation");
+const { User } = require("../../models/user");
 const { WorkspaceChats } = require("../../models/workspaceChats");
 const { safeJsonParse } = require("../http");
 const { USER_AGENT, WORKSPACE_AGENT } = require("./defaults");
@@ -523,29 +524,20 @@ class AgentHandler {
   async #loadAgents() {
     // Default User agent and workspace agent
     this.log(`Attaching user and default agent to Agent cluster.`);
-    this.aibitat.agent(USER_AGENT.name, await USER_AGENT.getDefinition());
-
-    // Get user object from invocation for variable expansion
     const user = this.invocation.user_id
-      ? { id: this.invocation.user_id }
+      ? await User.get({ id: Number(this.invocation.user_id) })
       : null;
-
-    this.aibitat.agent(
-      WORKSPACE_AGENT.name,
-      await WORKSPACE_AGENT.getDefinition(
-        this.provider,
-        this.invocation.workspace,
-        user
-      )
-    );
-
+    const userAgentDef = await USER_AGENT.getDefinition();
     const workspaceAgentDef = await WORKSPACE_AGENT.getDefinition(
       this.provider,
       this.invocation.workspace,
       user
     );
+
+    this.aibitat.agent(USER_AGENT.name, userAgentDef);
+    this.aibitat.agent(WORKSPACE_AGENT.name, workspaceAgentDef);
     this.#funcsToLoad = [
-      ...((await USER_AGENT.getDefinition())?.functions || []),
+      ...(userAgentDef?.functions || []),
       ...(workspaceAgentDef?.functions || []),
     ];
   }