add authentication support

src/auth/JLINCAuthBaseChatModel.js

@@ -0,0 +1,96 @@
+const { BaseChatModel } = require("@langchain/core/language_models/chat_models");
+const { JLINCTracer } = require("../tracer/index.js");
+const { authDecide, authInvoke, getDescription, getName } = require("./common.js");
+
+/**
+ * @typedef {import('./common').JLINCConfig} JLINCConfig
+ */
+
+/**
+ * @typedef {Object} JLINCAuthBaseChatModelFields
+ * @property {JLINCConfig} config
+ * @property {Tool} jlincAuthDecision
+ * @property {BaseChatModel} targetAuthorized
+ * @property {BaseChatModel|null} targetNotAuthorized
+ */
+
+class JLINCAuthBaseChatModel extends BaseChatModel {
+  /**
+   * @param {JLINCAuthBaseChatModelFields} fields
+   */
+  constructor({ config, jlincAuthDecision, targetAuthorized, targetNotAuthorized }) {
+    super({
+      name: "jlinc_auth_base_chat_model",
+      description: "", // overridden dynamically
+    });
+    Object.defineProperty(this, "name", {
+      enumerable: true,
+      configurable: true,
+      writable: true,
+      value: getName(targetAuthorized, targetNotAuthorized)
+    });
+    Object.defineProperty(this, "description", {
+      enumerable: true,
+      configurable: true,
+      writable: true,
+      value: getDescription(targetAuthorized, targetNotAuthorized, 'BaseChatModel')
+    });
+    /** @type {JLINCConfig} */
+    this.config = config;
+    /** @type {Tool} */
+    this.jlincAuthDecision = jlincAuthDecision;
+    /** @type {BaseChatModel} */
+    this.targetAuthorized = targetAuthorized;
+    /** @type {BaseChatModel|null} */
+    this.targetNotAuthorized = targetNotAuthorized;
+    /** @type {JLINCTracer} */
+    this.tracer = new JLINCTracer(config);
+    /** @type {string} */
+    this.authType = 'BaseChatModel';
+  }
+
+  /**
+   * @param {BindToolsInput[]} tools - The tools to bind to the underlying models.
+   * @param {Partial<CallOptions>} [kwargs] - Optional call options.
+   * @returns {Runnable<BaseLanguageModelInput, OutputMessageType, CallOptions>}
+   *   A new JLINCAuthBaseChatModel instance whose underlying models
+   *   have been bound with the provided tools.
+   */ 
+  bindTools(tools, kwargs = {}) {
+    if (this.config.debug)
+      console.log(`[JLINCAuth] Binding tools to BaseChatModels`);
+    const authorizedBound = this.targetAuthorized.bindTools(tools, kwargs);
+    const notAuthorizedBound = this.targetNotAuthorized
+      ? this.targetNotAuthorized.bindTools(tools, kwargs)
+      : null;
+    return new JLINCAuthBaseChatModel({
+      config: this.config,
+      jlincAuthDecision: this.jlincAuthDecision,
+      targetAuthorized: authorizedBound,
+      targetNotAuthorized: notAuthorizedBound,
+    });
+  }
+
+  /**
+   * @returns {string} - type
+   */
+  _llmType() {
+    const type = this.targetAuthorized?._llmType?.() ?? "unknown";
+    if (this.config?.debug)
+      console.log(`[JLINCAuth] Returning LLM type: ${type}`);
+    return type;
+  } 
+
+  /**
+   * @param {BaseLanguageModelInput} input
+   * @param {Partial<CallOptions>} [options] - Optional call options.
+   * @returns {Promise<OutputMessageType>} A promise that resolves with the output.
+   */
+  async invoke(input, runManager) {
+    return await authInvoke(this, input);
+  }
+}
+
+module.exports = {
+  JLINCAuthBaseChatModel,
+};

src/auth/JLINCAuthDecision.js

@@ -0,0 +1,79 @@
+const { Tool } = require("@langchain/core/tools");
+const axios = require("axios");
+
+/**
+ * @typedef {import('./common').JLINCConfig} JLINCConfig
+ */
+
+class JLINCAuthDecision extends Tool {
+  /**
+   * @param {JLINCConfig} fields
+   */
+  constructor(config) {
+    super({
+      name: "jlinc_auth",
+      description: "Authorization via AuthZEN format.",
+    });
+    Object.defineProperty(this, "name", {
+      enumerable: true,
+      configurable: true,
+      writable: true,
+      value: "jlinc-auth"
+    });
+    Object.defineProperty(this, "description", {
+      enumerable: true,
+      configurable: true,
+      writable: true,
+      value: "Authorization via AuthZEN format."
+    });
+    /** @type {JLINCConfig} */
+    this.config = config;
+    /** @type {Boolean} */
+    this.authenticated = false;
+  }
+
+  /**
+   * @param {any} payload
+   * @returns {Promise<any>}
+   */
+  async postToApi(auth) {
+    const response = await axios.post(
+      `${this.config.dataStoreApiUrl}/api/v1/auth`,
+      auth,
+      {
+        headers: {
+          'Authorization': `Bearer ${this.config.dataStoreApiKey}`,
+        }
+      }
+    );
+    return response.data;
+  }
+
+  /**
+   * @param {auth} auth
+   * @returns {Promise<boolean>} - authenticated
+   */
+  async evaluate(auth) {
+    this.authenticated = false;
+    try {
+      this.authenticated = (await this.postToApi(auth)).decision;
+    } catch (e) {
+      console.log(`[JLINCAuthTool] Error connecting to API, flagging as unauthorized`);
+      if (this.config.debug) console.error(e)
+    }
+    if (this.config.debug)
+      console.log(`[JLINCAuth] Got authorization of: ${this.authenticated}`);   
+    return this.authenticated;
+  }
+
+  /**
+   * @returns {boolean} - authenticated
+   */
+  getAuth() {
+    return this.authenticated;
+  }
+}
+
+module.exports = {
+  JLINCAuthDecision,
+};

src/auth/JLINCAuthTool.js

@@ -0,0 +1,68 @@
+const { Tool } = require("@langchain/core/tools");
+const { JLINCTracer } = require("../tracer/index.js");
+const { authDecide, authInvoke, getDescription, getName } = require("./common.js");
+
+/**
+ * @typedef {import('./common').JLINCConfig} JLINCConfig
+ */
+
+/**
+ * @typedef {Object} JLINCAuthToolFields
+ * @property {JLINCConfig} config
+ * @property {Tool} jlincAuthDecision
+ * @property {Tool} targetAuthorized
+ * @property {Tool|null} targetNotAuthorized
+ */
+
+class JLINCAuthTool extends Tool {
+  /**
+   * @param {JLINCAuthToolFields} fields
+   */
+  constructor({ config, jlincAuthDecision, targetAuthorized, targetNotAuthorized }) {
+    super({
+      name: "jlinc_auth_tool",
+      description: "", // overridden dynamically
+    });
+    Object.defineProperty(this, "name", {
+      enumerable: true,
+      configurable: true,
+      writable: true,
+      value: getName(targetAuthorized, targetNotAuthorized)
+    });
+    Object.defineProperty(this, "description", {
+      enumerable: true,
+      configurable: true,
+      writable: true,
+      value: getDescription(targetAuthorized, targetNotAuthorized, 'Tool')
+    });
+    /** @type {JLINCConfig} */
+    this.config = config;
+    /** @type {Tool} */
+    this.jlincAuthDecision = jlincAuthDecision;
+    /** @type {Tool} */
+    this.targetAuthorized = targetAuthorized;
+    /** @type {Tool|null} */
+    this.targetNotAuthorized = targetNotAuthorized;
+    /** @type {JLINCTracer} */
+    this.tracer = new JLINCTracer(config);
+    /** @type {string} */
+    this.authType = 'Tool';
+  }
+
+  /**
+   * @template TInput
+   * @template {undefined | ToolRunnableConfig} TConfig
+   * @template TOutput
+   *
+   * @param {TInput} input
+   * @param {TConfig} [config]
+   * @returns {Promise<ToolReturnType<TInput, TConfig, TOutput>>}
+   */
+  async invoke(input, runManager) {
+    return await authInvoke(this, input);
+  }
+}
+
+module.exports = {
+  JLINCAuthTool,
+};

src/auth/common.js

@@ -0,0 +1,85 @@
+const axios = require("axios");
+
+/**
+ * @typedef {Object} JLINCConfig
+ * @property {string} [dataStoreApiUrl]
+ * @property {string} [dataStoreApiKey]
+ * @property {string} [archiveApiUrl]
+ * @property {string} [archiveApiKey]
+ * @property {string} [systemPrefix]
+ * @property {string|null} [agreementId]
+ * @property {boolean} [debug]
+ */
+
+/**
+ * @param {any} - Authorized
+ * @param {any|null} - NotAuthorized
+ * @returns {string} - The description
+ */
+function getName(authorized, unauthorized) {
+  return unauthorized
+  ? `${authorized.name}-or-${unauthorized.name}`
+  : `authorized-${authorized.name}`
+}
+
+/**
+ * @param {any} - Authorized
+ * @param {any|null} - NotAuthorized
+ * @returns {string} - The description
+ */
+function getDescription(authorized, unauthorized, type) {
+  return unauthorized ? `
+    This is an authorization router ${type}. It decides whether
+    to route input to an "authorized" or "not authorized" ${type}.
+
+    Authorized ${type}: ${authorized.name}
+    Description: ${authorized.description}
+
+    Not Authorized ${type}: ${unauthorized.name}
+    Description: ${unauthorized.description}
+
+    The LLM can assume both ${type}s are available, but JLINC Auth will
+    only allow the Authorized ${type} to be called if you are authorized.
+  ` : `
+    This is an authorization router ${type}. It decides whether
+    to route input to an "authorized" ${type}.
+
+    Authorized ${type}: ${authorized.name}
+    Description: ${authorized.description}
+
+    The LLM can assume the ${type} is available, but JLINC Auth will
+    only allow the Authorized ${type} to be called if you are authorized.     
+  `
+}
+
+function getLogName(target) {
+  if (target.name) return target.name;
+  if (target.lc_kwargs?.bound?.model) return target.lc_kwargs.bound.model;
+  return 'no-name';
+}
+
+/**
+ * @param {any} settings
+ * @param {any} input
+ * @returns {Promise<any>} - The result of invoking the selected tool.
+ */
+async function authInvoke(settings, input) {
+  const authorized = await settings.jlincAuthDecision.getAuth();
+  let selectedTarget = settings.targetNotAuthorized
+  if (authorized) {
+    selectedTarget = settings.targetAuthorized
+  }
+  if (!selectedTarget)
+    throw new Error("No valid resource available");
+  if (settings.config.debug)
+    console.log(`[JLINCAuth] Invoking ${getLogName(selectedTarget)} (${settings.authType}|${authorized})`);
+  return await selectedTarget.invoke(input, {
+    callbacks: [settings.tracer],
+  });
+}
+
+module.exports = {
+  authInvoke,
+  getDescription,
+  getName,
+};

src/auth/index.js

@@ -0,0 +1,9 @@
+const { JLINCAuthDecision } = require("./JLINCAuthDecision.js");
+const { JLINCAuthBaseChatModel } = require("./JLINCAuthBaseChatModel.js");
+const { JLINCAuthTool } = require("./JLINCAuthTool.js");
+
+module.exports = {
+    JLINCAuthDecision,
+    JLINCAuthBaseChatModel,
+    JLINCAuthTool,
+};

src/index.js

@@ -0,0 +1,13 @@
+const {
+    JLINCAuthDecision,
+    JLINCAuthBaseChatModel,
+    JLINCAuthTool,
+} = require("./auth/index.js");
+const { JLINCTracer } = require("./tracer/index.js");
+
+module.exports = {
+    JLINCAuthDecision,
+    JLINCAuthBaseChatModel,
+    JLINCAuthTool,
+    JLINCTracer,
+};

src/tracer/index.js

@@ -235,7 +235,7 @@ class JLINCTracer extends LangChainTracer {
       console.error("[Tracer] Failed to log event:", error.message);
     }
   }
-  
+ 
 }
 
 module.exports = { JLINCTracer }