Adds managed Motorhead (#1561)

* add managed motorhead opt

* add example

* fix typings

* prettier

* update linting

* headers init

* fix specs

* Adds backcompat shim and docs

* Adds managed motorhead integration test

* Formatting

---------

Co-authored-by: James O'Dwyer <james@getmetal.io>

docs/docs/modules/indexes/retrievers/self_query/index.mdx

@@ -19,4 +19,4 @@ All Self Query retrievers require `peggy` as a peer dependency:
 
 ```bash npm2yarn
 npm install -S peggy
-```
+```

docs/docs/modules/memory/examples/motorhead_memory.mdx

@@ -0,0 +1,19 @@
+---
+hide_table_of_contents: true
+---
+
+import CodeBlock from "@theme/CodeBlock";
+
+# Motörhead Memory
+
+[Motörhead](https://github.com/getmetal/motorhead) is a memory server implemented in Rust. It automatically handles incremental summarization in the background and allows for stateless applications.
+
+## Setup
+
+See instructions at [Motörhead](https://github.com/getmetal/motorhead) for running the server locally, or https://getmetal.io to get API keys for the hosted version.
+
+## Usage
+
+import Example from "@examples/memory/motorhead.ts";
+
+<CodeBlock language="typescript">{Example}</CodeBlock>

examples/src/memory/motorhead.ts

@@ -0,0 +1,46 @@
+import { MotorheadMemory } from "langchain/memory";
+import { ChatOpenAI } from "langchain/chat_models/openai";
+import { ConversationChain } from "langchain/chains";
+
+// Managed Example (visit https://getmetal.io to get your keys)
+// const managedMemory = new MotorheadMemory({
+//   memoryKey: "chat_history",
+//   sessionId: "test",
+//   apiKey: "MY_API_KEY",
+//   clientId: "MY_CLIENT_ID",
+// });
+
+// Self Hosted Example
+const memory = new MotorheadMemory({
+  memoryKey: "chat_history",
+  sessionId: "test",
+  url: "localhost:8080", // Required for self hosted
+});
+
+const model = new ChatOpenAI({
+  modelName: "gpt-3.5-turbo",
+  temperature: 0,
+});
+
+const chain = new ConversationChain({ llm: model, memory });
+
+const res1 = await chain.call({ input: "Hi! I'm Jim." });
+console.log({ res1 });
+/*
+{
+  res1: {
+    text: "Hello Jim! It's nice to meet you. My name is AI. How may I assist you today?"
+  }
+}
+*/
+
+const res2 = await chain.call({ input: "What did I just say my name was?" });
+console.log({ res2 });
+
+/*
+{
+  res1: {
+    text: "You said your name was Jim."
+  }
+}
+*/

langchain/src/memory/motorhead_memory.ts

@@ -19,13 +19,19 @@ export interface MotorheadMemoryMessage {
 export type MotorheadMemoryInput = BaseChatMemoryInput &
   AsyncCallerParams & {
     sessionId: string;
+    /** @deprecated Use "url" instead. */
     motorheadURL?: string;
+    url?: string;
     memoryKey?: string;
     timeout?: number;
+    apiKey?: string;
+    clientId?: string;
   };
 
+const MANAGED_URL = "https://api.getmetal.io/v1/motorhead";
+
 export class MotorheadMemory extends BaseChatMemory {
-  motorheadURL = "localhost:8080";
+  url = MANAGED_URL;
 
   timeout = 3000;
 
@@ -37,40 +43,68 @@ export class MotorheadMemory extends BaseChatMemory {
 
   caller: AsyncCaller;
 
+  // Managed Params
+  apiKey?: string;
+
+  clientId?: string;
+
   constructor(fields: MotorheadMemoryInput) {
     const {
       sessionId,
+      url,
       motorheadURL,
       memoryKey,
       timeout,
       returnMessages,
       inputKey,
       outputKey,
       chatHistory,
+      apiKey,
+      clientId,
       ...rest
     } = fields;
     super({ returnMessages, inputKey, outputKey, chatHistory });
 
     this.caller = new AsyncCaller(rest);
     this.sessionId = sessionId;
-    this.motorheadURL = motorheadURL ?? this.motorheadURL;
+    this.url = url ?? motorheadURL ?? this.url;
     this.memoryKey = memoryKey ?? this.memoryKey;
     this.timeout = timeout ?? this.timeout;
+    this.apiKey = apiKey;
+    this.clientId = clientId;
   }
 
   get memoryKeys() {
     return [this.memoryKey];
   }
 
+  _getHeaders(): HeadersInit {
+    const isManaged = this.url === MANAGED_URL;
+
+    const headers: HeadersInit = {
+      "Content-Type": "application/json",
+    };
+
+    if (isManaged && !(this.apiKey && this.clientId)) {
+      throw new Error(
+        "apiKey and clientId are required for managed motorhead. Visit https://getmetal.io to get your keys."
+      );
+    }
+
+    if (isManaged && this.apiKey && this.clientId) {
+      headers["x-metal-api-key"] = this.apiKey;
+      headers["x-metal-client-id"] = this.clientId;
+    }
+    return headers;
+  }
+
   async init(): Promise<void> {
     const res = await this.caller.call(
       fetch,
-      `${this.motorheadURL}/sessions/${this.sessionId}/memory`,
+      `${this.url}/sessions/${this.sessionId}/memory`,
       {
         signal: this.timeout ? AbortSignal.timeout(this.timeout) : undefined,
-        headers: {
-          "Content-Type": "application/json",
-        },
+        headers: this._getHeaders(),
       }
     );
 
@@ -112,23 +146,17 @@ export class MotorheadMemory extends BaseChatMemory {
     const input = getInputValue(inputValues, this.inputKey);
     const output = getInputValue(outputValues, this.outputKey);
     await Promise.all([
-      this.caller.call(
-        fetch,
-        `${this.motorheadURL}/sessions/${this.sessionId}/memory`,
-        {
-          signal: this.timeout ? AbortSignal.timeout(this.timeout) : undefined,
-          method: "POST",
-          body: JSON.stringify({
-            messages: [
-              { role: "Human", content: `${input}` },
-              { role: "AI", content: `${output}` },
-            ],
-          }),
-          headers: {
-            "Content-Type": "application/json",
-          },
-        }
-      ),
+      this.caller.call(fetch, `${this.url}/sessions/${this.sessionId}/memory`, {
+        signal: this.timeout ? AbortSignal.timeout(this.timeout) : undefined,
+        method: "POST",
+        body: JSON.stringify({
+          messages: [
+            { role: "Human", content: `${input}` },
+            { role: "AI", content: `${output}` },
+          ],
+        }),
+        headers: this._getHeaders(),
+      }),
       super.saveContext(inputValues, outputValues),
     ]);
   }

langchain/src/memory/tests/motorhead_memory.int.test.ts

@@ -0,0 +1,22 @@
+/* eslint-disable no-process-env */
+/* eslint-disable @typescript-eslint/no-non-null-assertion */
+import { test, expect } from "@jest/globals";
+import { MotorheadMemory } from "../motorhead_memory.js";
+
+test("Test managed motörhead memory", async () => {
+  const memory = new MotorheadMemory({
+    sessionId: new Date().toISOString(),
+    apiKey: process.env.METAL_API_KEY!,
+    clientId: process.env.METAL_CLIENT_ID!,
+  });
+  const result1 = await memory.loadMemoryVariables({});
+  expect(result1).toStrictEqual({ history: "" });
+
+  await memory.saveContext(
+    { input: "Who is the best vocalist?" },
+    { response: "Ozzy Osbourne" }
+  );
+  const expectedString = "Human: Who is the best vocalist?\nAI: Ozzy Osbourne";
+  const result2 = await memory.loadMemoryVariables({});
+  expect(result2).toStrictEqual({ history: expectedString });
+});

langchain/src/memory/tests/motorhead_memory.test.ts

@@ -15,7 +15,7 @@ test("Test motörhead memory", async () => {
     } as Response)
   );
 
-  const memory = new MotorheadMemory({ sessionId: "1" });
+  const memory = new MotorheadMemory({ url: "localhost:8080", sessionId: "1" });
   const result1 = await memory.loadMemoryVariables({});
   expect(result1).toStrictEqual({ history: "" });
 
@@ -46,6 +46,7 @@ test("Test motörhead memory with pre-loaded history", async () => {
     } as Response)
   );
   const memory = new MotorheadMemory({
+    url: "localhost:8080",
     returnMessages: true,
     sessionId: "2",
   });