@maximai/maxim-js-langchain
Version:
This is the langchain wrapper built for (Maxim JS SDK)[https://www.npmjs.com/package/@maximai/maxim-js].
318 lines (317 loc) • 13.1 kB
JavaScript
"use strict";
Object.defineProperty(exports, "__esModule", { value: true });
exports.MaximLangchainTracer = void 0;
const base_1 = require("@langchain/core/callbacks/base");
const uuid_1 = require("uuid");
const models_1 = require("./models");
const utils_1 = require("./utils");
class MaximLangchainTracer extends base_1.BaseCallbackHandler {
constructor(logger, debug) {
super();
this.logger = logger;
this.name = "MaximLangchainTracer";
this.containers = new Map();
}
validateMaximMetadata(metadata) {
if (!metadata)
return;
const idKeys = ["sessionId", "traceId", "spanId"];
const presentKeys = idKeys.filter((key) => key in metadata);
if (presentKeys.length > 1) {
throw new Error(`Multiple keys found in metadata: ${presentKeys}. You can pass only one of these.`);
}
const validKeys = models_1.MetadataKeys;
const invalidKeys = Object.keys(metadata).filter((key) => !validKeys.includes(key));
if (invalidKeys.length > 0) {
throw new Error(`Invalid keys found in metadata: ${invalidKeys}. Valid keys are ${validKeys}`);
}
}
getMetadata(metadata) {
if (metadata && "maxim" in metadata) {
return new models_1.Metadata(metadata["maxim"]);
}
return null;
}
getContainerFromMetadata(runId) {
const maximMetadata = this.getMetadata();
let container = this.containers.get(runId);
if (maximMetadata) {
const spanId = maximMetadata.spanId;
if (spanId) {
container = new models_1.SpanContainer(this.logger, spanId, maximMetadata.spanName, undefined, true);
}
const traceId = maximMetadata.traceId;
if (traceId) {
container = new models_1.TraceContainer(this.logger, traceId, maximMetadata.traceName, undefined, true);
}
}
if (!container) {
const traceId = (0, uuid_1.v4)();
container = new models_1.TraceContainer(this.logger, traceId, "Trace");
container.create();
}
if (maximMetadata && container.type === "trace") {
const sessionId = maximMetadata.sessionId;
if (sessionId) {
container.parentId = sessionId;
}
}
return container;
}
getContainer(runId, parentRunId) {
let container;
if (!parentRunId) {
// This is the first activity in this run
container = this.getContainerFromMetadata(runId);
}
else {
container = this.containers.get(parentRunId);
}
return container;
}
async handleLLMStart(llm, prompts, runId, parentRunId, extraParams, tags, metadata, name) {
var _a, _b;
const maximMetadata = this.getMetadata(metadata);
let defaultGenerationName = "default-generation";
let generationTags = Object.assign({}, maximMetadata === null || maximMetadata === void 0 ? void 0 : maximMetadata.generationTags);
const [model, modelParameters] = (0, utils_1.parseLangchainModelAndParameters)(llm, metadata, extraParams);
const generationConfig = {
id: runId,
name: (_b = (_a = maximMetadata === null || maximMetadata === void 0 ? void 0 : maximMetadata.generationName) !== null && _a !== void 0 ? _a : name) !== null && _b !== void 0 ? _b : defaultGenerationName,
provider: (0, utils_1.determineProvider)(llm.id, metadata),
model,
messages: (0, utils_1.parseLangchainMessages)(prompts),
modelParameters,
tags: generationTags,
};
const container = this.getContainer(runId, parentRunId);
if (!container) {
console.error("[MaximSDK] Couldn't find a container for generation");
return;
}
if (!container.isCreated) {
container.create();
}
container.addGeneration(generationConfig);
}
handleLLMEnd(output, runId, parentRunId, tags) {
const result = (0, utils_1.convertLLMResultToCompletionResult)(output);
const container = this.getContainer(runId, parentRunId);
if (!container) {
console.error("[MaximSDK] Couldn't find a container for generation");
return;
}
this.logger.generationResult(runId, result);
if (!container.parentId) {
container.end();
this.containers.delete(runId);
}
}
handleLLMError(err, runId, parentRunId, tags) {
const container = this.getContainer(runId, parentRunId);
if (!container) {
console.error("[MaximSDK] Couldn't find a container for generation");
return;
}
const generationError = (0, utils_1.parseLangchainLLMError)(err);
this.logger.generationError(runId, generationError);
if (!container.parentId) {
container.end();
this.containers.delete(runId);
}
}
async handleChatModelStart(llm, messages, runId, parentRunId, extraParams, tags, metadata, runName) {
var _a, _b;
const [model, modelParameters] = (0, utils_1.parseLangchainModelAndParameters)(llm, metadata, extraParams);
const maximMetadata = this.getMetadata(metadata);
let defaultGenerationName = "default-generation";
let generationTags = Object.assign({}, maximMetadata === null || maximMetadata === void 0 ? void 0 : maximMetadata.generationTags);
const generationConfig = {
id: runId,
name: (_b = (_a = maximMetadata === null || maximMetadata === void 0 ? void 0 : maximMetadata.generationName) !== null && _a !== void 0 ? _a : runName) !== null && _b !== void 0 ? _b : defaultGenerationName,
provider: (0, utils_1.determineProvider)(llm.id, metadata),
model,
messages: (0, utils_1.parseLangchainMessages)(messages),
modelParameters,
tags: generationTags,
};
const container = this.getContainer(runId, parentRunId);
if (!container) {
console.error("[MaximSDK] Couldn't find a container for generation");
return;
}
if (!container.isCreated) {
container.create();
}
container.addGeneration(generationConfig);
if (!container.parentId) {
this.containers.set(runId, container);
}
}
handleRetrieverStart(retriever, query, runId, parentRunId, tags, metadata, name) {
var _a, _b;
const retrievalId = runId;
const defaultRetrievalName = "default-retrieval";
const maximMetadata = this.getMetadata(metadata);
const retrievalConfig = {
id: retrievalId,
name: (_b = (_a = maximMetadata === null || maximMetadata === void 0 ? void 0 : maximMetadata.retrievalName) !== null && _a !== void 0 ? _a : name) !== null && _b !== void 0 ? _b : defaultRetrievalName,
};
const container = this.getContainer(runId, parentRunId);
if (!container) {
console.error("[MaximSDK] Couldn't find a container for retrieval");
return;
}
if (!container.isCreated) {
container.create();
}
const retrieval = container.addRetrieval(retrievalConfig);
retrieval.input(query);
}
handleRetrieverEnd(documents, runId, parentRunId, tags) {
const container = this.getContainer(runId, parentRunId);
if (!container) {
console.error("[MaximSDK] Couldn't find a container for retrieval");
return;
}
this.logger.retrievalOutput(runId, JSON.stringify(documents));
if (!container.parentId) {
container.end();
this.containers.delete(runId);
}
}
handleChainStart(chain, inputs, runId, parentRunId, tags, metadata, runType, runName) {
try {
const name = runName;
const maximMetadata = this.getMetadata(metadata);
const chainTags = Object.assign({}, maximMetadata === null || maximMetadata === void 0 ? void 0 : maximMetadata.chainTags);
const container = this.getContainer(runId, parentRunId);
if (!container) {
console.error("[MaximSDK] Couldn't find a container for chain");
return;
}
if (!container.isCreated()) {
container.create();
}
// Add run_id and parent_run_id to tags
chainTags["run_id"] = runId;
if (parentRunId) {
chainTags["parent_run_id"] = parentRunId;
}
// Add custom tags if provided
if (tags) {
tags.forEach((tag) => {
const [key, value] = tag.split(":", 2);
if (key && value) {
chainTags[key.trim()] = value.trim();
}
});
}
// Add metadata as tags if provided
if (metadata) {
Object.entries(metadata).forEach(([key, value]) => {
chainTags[key.trim()] = typeof value === "string" ? value.trim() : JSON.stringify(value);
});
}
const spanConfig = {
id: runId,
name: name,
tags: chainTags,
};
container.addSpan(spanConfig);
this.containers.set(runId, new models_1.SpanContainer(this.logger, runId, name, parentRunId, true));
}
catch (e) {
console.error("[MaximSDK] Error while processing chain_start", e);
}
}
handleChainEnd(outputs, runId, parentRunId, tags, kwargs) {
try {
const chainTags = {};
const container = this.getContainer(runId, parentRunId);
if (!container) {
console.error("[MaximSDK] Couldn't find a container for chain");
return;
}
// Process tags if provided
if (tags) {
tags.forEach((tag) => {
const [key, value] = tag.split(":", 2);
if (key && value) {
chainTags[key.trim()] = value.trim();
}
});
}
container.addTags(chainTags);
container.end();
this.containers.delete(runId);
}
catch (e) {
console.error("[MaximSDK] Failed to parse chain-end:", e);
}
}
handleToolStart(tool, input, runId, parentRunId, tags, metadata, runName) {
var _a;
try {
const name = (_a = tool.name) !== null && _a !== void 0 ? _a : "tool";
const description = "description" in tool ? tool.description : "";
const container = this.getContainer(runId, parentRunId);
if (!container) {
console.error("[MaximSDK] Couldn't find a container for this tool call");
return;
}
if (!container.isCreated()) {
container.create();
}
container.addToolCall({ id: runId, name, description, args: input });
if (!container.parentId) {
this.containers.set(runId, container);
}
}
catch (e) {
console.error("[MaximSDK] Failed to parse tool-start:", e);
}
}
handleToolEnd(output, runId, parentRunId, tags) {
try {
const container = this.getContainer(runId, parentRunId);
if (!container) {
console.error("[MaximSDK] Couldn't find a container for this tool call");
return;
}
if (output && "status" in output) {
if (output.status === "success") {
this.logger.toolCallResult(runId, output.content);
}
else if (output.status === "error") {
this.logger.toolCallError(runId, output.content);
}
}
if (!container.parentId) {
container.end();
this.containers.delete(runId);
}
}
catch (e) {
console.error("[MaximSDK] Failed to parse tool-end:", e);
}
}
handleToolError(error, runId, parentRunId, tags) {
try {
const container = this.getContainer(runId, parentRunId);
if (!container) {
console.error("[MaximSDK] Couldn't find a container for this tool call");
return;
}
this.logger.toolCallError(runId, error);
if (!container.parentId) {
container.end();
this.containers.delete(runId);
}
}
catch (e) {
console.error("[MaximSDK] Failed to parse tool-end:", e);
}
}
}
exports.MaximLangchainTracer = MaximLangchainTracer;