Skip to content

Commit 72132b8

Browse files
cleanup
1 parent 3f24a6a commit 72132b8

File tree

5 files changed

+54
-43
lines changed

5 files changed

+54
-43
lines changed

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

Lines changed: 4 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -39,14 +39,13 @@ describe("WORKSPACE_AGENT.getDefinition", () => {
3939
};
4040
const user = { id: 1 };
4141
const provider = "openai";
42-
42+
const expectedPrompt = await Provider.systemPrompt({ provider, workspace, user });
4343
const definition = await WORKSPACE_AGENT.getDefinition(
4444
provider,
4545
workspace,
4646
user
4747
);
48-
49-
expect(definition.role).toBe(Provider.systemPrompt(provider));
48+
expect(definition.role).toBe(expectedPrompt);
5049
expect(SystemPromptVariables.expandSystemPromptVariables).not.toHaveBeenCalled();
5150
});
5251

@@ -84,7 +83,6 @@ describe("WORKSPACE_AGENT.getDefinition", () => {
8483
};
8584
const user = null;
8685
const provider = "lmstudio";
87-
8886
const expandedPrompt = "You are a helpful assistant. Today is January 1, 2024.";
8987
SystemPromptVariables.expandSystemPromptVariables.mockResolvedValue(expandedPrompt);
9088

@@ -118,15 +116,15 @@ describe("WORKSPACE_AGENT.getDefinition", () => {
118116

119117
it("should use LMStudio specific prompt when workspace has no openAiPrompt", async () => {
120118
const workspace = { id: 1, openAiPrompt: null };
119+
const user = null;
121120
const provider = "lmstudio";
122-
123121
const definition = await WORKSPACE_AGENT.getDefinition(
124122
provider,
125123
workspace,
126124
null
127125
);
128126

129-
expect(definition.role).toBe(Provider.systemPrompt(provider));
127+
expect(definition.role).toBe(await Provider.systemPrompt({ provider, workspace, user }));
130128
expect(definition.role).toContain("helpful ai assistant");
131129
});
132130
});

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

Lines changed: 25 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,9 @@ const { toValidNumber, safeJsonParse } = require("../../../http");
1919
const { getLLMProviderClass } = require("../../../helpers");
2020
const { parseLMStudioBasePath } = require("../../../AiProviders/lmStudio");
2121
const { parseFoundryBasePath } = require("../../../AiProviders/foundry");
22+
const {
23+
SystemPromptVariables,
24+
} = require("../../../../models/systemPromptVariables");
2225

2326
const DEFAULT_WORKSPACE_PROMPT =
2427
"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 {
288291
return llm.promptWindowLimit(modelName);
289292
}
290293

291-
// For some providers we may want to override the system prompt to be more verbose.
292-
// Currently we only do this for lmstudio, but we probably will want to expand this even more
293-
// to any Untooled LLM.
294-
static systemPrompt(provider = null) {
294+
static defaultSystemPromptForProvider(provider = null) {
295295
switch (provider) {
296296
case "lmstudio":
297297
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 {
300300
}
301301
}
302302

303+
/**
304+
* Get the system prompt for a provider.
305+
* @param {string} provider
306+
* @param {import("@prisma/client").workspaces | null} workspace
307+
* @param {import("@prisma/client").users | null} user
308+
* @returns {Promise<string>}
309+
*/
310+
static async systemPrompt({
311+
provider = null,
312+
workspace = null,
313+
user = null,
314+
}) {
315+
if (!workspace?.openAiPrompt)
316+
return Provider.defaultSystemPromptForProvider(provider);
317+
return await SystemPromptVariables.expandSystemPromptVariables(
318+
workspace.openAiPrompt,
319+
user?.id || null,
320+
workspace.id
321+
);
322+
}
323+
303324
/**
304325
* Whether the provider supports agent streaming.
305326
* Disabled by default and needs to be explicitly enabled in the provider

server/utils/agents/defaults.js

Lines changed: 9 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,7 @@ const DEFAULT_SKILLS = [
1616

1717
const USER_AGENT = {
1818
name: "USER",
19-
getDefinition: async () => {
19+
getDefinition: () => {
2020
return {
2121
interrupt: "ALWAYS",
2222
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 = {
2626

2727
const WORKSPACE_AGENT = {
2828
name: "@agent",
29+
/**
30+
* Get the definition for the workspace agent with its role (prompt) and functions in Aibitat format
31+
* @param {string} provider
32+
* @param {import("@prisma/client").workspaces | null} workspace
33+
* @param {import("@prisma/client").users | null} user
34+
* @returns {Promise<{ role: string, functions: object[] }>}
35+
*/
2936
getDefinition: async (provider = null, workspace = null, user = null) => {
30-
// If workspace has a system prompt, use it with variable expansion
31-
// Otherwise fall back to provider default
32-
let role = Provider.systemPrompt(provider);
33-
if (workspace?.openAiPrompt) {
34-
role = await SystemPromptVariables.expandSystemPromptVariables(
35-
workspace.openAiPrompt,
36-
user?.id || null,
37-
workspace.id
38-
);
39-
}
40-
4137
return {
42-
role,
38+
role: await Provider.systemPrompt({ provider, workspace, user }),
4339
functions: [
4440
...(await agentSkillsFromSystemSettings()),
4541
...ImportedPlugin.activeImportedPlugins(),

server/utils/agents/ephemeral.js

Lines changed: 9 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@ const ImportedPlugin = require("./imported");
44
const MCPCompatibilityLayer = require("../MCP");
55
const { AgentFlows } = require("../agentFlows");
66
const { httpSocket } = require("./aibitat/plugins/http-socket.js");
7+
const { User } = require("../../models/user");
78
const { WorkspaceChats } = require("../../models/workspaceChats");
89
const { safeJsonParse } = require("../http");
910
const {
@@ -26,7 +27,7 @@ class EphemeralAgentHandler extends AgentHandler {
2627
#invocationUUID = null;
2728
/** @type {import("@prisma/client").workspaces|null} the workspace to use for the agent */
2829
#workspace = null;
29-
/** @type {import("@prisma/client").users|null} the user id to use for the agent */
30+
/** @type {import("@prisma/client").users["id"]|null} the user id to use for the agent */
3031
#userId = null;
3132
/** @type {import("@prisma/client").workspace_threads|null} the workspace thread id to use for the agent */
3233
#threadId = null;
@@ -69,6 +70,9 @@ class EphemeralAgentHandler extends AgentHandler {
6970
this.#workspace = workspace;
7071
this.#prompt = prompt;
7172

73+
// Note: userId for ephemeral agent is only available
74+
// via the workspace-thread chat endpoints for the API
75+
// since workspaces can belong to multiple users.
7276
this.#userId = userId;
7377
this.#threadId = threadId;
7478
this.#sessionId = sessionId;
@@ -319,10 +323,10 @@ class EphemeralAgentHandler extends AgentHandler {
319323
async #loadAgents() {
320324
// Default User agent and workspace agent
321325
this.log(`Attaching user and default agent to Agent cluster.`);
322-
this.aibitat.agent(USER_AGENT.name, await USER_AGENT.getDefinition());
323-
324-
// Get user object from invocation for variable expansion
325-
const user = this.#userId ? { id: this.#userId } : null;
326+
this.aibitat.agent(USER_AGENT.name, USER_AGENT.getDefinition());
327+
const user = this.#userId
328+
? await User.get({ id: Number(this.#userId) })
329+
: null;
326330

327331
this.aibitat.agent(
328332
WORKSPACE_AGENT.name,

server/utils/agents/index.js

Lines changed: 7 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@ const AgentPlugins = require("./aibitat/plugins");
33
const {
44
WorkspaceAgentInvocation,
55
} = require("../../models/workspaceAgentInvocation");
6+
const { User } = require("../../models/user");
67
const { WorkspaceChats } = require("../../models/workspaceChats");
78
const { safeJsonParse } = require("../http");
89
const { USER_AGENT, WORKSPACE_AGENT } = require("./defaults");
@@ -523,29 +524,20 @@ class AgentHandler {
523524
async #loadAgents() {
524525
// Default User agent and workspace agent
525526
this.log(`Attaching user and default agent to Agent cluster.`);
526-
this.aibitat.agent(USER_AGENT.name, await USER_AGENT.getDefinition());
527-
528-
// Get user object from invocation for variable expansion
529527
const user = this.invocation.user_id
530-
? { id: this.invocation.user_id }
528+
? await User.get({ id: Number(this.invocation.user_id) })
531529
: null;
532-
533-
this.aibitat.agent(
534-
WORKSPACE_AGENT.name,
535-
await WORKSPACE_AGENT.getDefinition(
536-
this.provider,
537-
this.invocation.workspace,
538-
user
539-
)
540-
);
541-
530+
const userAgentDef = await USER_AGENT.getDefinition();
542531
const workspaceAgentDef = await WORKSPACE_AGENT.getDefinition(
543532
this.provider,
544533
this.invocation.workspace,
545534
user
546535
);
536+
537+
this.aibitat.agent(USER_AGENT.name, userAgentDef);
538+
this.aibitat.agent(WORKSPACE_AGENT.name, workspaceAgentDef);
547539
this.#funcsToLoad = [
548-
...((await USER_AGENT.getDefinition())?.functions || []),
540+
...(userAgentDef?.functions || []),
549541
...(workspaceAgentDef?.functions || []),
550542
];
551543
}

0 commit comments

Comments
 (0)