Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Use @inngest/ai and optional step tooling #43

Open
wants to merge 6 commits into
base: main
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 2 additions & 1 deletion package.json
Original file line number Diff line number Diff line change
Expand Up @@ -46,10 +46,11 @@
},
"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",
"inngest": "^3.29.0",
"inngest": "^3.30.0",
"openai-zod-to-json-schema": "^1.0.3",
"zod": "^3.23.8"
},
Expand Down
44 changes: 35 additions & 9 deletions pnpm-lock.yaml

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

2 changes: 1 addition & 1 deletion src/adapters/anthropic.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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";
Expand Down
2 changes: 1 addition & 1 deletion 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";
Expand Down
2 changes: 1 addition & 1 deletion src/adapters/openai.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down
42 changes: 20 additions & 22 deletions 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 {
Expand All @@ -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.
Expand Down Expand Up @@ -284,7 +282,7 @@ export class Agent {
const result = await found.handler(tool.input, {
agent: this,
network,
step: await getStepTools(),
// step: await getStepTools(),
});

// TODO: handle error and send them back to the LLM
Expand Down Expand Up @@ -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;
},
});
Expand Down
46 changes: 41 additions & 5 deletions 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";
Expand Down Expand Up @@ -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>;
jpwilliams marked this conversation as resolved.
Show resolved Hide resolved

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 };
}
Expand Down
2 changes: 1 addition & 1 deletion 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 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,
Expand Down
21 changes: 1 addition & 20 deletions 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 } from "./util";

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

/**
* 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>;
14 changes: 6 additions & 8 deletions 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>;
Expand Down Expand Up @@ -29,16 +29,14 @@ 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;
};

/**
Expand Down
Loading