Skip to content

Commit

Permalink
Use @inngest/ai and optional step tooling (#43)
Browse files Browse the repository at this point in the history
* Use `@inngest/ai` and optional step tooling

* Make sure to get JSON response of `infer()` call

* Use set `@inngest/ai` versions

`types` type needs loosening.

* Await `fetch().json()`

* Make sure we call `Model#onCall()` if we're not using `step.ai.infer()`

* Remove comment

* Create seven-beers-press.md

* Provide optional access to step tooling in agent tool handlers
  • Loading branch information
jpwilliams authored Jan 20, 2025

Verified

This commit was created on GitHub.com and signed with GitHub’s verified signature.
1 parent dd54d17 commit ae56867
Showing 13 changed files with 85 additions and 60 deletions.
5 changes: 5 additions & 0 deletions .changeset/seven-beers-press.md
Original file line number Diff line number Diff line change
@@ -0,0 +1,5 @@
---
"@inngest/agent-kit": patch
---

Use `@inngest/ai` and only optionally use step tooling
2 changes: 1 addition & 1 deletion examples/swebench/agents/planner.ts
Original file line number Diff line number Diff line change
@@ -32,7 +32,7 @@ export const planningAgent = createAgent({

handler: async (plan, opts) => {
// Store this in the function state for introspection in tracing.
await opts.step.run("plan created", () => plan);
await opts.step?.run("plan created", () => plan);
opts.network?.state.kv.set("plan", plan);
},
}),
1 change: 1 addition & 0 deletions packages/agent-kit/package.json
Original file line number Diff line number Diff line change
@@ -46,6 +46,7 @@
},
"dependencies": {
"@dmitryrechkin/json-schema-to-zod": "^1.0.0",
"@inngest/ai": "^0.0.0",
"@modelcontextprotocol/sdk": "^1.1.1",
"eventsource": "^3.0.2",
"express": "^4.21.1",
2 changes: 1 addition & 1 deletion packages/agent-kit/src/adapters/anthropic.ts
Original file line number Diff line number Diff line change
@@ -7,7 +7,7 @@ import {
type AiAdapter,
type Anthropic,
type AnthropicAiAdapter,
} from "inngest";
} from "@inngest/ai";
import { zodToJsonSchema } from "openai-zod-to-json-schema";
import { z } from "zod";
import { type AgenticModel } from "../model";
2 changes: 1 addition & 1 deletion packages/agent-kit/src/adapters/index.ts
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
import { type AiAdapter, type AiAdapters } from "inngest";
import { type AiAdapter, type AiAdapters } from "@inngest/ai";
import { type AgenticModel } from "../model";
import * as anthropic from "./anthropic";
import * as openai from "./openai";
2 changes: 1 addition & 1 deletion packages/agent-kit/src/adapters/openai.ts
Original file line number Diff line number Diff line change
@@ -4,7 +4,7 @@
* @module
*/

import { type AiAdapter, type OpenAi } from "inngest";
import { type AiAdapter, type OpenAi } from "@inngest/ai";
import { zodToJsonSchema } from "openai-zod-to-json-schema";
import { type AgenticModel } from "../model";
import {
40 changes: 19 additions & 21 deletions packages/agent-kit/src/agent.ts
Original file line number Diff line number Diff line change
@@ -1,5 +1,15 @@
import { type AiAdapter } from "inngest";
import {
JSONSchemaToZod,
type JSONSchema,
} from "@dmitryrechkin/json-schema-to-zod";
import { type AiAdapter } from "@inngest/ai";
import { Client as MCPClient } from "@modelcontextprotocol/sdk/client/index.js";
import { SSEClientTransport } from "@modelcontextprotocol/sdk/client/sse.js";
import { WebSocketClientTransport } from "@modelcontextprotocol/sdk/client/websocket.js";
import { type Transport } from "@modelcontextprotocol/sdk/shared/transport";
import { ListToolsResultSchema } from "@modelcontextprotocol/sdk/types.js";
import { EventSource } from "eventsource";
import type { ZodType } from "zod";
import { createAgenticModelFromAiAdapter, type AgenticModel } from "./model";
import { NetworkRun } from "./networkRun";
import {
@@ -8,20 +18,8 @@ import {
type Message,
type ToolResultMessage,
} from "./state";
import { type Tool, type MCP } from "./types";
import { type MCP, type Tool } from "./types";
import { getStepTools, type AnyZodType, type MaybePromise } from "./util";
// MCP
import { Client as MCPClient } from "@modelcontextprotocol/sdk/client/index.js";
import { SSEClientTransport } from "@modelcontextprotocol/sdk/client/sse.js";
import { WebSocketClientTransport } from "@modelcontextprotocol/sdk/client/websocket.js";
import { ListToolsResultSchema } from "@modelcontextprotocol/sdk/types.js";
import {
type JSONSchema,
JSONSchemaToZod,
} from "@dmitryrechkin/json-schema-to-zod";
import type { ZodType } from "zod";
import { EventSource } from "eventsource";

/**
* createTool is a helper that properly types the input argument for a handler
* based off of the Zod parameter types.
@@ -392,16 +390,16 @@ export class Agent {
server,
tool: t,
},
handler: async (
input: { [x: string]: unknown } | undefined,
opts
) => {
const result = await opts.step.run(name, async () => {
return await client.callTool({
handler: async (input: { [x: string]: unknown } | undefined) => {
const fn = () =>
client.callTool({
name: t.name,
arguments: input,
});
});

const step = await getStepTools();
const result = await (step?.run(name, fn) ?? fn());

return result.content;
},
});
46 changes: 41 additions & 5 deletions packages/agent-kit/src/model.ts
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
import { type AiAdapter } from "inngest";
import { type AiAdapter } from "@inngest/ai";
import { adapters } from "./adapters";
import { type Message } from "./state";
import { type Tool } from "./types";
@@ -41,12 +41,48 @@ export class AgenticModel<TAiAdapter extends AiAdapter.Any> {
tools: Tool.Any[],
tool_choice: Tool.Choice
): Promise<AgenticModel.InferenceResponse> {
const body = this.requestParser(this.#model, input, tools, tool_choice);
let result: AiAdapter.Input<TAiAdapter>;

const step = await getStepTools();

const result = (await step.ai.infer(stepID, {
model: this.#model,
body: this.requestParser(this.#model, input, tools, tool_choice),
})) as AiAdapter.Input<TAiAdapter>;
if (step) {
result = (await step.ai.infer(stepID, {
model: this.#model,
body,
})) as AiAdapter.Input<TAiAdapter>;
} else {
// Allow the model to mutate options and body for this call
const modelCopy = { ...this.#model };
this.#model.onCall?.(modelCopy, body);

const url = new URL(modelCopy.url || "");

const headers: Record<string, string> = {
"Content-Type": "application/json",
};

// Make sure we handle every known format in `@inngest/ai`.
const formatHandlers: Record<AiAdapter.Format, () => void> = {
"openai-chat": () => {
headers["Authorization"] = `Bearer ${modelCopy.authKey}`;
},
anthropic: () => {
headers["x-api-key"] = modelCopy.authKey;
},
};

formatHandlers[modelCopy.format as AiAdapter.Format]();

// eslint-disable-next-line @typescript-eslint/no-unsafe-assignment
result = await (
await fetch(url, {
method: "POST",
headers,
body,
})
).json();
}

return { output: this.responseParser(result), raw: result };
}
2 changes: 1 addition & 1 deletion packages/agent-kit/src/models.ts
Original file line number Diff line number Diff line change
@@ -1 +1 @@
export { openai, anthropic, gemini } from "inngest";
export { anthropic, gemini, openai } from "@inngest/ai";
2 changes: 1 addition & 1 deletion packages/agent-kit/src/network.ts
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
import { type AiAdapter } from "inngest";
import { type AiAdapter } from "@inngest/ai";
import { z } from "zod";
import {
createRoutingAgent,
22 changes: 2 additions & 20 deletions packages/agent-kit/src/types.ts
Original file line number Diff line number Diff line change
@@ -1,13 +1,7 @@
import { type GetStepTools, type Inngest } from "inngest";
import { type output as ZodOutput } from "zod";
import { type Agent } from "./agent";
import { type NetworkRun } from "./networkRun";
import {
type AnyZodType,
type GenericizeFunctionsInObject,
type MaybePromise,
type SimplifyDeep,
} from "./util";
import { type AnyZodType, type MaybePromise, type StepTools } from "./util";

export type Tool<T extends AnyZodType> = {
name: string;
@@ -69,17 +63,5 @@ export namespace MCP {
export type ToolHandlerArgs = {
agent: Agent;
network?: NetworkRun;
step: GetStepTools<Inngest.Any>;
step?: StepTools;
};

/**
* Represents step tooling from an Inngest client, purposefully genericized to
* allow for more flexible usage.
*
* Prefer use of `GetStepTools` in most cases, especially when you have access
* to a client.
*/
export type AnyStepTools = SimplifyDeep<
GenericizeFunctionsInObject<GetStepTools<Inngest.Any>>
> &
Record<string, unknown>;
16 changes: 8 additions & 8 deletions packages/agent-kit/src/util.ts
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
import { getAsyncCtx } from "inngest/experimental";
import { type AsyncContext, getAsyncCtx } from "inngest/experimental";
import { type ZodType } from "zod";

export type MaybePromise<T> = T | Promise<T>;
@@ -29,18 +29,18 @@ export const stringifyError = (e: unknown): string => {
};

/**
* Attempts to retrieve the step tools from the async context. If the context is
* not found, an error is thrown.
* Attempts to retrieve the step tools from the async context.
*/
export const getStepTools = async () => {
export const getStepTools = async (): Promise<
AsyncContext["ctx"]["step"] | undefined
> => {
const asyncCtx = await getAsyncCtx();
if (!asyncCtx) {
throw new Error("Could not find Inngest step tooling in async context");
}

return asyncCtx.ctx.step;
return asyncCtx?.ctx.step;
};

export type StepTools = Awaited<ReturnType<typeof getStepTools>>;

/**
* Given an object `T`, return a new object where all keys with function types
* as values are genericized. If the value is an object, recursively apply this
3 changes: 3 additions & 0 deletions pnpm-lock.yaml

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

0 comments on commit ae56867

Please sign in to comment.