Skip to content

Commit 3f24a6a

Browse files
shatfield4shatfield4
committed
agent workspace system prompt with variable expansion
1 parent 6270a0a commit 3f24a6a

File tree

4 files changed

+169
-5
lines changed

4 files changed

+169
-5
lines changed

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

Lines changed: 133 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,133 @@
1+
// Set required env vars before requiring modules
2+
process.env.STORAGE_DIR = __dirname;
3+
process.env.NODE_ENV = "test";
4+
5+
const { SystemPromptVariables } = require("../../../models/systemPromptVariables");
6+
const Provider = require("../../../utils/agents/aibitat/providers/ai-provider");
7+
8+
jest.mock("../../../models/systemPromptVariables");
9+
jest.mock("../../../models/systemSettings");
10+
jest.mock("../../../utils/agents/imported", () => ({
11+
activeImportedPlugins: jest.fn().mockReturnValue([]),
12+
}));
13+
jest.mock("../../../utils/agentFlows", () => ({
14+
AgentFlows: {
15+
activeFlowPlugins: jest.fn().mockReturnValue([]),
16+
},
17+
}));
18+
jest.mock("../../../utils/MCP", () => {
19+
return jest.fn().mockImplementation(() => ({
20+
activeMCPServers: jest.fn().mockResolvedValue([]),
21+
}));
22+
});
23+
24+
const { WORKSPACE_AGENT } = require("../../../utils/agents/defaults");
25+
26+
describe("WORKSPACE_AGENT.getDefinition", () => {
27+
beforeEach(() => {
28+
jest.clearAllMocks();
29+
// Mock SystemSettings to return empty arrays for agent skills
30+
const { SystemSettings } = require("../../../models/systemSettings");
31+
SystemSettings.getValueOrFallback = jest.fn().mockResolvedValue("[]");
32+
});
33+
34+
it("should use provider default system prompt when workspace has no openAiPrompt", async () => {
35+
const workspace = {
36+
id: 1,
37+
name: "Test Workspace",
38+
openAiPrompt: null,
39+
};
40+
const user = { id: 1 };
41+
const provider = "openai";
42+
43+
const definition = await WORKSPACE_AGENT.getDefinition(
44+
provider,
45+
workspace,
46+
user
47+
);
48+
49+
expect(definition.role).toBe(Provider.systemPrompt(provider));
50+
expect(SystemPromptVariables.expandSystemPromptVariables).not.toHaveBeenCalled();
51+
});
52+
53+
it("should use workspace system prompt with variable expansion when openAiPrompt exists", async () => {
54+
const workspace = {
55+
id: 1,
56+
name: "Test Workspace",
57+
openAiPrompt: "You are a helpful assistant for {workspace.name}. The current user is {user.name}.",
58+
};
59+
const user = { id: 1 };
60+
const provider = "openai";
61+
62+
const expandedPrompt = "You are a helpful assistant for Test Workspace. The current user is John Doe.";
63+
SystemPromptVariables.expandSystemPromptVariables.mockResolvedValue(expandedPrompt);
64+
65+
const definition = await WORKSPACE_AGENT.getDefinition(
66+
provider,
67+
workspace,
68+
user
69+
);
70+
71+
expect(SystemPromptVariables.expandSystemPromptVariables).toHaveBeenCalledWith(
72+
workspace.openAiPrompt,
73+
user.id,
74+
workspace.id
75+
);
76+
expect(definition.role).toBe(expandedPrompt);
77+
});
78+
79+
it("should handle workspace system prompt without user context", async () => {
80+
const workspace = {
81+
id: 1,
82+
name: "Test Workspace",
83+
openAiPrompt: "You are a helpful assistant. Today is {date}.",
84+
};
85+
const user = null;
86+
const provider = "lmstudio";
87+
88+
const expandedPrompt = "You are a helpful assistant. Today is January 1, 2024.";
89+
SystemPromptVariables.expandSystemPromptVariables.mockResolvedValue(expandedPrompt);
90+
91+
const definition = await WORKSPACE_AGENT.getDefinition(
92+
provider,
93+
workspace,
94+
user
95+
);
96+
97+
expect(SystemPromptVariables.expandSystemPromptVariables).toHaveBeenCalledWith(
98+
workspace.openAiPrompt,
99+
null,
100+
workspace.id
101+
);
102+
expect(definition.role).toBe(expandedPrompt);
103+
});
104+
105+
it("should return functions array in definition", async () => {
106+
const workspace = { id: 1, openAiPrompt: null };
107+
const provider = "openai";
108+
109+
const definition = await WORKSPACE_AGENT.getDefinition(
110+
provider,
111+
workspace,
112+
null
113+
);
114+
115+
expect(definition).toHaveProperty("functions");
116+
expect(Array.isArray(definition.functions)).toBe(true);
117+
});
118+
119+
it("should use LMStudio specific prompt when workspace has no openAiPrompt", async () => {
120+
const workspace = { id: 1, openAiPrompt: null };
121+
const provider = "lmstudio";
122+
123+
const definition = await WORKSPACE_AGENT.getDefinition(
124+
provider,
125+
workspace,
126+
null
127+
);
128+
129+
expect(definition.role).toBe(Provider.systemPrompt(provider));
130+
expect(definition.role).toContain("helpful ai assistant");
131+
});
132+
});
133+

‎server/utils/agents/defaults.js‎

Lines changed: 14 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@ const Provider = require("./aibitat/providers/ai-provider");
55
const ImportedPlugin = require("./imported");
66
const { AgentFlows } = require("../agentFlows");
77
const MCPCompatibilityLayer = require("../MCP");
8+
const { SystemPromptVariables } = require("../../models/systemPromptVariables");
89

910
// This is a list of skills that are built-in and default enabled.
1011
const DEFAULT_SKILLS = [
@@ -25,9 +26,20 @@ const USER_AGENT = {
2526

2627
const WORKSPACE_AGENT = {
2728
name: "@agent",
28-
getDefinition: async (provider = null) => {
29+
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+
2941
return {
30-
role: Provider.systemPrompt(provider),
42+
role,
3143
functions: [
3244
...(await agentSkillsFromSystemSettings()),
3345
...ImportedPlugin.activeImportedPlugins(),

‎server/utils/agents/ephemeral.js‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -320,9 +320,13 @@ class EphemeralAgentHandler extends AgentHandler {
320320
// Default User agent and workspace agent
321321
this.log(`Attaching user and default agent to Agent cluster.`);
322322
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+
323327
this.aibitat.agent(
324328
WORKSPACE_AGENT.name,
325-
await WORKSPACE_AGENT.getDefinition(this.provider)
329+
await WORKSPACE_AGENT.getDefinition(this.provider, this.#workspace, user)
326330
);
327331

328332
this.#funcsToLoad = [

‎server/utils/agents/index.js‎

Lines changed: 17 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -524,14 +524,29 @@ class AgentHandler {
524524
// Default User agent and workspace agent
525525
this.log(`Attaching user and default agent to Agent cluster.`);
526526
this.aibitat.agent(USER_AGENT.name, await USER_AGENT.getDefinition());
527+
528+
// Get user object from invocation for variable expansion
529+
const user = this.invocation.user_id
530+
? { id: this.invocation.user_id }
531+
: null;
532+
527533
this.aibitat.agent(
528534
WORKSPACE_AGENT.name,
529-
await WORKSPACE_AGENT.getDefinition(this.provider)
535+
await WORKSPACE_AGENT.getDefinition(
536+
this.provider,
537+
this.invocation.workspace,
538+
user
539+
)
530540
);
531541

542+
const workspaceAgentDef = await WORKSPACE_AGENT.getDefinition(
543+
this.provider,
544+
this.invocation.workspace,
545+
user
546+
);
532547
this.#funcsToLoad = [
533548
...((await USER_AGENT.getDefinition())?.functions || []),
534-
...((await WORKSPACE_AGENT.getDefinition())?.functions || []),
549+
...(workspaceAgentDef?.functions || []),
535550
];
536551
}
537552

0 commit comments

Comments
 (0)