fix: add apiKey to model (#60)

Author: CahidArda

src/agents/adapters.ts

@@ -9,20 +9,21 @@ import { tool } from "ai";
 import { AISDKTool, LangchainTool } from "./types";
 import { AGENT_NAME_HEADER } from "./constants";
 
-type ModelParams = Parameters<ReturnType<typeof createWorkflowOpenAI>>;
-type ModelSettingsWithBaseURL = ModelParams["1"] & { baseURL?: string };
-export type ModelParamsWithBaseURL = [ModelParams[0], ModelSettingsWithBaseURL?];
-
 /**
  * creates an AI SDK openai client with a custom
  * fetch implementation which uses context.call.
  *
  * @param context workflow context
  * @returns ai sdk openai
  */
-export const createWorkflowOpenAI = (context: WorkflowContext, baseURL?: string) => {
+export const createWorkflowOpenAI = (
+  context: WorkflowContext,
+  config?: { baseURL?: string; apiKey?: string }
+) => {
+  const { baseURL, apiKey } = config ?? {};
   return createOpenAI({
     baseURL,
+    apiKey,
     compatibility: "strict",
     fetch: async (input, init) => {
       try {

src/agents/index.ts

@@ -1,18 +1,16 @@
 import { WorkflowContext } from "../context";
-import { createWorkflowOpenAI, ModelParamsWithBaseURL, wrapTools } from "./adapters";
+import { createWorkflowOpenAI, wrapTools } from "./adapters";
 import { Agent } from "./agent";
 import { Task } from "./task";
 import {
   AgentParameters,
   AISDKTool,
+  CustomModelParams,
   LangchainTool,
   MultiAgentTaskParams,
   SingleAgentTaskParams,
 } from "./types";
 
-export { createWorkflowOpenAI } from "./adapters";
-export { Agent } from "./agent";
-
 /**
  * Workflow Agents API
  *
@@ -94,10 +92,10 @@ export class WorkflowAgents {
   /**
    * creates an openai model for agents
    */
-  public openai(...params: ModelParamsWithBaseURL) {
+  public openai(...params: CustomModelParams) {
     const [model, settings] = params;
-    const { baseURL, ...otherSettings } = settings ?? {};
-    const openai = createWorkflowOpenAI(this.context, baseURL);
+    const { baseURL, apiKey, ...otherSettings } = settings ?? {};
+    const openai = createWorkflowOpenAI(this.context, { baseURL, apiKey });
     return openai(model, otherSettings);
   }
 }

src/agents/task.test.ts

@@ -9,8 +9,8 @@ import { z } from "zod";
 import { DisabledWorkflowContext } from "../serve/authorization";
 
 export const getAgentsApi = ({ disabledContext }: { disabledContext: boolean }) => {
-  const token = getWorkflowRunId();
-  const workflowRunId = nanoid();
+  const workflowRunId = getWorkflowRunId();
+  const token = nanoid();
 
   let context: WorkflowContext;
   if (disabledContext) {
@@ -130,9 +130,12 @@ describe("tasks", () => {
 
   test("multi agent with baseURL", async () => {
     const { agentsApi, agent, token, workflowRunId } = getAgentsApi({ disabledContext: false });
+
+    const customURL = "https://api.deepseek.com/v1";
+    const customApiKey = nanoid();
     const task = agentsApi.task({
       agents: [agent],
-      model: agentsApi.openai("gpt-3.5-turbo", { baseURL: "https://api.deepseek.com/v1" }),
+      model: agentsApi.openai("gpt-3.5-turbo", { baseURL: customURL, apiKey: customApiKey }),
       maxSteps: 2,
       prompt: "hello world!",
     });
@@ -173,7 +176,7 @@ describe("tasks", () => {
               "upstash-callback-workflow-url": "https://requestcatcher.com/api",
               "upstash-failure-callback-retries": "3",
               "upstash-feature-set": "WF_NoDelete,InitialBody",
-              "upstash-forward-authorization": `Bearer ${openaiToken}`,
+              "upstash-forward-authorization": `Bearer ${customApiKey}`,
               "upstash-forward-content-type": "application/json",
               "upstash-forward-upstash-agent-name": "Manager LLM",
               "upstash-method": "POST",

src/agents/types.ts

@@ -1,5 +1,6 @@
 import type { CoreTool, generateText } from "ai";
 import { Agent } from "./agent";
+import { createWorkflowOpenAI } from "./adapters";
 
 export type AISDKTool = CoreTool;
 export type LangchainTool = {
@@ -88,3 +89,7 @@ export type ManagerAgentParameters = {
   model: Model;
 } & Pick<Partial<AgentParameters>, "name" | "background"> &
   Pick<AgentParameters, "maxSteps">;
+
+type ModelParams = Parameters<ReturnType<typeof createWorkflowOpenAI>>;
+type CustomModelSettings = ModelParams["1"] & { baseURL?: string; apiKey?: string };
+export type CustomModelParams = [ModelParams[0], CustomModelSettings?];