fix weird type thing

patrykkopycinski · Jul 22, 2024 · 50d73f1 · 50d73f1
1 parent 4db15dc
commit 50d73f1
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 24 deletions.
diff --git a/...ck/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/graph.ts b/...ck/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/graph.ts
@@ -41,7 +41,6 @@ interface GetDefaultAssistantGraphParams {
   agentRunnable: AgentRunnableSequence;
   dataClients?: AssistantDataClients;
   conversationId?: string;
-  getLlmInstance: () => BaseChatModel;
   llm: BaseChatModel;
   logger: Logger;
   tools: StructuredTool[];
@@ -61,7 +60,6 @@ export const getDefaultAssistantGraph = ({
   agentRunnable,
   conversationId,
   dataClients,
-  getLlmInstance,
   llm,
   logger,
   responseLanguage,
@@ -154,7 +152,7 @@ export const getDefaultAssistantGraph = ({
     const respondNode = (state: AgentState) =>
       respond({
         ...nodeParams,
-        llm: getLlmInstance(),
+        llm,
         state,
       });
     const shouldContinueEdge = (state: AgentState) => shouldContinue({ ...nodeParams, state });

diff --git a/...ck/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/index.ts b/...ck/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/index.ts
@@ -53,26 +53,23 @@ export const callAssistantGraph: AgentExecutor<true | false> = async ({
   const logger = parentLogger.get('defaultAssistantGraph');
   const isOpenAI = llmType === 'openai';
   const llmClass = getLlmClass(llmType, bedrockChatEnabled);
-  const getLlmInstance = () =>
-    new llmClass({
-      actionsClient,
-      connectorId,
-      llmType,
-      logger,
-      // possible client model override,
-      // let this be undefined otherwise so the connector handles the model
-      model: request.body.model,
-      // ensure this is defined because we default to it in the language_models
-      // This is where the LangSmith logs (Metadata > Invocation Params) are set
-      temperature: getDefaultArguments(llmType).temperature,
-      signal: abortSignal,
-      streaming: isStream,
-      // prevents the agent from retrying on failure
-      // failure could be due to bad connector, we should deliver that result to the client asap
-      maxRetries: 0,
-    });
-
-  const llm = getLlmInstance();
+  const llm = new llmClass({
+    actionsClient,
+    connectorId,
+    llmType,
+    logger,
+    // possible client model override,
+    // let this be undefined otherwise so the connector handles the model
+    model: request.body.model,
+    // ensure this is defined because we default to it in the language_models
+    // This is where the LangSmith logs (Metadata > Invocation Params) are set
+    temperature: getDefaultArguments(llmType).temperature,
+    signal: abortSignal,
+    streaming: isStream,
+    // prevents the agent from retrying on failure
+    // failure could be due to bad connector, we should deliver that result to the client asap
+    maxRetries: 0,
+  });
 
   const anonymizationFieldsRes =
     await dataClients?.anonymizationFieldsDataClient?.findDocuments<EsAnonymizationFieldsSchema>({
@@ -152,7 +149,6 @@ export const callAssistantGraph: AgentExecutor<true | false> = async ({
     conversationId,
     dataClients,
     llm,
-    getLlmInstance,
     logger,
     tools,
     responseLanguage,