Enhance system prompt context handling by adding support for environment-specific information such as PostHog project ID, GitHub token, and Vercel token. Update related interfaces and functions to accommodate the new context parameter, improving the agent's ability to utilize contextual data during task execution.

Rhys Sullivan committed 9 months ago

Commit 66259751 parent 7a5c9454 files · +55 −7Browse files

Unified · Split

4 files

apps/server/src/index.ts +14 −1
apps/server/src/task-runner.ts +3 −0
packages/core/src/agent.ts +29 −3
packages/core/src/discovery.ts +9 −3

@@ -189,7 +189,20 @@ console.log("Loading tools..."); · 188 unmodified lines
189	const tools = await loadTools();	189	const tools = await loadTools();
190	const model = createPiAiModel({ apiKey });	190	const model = createPiAiModel({ apiKey });
191		191
192	const app = createApp({ tools, model });	192	// Build context for the system prompt
		193	const contextLines: string[] = [];
		194	if (POSTHOG_PROJECT_ID) {
		195	contextLines.push(`- PostHog project ID: ${POSTHOG_PROJECT_ID} (use this for PostHog API calls that require a project_id)`);
		196	}
		197	if (GITHUB_TOKEN) {
		198	contextLines.push(`- GitHub: authenticated via personal access token`);
		199	}
		200	if (VERCEL_TOKEN) {
		201	contextLines.push(`- Vercel: authenticated via personal token`);
		202	}
		203	const context = contextLines.length > 0 ? contextLines.join("\n") : undefined;
		204
		205	const app = createApp({ tools, model, context });
193		206
194	app.listen(PORT);	207	app.listen(PORT);
195		208

apps/server/src/task-runner.ts+3−0

 export interface TaskRunnerOptions {
   readonly tools: ToolTree;
   readonly model: LanguageModel;
+  /** Additional context for the system prompt (project IDs, org info, etc.) */
+  readonly context?: string | undefined;
 }
 /**
   const agent = createAgent({
     tools: options.tools,
     model: options.model,
+    context: options.context,
     requestApproval: (request: ApprovalRequest): Promise<ApprovalDecision> => {
       // Emit the approval_request event so SSE subscribers see it
       emitTaskEvent(taskId, {

packages/core/src/agent.ts+29−3

@@ -79,6 +79,12 @@ export interface AgentOptions { · 78 unmodified lines
79	* Defaults to 50.	79	* Defaults to 50.
80	*/	80	*/
81	readonly discoveryThreshold?: number \| undefined;	81	readonly discoveryThreshold?: number \| undefined;
		82	/**
		83	* Additional context to include in the system prompt.
		84	* Use this to provide environment-specific info like project IDs,
		85	* usernames, organization names, etc. that the agent may need.
		86	*/
		87	readonly context?: string \| undefined;
82	}	88	}
83		89
84	export interface AgentResult {	90	export interface AgentResult {
@@ -181,7 +187,7 @@ function prepareToolsForAgent( · 96 unmodified lines
181	// System prompt	187	// System prompt
182	// ---------------------------------------------------------------------------	188	// ---------------------------------------------------------------------------
183		189
184	function buildSystemPrompt(guidance: string, discoveryMode: boolean): string {	190	function buildSystemPrompt(guidance: string, discoveryMode: boolean, context?: string): string {
185	const discoveryInstructions = discoveryMode	191	const discoveryInstructions = discoveryMode
186	? `	192	? `
187	## Tool Discovery	193	## Tool Discovery
@@ -222,13 +228,21 @@ for (const issue of issues) { · 34 unmodified lines
222		228
223	return { closed: results.length, issues: results };	229	return { closed: results.length, issues: results };
224	\`\`\`	230	\`\`\`
		231	`
		232	: "";
		233
		234	const contextSection = context
		235	? `
		236	## Context
		237
		238	${context}
225	`	239	`
226	: "";	240	: "";
227		241
228	return `You are an AI assistant that executes tasks by generating TypeScript code.	242	return `You are an AI assistant that executes tasks by generating TypeScript code.
229		243
230	You have access to a set of tools via the \`tools\` object. When the user asks you to do something, generate TypeScript code that calls these tools to accomplish the task.	244	You have access to a set of tools via the \`tools\` object. When the user asks you to do something, generate TypeScript code that calls these tools to accomplish the task.
231		245	${contextSection}
232	## Available Tools	246	## Available Tools
233		247
234	${guidance}	248	${guidance}
@@ -241,6 +255,17 @@ ${discoveryInstructions} · 6 unmodified lines
241	- No \`fetch\`, \`process\`, \`require\`, or \`import\` — use tools for all external interactions	255	- No \`fetch\`, \`process\`, \`require\`, or \`import\` — use tools for all external interactions
242	- Write clean, straightforward TypeScript — use for loops, map, filter, Promise.all as needed	256	- Write clean, straightforward TypeScript — use for loops, map, filter, Promise.all as needed
243	- Handle errors gracefully with try/catch — if a tool call fails or is denied, continue with remaining work	257	- Handle errors gracefully with try/catch — if a tool call fails or is denied, continue with remaining work
		258	- Pagination: If a list endpoint's return type has \`pagination\`, \`next\`, \`cursor\`, or similar fields, the results may be paginated. Check \`pagination.next\` (or equivalent) in the response — if it's not null, there are more pages. Paginate when the user needs complete results. Example pattern:
		259	\`\`\`ts
		260	let allItems = [];
		261	let cursor = undefined;
		262	while (true) {
		263	const resp = await tools.api.list({ limit: 100, since: cursor });
		264	allItems.push(...resp.items);
		265	if (!resp.pagination?.next) break;
		266	cursor = resp.pagination.next;
		267	}
		268	\`\`\`
244	- Return a structured result from your script, then summarize what happened	269	- Return a structured result from your script, then summarize what happened
245	- Be concise and accurate — base your response on actual tool results, not assumptions`;	270	- Be concise and accurate — base your response on actual tool results, not assumptions`;
246	}	271	}
@@ -261,6 +286,7 @@ export function createAgent(options: AgentOptions): { · 14 unmodified lines
261	maxCodeRuns = 100,	286	maxCodeRuns = 100,
262	timeoutMs = 30_000,	287	timeoutMs = 30_000,
263	discoveryThreshold = DISCOVERY_THRESHOLD,	288	discoveryThreshold = DISCOVERY_THRESHOLD,
		289	context,
264	} = options;	290	} = options;
265		291
266	const { sandboxTools, promptGuidance, toolDeclarations, discoveryMode } =	292	const { sandboxTools, promptGuidance, toolDeclarations, discoveryMode } =
@@ -278,7 +304,7 @@ export function createAgent(options: AgentOptions): { · 11 unmodified lines
278		304
279	return {	305	return {
280	async run(prompt: string): Promise<AgentResult> {	306	async run(prompt: string): Promise<AgentResult> {
281	const systemPrompt = buildSystemPrompt(promptGuidance, discoveryMode);	307	const systemPrompt = buildSystemPrompt(promptGuidance, discoveryMode, context);
282	const runs: CodeRun[] = [];	308	const runs: CodeRun[] = [];
283	const allReceipts: ToolCallReceipt[] = [];	309	const allReceipts: ToolCallReceipt[] = [];
284		310

packages/core/src/discovery.ts+9−3

       const path = prefix ? `${prefix}.${key}` : key;
       if (isToolDefinition(value)) {
         const tool = value as ToolDefinition;
+        const returnsType = getReturnsTypeString(tool);
+        // Detect paginated endpoints from return type
+        const isPaginated = /\bpagination\b|\bnext_cursor\b|\bhas_more\b|\bnextPage\b/.test(returnsType);
+        const description = isPaginated
+          ? `${tool.description} (⚠️ paginated — loop until pagination.next is null to get all results)`
+          : tool.description;
         entries.push({
           path,
-          description: tool.description,
+          description,
           approval: tool.approval,
           argsType: getArgsTypeString(tool),
-          returnsType: getReturnsTypeString(tool),
+          returnsType,
-          searchText: `${path} ${tool.description}`.toLowerCase(),
+          searchText: `${path} ${description}`.toLowerCase(),
         });
       } else {
         walk(value as ToolTree, path);