Skip to content

Commit

Permalink
feat: working stage of tracing
Browse files Browse the repository at this point in the history
  • Loading branch information
5war00p committed Jan 18, 2024
1 parent 2cd57ef commit d15d635
Show file tree
Hide file tree
Showing 3 changed files with 100 additions and 58 deletions.
99 changes: 70 additions & 29 deletions packages/instrumentation-vertexai/src/instrumentation.ts
Original file line number Diff line number Diff line change
Expand Up @@ -29,14 +29,9 @@ import {
} from "@opentelemetry/instrumentation";
import { VertexAIInstrumentationConfig } from "./types";
import { SpanAttributes } from "@traceloop/ai-semantic-conventions";
import {
GenerateContentRequest,
GenerateContentResponse,
Part,
VertexAI,
} from "@google-cloud/vertexai";
import * as vertexAI from "@google-cloud/vertexai";

export class VertexAIInstrumentation extends InstrumentationBase {
export class VertexAIInstrumentation extends InstrumentationBase<any> {
protected override _config!: VertexAIInstrumentationConfig;

constructor(config: VertexAIInstrumentationConfig = {}) {
Expand All @@ -47,33 +42,66 @@ export class VertexAIInstrumentation extends InstrumentationBase {
super.setConfig(config);
}

protected init(): InstrumentationModuleDefinition<unknown> {
const module = new InstrumentationNodeModuleDefinition<unknown>(
protected init(): InstrumentationModuleDefinition<any> {
const module = new InstrumentationNodeModuleDefinition<any>(
"@google-cloud/vertexai",
[">=0.2.1"],
this.patch.bind(this),
this.unpatch.bind(this),
);
return module;
}

public manuallyInstrument(module: typeof VertexAI) {
public manuallyInstrument(module: typeof vertexAI) {
// this._wrap(
// module.VertexAI_Preview.prototype,
// "getGenerativeModel",
// this.patchVertexAI(),
// );
this._wrap(
module.GenerativeModel.prototype,
"generateContent",
this.patchVertexAI(),
);
}

private patch(moduleExports: typeof vertexAI) {
// this._wrap(
// moduleExports.VertexAI_Preview.prototype,
// "getGenerativeModel",
// this.patchVertexAI(),
// );
this._wrap(
module.prototype.preview.getGenerativeModel.prototype,
moduleExports.GenerativeModel.prototype,
"generateContent",
this.patchVertexAI(),
);
return moduleExports;
}

private unpatch(moduleExports: typeof vertexAI): void {
// this._unwrap(
// moduleExports.VertexAI_Preview.prototype,
// "getGenerativeModel",
// );
this._unwrap(moduleExports.GenerativeModel.prototype, "generateContent");
}

private patchVertexAI() {
// eslint-disable-next-line @typescript-eslint/no-this-alias
const plugin = this;
// eslint-disable-next-line @typescript-eslint/ban-types
return (original: Function) => {
return function method(this: any, ...args: GenerateContentRequest[]) {
return function method(
this: any,
...args: vertexAI.GenerateContentRequest[]
) {
const span = plugin._startSpan({
params: args[0],
});

const execContext = trace.setSpan(context.active(), span);

const execPromise = safeExecuteInTheMiddle(
() => {
return context.with(execContext, () => {
Expand All @@ -84,20 +112,30 @@ export class VertexAIInstrumentation extends InstrumentationBase {
() => {},
);

// Its to get the model name
// if (args[0].model) {
// this.model = args[0].model;
// return context.bind(execContext, execPromise);
// }

const wrappedPromise = plugin._wrapPromise(span, execPromise);

return context.bind(execContext, wrappedPromise as any);
};
};
}

private _startSpan({ params }: { params: GenerateContentRequest }): Span {
private _startSpan({
params,
}: {
params: vertexAI.GenerateContentRequest;
}): Span {
const attributes: Attributes = {
[SpanAttributes.LLM_VENDOR]: "VertexAI",
[SpanAttributes.LLM_REQUEST_TYPE]: "completion",
};

attributes[SpanAttributes.LLM_REQUEST_MODEL] = "Test Model";
// attributes[SpanAttributes.LLM_REQUEST_MODEL] = "";

if (
params.generation_config !== undefined &&
Expand Down Expand Up @@ -125,7 +163,6 @@ export class VertexAIInstrumentation extends InstrumentationBase {
attributes[`${SpanAttributes.LLM_PROMPTS}.0.content`] =
this._formatPartsData(params.contents[0].parts);
}

return this.tracer.startSpan(`vertexai.completion`, {
kind: SpanKind.CLIENT,
attributes,
Expand All @@ -136,7 +173,10 @@ export class VertexAIInstrumentation extends InstrumentationBase {
return promise
.then((result) => {
return new Promise<T>((resolve) => {
this._endSpan({ span, result: result as GenerateContentResponse });
this._endSpan({
span,
result: result as vertexAI.GenerateContentResult,
});
resolve(result);
});
})
Expand All @@ -159,31 +199,32 @@ export class VertexAIInstrumentation extends InstrumentationBase {
result,
}: {
span: Span;
result: GenerateContentResponse;
result: vertexAI.GenerateContentResult;
}) {
span.setAttribute(SpanAttributes.LLM_RESPONSE_MODEL, "Test Model");
if (result.usageMetadata) {
if (result.usageMetadata.totalTokenCount !== undefined)
// span.setAttribute(SpanAttributes.LLM_RESPONSE_MODEL, '');

if (result.response.usageMetadata) {
if (result.response.usageMetadata.totalTokenCount !== undefined)
span.setAttribute(
SpanAttributes.LLM_USAGE_TOTAL_TOKENS,
result.usageMetadata.totalTokenCount,
result.response.usageMetadata.totalTokenCount,
);

if (result.usageMetadata.candidates_token_count)
if (result.response.usageMetadata.candidates_token_count)
span.setAttribute(
SpanAttributes.LLM_USAGE_COMPLETION_TOKENS,
result.usageMetadata.candidates_token_count,
result.response.usageMetadata.candidates_token_count,
);

if (result.usageMetadata.prompt_token_count)
if (result.response.usageMetadata.prompt_token_count)
span.setAttribute(
SpanAttributes.LLM_USAGE_PROMPT_TOKENS,
result.usageMetadata?.prompt_token_count,
result.response.usageMetadata?.prompt_token_count,
);
}

if (this._shouldSendPrompts()) {
result.candidates.forEach((candidate, index) => {
if (this._shouldSendPrompts() && result.response) {
result.response.candidates.forEach((candidate, index) => {
if (candidate.finishReason)
span.setAttribute(
`${SpanAttributes.LLM_COMPLETIONS}.${index}.finish_reason`,
Expand All @@ -207,7 +248,7 @@ export class VertexAIInstrumentation extends InstrumentationBase {
span.end();
}

private _formatPartsData(parts: Part[]): Array<string> {
private _formatPartsData(parts: vertexAI.Part[]): string {
const result = parts
.map((part) => {
if (part.text) return part.text;
Expand All @@ -219,7 +260,7 @@ export class VertexAIInstrumentation extends InstrumentationBase {
})
.filter(Boolean);

return result;
return result.join("\n");
}

private _shouldSendPrompts() {
Expand Down
55 changes: 28 additions & 27 deletions packages/sample-app/src/vertexai/gemini_generate.ts
Original file line number Diff line number Diff line change
@@ -1,45 +1,46 @@
import * as traceloop from "@traceloop/node-server-sdk";
import { VertexAI } from "@google-cloud/vertexai";
// import { ConsoleSpanExporter } from "@opentelemetry/sdk-trace-node";

traceloop.initialize({
appName: "sample_vertexai",
apiKey: process.env.TRACELOOP_API_KEY,
disableBatch: true,
// exporter: new ConsoleSpanExporter(),
});

async function createNonStreamingContent() {
// Initialize Vertex with your Cloud project and location
const vertexAI = new VertexAI({
project: process.env.VERTEXAI_PROJECT_ID ?? "",
location: process.env.VERTEXAI_LOCATION ?? "",
});

// Instantiate the model
const generativeModel = vertexAI.preview.getGenerativeModel({
model: "gemini-pro-vision",
});
// Initialize Vertex with your Cloud project and location
const vertexAI = new VertexAI({
project: process.env.VERTEXAI_PROJECT_ID ?? "",
location: process.env.VERTEXAI_LOCATION ?? "",
});

const request = {
contents: [{ role: "user", parts: [{ text: "What is Node.js?" }] }],
};
async function createNonStreamingContent() {
return await traceloop.withWorkflow("sample_completion", {}, async () => {
// Instantiate the model
const generativeModel = vertexAI.preview.getGenerativeModel({
model: "gemini-pro-vision",
});

console.log("Prompt:");
console.log(request.contents[0].parts[0].text);
console.log("Non-Streaming Response Text:");
const request = {
contents: [{ role: "user", parts: [{ text: "What is Node.js?" }] }],
};

// Create the response stream
const responseStream = await generativeModel.generateContent(request);
// Create the response stream
const responseStream = await generativeModel.generateContent(request);

// Wait for the response stream to complete
const aggregatedResponse = await responseStream.response;
// Wait for the response stream to complete
const aggregatedResponse = await responseStream.response;

// Select the text from the response
const fullTextResponse =
aggregatedResponse.candidates[0].content.parts[0].text;
// Select the text from the response
const fullTextResponse =
aggregatedResponse.candidates[0].content.parts[0].text;

console.log(fullTextResponse);
return fullTextResponse;
});
}

createNonStreamingContent().catch((err) => {
console.error(err.message);
traceloop.withAssociationProperties({}, async () => {
const completionResponse = await createNonStreamingContent();
console.log(completionResponse);
});
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
import { SpanExporter } from "@opentelemetry/sdk-trace-base";
import * as openai from "openai";
import * as llamaIndex from "llamaindex";
import * as vertexAI from "@google-cloud/vertexai";
import { VertexAI } from "@google-cloud/vertexai";

/**
* Options for initializing the Traceloop SDK.
Expand Down Expand Up @@ -56,7 +56,7 @@ export interface InitializeOptions {
instrumentModules?: {
openAI: typeof openai.OpenAI;
llamaIndex: typeof llamaIndex;
vertexAI: typeof vertexAI.VertexAI;
vertexAI: VertexAI;
};

/**
Expand Down

0 comments on commit d15d635

Please sign in to comment.