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
1111 import { defineTool, mergeToolTrees, type ToolTree } from "@openassistant/core/tools";
1212 import { generateMcpTools } from "@openassistant/tool-gen/mcp";
1313 import type { McpToolSource } from "@openassistant/tool-gen/mcp";
14+import { generateOpenApiTools } from "@openassistant/tool-gen/openapi";
15+import type { OpenApiToolSource } from "@openassistant/tool-gen/openapi";
1416 import { z } from "zod";
1517
1618 // ---------------------------------------------------------------------------
@@ -54,6 +56,28 @@ const mcpSources: McpToolSource[] = [ · 37 unmodified lines
5456 },
5557 ];
5658
59+const POSTHOG_API_KEY = process.env["POSTHOG_PERSONAL_API_KEY"];
60+const POSTHOG_PROJECT_ID = process.env["POSTHOG_PROJECT_ID"];
61+const GITHUB_TOKEN = process.env["OPENASSISTANT_GITHUB_TOKEN"];
62+
63+const 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+
5781 // ---------------------------------------------------------------------------
5882 // Hand-written tools (approval testing)
5983 // ---------------------------------------------------------------------------
@@ -111,6 +135,7 @@ const builtinTools: ToolTree = { · 51 unmodified lines
111135 async function loadTools(): Promise<ToolTree> {
112136 const trees: ToolTree[] = [builtinTools];
113137
138+ // Load MCP tools
114139 for (const source of mcpSources) {
115140 try {
116141 console.log(`Loading MCP tools from ${source.name} (${source.url})...`);
@@ -122,6 +147,22 @@ async function loadTools(): Promise<ToolTree> { · 5 unmodified lines
122147 }
123148 }
124149
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+
125166 return mergeToolTrees(...trees);
126167 }
127168
packages/core/package.json+2−1
@@ -8,7 +8,8 @@ · 7 unmodified lines
88 "./runner": "./src/runner.ts",
99 "./typechecker": "./src/typechecker.ts",
1010 "./agent": "./src/agent.ts",
11− "./events": "./src/events.ts"
11+ "./events": "./src/events.ts",
12+ "./discovery": "./src/discovery.ts"
1213 },
1314 "scripts": {
1415 "test": "bun test",
packages/core/src/agent.ts+127−7
@@ -9,12 +9,17 @@ · 8 unmodified lines
99 * 6. Execute in sandbox
1010 * 7. Collect receipts, feed back to Claude
1111 * 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.
1216 */
1317
1418 import type { ToolTree, ApprovalDecision, ApprovalRequest, ToolCallReceipt } from "./tools.js";
1519 import type { RunResult } from "./runner.js";
1620 import { createRunner } from "./runner.js";
1721 import { generateToolDeclarations, generatePromptGuidance, typecheckCode } from "./typechecker.js";
22+import { countTools, createDiscoverTool } from "./discovery.js";
1823 import type { TaskEvent } from "./events.js";
1924
2025 // ---------------------------------------------------------------------------
@@ -67,6 +72,13 @@ export interface AgentOptions { · 46 unmodified lines
6772 readonly maxCodeRuns?: number | undefined;
6873 /** Execution timeout per code run in ms. Defaults to 30_000. */
6974 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;
7082 }
7183
7284 export interface AgentResult {
@@ -81,12 +93,117 @@ export interface CodeRun { · 8 unmodified lines
8193 readonly result: RunResult;
8294 }
8395
96+// ---------------------------------------------------------------------------
97+// Discovery mode
98+// ---------------------------------------------------------------------------
99+
100+const 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+ */
110+function 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+
84180 // ---------------------------------------------------------------------------
85181 // System prompt
86182 // ---------------------------------------------------------------------------
87183
88−function buildSystemPrompt(tools: ToolTree): string {
89− const guidance = generatePromptGuidance(tools);
184+function buildSystemPrompt(guidance: string, discoveryMode: boolean): string {
185+ const discoveryInstructions = discoveryMode
186+ ? `
187+## Tool Discovery
188+
189+Some tool namespaces are too large to list here. Use \`tools.discover({ query: "..." })\` to search for tools by keyword.
190+The discover tool returns an array of matching tools with their paths, descriptions, and TypeScript signatures.
191+
192+**Workflow for large APIs:**
193+1. Call \`tools.discover({ query: "relevant keywords" })\` to find relevant tools
194+2. Read the returned signatures to understand the input/output types
195+3. Call the discovered tools using the path from the results, e.g. \`tools.posthog.feature_flags.feature_flags_list({ project_id: "123" })\`
196+4. You can call discover multiple times with different queries
197+
198+Example:
199+\`\`\`ts
200+const found = await tools.discover({ query: "feature flags" });
201+// found.results = [{ path: "posthog.feature_flags.feature_flags_list", signature: "...", ... }]
202+const flags = await tools.posthog.feature_flags.feature_flags_list({ project_id: "123" });
203+\`\`\`
204+`
205+ : "";
206+
90207 return `You are an AI assistant that executes tasks by generating TypeScript code.
91208
92209 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.
@@ -94,7 +211,7 @@ You have access to a set of tools via the \`tools\` object. When the user asks y · 1 unmodified line
94211 ## Available Tools
95212
96213 ${guidance}
97−
214+${discoveryInstructions}
98215 ## Instructions
99216
100217 - Use the \`run_code\` tool to execute TypeScript code
@@ -121,12 +238,14 @@ export function createAgent(options: AgentOptions): { · 20 unmodified lines
121238 maxTypecheckRetries = 3,
122239 maxCodeRuns = 10,
123240 timeoutMs = 30_000,
241+ discoveryThreshold = DISCOVERY_THRESHOLD,
124242 } = options;
125243
126− const toolDeclarations = generateToolDeclarations(tools);
244+ const { sandboxTools, promptGuidance, toolDeclarations, discoveryMode } =
245+ prepareToolsForAgent(tools, discoveryThreshold);
127246
128247 const runner = createRunner({
129− tools,
248+ tools: sandboxTools,
130249 requestApproval,
131250 timeoutMs,
132251 });
@@ -137,7 +256,7 @@ export function createAgent(options: AgentOptions): { · 4 unmodified lines
137256
138257 return {
139258 async run(prompt: string): Promise<AgentResult> {
140− const systemPrompt = buildSystemPrompt(tools);
259+ const systemPrompt = buildSystemPrompt(promptGuidance, discoveryMode);
141260 const runs: CodeRun[] = [];
142261 const allReceipts: ToolCallReceipt[] = [];
143262
@@ -194,6 +313,7 @@ export function createAgent(options: AgentOptions): { · 50 unmodified lines
194313 typecheckOk = true;
195314 break;
196315 }
316+
197317 typecheckErrors = check.errors;
198318
199319 if (attempt < maxTypecheckRetries) {
@@ -270,7 +390,7 @@ export function createAgent(options: AgentOptions): { · 70 unmodified lines
270390 }
271391
272392 // 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.";
274394 emit({ type: "agent_message", text });
275395 emit({ type: "completed", receipts: allReceipts });
276396 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+
14+import { defineTool, isToolDefinition, type ToolDefinition, type ToolTree } from "./tools.js";
15+import { zodToTypeString } from "./typechecker.js";
16+import { z } from "zod";
17+
18+// ---------------------------------------------------------------------------
19+// Index entry
20+// ---------------------------------------------------------------------------
21+
22+export 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+
39+export 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+ */
74+function 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+
106+export 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+ */
119+export 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+ */
159+export 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
5151 type PiAiModelOptions,
5252 } from "./pi-ai-adapter.js";
5353
54+// Discovery
55+export {
56+ buildToolIndex,
57+ createDiscoverTool,
58+ countTools,
59+ type ToolIndexEntry,
60+ type DiscoverToolOptions,
61+} from "./discovery.js";
62+
5463 // Events
5564 export type {
5665 TaskEvent,
packages/tool-gen/src/openapi.ts+24−7
@@ -30,6 +30,8 @@ export interface OpenApiToolSource { · 29 unmodified lines
3030 readonly defaultReadApproval?: ApprovalMode | undefined;
3131 /** Default approval mode for write operations (POST/PUT/DELETE/PATCH). Defaults to "required". */
3232 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;
3335 }
3436
3537 export type OpenApiAuth =
@@ -237,9 +239,9 @@ export async function generateOpenApiTools( · 201 unmodified lines
237239 // Parse and dereference the spec
238240 const api = await SwaggerParser.dereference(source.spec as string) as Record<string, unknown>;
239241
240− // Get base URL
242+ // Get base URL (config override > spec servers > empty)
241243 const servers = (api["servers"] as Array<{ url: string }> | undefined) ?? [];
242− const baseUrl = servers[0]?.url ?? "";
244+ const baseUrl = source.baseUrl ?? servers[0]?.url ?? "";
243245
244246 const operations = parseOperations(api);
245247 const authHeaders = buildAuthHeaders(source.auth);
@@ -302,11 +304,26 @@ export async function generateOpenApiTools( · 56 unmodified lines
302304 required: inputRequired.length > 0 ? inputRequired : undefined,
303305 };
304306
305− const argsZod = jsonSchemaToZod(combinedInputSchema);
306− const argsTypeString = jsonSchemaToTypeString(combinedInputSchema);
307− const returnsTypeString = op.responseSchema
308− ? jsonSchemaToTypeString(op.responseSchema)
309− : "any";
307+ let argsZod: z.ZodType;
308+ let argsTypeString: string;
309+ let returnsTypeString: string;
310+
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+ }
310327
311328 const description = op.summary || op.description || `${op.method.toUpperCase()} ${op.path}`;
312329