feat: allow for model to be configurable

.env.example

@@ -8,3 +8,5 @@ LANGSMITH_TEST_TRACKING="true"
 
 # LLMs
 ANTHROPIC_API_KEY=""
+# Only used for generating eval data
+# OPENAI_API_KEY=""

src/llmanager/nodes/final-answer.ts

@@ -1,7 +1,8 @@
 import { z } from "zod";
 import { AgentState, AgentUpdate } from "../types.js";
-import { ChatAnthropic } from "@langchain/anthropic";
 import { findQueryStringOrThrow } from "../../utils/query.js";
+import { LangGraphRunnableConfig } from "@langchain/langgraph";
+import { loadModelFromConfig } from "../../utils/model.js";
 
 const FINAL_ANSWER_PROMPT = `You're a highly advanced AI manager, tasked with approving or rejecting one of your employees requests.
 
@@ -26,7 +27,10 @@ Here is the users request:
 Ensure your answer is accurate, and accounts for all of the context provided above.
 `;
 
-export async function finalAnswer(state: AgentState): Promise<AgentUpdate> {
+export async function finalAnswer(
+  state: AgentState,
+  config: LangGraphRunnableConfig,
+): Promise<AgentUpdate> {
   const query = findQueryStringOrThrow(state.messages);
 
   const finalAnswerSchema = z.object({
@@ -49,10 +53,10 @@ export async function finalAnswer(state: AgentState): Promise<AgentUpdate> {
     .replace("{REASONING}", state.generatedReasoning)
     .replace("{REQUEST}", query);
 
-  const model = new ChatAnthropic({
-    model: "claude-3-7-sonnet-latest",
+  const model = await loadModelFromConfig(config, {
     temperature: 0,
-  }).bindTools([
+  });
+  const modelWithTools = model.bindTools([
     {
       name: "finalAnswer",
       schema: finalAnswerSchema,
@@ -61,7 +65,7 @@ export async function finalAnswer(state: AgentState): Promise<AgentUpdate> {
     },
   ]);
 
-  const response = await model.invoke([
+  const response = await modelWithTools.invoke([
     {
       role: "user",
       content: formattedPrompt,

src/reasoning/nodes/initial-reasoning.ts

@@ -1,11 +1,11 @@
 import { ReasoningState, ReasoningUpdate } from "../types.js";
 import { LangGraphRunnableConfig } from "@langchain/langgraph";
-import { ChatAnthropic } from "@langchain/anthropic";
 import {
   buildContext,
   formatContextPrompt,
 } from "../../utils/build-context.js";
 import { findQueryStringOrThrow } from "../../utils/query.js";
+import { loadModelFromConfig } from "../../utils/model.js";
 
 const INITIAL_REASONING_PROMPT = `You're an AI manager tasked with analyzing and reasoning about a request one of your employees has made.
 Your task is to analyze the request from one of your employees, and reason about whether it should be approved or rejected.
@@ -48,8 +48,7 @@ export async function initialReasoning(
     }),
   );
 
-  const model = new ChatAnthropic({
-    model: "claude-3-7-sonnet-latest",
+  const model = await loadModelFromConfig(config, {
     temperature: 0,
   });
 

src/reflection/nodes/explanation-reflection.ts

@@ -1,8 +1,8 @@
 import { LangGraphRunnableConfig } from "@langchain/langgraph";
 import { getReflections, putReflections } from "../../stores/reflection.js";
 import { ReflectionState, ReflectionUpdate } from "../types.js";
-import { ChatAnthropic } from "@langchain/anthropic";
 import { z } from "zod";
+import { loadModelFromConfig } from "../../utils/model.js";
 
 const EXPLANATION_REFLECTION_PROMPT = `You're an advanced AI assistant tasked with generating reflections on an incorrect explanation for a user request, even though the final answer was correct.
 A human manually reviewed the explanation and determined it was incorrect, despite the answer being correct.
@@ -82,14 +82,15 @@ export async function explanationReflection(
     reflections: reflections,
   });
 
-  const model = new ChatAnthropic({
-    model: "claude-3-7-sonnet-latest",
+  const model = await loadModelFromConfig(config, {
     temperature: 0,
     thinking: {
       type: "enabled",
       budget_tokens: 3072,
     },
-  }).bindTools(
+  });
+
+  const modelWithTools = model.bindTools(
     [
       {
         name: "generate_reflections",
@@ -109,7 +110,9 @@ export async function explanationReflection(
     },
   );
 
-  const response = await model.invoke([{ role: "human", content: prompt }]);
+  const response = await modelWithTools.invoke([
+    { role: "human", content: prompt },
+  ]);
 
   const newReflections = response.tool_calls?.[0]?.args?.reflections as
     | string[]

src/reflection/nodes/full-reflection.ts

@@ -3,6 +3,7 @@ import { getReflections, putReflections } from "../../stores/reflection.js";
 import { ReflectionState, ReflectionUpdate } from "../types.js";
 import { ChatAnthropic } from "@langchain/anthropic";
 import { z } from "zod";
+import { loadModelFromConfig } from "../../utils/model.js";
 
 const FULL_REFLECTION_PROMPT = `You're an advanced AI assistant tasked with generating reflections on a an incorrect answer to a user request.
 A human manually reviewed the answer, and the explanation and determined they were both incorrect.
@@ -86,14 +87,15 @@ export async function fullReflection(
     reflections: reflections,
   });
 
-  const model = new ChatAnthropic({
-    model: "claude-3-7-sonnet-latest",
+  const model = await loadModelFromConfig(config, {
     temperature: 0,
     thinking: {
       type: "enabled",
       budget_tokens: 3072,
     },
-  }).bindTools(
+  });
+
+  const modelWithTools = model.bindTools(
     [
       {
         name: "generate_reflections",
@@ -113,7 +115,7 @@ export async function fullReflection(
     },
   );
 
-  const response = await model.invoke([
+  const response = await modelWithTools.invoke([
     {
       role: "human",
       content: prompt,

src/utils/model.ts

@@ -0,0 +1,41 @@
+import {
+  BaseChatModel,
+  BaseChatModelCallOptions,
+  BindToolsInput,
+} from "@langchain/core/language_models/chat_models";
+import { LangGraphRunnableConfig } from "@langchain/langgraph";
+import { initChatModel } from "langchain/chat_models/universal";
+import { Runnable, RunnableConfig } from "@langchain/core/runnables";
+import { BaseLanguageModelInput } from "@langchain/core/language_models/base";
+import { AIMessageChunk } from "@langchain/core/messages";
+
+type BaseChatModelWithBindTools = BaseChatModel & {
+  bindTools: (
+    tools: BindToolsInput[],
+    kwargs?: Partial<RunnableConfig>,
+  ) => Runnable<
+    BaseLanguageModelInput,
+    AIMessageChunk,
+    BaseChatModelCallOptions
+  >;
+};
+/**
+ * Loads a chat model from the configuration object. Defaults to
+ * `anthropic/claude-3-7-sonnet-latest`.
+ *
+ * @param config The configuration object containing the model ID.
+ * @returns The loaded chat model.
+ */
+export async function loadModelFromConfig(
+  config: LangGraphRunnableConfig,
+  modelConfig?: Record<string, unknown>,
+): Promise<BaseChatModelWithBindTools> {
+  const modelId =
+    config.configurable?.modelId ?? "anthropic/claude-3-7-sonnet-latest";
+
+  const model = await initChatModel(modelId, modelConfig);
+  if (!model.bindTools) {
+    throw new Error("Model does not support binding tools");
+  }
+  return model;
+}