Skip to content

Commit

Permalink
fix(pinecone): handle exceptions
Browse files Browse the repository at this point in the history
  • Loading branch information
@nirga
nirga committed Apr 22, 2024
1 parent da05d15 commit 86b8a1b
Show file tree
Hide file tree
Showing 3 changed files with 108 additions and 82 deletions.
174 changes: 95 additions & 79 deletions packages/instrumentation-pinecone/src/instrumentation.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,6 @@ import type * as pinecone from "@pinecone-database/pinecone";
import { context, trace, Tracer, SpanStatusCode } from "@opentelemetry/api";
import {
InstrumentationBase,
InstrumentationConfig,
InstrumentationModuleDefinition,
InstrumentationNodeModuleDefinition,
safeExecuteInTheMiddle,
Expand All @@ -28,12 +27,17 @@ import {
EventAttributes,
} from "@traceloop/ai-semantic-conventions";
import { version } from "../package.json";
import { PineconeInstrumentationConfig } from "./types";

export class PineconeInstrumentation extends InstrumentationBase<any> {
constructor(config: InstrumentationConfig = {}) {
constructor(config: PineconeInstrumentationConfig = {}) {
super("@traceloop/instrumentation-pinecone", version, config);
}

public override setConfig(config: PineconeInstrumentationConfig = {}) {
super.setConfig(config);
}

public manuallyInstrument(module: typeof pinecone) {
this.patch(module);
}
Expand Down Expand Up @@ -145,33 +149,38 @@ export class PineconeInstrumentation extends InstrumentationBase<any> {
return function method(this: any, ...args: unknown[]) {
const span = tracer.startSpan(`pinecone.query`);
const execContext = trace.setSpan(context.active(), span);
const options = args[0] as pinecone.QueryOptions;
span.setAttribute(SpanAttributes.VECTOR_DB_VENDOR, "Pinecone");
const query_request_event = span.addEvent("pinecone.query.request");
query_request_event.setAttribute(
EventAttributes.VECTOR_DB_QUERY_TOP_K,
options.topK,
);
query_request_event.setAttribute(
EventAttributes.VECTOR_DB_QUERY_INCLUDE_VALUES,
options.includeValues || false,
);
query_request_event.setAttribute(
EventAttributes.VECTOR_DB_QUERY_INCLUDE_METADATA,
options.includeMetadata || false,
);
query_request_event.setAttribute(
EventAttributes.VECTOR_DB_QUERY_ID,
(options as pinecone.QueryByRecordId).id,
);
query_request_event.setAttribute(
EventAttributes.VECTOR_DB_QUERY_EMBEDDINGS_VECTOR,
(options as pinecone.QueryByVectorValues).vector,
);
query_request_event.setAttribute(
EventAttributes.VECTOR_DB_QUERY_METADATA_FILTER,
JSON.stringify(options.filter ? options.filter : {}),
);
try {
const options = args[0] as pinecone.QueryOptions;
span.setAttribute(SpanAttributes.VECTOR_DB_VENDOR, "Pinecone");
const query_request_event = span.addEvent("pinecone.query.request");
query_request_event.setAttribute(
EventAttributes.VECTOR_DB_QUERY_TOP_K,
options.topK,
);
query_request_event.setAttribute(
EventAttributes.VECTOR_DB_QUERY_INCLUDE_VALUES,
options.includeValues || false,
);
query_request_event.setAttribute(
EventAttributes.VECTOR_DB_QUERY_INCLUDE_METADATA,
options.includeMetadata || false,
);
query_request_event.setAttribute(
EventAttributes.VECTOR_DB_QUERY_ID,
(options as pinecone.QueryByRecordId).id,
);
query_request_event.setAttribute(
EventAttributes.VECTOR_DB_QUERY_EMBEDDINGS_VECTOR,
(options as pinecone.QueryByVectorValues).vector,
);
query_request_event.setAttribute(
EventAttributes.VECTOR_DB_QUERY_METADATA_FILTER,
JSON.stringify(options.filter ? options.filter : {}),
);
} catch (e) {
this._diag.warn(e);
this._config.exceptionLogger?.(e);
}

const execPromise = safeExecuteInTheMiddle(
() => {
Expand All @@ -189,71 +198,78 @@ export class PineconeInstrumentation extends InstrumentationBase<any> {
.then((result: any) => {
return new Promise((resolve) => {
span.setStatus({ code: SpanStatusCode.OK });
const result_obj =
result as pinecone.QueryResponse<pinecone.RecordMetadata>;
const query_result_event = span.addEvent("pinecone.query.result");
query_result_event.setAttribute(
EventAttributes.VECTOR_DB_QUERY_RESULT_NAMESPACE,
result_obj.namespace,
);
if (result_obj.usage?.readUnits !== undefined) {
query_result_event.setAttribute(
EventAttributes.VECTOR_DB_QUERY_RESULT_READ_UNITS_CONSUMED,
result_obj.usage?.readUnits,
try {
const result_obj =
result as pinecone.QueryResponse<pinecone.RecordMetadata>;
const query_result_event = span.addEvent(
"pinecone.query.result",
);
}
query_result_event.setAttribute(
EventAttributes.VECTOR_DB_QUERY_RESULT_MATCHES_LENGTH,
result_obj.matches.length,
);
for (let i = 0; i < result_obj.matches.length; i++) {
const match = result_obj.matches[i];
const query_result_match_event = query_result_event.addEvent(
`pinecone.query.result.${i}`,
query_result_event.setAttribute(
EventAttributes.VECTOR_DB_QUERY_RESULT_NAMESPACE,
result_obj.namespace,
);
if (match.score !== undefined) {
query_result_match_event.setAttribute(
EventAttributes.VECTOR_DB_QUERY_RESULT_SCORE.replace(
"{i}",
i.toString(),
),
match.score,
if (result_obj.usage?.readUnits !== undefined) {
query_result_event.setAttribute(
EventAttributes.VECTOR_DB_QUERY_RESULT_READ_UNITS_CONSUMED,
result_obj.usage?.readUnits,
);
}
if (match.sparseValues !== undefined) {
query_result_event.setAttribute(
EventAttributes.VECTOR_DB_QUERY_RESULT_MATCHES_LENGTH,
result_obj.matches.length,
);
for (let i = 0; i < result_obj.matches.length; i++) {
const match = result_obj.matches[i];
const query_result_match_event = query_result_event.addEvent(
`pinecone.query.result.${i}`,
);
if (match.score !== undefined) {
query_result_match_event.setAttribute(
EventAttributes.VECTOR_DB_QUERY_RESULT_SCORE.replace(
"{i}",
i.toString(),
),
match.score,
);
}
if (match.sparseValues !== undefined) {
query_result_match_event.setAttribute(
EventAttributes.VECTOR_DB_QUERY_RESULT_SPARSE_INDICES.replace(
"{i}",
i.toString(),
),
match.sparseValues?.indices,
);
query_result_match_event.setAttribute(
EventAttributes.VECTOR_DB_QUERY_RESULT_SPARSE_VALUES.replace(
"{i}",
i.toString(),
),
match.sparseValues?.values,
);
}
query_result_match_event.setAttribute(
EventAttributes.VECTOR_DB_QUERY_RESULT_SPARSE_INDICES.replace(
EventAttributes.VECTOR_DB_QUERY_RESULT_ID.replace(
"{i}",
i.toString(),
),
match.sparseValues?.indices,
match.id,
);
query_result_match_event.setAttribute(
EventAttributes.VECTOR_DB_QUERY_RESULT_SPARSE_VALUES.replace(
EventAttributes.VECTOR_DB_QUERY_RESULT_VALUES.replace(
"{i}",
i.toString(),
),
match.sparseValues?.values,
match.values,
);
query_result_match_event.addEvent(
`pinecone.query.result.${i}.metadata`,
match.metadata,
);
}
query_result_match_event.setAttribute(
EventAttributes.VECTOR_DB_QUERY_RESULT_ID.replace(
"{i}",
i.toString(),
),
match.id,
);
query_result_match_event.setAttribute(
EventAttributes.VECTOR_DB_QUERY_RESULT_VALUES.replace(
"{i}",
i.toString(),
),
match.values,
);
query_result_match_event.addEvent(
`pinecone.query.result.${i}.metadata`,
match.metadata,
);
} catch (e) {
this._diag.warn(e);
this._config.exceptionLogger?.(e);
}
span.end();
resolve(result);
Expand Down
8 changes: 8 additions & 0 deletions packages/instrumentation-pinecone/src/types.ts
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,8 @@
import { InstrumentationConfig } from "@opentelemetry/instrumentation";

export interface PineconeInstrumentationConfig extends InstrumentationConfig {
/**
* A custom logger to log any exceptions that happen during span creation.
*/
exceptionLogger?: (e: Error) => void;
}
8 changes: 5 additions & 3 deletions packages/traceloop-sdk/src/lib/tracing/index.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -78,7 +78,7 @@ export const initInstrumentations = () => {
bedrockInstrumentation = new BedrockInstrumentation({ exceptionLogger });
instrumentations.push(bedrockInstrumentation);

pineconeInstrumentation = new PineconeInstrumentation();
pineconeInstrumentation = new PineconeInstrumentation({ exceptionLogger });
instrumentations.push(pineconeInstrumentation);

langchainInstrumentation = new LangChainInstrumentation({ exceptionLogger });
Expand Down Expand Up @@ -155,7 +155,7 @@ export const manuallyInitInstrumentations = (
}

if (instrumentModules?.pinecone) {
const instrumentation = new PineconeInstrumentation();
const instrumentation = new PineconeInstrumentation({ exceptionLogger });
instrumentations.push(instrumentation as Instrumentation);
instrumentation.manuallyInstrument(instrumentModules.pinecone);
}
Expand All @@ -169,7 +169,9 @@ export const manuallyInitInstrumentations = (
}

if (instrumentModules?.llamaIndex) {
llamaIndexInstrumentation = new LlamaIndexInstrumentation();
llamaIndexInstrumentation = new LlamaIndexInstrumentation({
exceptionLogger,
});
instrumentations.push(llamaIndexInstrumentation);
llamaIndexInstrumentation.manuallyInstrument(instrumentModules.llamaIndex);
}
Expand Down

0 comments on commit 86b8a1b

Please sign in to comment.