History

posthog and github

Rhys Sullivan committed 9 months ago

Commit b9fea141 parent 98faa996 files · +372 −15Browse files

Unified · Split

6 files
  1. apps/server/src/index.ts +41 −0
  2. packages/core/package.json +2 −1
  3. packages/core/src/agent.ts +127 −7
  4. packages/core/src/discovery.ts +169 −0
  5. packages/core/src/index.ts +9 −0
  6. packages/tool-gen/src/openapi.ts +24 −7
apps/server/src/index.ts+41−0
@@ -11,6 +11,8 @@ import { createPiAiModel } from "@openassistant/core"; · 10 unmodified lines
11import { defineTool, mergeToolTrees, type ToolTree } from "@openassistant/core/tools";11import { defineTool, mergeToolTrees, type ToolTree } from "@openassistant/core/tools";
12import { generateMcpTools } from "@openassistant/tool-gen/mcp";12import { generateMcpTools } from "@openassistant/tool-gen/mcp";
13import type { McpToolSource } from "@openassistant/tool-gen/mcp";13import type { McpToolSource } from "@openassistant/tool-gen/mcp";
14import { generateOpenApiTools } from "@openassistant/tool-gen/openapi";
15import type { OpenApiToolSource } from "@openassistant/tool-gen/openapi";
14import { z } from "zod";16import { z } from "zod";
15 17
16// ---------------------------------------------------------------------------18// ---------------------------------------------------------------------------
@@ -54,6 +56,28 @@ const mcpSources: McpToolSource[] = [ · 37 unmodified lines
54 },56 },
55];57];
56 58
59const POSTHOG_API_KEY = process.env["POSTHOG_PERSONAL_API_KEY"];
60const POSTHOG_PROJECT_ID = process.env["POSTHOG_PROJECT_ID"];
61const GITHUB_TOKEN = process.env["OPENASSISTANT_GITHUB_TOKEN"];
62
63const openApiSources: OpenApiToolSource[] = [
64 ...(POSTHOG_API_KEY ? [{
65 name: "posthog",
66 spec: "https://app.posthog.com/api/schema/?format=json",
67 baseUrl: "https://app.posthog.com",
68 auth: { type: "bearer" as const, token: POSTHOG_API_KEY },
69 defaultReadApproval: "auto" as const,
70 defaultWriteApproval: "required" as const,
71 }] : []),
72 ...(GITHUB_TOKEN ? [{
73 name: "github",
74 spec: "https://raw.githubusercontent.com/github/rest-api-description/main/descriptions/api.github.com/api.github.com.json",
75 auth: { type: "bearer" as const, token: GITHUB_TOKEN },
76 defaultReadApproval: "auto" as const,
77 defaultWriteApproval: "required" as const,
78 }] : []),
79];
80
57// ---------------------------------------------------------------------------81// ---------------------------------------------------------------------------
58// Hand-written tools (approval testing)82// Hand-written tools (approval testing)
59// ---------------------------------------------------------------------------83// ---------------------------------------------------------------------------
@@ -111,6 +135,7 @@ const builtinTools: ToolTree = { · 51 unmodified lines
111async function loadTools(): Promise<ToolTree> {135async function loadTools(): Promise<ToolTree> {
112 const trees: ToolTree[] = [builtinTools];136 const trees: ToolTree[] = [builtinTools];
113 137
138 // Load MCP tools
114 for (const source of mcpSources) {139 for (const source of mcpSources) {
115 try {140 try {
116 console.log(`Loading MCP tools from ${source.name} (${source.url})...`);141 console.log(`Loading MCP tools from ${source.name} (${source.url})...`);
@@ -122,6 +147,22 @@ async function loadTools(): Promise<ToolTree> { · 5 unmodified lines
122 }147 }
123 }148 }
124 149
150 // Load OpenAPI tools
151 for (const source of openApiSources) {
152 try {
153 console.log(`Loading OpenAPI tools from ${source.name}...`);
154 const result = await generateOpenApiTools(source);
155 trees.push(result.tools);
156 const nsTools = result.tools[source.name];
157 const count = nsTools ? Object.values(nsTools).reduce((sum, tag) => {
158 return sum + (typeof tag === "object" && tag !== null ? Object.keys(tag).length : 1);
159 }, 0) : 0;
160 console.log(` Loaded ~${count} tools from ${source.name}`);
161 } catch (error) {
162 console.error(` Failed to load ${source.name}:`, error instanceof Error ? error.message : error);
163 }
164 }
165
125 return mergeToolTrees(...trees);166 return mergeToolTrees(...trees);
126}167}
127 168
packages/core/package.json+2−1
@@ -8,7 +8,8 @@ · 7 unmodified lines
8 "./runner": "./src/runner.ts",8 "./runner": "./src/runner.ts",
9 "./typechecker": "./src/typechecker.ts",9 "./typechecker": "./src/typechecker.ts",
10 "./agent": "./src/agent.ts",10 "./agent": "./src/agent.ts",
11 "./events": "./src/events.ts"11 "./events": "./src/events.ts",
12 "./discovery": "./src/discovery.ts"
12 },13 },
13 "scripts": {14 "scripts": {
14 "test": "bun test",15 "test": "bun test",
packages/core/src/agent.ts+127−7
@@ -9,12 +9,17 @@ · 8 unmodified lines
9 * 6. Execute in sandbox9 * 6. Execute in sandbox
10 * 7. Collect receipts, feed back to Claude10 * 7. Collect receipts, feed back to Claude
11 * 8. Claude may call run_code again or produce final text11 * 8. Claude may call run_code again or produce final text
12 *
13 * When the tool count exceeds DISCOVERY_THRESHOLD, the agent injects
14 * a `tools.discover()` meta-tool so the LLM can search for relevant
15 * tools by keyword instead of having all tools described in the prompt.
12 */16 */
13 17
14import type { ToolTree, ApprovalDecision, ApprovalRequest, ToolCallReceipt } from "./tools.js";18import type { ToolTree, ApprovalDecision, ApprovalRequest, ToolCallReceipt } from "./tools.js";
15import type { RunResult } from "./runner.js";19import type { RunResult } from "./runner.js";
16import { createRunner } from "./runner.js";20import { createRunner } from "./runner.js";
17import { generateToolDeclarations, generatePromptGuidance, typecheckCode } from "./typechecker.js";21import { generateToolDeclarations, generatePromptGuidance, typecheckCode } from "./typechecker.js";
22import { countTools, createDiscoverTool } from "./discovery.js";
18import type { TaskEvent } from "./events.js";23import type { TaskEvent } from "./events.js";
19 24
20// ---------------------------------------------------------------------------25// ---------------------------------------------------------------------------
@@ -67,6 +72,13 @@ export interface AgentOptions { · 46 unmodified lines
67 readonly maxCodeRuns?: number | undefined;72 readonly maxCodeRuns?: number | undefined;
68 /** Execution timeout per code run in ms. Defaults to 30_000. */73 /** Execution timeout per code run in ms. Defaults to 30_000. */
69 readonly timeoutMs?: number | undefined;74 readonly timeoutMs?: number | undefined;
75 /**
76 * Tool count threshold for enabling discovery mode.
77 * When total tools exceed this, inject `tools.discover()` and
78 * only describe small tool sources in the system prompt.
79 * Defaults to 50.
80 */
81 readonly discoveryThreshold?: number | undefined;
70}82}
71 83
72export interface AgentResult {84export interface AgentResult {
@@ -81,12 +93,117 @@ export interface CodeRun { · 8 unmodified lines
81 readonly result: RunResult;93 readonly result: RunResult;
82}94}
83 95
96// ---------------------------------------------------------------------------
97// Discovery mode
98// ---------------------------------------------------------------------------
99
100const DISCOVERY_THRESHOLD = 50;
101
102/**
103 * When discovery mode is active, separate tools into:
104 * - "small" tools that fit in the prompt (under the threshold)
105 * - "large" tools that are only available via discover()
106 *
107 * The discover tool searches ALL tools. Small tools are described
108 * in the prompt AND available via discover.
109 */
110function prepareToolsForAgent(
111 tools: ToolTree,
112 threshold: number,
113): {
114 /** Tools actually wired in the sandbox (always all of them + discover if needed) */
115 sandboxTools: ToolTree;
116 /** Prompt guidance (subset or all) */
117 promptGuidance: string;
118 /** Tool declarations for typechecker (subset or all) */
119 toolDeclarations: string;
120 /** Whether discovery mode is active */
121 discoveryMode: boolean;
122} {
123 const totalCount = countTools(tools);
124
125 if (totalCount <= threshold) {
126 // Small enough — describe everything in the prompt
127 return {
128 sandboxTools: tools,
129 promptGuidance: generatePromptGuidance(tools),
130 toolDeclarations: generateToolDeclarations(tools),
131 discoveryMode: false,
132 };
133 }
134
135 // Discovery mode: inject discover tool, only describe small namespaces in prompt
136 const discoverTool = createDiscoverTool(tools);
137
138 // Separate namespaces into small (described) and large (discover-only)
139 const described: Record<string, ToolTree | import("./tools.js").ToolDefinition> = { discover: discoverTool };
140 const largeNamespaces: string[] = [];
141
142 for (const [key, value] of Object.entries(tools)) {
143 const nsCount = countTools({ [key]: value });
144 if (nsCount <= threshold) {
145 described[key] = value;
146 } else {
147 largeNamespaces.push(`${key} (${nsCount} tools)`);
148 }
149 }
150
151 const describedTree = described as ToolTree;
152
153 // Sandbox gets ALL tools + discover
154 const sandboxTools = { ...tools, discover: discoverTool } as ToolTree;
155
156 // Prompt only describes the small namespaces + discover
157 const promptGuidance = generatePromptGuidance(describedTree);
158
159 // Typechecker: small namespaces are fully typed, large namespaces
160 // are declared as `Record<string, any>` so the LLM can call
161 // discovered tools without typecheck errors.
162 const largeNsDeclarations = Object.entries(tools)
163 .filter(([key]) => !described[key])
164 .map(([key]) => ` ${key}: Record<string, Record<string, (...args: any[]) => Promise<any>>>;`)
165 .join("\n");
166
167 const toolDeclarations = generateToolDeclarations(describedTree).replace(
168 /\};$/,
169 largeNsDeclarations ? `${largeNsDeclarations}\n};` : "};",
170 );
171
172 return {
173 sandboxTools,
174 promptGuidance,
175 toolDeclarations,
176 discoveryMode: true,
177 };
178}
179
84// ---------------------------------------------------------------------------180// ---------------------------------------------------------------------------
85// System prompt181// System prompt
86// ---------------------------------------------------------------------------182// ---------------------------------------------------------------------------
87 183
88function buildSystemPrompt(tools: ToolTree): string {184function buildSystemPrompt(guidance: string, discoveryMode: boolean): string {
89 const guidance = generatePromptGuidance(tools);185 const discoveryInstructions = discoveryMode
186 ? `
187## Tool Discovery
188
189Some tool namespaces are too large to list here. Use \`tools.discover({ query: "..." })\` to search for tools by keyword.
190The discover tool returns an array of matching tools with their paths, descriptions, and TypeScript signatures.
191
192**Workflow for large APIs:**
1931. Call \`tools.discover({ query: "relevant keywords" })\` to find relevant tools
1942. Read the returned signatures to understand the input/output types
1953. Call the discovered tools using the path from the results, e.g. \`tools.posthog.feature_flags.feature_flags_list({ project_id: "123" })\`
1964. You can call discover multiple times with different queries
197
198Example:
199\`\`\`ts
200const found = await tools.discover({ query: "feature flags" });
201// found.results = [{ path: "posthog.feature_flags.feature_flags_list", signature: "...", ... }]
202const flags = await tools.posthog.feature_flags.feature_flags_list({ project_id: "123" });
203\`\`\`
204`
205 : "";
206
90 return `You are an AI assistant that executes tasks by generating TypeScript code.207 return `You are an AI assistant that executes tasks by generating TypeScript code.
91 208
92You 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.209You 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.
@@ -94,7 +211,7 @@ You have access to a set of tools via the \`tools\` object. When the user asks y · 1 unmodified line
94## Available Tools211## Available Tools
95 212
96${guidance}213${guidance}
97 214${discoveryInstructions}
98## Instructions215## Instructions
99 216
100- Use the \`run_code\` tool to execute TypeScript code217- Use the \`run_code\` tool to execute TypeScript code
@@ -121,12 +238,14 @@ export function createAgent(options: AgentOptions): { · 20 unmodified lines
121 maxTypecheckRetries = 3,238 maxTypecheckRetries = 3,
122 maxCodeRuns = 10,239 maxCodeRuns = 10,
123 timeoutMs = 30_000,240 timeoutMs = 30_000,
241 discoveryThreshold = DISCOVERY_THRESHOLD,
124 } = options;242 } = options;
125 243
126 const toolDeclarations = generateToolDeclarations(tools);244 const { sandboxTools, promptGuidance, toolDeclarations, discoveryMode } =
245 prepareToolsForAgent(tools, discoveryThreshold);
127 246
128 const runner = createRunner({247 const runner = createRunner({
129 tools,248 tools: sandboxTools,
130 requestApproval,249 requestApproval,
131 timeoutMs,250 timeoutMs,
132 });251 });
@@ -137,7 +256,7 @@ export function createAgent(options: AgentOptions): { · 4 unmodified lines
137 256
138 return {257 return {
139 async run(prompt: string): Promise<AgentResult> {258 async run(prompt: string): Promise<AgentResult> {
140 const systemPrompt = buildSystemPrompt(tools);259 const systemPrompt = buildSystemPrompt(promptGuidance, discoveryMode);
141 const runs: CodeRun[] = [];260 const runs: CodeRun[] = [];
142 const allReceipts: ToolCallReceipt[] = [];261 const allReceipts: ToolCallReceipt[] = [];
143 262
@@ -194,6 +313,7 @@ export function createAgent(options: AgentOptions): { · 50 unmodified lines
194 typecheckOk = true;313 typecheckOk = true;
195 break;314 break;
196 }315 }
316
197 typecheckErrors = check.errors;317 typecheckErrors = check.errors;
198 318
199 if (attempt < maxTypecheckRetries) {319 if (attempt < maxTypecheckRetries) {
@@ -270,7 +390,7 @@ export function createAgent(options: AgentOptions): { · 70 unmodified lines
270 }390 }
271 391
272 // Hit max code runs392 // Hit max code runs
273 const text = "Reached maximum number of code executions.";393 const text = "I've reached the maximum number of code executions for this task.";
274 emit({ type: "agent_message", text });394 emit({ type: "agent_message", text });
275 emit({ type: "completed", receipts: allReceipts });395 emit({ type: "completed", receipts: allReceipts });
276 return { text, runs, allReceipts };396 return { text, runs, allReceipts };
packages/core/src/discovery.ts+169−0
@@ -0,0 +1,169 @@
1/**
2 * Tool discovery — searchable index of tools for large APIs.
3 *
4 * When there are too many tools to fit in the LLM context window,
5 * the agent uses `tools.discover({ query })` to search for relevant
6 * tools by name, path, and description. Returns type signatures so
7 * the LLM knows how to call what it finds.
8 *
9 * The discovered tools are already wired in the sandbox — the LLM
10 * just doesn't get told about all of them in the system prompt.
11 * `discover` fills that gap on demand.
12 */
13
14import { defineTool, isToolDefinition, type ToolDefinition, type ToolTree } from "./tools.js";
15import { zodToTypeString } from "./typechecker.js";
16import { z } from "zod";
17
18// ---------------------------------------------------------------------------
19// Index entry
20// ---------------------------------------------------------------------------
21
22export interface ToolIndexEntry {
23 /** Full dotted path: "posthog.feature_flags.list" */
24 readonly path: string;
25 /** Tool description */
26 readonly description: string;
27 /** Approval mode */
28 readonly approval: "auto" | "required";
29 /** TypeScript signature: "(input: { ... }): Promise<{ ... }>" */
30 readonly signature: string;
31 /** Searchable text (lowercase): path + description */
32 readonly searchText: string;
33}
34
35// ---------------------------------------------------------------------------
36// Build index
37// ---------------------------------------------------------------------------
38
39export function buildToolIndex(tree: ToolTree): ToolIndexEntry[] {
40 const entries: ToolIndexEntry[] = [];
41
42 function walk(node: ToolTree, prefix: string) {
43 for (const [key, value] of Object.entries(node)) {
44 const path = prefix ? `${prefix}.${key}` : key;
45 if (isToolDefinition(value)) {
46 const tool = value as ToolDefinition;
47 const argsType = zodToTypeString(tool.args);
48 const returnsType = zodToTypeString(tool.returns);
49 entries.push({
50 path,
51 description: tool.description,
52 approval: tool.approval,
53 signature: `(input: ${argsType}): Promise<${returnsType}>`,
54 searchText: `${path} ${tool.description}`.toLowerCase(),
55 });
56 } else {
57 walk(value as ToolTree, path);
58 }
59 }
60 }
61
62 walk(tree, "");
63 return entries;
64}
65
66// ---------------------------------------------------------------------------
67// Search
68// ---------------------------------------------------------------------------
69
70/**
71 * Simple keyword search. Matches entries where ALL query terms
72 * appear somewhere in the path + description.
73 */
74function searchIndex(
75 index: ToolIndexEntry[],
76 query: string,
77 limit: number,
78): ToolIndexEntry[] {
79 const terms = query.toLowerCase().split(/\s+/).filter(Boolean);
80 if (terms.length === 0) return index.slice(0, limit);
81
82 const scored = index
83 .map((entry) => {
84 let score = 0;
85 for (const term of terms) {
86 if (!entry.searchText.includes(term)) return { entry, score: -1 };
87 // Bonus for path match (more specific)
88 if (entry.path.toLowerCase().includes(term)) score += 2;
89 // Bonus for exact word boundary match
90 if (entry.searchText.includes(` ${term}`) || entry.searchText.startsWith(term)) score += 1;
91 score += 1;
92 }
93 return { entry, score };
94 })
95 .filter((s) => s.score > 0)
96 .sort((a, b) => b.score - a.score)
97 .slice(0, limit);
98
99 return scored.map((s) => s.entry);
100}
101
102// ---------------------------------------------------------------------------
103// Create discover tool
104// ---------------------------------------------------------------------------
105
106export interface DiscoverToolOptions {
107 /** Max results to return per query. Defaults to 20. */
108 readonly maxResults?: number;
109}
110
111/**
112 * Create a `discover` tool that searches the tool index.
113 *
114 * Usage in LLM-generated code:
115 * const results = await tools.discover({ query: "feature flags" });
116 * // results is an array of { path, description, approval, signature }
117 * // Then call: await tools.posthog.feature_flags.list({ ... })
118 */
119export function createDiscoverTool(
120 tree: ToolTree,
121 options?: DiscoverToolOptions,
122) {
123 const index = buildToolIndex(tree);
124 const maxResults = options?.maxResults ?? 20;
125
126 return defineTool({
127 description: `Search for available tools by keyword. Returns matching tool paths, descriptions, and TypeScript signatures. Use this when you need to find tools for a specific task. There are ${index.length} tools available across all connected services.`,
128 approval: "auto" as const,
129 args: z.object({
130 query: z.string().describe("Search keywords (e.g. 'feature flags', 'create user', 'analytics events')"),
131 }),
132 returns: z.object({
133 results: z.array(z.object({
134 path: z.string(),
135 description: z.string(),
136 approval: z.string(),
137 signature: z.string(),
138 })),
139 total: z.number(),
140 }),
141 run: async (input) => {
142 const results = searchIndex(index, input.query, maxResults);
143 return {
144 results: results.map((e) => ({
145 path: e.path,
146 description: e.description,
147 approval: e.approval,
148 signature: e.signature,
149 })),
150 total: results.length,
151 };
152 },
153 });
154}
155
156/**
157 * Count the total number of tool definitions in a tree.
158 */
159export function countTools(tree: ToolTree): number {
160 let count = 0;
161 for (const value of Object.values(tree)) {
162 if (isToolDefinition(value)) {
163 count++;
164 } else {
165 count += countTools(value as ToolTree);
166 }
167 }
168 return count;
169}
packages/core/src/index.ts+9−0
@@ -51,6 +51,15 @@ export { · 50 unmodified lines
51 type PiAiModelOptions,51 type PiAiModelOptions,
52} from "./pi-ai-adapter.js";52} from "./pi-ai-adapter.js";
53 53
54// Discovery
55export {
56 buildToolIndex,
57 createDiscoverTool,
58 countTools,
59 type ToolIndexEntry,
60 type DiscoverToolOptions,
61} from "./discovery.js";
62
54// Events63// Events
55export type {64export type {
56 TaskEvent,65 TaskEvent,
packages/tool-gen/src/openapi.ts+24−7
@@ -30,6 +30,8 @@ export interface OpenApiToolSource { · 29 unmodified lines
30 readonly defaultReadApproval?: ApprovalMode | undefined;30 readonly defaultReadApproval?: ApprovalMode | undefined;
31 /** Default approval mode for write operations (POST/PUT/DELETE/PATCH). Defaults to "required". */31 /** Default approval mode for write operations (POST/PUT/DELETE/PATCH). Defaults to "required". */
32 readonly defaultWriteApproval?: ApprovalMode | undefined;32 readonly defaultWriteApproval?: ApprovalMode | undefined;
33 /** Override the base URL from the spec. Useful when the spec has no `servers` or you're targeting a different environment. */
34 readonly baseUrl?: string | undefined;
33}35}
34 36
35export type OpenApiAuth =37export type OpenApiAuth =
@@ -237,9 +239,9 @@ export async function generateOpenApiTools( · 201 unmodified lines
237 // Parse and dereference the spec239 // Parse and dereference the spec
238 const api = await SwaggerParser.dereference(source.spec as string) as Record<string, unknown>;240 const api = await SwaggerParser.dereference(source.spec as string) as Record<string, unknown>;
239 241
240 // Get base URL242 // Get base URL (config override > spec servers > empty)
241 const servers = (api["servers"] as Array<{ url: string }> | undefined) ?? [];243 const servers = (api["servers"] as Array<{ url: string }> | undefined) ?? [];
242 const baseUrl = servers[0]?.url ?? "";244 const baseUrl = source.baseUrl ?? servers[0]?.url ?? "";
243 245
244 const operations = parseOperations(api);246 const operations = parseOperations(api);
245 const authHeaders = buildAuthHeaders(source.auth);247 const authHeaders = buildAuthHeaders(source.auth);
@@ -302,11 +304,26 @@ export async function generateOpenApiTools( · 56 unmodified lines
302 required: inputRequired.length > 0 ? inputRequired : undefined,304 required: inputRequired.length > 0 ? inputRequired : undefined,
303 };305 };
304 306
305 const argsZod = jsonSchemaToZod(combinedInputSchema);307 let argsZod: z.ZodType;
306 const argsTypeString = jsonSchemaToTypeString(combinedInputSchema);308 let argsTypeString: string;
307 const returnsTypeString = op.responseSchema309 let returnsTypeString: string;
308 ? jsonSchemaToTypeString(op.responseSchema)310
309 : "any";311 try {
312 argsZod = jsonSchemaToZod(combinedInputSchema);
313 argsTypeString = jsonSchemaToTypeString(combinedInputSchema);
314 } catch {
315 // Recursive/circular schemas — fall back to z.any()
316 argsZod = z.any();
317 argsTypeString = "any";
318 }
319
320 try {
321 returnsTypeString = op.responseSchema
322 ? jsonSchemaToTypeString(op.responseSchema)
323 : "any";
324 } catch {
325 returnsTypeString = "any";
326 }
310 327
311 const description = op.summary || op.description || `${op.method.toUpperCase()} ${op.path}`;328 const description = op.summary || op.description || `${op.method.toUpperCase()} ${op.path}`;
312 329