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	189	const tools = await loadTools();
190	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	207	app.listen(PORT);
195	208

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

@@ -17,6 +17,8 @@ import { · 16 unmodified lines
17	17	export interface TaskRunnerOptions {
18	18	readonly tools: ToolTree;
19	19	readonly model: LanguageModel;
	20	+ /** Additional context for the system prompt (project IDs, org info, etc.) */
	21	+ readonly context?: string \| undefined;
20	22	}
21	23
22	24	/**
@@ -36,6 +38,7 @@ export async function runTask( · 13 unmodified lines
36	38	const agent = createAgent({
37	39	tools: options.tools,
38	40	model: options.model,
	41	+ context: options.context,
39	42	requestApproval: (request: ApprovalRequest): Promise<ApprovalDecision> => {
40	43	// Emit the approval_request event so SSE subscribers see it
41	44	emitTaskEvent(taskId, {

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

@@ -79,6 +79,12 @@ export interface AgentOptions { · 78 unmodified lines
79	79	* Defaults to 50.
80	80	*/
81	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	90	export interface AgentResult {
@@ -181,7 +187,7 @@ function prepareToolsForAgent( · 96 unmodified lines
181	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	191	const discoveryInstructions = discoveryMode
186	192	? `
187	193	## Tool Discovery
@@ -222,13 +228,21 @@ for (const issue of issues) { · 34 unmodified lines
222	228
223	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	242	return `You are an AI assistant that executes tasks by generating TypeScript code.
229	243
230	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	246	## Available Tools
233	247
234	248	${guidance}
@@ -241,6 +255,17 @@ ${discoveryInstructions} · 6 unmodified lines
241	255	- No \`fetch\`, \`process\`, \`require\`, or \`import\` — use tools for all external interactions
242	256	- Write clean, straightforward TypeScript — use for loops, map, filter, Promise.all as needed
243	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	269	- Return a structured result from your script, then summarize what happened
245	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	286	maxCodeRuns = 100,
262	287	timeoutMs = 30_000,
263	288	discoveryThreshold = DISCOVERY_THRESHOLD,
	289	+ context,
264	290	} = options;
265	291
266	292	const { sandboxTools, promptGuidance, toolDeclarations, discoveryMode } =
@@ -278,7 +304,7 @@ export function createAgent(options: AgentOptions): { · 11 unmodified lines
278	304
279	305	return {
280	306	async run(prompt: string): Promise<AgentResult> {
281		− const systemPrompt = buildSystemPrompt(promptGuidance, discoveryMode);
	307	+ const systemPrompt = buildSystemPrompt(promptGuidance, discoveryMode, context);
282	308	const runs: CodeRun[] = [];
283	309	const allReceipts: ToolCallReceipt[] = [];
284	310

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

@@ -48,13 +48,19 @@ export function buildToolIndex(tree: ToolTree): ToolIndexEntry[] { · 47 unmodified lines
48	48	const path = prefix ? `${prefix}.${key}` : key;
49	49	if (isToolDefinition(value)) {
50	50	const tool = value as ToolDefinition;
	51	+ const returnsType = getReturnsTypeString(tool);
	52	+ // Detect paginated endpoints from return type
	53	+ const isPaginated = /\bpagination\b\|\bnext_cursor\b\|\bhas_more\b\|\bnextPage\b/.test(returnsType);
	54	+ const description = isPaginated
	55	+ ? `${tool.description} (⚠️ paginated — loop until pagination.next is null to get all results)`
	56	+ : tool.description;
51	57	entries.push({
52	58	path,
53		− description: tool.description,
	59	+ description,
54	60	approval: tool.approval,
55	61	argsType: getArgsTypeString(tool),
56		− returnsType: getReturnsTypeString(tool),
57		− searchText: `${path} ${tool.description}`.toLowerCase(),
	62	+ returnsType,
	63	+ searchText: `${path} ${description}`.toLowerCase(),
58	64	});
59	65	} else {
60	66	walk(value as ToolTree, path);