@traceloop/instrumentation-together
Version:
Together Instrumentation
399 lines (395 loc) • 21.8 kB
JavaScript
import { __asyncGenerator, __await, __asyncValues } from 'tslib';
import { trace, context, SpanKind } from '@opentelemetry/api';
import { InstrumentationBase, InstrumentationNodeModuleDefinition, safeExecuteInTheMiddle } from '@opentelemetry/instrumentation';
import { SpanAttributes, CONTEXT_KEY_ALLOW_TRACE_CONTENT } from '@traceloop/ai-semantic-conventions';
var version = "0.14.6";
class TogetherInstrumentation extends InstrumentationBase {
constructor(config = {}) {
super("@traceloop/instrumentation-together-ai", version, Object.assign({ enrichTokens: true }, config));
}
setConfig(config = {}) {
super.setConfig(config);
}
manuallyInstrument(module) {
this._diag.debug(`Manually instrumenting togetherai`);
this._wrap(module.Chat.Completions.prototype, "create", this.patchTogether("chat"));
this._wrap(module.Completions.prototype, "create", this.patchTogether("completion"));
}
init() {
const module = new InstrumentationNodeModuleDefinition("together-ai", [">=0.13.0"], this.patch.bind(this), this.unpatch.bind(this));
return module;
}
patch(moduleExports, moduleVersion) {
this._diag.debug(`Patching togetherai@${moduleVersion}`);
this._wrap(moduleExports.Together.Chat.Completions.prototype, "create", this.patchTogether("chat"));
this._wrap(moduleExports.Together.Completions.prototype, "create", this.patchTogether("completion"));
return moduleExports;
}
unpatch(moduleExports, moduleVersion) {
this._diag.debug(`Unpatching togetherai@${moduleVersion}`);
this._unwrap(moduleExports.Together.Chat.Completions.prototype, "create");
this._unwrap(moduleExports.Together.Completions.prototype, "create");
}
patchTogether(type) {
// eslint-disable-next-line @typescript-eslint/no-this-alias
const plugin = this;
// eslint-disable-next-line
return (original) => {
return function method(...args) {
const span = type === "chat"
? plugin.startSpan({
type,
params: args[0],
})
: plugin.startSpan({
type,
params: args[0],
});
const execContext = trace.setSpan(context.active(), span);
const execPromise = safeExecuteInTheMiddle(() => {
return context.with(execContext, () => {
var _a;
if ((_a = args === null || args === void 0 ? void 0 : args[0]) === null || _a === void 0 ? void 0 : _a.extraAttributes) {
delete args[0].extraAttributes;
}
return original.apply(this, args);
});
}, (e) => {
if (e) {
console.error(`Error in ${type} instrumentation:`, e);
plugin._diag.error("Together instrumentation: error", e);
}
});
if (args[0].stream) {
return context.bind(execContext, plugin._streamingWrapPromise({
span,
type,
params: args[0],
promise: execPromise,
}));
}
const wrappedPromise = plugin._wrapPromise(type, span, execPromise);
return context.bind(execContext, wrappedPromise);
};
};
}
startSpan({ type, params, }) {
var _a, _b, _c, _d;
const attributes = {
[SpanAttributes.LLM_SYSTEM]: "TogetherAI",
[SpanAttributes.LLM_REQUEST_TYPE]: type,
};
try {
attributes[SpanAttributes.LLM_REQUEST_MODEL] = params.model;
if (params.max_tokens) {
attributes[SpanAttributes.LLM_REQUEST_MAX_TOKENS] = params.max_tokens;
}
if (params.temperature) {
attributes[SpanAttributes.LLM_REQUEST_TEMPERATURE] = params.temperature;
}
if (params.top_p) {
attributes[SpanAttributes.LLM_REQUEST_TOP_P] = params.top_p;
}
if (params.frequency_penalty) {
attributes[SpanAttributes.LLM_FREQUENCY_PENALTY] =
params.frequency_penalty;
}
if (params.presence_penalty) {
attributes[SpanAttributes.LLM_PRESENCE_PENALTY] =
params.presence_penalty;
}
if (params.extraAttributes !== undefined &&
typeof params.extraAttributes === "object") {
Object.keys(params.extraAttributes).forEach((key) => {
attributes[key] = params.extraAttributes[key];
});
}
if (this._shouldSendPrompts()) {
if (type === "chat") {
params.messages.forEach((message, index) => {
attributes[`${SpanAttributes.LLM_PROMPTS}.${index}.role`] =
message.role;
if (typeof message.content === "string") {
attributes[`${SpanAttributes.LLM_PROMPTS}.${index}.content`] =
message.content || "";
}
else {
attributes[`${SpanAttributes.LLM_PROMPTS}.${index}.content`] =
JSON.stringify(message.content);
}
});
(_a = params.tools) === null || _a === void 0 ? void 0 : _a.forEach((func, index) => {
var _a, _b, _c;
attributes[`${SpanAttributes.LLM_REQUEST_FUNCTIONS}.${index}.name`] = (_a = func.function) === null || _a === void 0 ? void 0 : _a.name;
attributes[`${SpanAttributes.LLM_REQUEST_FUNCTIONS}.${index}.description`] = (_b = func.function) === null || _b === void 0 ? void 0 : _b.description;
attributes[`${SpanAttributes.LLM_REQUEST_FUNCTIONS}.${index}.arguments`] = JSON.stringify((_c = func.function) === null || _c === void 0 ? void 0 : _c.parameters);
});
(_b = params.tools) === null || _b === void 0 ? void 0 : _b.forEach((tool, index) => {
if (!tool.function) {
return;
}
attributes[`${SpanAttributes.LLM_REQUEST_FUNCTIONS}.${index}.name`] = tool.function.name;
attributes[`${SpanAttributes.LLM_REQUEST_FUNCTIONS}.${index}.description`] = tool.function.description;
attributes[`${SpanAttributes.LLM_REQUEST_FUNCTIONS}.${index}.arguments`] = JSON.stringify(tool.function.parameters);
});
}
else {
attributes[`${SpanAttributes.LLM_PROMPTS}.0.role`] = "user";
if (typeof params.prompt === "string") {
attributes[`${SpanAttributes.LLM_PROMPTS}.0.content`] =
params.prompt;
}
else {
attributes[`${SpanAttributes.LLM_PROMPTS}.0.content`] =
JSON.stringify(params.prompt);
}
}
}
}
catch (e) {
console.error("Error setting span attributes:", e);
this._diag.debug(e);
(_d = (_c = this._config).exceptionLogger) === null || _d === void 0 ? void 0 : _d.call(_c, e);
}
return this.tracer.startSpan(`together.${type}`, {
kind: SpanKind.CLIENT,
attributes,
});
}
_streamingWrapPromise(_a) {
return __asyncGenerator(this, arguments, function* _streamingWrapPromise_1({ span, type, params, promise, }) {
var _b, e_1, _c, _d, _e, e_2, _f, _g;
var _h, _j, _k, _l, _m, _o, _p, _q, _r, _s, _t, _u, _v, _w, _x, _y, _z, _0, _1;
if (type === "chat") {
const result = {
id: "0",
created: -1,
model: "",
choices: [
{
index: 0,
logprobs: null,
finish_reason: "stop",
message: {
role: "assistant",
content: "",
tool_calls: [],
},
},
],
object: "chat.completion",
};
const stream = yield __await(promise);
try {
for (var _2 = true, stream_1 = __asyncValues(stream), stream_1_1; stream_1_1 = yield __await(stream_1.next()), _b = stream_1_1.done, !_b; _2 = true) {
_d = stream_1_1.value;
_2 = false;
const chunk = _d;
yield yield __await(chunk);
result.id = chunk.id;
result.created = chunk.created;
result.model = chunk.model;
if ((_h = chunk.choices[0]) === null || _h === void 0 ? void 0 : _h.finish_reason) {
result.choices[0].finish_reason = chunk.choices[0].finish_reason;
}
if (result.choices[0].message && ((_j = chunk.choices[0]) === null || _j === void 0 ? void 0 : _j.delta.content)) {
result.choices[0].message.content += chunk.choices[0].delta.content;
}
if (result.choices[0].message &&
((_k = chunk.choices[0]) === null || _k === void 0 ? void 0 : _k.delta.function_call) &&
((_l = chunk.choices[0]) === null || _l === void 0 ? void 0 : _l.delta.function_call.arguments) &&
((_m = chunk.choices[0]) === null || _m === void 0 ? void 0 : _m.delta.function_call.name)) {
result.choices[0].message.function_call = {
name: chunk.choices[0].delta.function_call.name,
arguments: chunk.choices[0].delta.function_call.arguments,
};
}
for (const toolCall of (_q = (_p = (_o = chunk.choices[0]) === null || _o === void 0 ? void 0 : _o.delta) === null || _p === void 0 ? void 0 : _p.tool_calls) !== null && _q !== void 0 ? _q : []) {
if (result.choices[0].message &&
((_t = (_s = (_r = result.choices[0].message) === null || _r === void 0 ? void 0 : _r.tool_calls) === null || _s === void 0 ? void 0 : _s.length) !== null && _t !== void 0 ? _t : 0) <
toolCall.index + 1) {
(_u = result.choices[0].message.tool_calls) === null || _u === void 0 ? void 0 : _u.push({
function: {
name: "",
arguments: "",
},
id: "",
type: "function",
index: toolCall.index,
});
}
if (result.choices[0].message &&
result.choices[0].message.tool_calls) {
if (toolCall.id) {
result.choices[0].message.tool_calls[toolCall.index].id +=
toolCall.id;
}
if (result.choices[0].message && toolCall.type) {
result.choices[0].message.tool_calls[toolCall.index].type +=
toolCall.type;
}
if (result.choices[0].message && ((_v = toolCall.function) === null || _v === void 0 ? void 0 : _v.name)) {
result.choices[0].message.tool_calls[toolCall.index].function.name += toolCall.function.name;
}
if (result.choices[0].message && ((_w = toolCall.function) === null || _w === void 0 ? void 0 : _w.arguments)) {
result.choices[0].message.tool_calls[toolCall.index].function.arguments += toolCall.function.arguments;
}
}
}
// Use token usage information from the final chunk if available
if (chunk.usage) {
result.usage = chunk.usage;
}
}
}
catch (e_1_1) { e_1 = { error: e_1_1 }; }
finally {
try {
if (!_2 && !_b && (_c = stream_1.return)) yield __await(_c.call(stream_1));
}
finally { if (e_1) throw e_1.error; }
}
this._endSpan({ span, type, result });
}
else {
const result = {
id: "0",
created: -1,
model: "",
choices: [
{
finish_reason: "stop",
text: "",
},
],
object: "text.completion",
usage: {
prompt_tokens: 0,
completion_tokens: 0,
total_tokens: 0,
},
};
try {
for (var _3 = true, _4 = __asyncValues(yield __await(promise)), _5; _5 = yield __await(_4.next()), _e = _5.done, !_e; _3 = true) {
_g = _5.value;
_3 = false;
const chunk = _g;
yield yield __await(chunk);
try {
result.id = chunk.id;
result.created = chunk.created;
result.model = chunk.model;
if ((_x = chunk.choices[0]) === null || _x === void 0 ? void 0 : _x.finish_reason) {
result.choices[0].finish_reason = chunk.choices[0].finish_reason;
}
if ((_y = chunk.choices[0]) === null || _y === void 0 ? void 0 : _y.logprobs) {
result.choices[0].logprobs = chunk.choices[0].logprobs;
}
if ((_z = chunk.choices[0]) === null || _z === void 0 ? void 0 : _z.text) {
result.choices[0].text += chunk.choices[0].text;
}
// Log any usage info from chunk
if (chunk.usage) {
result.usage = chunk.usage;
}
}
catch (e) {
console.error("Error processing chunk:", e);
this._diag.debug(e);
(_1 = (_0 = this._config).exceptionLogger) === null || _1 === void 0 ? void 0 : _1.call(_0, e);
}
}
}
catch (e_2_1) { e_2 = { error: e_2_1 }; }
finally {
try {
if (!_3 && !_e && (_f = _4.return)) yield __await(_f.call(_4));
}
finally { if (e_2) throw e_2.error; }
}
this._endSpan({ span, type, result });
}
});
}
_wrapPromise(type, span, promise) {
return promise._thenUnwrap((result) => {
if (type === "chat") {
this._endSpan({
type,
span,
result: result,
});
}
else {
this._endSpan({
type,
span,
result: result,
});
}
return result;
});
}
_shouldSendPrompts() {
const contextShouldSendPrompts = context
.active()
.getValue(CONTEXT_KEY_ALLOW_TRACE_CONTENT);
if (contextShouldSendPrompts !== undefined) {
return contextShouldSendPrompts;
}
return this._config.traceContent !== undefined
? this._config.traceContent
: true;
}
_endSpan({ span, type, result, }) {
var _a, _b, _c, _d, _e;
try {
span.setAttribute(SpanAttributes.LLM_RESPONSE_MODEL, result.model);
if (result.usage) {
span.setAttribute(SpanAttributes.LLM_USAGE_TOTAL_TOKENS, (_a = result.usage) === null || _a === void 0 ? void 0 : _a.total_tokens);
span.setAttribute(SpanAttributes.LLM_USAGE_COMPLETION_TOKENS, (_b = result.usage) === null || _b === void 0 ? void 0 : _b.completion_tokens);
span.setAttribute(SpanAttributes.LLM_USAGE_PROMPT_TOKENS, (_c = result.usage) === null || _c === void 0 ? void 0 : _c.prompt_tokens);
}
if (this._shouldSendPrompts()) {
if (type === "chat") {
result.choices.forEach((choice, index) => {
var _a, _b, _c, _d, _e, _f, _g, _h, _j, _k, _l, _m, _o, _p, _q, _r, _s;
span.setAttribute(`${SpanAttributes.LLM_COMPLETIONS}.${index}.finish_reason`, (_a = choice.finish_reason) !== null && _a !== void 0 ? _a : "");
span.setAttribute(`${SpanAttributes.LLM_COMPLETIONS}.${index}.role`, (_c = (_b = choice.message) === null || _b === void 0 ? void 0 : _b.role) !== null && _c !== void 0 ? _c : "");
span.setAttribute(`${SpanAttributes.LLM_COMPLETIONS}.${index}.content`, (_e = (_d = choice.message) === null || _d === void 0 ? void 0 : _d.content) !== null && _e !== void 0 ? _e : "");
if ((_f = choice.message) === null || _f === void 0 ? void 0 : _f.function_call) {
span.setAttribute(`${SpanAttributes.LLM_COMPLETIONS}.${index}.function_call.name`, (_j = (_h = (_g = choice.message) === null || _g === void 0 ? void 0 : _g.function_call) === null || _h === void 0 ? void 0 : _h.name) !== null && _j !== void 0 ? _j : "");
span.setAttribute(`${SpanAttributes.LLM_COMPLETIONS}.${index}.function_call.arguments`, (_m = (_l = (_k = choice.message) === null || _k === void 0 ? void 0 : _k.function_call) === null || _l === void 0 ? void 0 : _l.arguments) !== null && _m !== void 0 ? _m : "");
}
else if ((_p = (_o = choice.message) === null || _o === void 0 ? void 0 : _o.tool_calls) === null || _p === void 0 ? void 0 : _p[0]) {
span.setAttribute(`${SpanAttributes.LLM_COMPLETIONS}.${index}.function_call.name`, (_q = choice.message.tool_calls[0].function.name) !== null && _q !== void 0 ? _q : "");
span.setAttribute(`${SpanAttributes.LLM_COMPLETIONS}.${index}.function_call.arguments`, (_r = choice.message.tool_calls[0].function.arguments) !== null && _r !== void 0 ? _r : "");
}
if ((_s = choice === null || choice === void 0 ? void 0 : choice.message) === null || _s === void 0 ? void 0 : _s.tool_calls) {
for (const [toolIndex, toolCall,] of choice.message.tool_calls.entries()) {
span.setAttribute(`${SpanAttributes.LLM_COMPLETIONS}.${index}.tool_calls.${toolIndex}.name`, toolCall.function.name);
span.setAttribute(`${SpanAttributes.LLM_COMPLETIONS}.${index}.tool_calls.${toolIndex}.arguments`, toolCall.function.arguments);
}
}
});
}
else {
result.choices.forEach((choice, index) => {
var _a, _b;
span.setAttribute(`${SpanAttributes.LLM_COMPLETIONS}.${index}.finish_reason`, (_a = choice.finish_reason) !== null && _a !== void 0 ? _a : "");
span.setAttribute(`${SpanAttributes.LLM_COMPLETIONS}.${index}.role`, "assistant");
span.setAttribute(`${SpanAttributes.LLM_COMPLETIONS}.${index}.content`, (_b = choice.text) !== null && _b !== void 0 ? _b : "");
});
}
}
}
catch (e) {
console.error("Error ending span:", e);
this._diag.debug(e);
(_e = (_d = this._config).exceptionLogger) === null || _e === void 0 ? void 0 : _e.call(_d, e);
}
span.end();
}
}
export { TogetherInstrumentation };
//# sourceMappingURL=index.mjs.map