Working, both static and streaming

node/src/chat/defaults.js

@@ -0,0 +1,41 @@
+module.exports = {
+  DefaultIndex: "Work",
+  DefaultKey: "title",
+  DefaultAttributes: [
+    "title",
+    "accession_number",
+    "alternate_title",
+    "api_model",
+    "catalog_key",
+    "collection",
+    "contributor",
+    "create_date",
+    "creator",
+    "date_created",
+    "description",
+    "genre",
+    "identifier_descriptive",
+    "keywords",
+    "language",
+    "library_unit",
+    "location",
+    "physical_description_material",
+    "physical_description_size",
+    "preservation_level",
+    "published",
+    "related_material",
+    "related_url",
+    "rights_holder",
+    "rights_statement",
+    "scope_and_contents",
+    "series",
+    "source",
+    "status",
+    "style_period",
+    "subject",
+    "table_of_contents",
+    "technique",
+    "visibility",
+    "work_type",
+  ],
+};

node/src/chat/multipart.js

@@ -0,0 +1,41 @@
+const { HttpResponseStream } = global.awslambda;
+
+class MultiPartStream {
+  constructor(stream) {
+    if (stream) {
+      this.boundary = `--stream-boundary-${Number(new Date())}`;
+      this.stream = HttpResponseStream.from(stream, {
+        statusCode: 200,
+        headers: {
+          "Access-Control-Allow-Origin": "*",
+          "Access-Control-Expose-Headers": "*",
+          "Access-Control-Allow-Headers": "*",
+          "Access-Control-Max-Age": "3600",
+          "Content-Type": `multipart/mixed; boundary='${this.boundary}'`,
+        },
+      });
+    }
+  }
+
+  newPart(headers) {
+    if (!this.stream) return;
+    const head = [this.boundary];
+    for (const header in headers) {
+      head.push(`${header}: ${headers[header]}`);
+    }
+    head.push("");
+    this.write(head.join("\n"));
+  }
+
+  write(content) {
+    if (!this.stream) return;
+    this.stream.write(`${content.length.toString(16)}\r\n${content}\r\n`);
+  }
+
+  end() {
+    if (!this.stream) return;
+    this.stream.end();
+  }
+}
+
+module.exports = MultiPartStream;

node/src/chat/setup.js

@@ -0,0 +1,61 @@
+const { ChatOpenAI } = require("langchain/chat_models/openai");
+const { default: Weaviate } = require("weaviate-ts-client");
+const { EmbeddedWeaviateStore } = require("../chat/embedded-weaviate-store");
+const { BaseTracer } = require("langchain/callbacks");
+
+class StreamingMultipartCallbackHandler extends BaseTracer {
+  name = "streaming_stdout_callback_handler";
+
+  persistRun(_) {
+    return Promise.resolve();
+  }
+
+  constructor({ multipart, ...args }) {
+    super(args);
+    this.multipart = multipart;
+  }
+
+  onLLMNewToken(run) {
+    const { token } = run.events[run.events.length - 1].kwargs;
+    this.multipart.write(token);
+  }
+
+  onLLMEnd(_run) {
+    this.multipart.end();
+  }
+}
+
+function openAIClient(deployment, { stream }) {
+  const callbacks = stream
+    ? [new StreamingMultipartCallbackHandler({ stream })]
+    : [];
+
+  return new ChatOpenAI({
+    azureOpenAIApiKey: process.env.AZURE_OPENAI_API_KEY,
+    azureOpenAIApiInstanceName: process.env.AZURE_OPENAI_RESOURCE_NAME,
+    azureOpenAIApiDeploymentName: deployment,
+    azureOpenAIApiVersion: "2023-07-01-preview",
+    maxTokens: -1,
+    modelName: "gpt-4",
+    streaming: true,
+    callbacks,
+  });
+}
+
+function weaviateVectorStore({ indexName, textKey, attributes }) {
+  const url = new URL(process.env.WEAVIATE_URL);
+  const client = new Weaviate.client({
+    scheme: url.protocol.slice(0, -1),
+    host: url.hostname,
+    apiKey: new Weaviate.ApiKey(process.env.WEAVIATE_API_KEY),
+  });
+
+  return EmbeddedWeaviateStore.fromExistingIndex(null, {
+    client,
+    indexName,
+    textKey,
+    metadataKeys: attributes,
+  });
+}
+
+module.exports = { openAIClient, weaviateVectorStore };