Skip to content

Commit

Permalink
feat(instrumentation): add fw metrics
Browse files Browse the repository at this point in the history
Signed-off-by: GALLLASMILAN <[email protected]>
  • Loading branch information
GALLLASMILAN committed Nov 20, 2024
1 parent d7c4060 commit 89064aa
Show file tree
Hide file tree
Showing 14 changed files with 458 additions and 98 deletions.
4 changes: 4 additions & 0 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -152,6 +152,10 @@ This project and everyone participating in it are governed by the [Code of Condu

All content in these repositories including code has been provided by IBM under the associated open source software license and IBM is under no obligation to provide enhancements, updates, or support. IBM developers produced this code as an open source project (not as an IBM product), and IBM makes no assertions as to the level of quality nor security, and will not be maintaining this code going forward.

## Telemetry

Some metrics are collected by default. See the [Native Telemetry](./docs/native-telemetry.md) for more detail.

## Contributors

Special thanks to our contributors for helping us improve Bee Agent Framework.
Expand Down
26 changes: 26 additions & 0 deletions docs/native-telemetry.md
Original file line number Diff line number Diff line change
@@ -0,0 +1,26 @@
# Native Telemetry in Bee Agent Framework

The Bee Agent Framework comes with built-in telemetry capabilities to help users monitor and optimize their applications. This document provides an overview of the telemetry feature, including what data is collected and how to disable it if necessary.

## Overview

The telemetry functionality in the Bee Agent Framework collects performance metrics and operational data to provide insights into how the framework operates in real-world environments. This feature helps us:

- Identify performance bottlenecks.
- Improve framework stability and reliability.
- Enhance user experience by understanding usage patterns.

## Data Collected

We value your privacy and ensure that **no sensitive data** is collected through telemetry. The following types of information are gathered:

- Framework version and runtime environment details.
- Anonymized usage statistics for built-in features.

## Disabling Telemetry

We understand that not all users want to send telemetry data. You can easily disable this feature by setting an environment variable:

```bash
BEE_FRAMEWORK_INSTRUMENTATION_METRICS_DISABLED=true
```
20 changes: 11 additions & 9 deletions package.json
Original file line number Diff line number Diff line change
Expand Up @@ -147,13 +147,13 @@
"lint:fix": "yarn eslint --fix",
"format": "yarn prettier --check .",
"format:fix": "yarn prettier --write .",
"test:unit": "vitest run src",
"test:unit:watch": "vitest run src",
"test:e2e": "vitest run tests/e2e",
"test:e2e:watch": "vitest watch tests/e2e",
"test:unit": "BEE_FRAMEWORK_INSTRUMENTATION_METRICS_DISABLED=true vitest run src",
"test:unit:watch": "BEE_FRAMEWORK_INSTRUMENTATION_METRICS_DISABLED=true vitest run src",
"test:e2e": "BEE_FRAMEWORK_INSTRUMENTATION_METRICS_DISABLED=true vitest run tests/e2e",
"test:e2e:watch": "BEE_FRAMEWORK_INSTRUMENTATION_METRICS_DISABLED=true vitest watch tests/e2e",
"test:all": "yarn test:unit && yarn test:e2e && yarn test:examples",
"test:examples": "vitest --config ./examples/vitest.examples.config.ts run tests/examples",
"test:examples:watch": "vitest --config ./examples/vitest.examples.config.ts watch tests/examples",
"test:examples": "BEE_FRAMEWORK_INSTRUMENTATION_METRICS_DISABLED=true vitest --config ./examples/vitest.examples.config.ts run tests/examples",
"test:examples:watch": "BEE_FRAMEWORK_INSTRUMENTATION_METRICS_DISABLED=true vitest --config ./examples/vitest.examples.config.ts watch tests/examples",
"prepare": "husky",
"copyright": "./scripts/copyright.sh",
"release": "release-it",
Expand All @@ -165,6 +165,11 @@
"@connectrpc/connect": "^1.6.1",
"@connectrpc/connect-node": "^1.6.1",
"@opentelemetry/api": "^1.9.0",
"@opentelemetry/exporter-metrics-otlp-http": "^0.54.2",
"@opentelemetry/resources": "^1.27.0",
"@opentelemetry/sdk-metrics": "^1.27.0",
"@opentelemetry/sdk-node": "^0.54.2",
"@opentelemetry/semantic-conventions": "^1.27.0",
"@streamparser/json": "^0.0.21",
"ajv": "^8.17.1",
"ajv-formats": "^3.0.1",
Expand Down Expand Up @@ -262,10 +267,7 @@
"@langchain/community": "~0.3.12",
"@langchain/core": "~0.3.17",
"@opentelemetry/instrumentation": "^0.54.0",
"@opentelemetry/resources": "^1.27.0",
"@opentelemetry/sdk-node": "^0.54.0",
"@opentelemetry/sdk-trace-node": "^1.27.0",
"@opentelemetry/semantic-conventions": "^1.27.0",
"@release-it/conventional-changelog": "^8.0.2",
"@rollup/plugin-commonjs": "^28.0.1",
"@swc/core": "^1.7.36",
Expand Down
12 changes: 10 additions & 2 deletions src/agents/base.ts
Original file line number Diff line number Diff line change
Expand Up @@ -21,8 +21,12 @@ import { GetRunContext, RunContext } from "@/context.js";
import { Emitter } from "@/emitter/emitter.js";
import { BaseMemory } from "@/memory/base.js";
import { createTelemetryMiddleware } from "@/instrumentation/create-telemetry-middleware.js";
import { INSTRUMENTATION_ENABLED } from "@/instrumentation/config.js";
import {
INSTRUMENTATION_ENABLED,
INSTRUMENTATION_METRICS_ENABLED,
} from "@/instrumentation/config.js";
import { doNothing } from "remeda";
import { createTelemetryMetricsMiddleware } from "@/instrumentation/create-telemetry-metrics-middleware.js";

export class AgentError extends FrameworkError {}

Expand Down Expand Up @@ -65,7 +69,11 @@ export abstract class BaseAgent<
this.isRunning = false;
}
},
).middleware(INSTRUMENTATION_ENABLED ? createTelemetryMiddleware() : doNothing());
)
.middleware(INSTRUMENTATION_ENABLED ? createTelemetryMiddleware() : doNothing())
.middleware(
INSTRUMENTATION_METRICS_ENABLED ? createTelemetryMetricsMiddleware() : doNothing(),
);
}

protected abstract _run(
Expand Down
3 changes: 3 additions & 0 deletions src/instrumentation/config.ts
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,9 @@ import { parseEnv } from "@/internals/env.js";
import { z } from "zod";

export const INSTRUMENTATION_ENABLED = parseEnv.asBoolean("BEE_FRAMEWORK_INSTRUMENTATION_ENABLED");
export const INSTRUMENTATION_METRICS_ENABLED = !parseEnv.asBoolean(
"BEE_FRAMEWORK_INSTRUMENTATION_METRICS_DISABLED",
);

export const INSTRUMENTATION_IGNORED_KEYS = parseEnv(
"BEE_FRAMEWORK_INSTRUMENTATION_IGNORED_KEYS",
Expand Down
59 changes: 59 additions & 0 deletions src/instrumentation/create-telemetry-metrics-middleware.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,59 @@
/**
* Copyright 2024 IBM Corp.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/

import { GetRunContext, RunInstance } from "@/context.js";
import { FrameworkError } from "@/errors.js";
import { buildModuleUsageMetric, isMeasurementedInstance } from "./opentelemetry.js";
import { GenerateCallbacks } from "@/llms/base.js";
import { createFullPath } from "@/emitter/utils.js";
import { metricReader } from "./sdk.js";

export const activeTracesMap = new Map<string, string>();

/**
* This middleware collects the usage metrics from framework entities runs and sends them to the central collector
* @returns
*/
export function createTelemetryMetricsMiddleware() {
return (context: GetRunContext<RunInstance, unknown>) => {
const traceId = context.emitter?.trace?.id;
if (!traceId) {
throw new FrameworkError(`Fatal error. Missing traceId`, [], { context });
}
if (activeTracesMap.has(traceId)) {
return;
}
activeTracesMap.set(traceId, context.instance.constructor.name);

const { emitter } = context;
const basePath = createFullPath(emitter.namespace, "");

const startEventName: keyof GenerateCallbacks = `start`;
const finishEventName: keyof GenerateCallbacks = `finish`;

// collect module_usage metric for llm|tool|agent start event
emitter.match(
(event) => event.name === startEventName && isMeasurementedInstance(event.creator),
(_, meta) => buildModuleUsageMetric({ traceId, instance: meta.creator, eventId: meta.id }),
);

// send metrics to the public collector
emitter.match(
(event) => event.path === `${basePath}.run.${finishEventName}`,
async () => await metricReader.forceFlush(),
);
};
}
4 changes: 3 additions & 1 deletion src/instrumentation/create-telemetry-middleware.ts
Original file line number Diff line number Diff line change
Expand Up @@ -27,7 +27,7 @@ import { Version } from "@/version.js";
import { Role } from "@/llms/primitives/message.js";
import type { GetRunContext, RunInstance } from "@/context.js";
import type { GeneratedResponse, FrameworkSpan } from "./types.js";
import { activeTracesMap, buildTraceTree } from "./tracer.js";
import { buildTraceTree } from "./opentelemetry.js";
import { traceSerializer } from "./helpers/trace-serializer.js";
import { INSTRUMENTATION_IGNORED_KEYS } from "./config.js";
import { createFullPath } from "@/emitter/utils.js";
Expand All @@ -36,6 +36,8 @@ import { instrumentationLogger } from "./logger.js";
import { BaseAgent } from "@/agents/base.js";
import { assertLLMWithMessagesToPromptFn } from "./helpers/utils.js";

export const activeTracesMap = new Map<string, string>();

export function createTelemetryMiddleware() {
return (context: GetRunContext<RunInstance, unknown>) => {
if (!context.emitter?.trace?.id) {
Expand Down
108 changes: 108 additions & 0 deletions src/instrumentation/opentelemetry.spec.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,108 @@
/**
* Copyright 2024 IBM Corp.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/

import { BeeAgent } from "@/agents/bee/agent.js";
import { ElasticSearchTool } from "@/tools/database/elasticsearch.js";
import { SQLTool } from "@/tools/database/sql.js";
import { GoogleSearchTool } from "@/tools/search/googleSearch.js";
import { OpenMeteoTool } from "@/tools/weather/openMeteo.js";
import { expect, describe } from "vitest";
import { isMeasurementedInstance } from "./opentelemetry.js";
import { DuckDuckGoSearchTool } from "@/tools/search/duckDuckGoSearch.js";
import { WebCrawlerTool } from "@/tools/web/webCrawler.js";
import { ArXivTool } from "@/tools/arxiv.js";
import { CalculatorTool } from "@/tools/calculator.js";
import { LLMTool } from "@/tools/llm.js";
import { OllamaLLM } from "@/adapters/ollama/llm.js";
import { BAMLLM } from "@/adapters/bam/llm.js";
import { WatsonXLLM } from "@/adapters/watsonx/llm.js";
import { LLM } from "@/llms/llm.js";
import { Emitter } from "@/emitter/emitter.js";
import { GenerateCallbacks, LLMMeta, BaseLLMTokenizeOutput, AsyncStream } from "@/llms/base.js";
import { OllamaChatLLM } from "@/adapters/ollama/chat.js";
import { TokenMemory } from "@/memory/tokenMemory.js";
import { GraniteBeeAgent } from "@/agents/granite/agent.js";
import { SlidingCache } from "@/cache/slidingCache.js";

export class CustomLLM extends LLM<any> {
public readonly emitter = Emitter.root.child<GenerateCallbacks>({
namespace: ["bam", "llm"],
creator: this,
});
meta(): Promise<LLMMeta> {
throw new Error("Method not implemented.");
}
tokenize(): Promise<BaseLLMTokenizeOutput> {
throw new Error("Method not implemented.");
}
protected _generate(): Promise<any> {
throw new Error("Method not implemented.");
}
protected _stream(): AsyncStream<any, void> {
throw new Error("Method not implemented.");
}
}

const llm = new OllamaChatLLM({ modelId: "llama3.1" });
const memory = new TokenMemory({ llm });

describe("opentelemetry", () => {
describe("isMeasurementedInstance", () => {
it.each([
// tool
new OpenMeteoTool(),
new GoogleSearchTool({ apiKey: "xx", cseId: "xx", maxResults: 10 }),
new ElasticSearchTool({ connection: { cloud: { id: "" } } }),
new SQLTool({ connection: { dialect: "mariadb" }, provider: "mysql" }),
new DuckDuckGoSearchTool(),
new OpenMeteoTool(),
new WebCrawlerTool(),
new ArXivTool(),
new CalculatorTool(),
new LLMTool({ llm: new OllamaLLM({ modelId: "llama3.1" }) }),
// llm
new OllamaLLM({ modelId: "llama3.1" }),
new BAMLLM({ modelId: "llama3.1" }),
new WatsonXLLM({ modelId: "llama3.1", apiKey: "xx" }),
new CustomLLM("llama3.1"),
new OllamaChatLLM({ modelId: "llama3.1" }),
// agent
new BeeAgent({ llm, memory, tools: [] }),
new GraniteBeeAgent({ llm, memory, tools: [] }),
])("Should return true for '%s'", (value) => {
expect(isMeasurementedInstance(value)).toBeTruthy();
});

it.each([
null,
undefined,
"",
0,
"string",
{},
memory,
new SlidingCache({
size: 50,
}),
new Emitter({
namespace: ["app"],
creator: this,
}),
])("Should return false for '%s'", (value) => {
expect(isMeasurementedInstance(value)).toBeFalsy();
});
});
});
Original file line number Diff line number Diff line change
Expand Up @@ -17,10 +17,16 @@
import { Version } from "@/version.js";
import opentelemetry, { SpanStatusCode, TimeInput } from "@opentelemetry/api";
import { FrameworkSpan, GeneratedResponse } from "./types.js";
import { BaseAgent } from "@/agents/base.js";
import { Tool } from "@/tools/base.js";
import { BaseLLM } from "@/llms/base.js";
import os from "os";

export const tracer = opentelemetry.trace.getTracer("bee-agent-framework", Version);
const name = "bee-agent-framework";

export const activeTracesMap = new Map<string, string>();
export const meter = opentelemetry.metrics.getMeter(name, Version);
export const tracer = opentelemetry.trace.getTracer(name, Version);
const moduleUsageGauge = meter.createGauge("module_usage");

interface ComputeTreeProps {
prompt?: string | null;
Expand All @@ -40,6 +46,12 @@ interface BuildSpansForParentProps {
parentId: string | undefined;
}

interface BuildModuleUsageMetricProps {
instance: object;
traceId: string;
eventId: string;
}

function buildSpansForParent({ spans, parentId }: BuildSpansForParentProps) {
spans
.filter((fwSpan) => fwSpan.parent_id === parentId)
Expand Down Expand Up @@ -114,3 +126,26 @@ export function buildTraceTree({
},
);
}

export const isMeasurementedInstance = (instance: any) =>
Boolean(
instance &&
(instance instanceof BaseAgent || instance instanceof Tool || instance instanceof BaseLLM),
);

export function buildModuleUsageMetric({
instance,
traceId,
eventId,
}: BuildModuleUsageMetricProps) {
moduleUsageGauge.record(1, {
source: instance.constructor.name,
type: instance instanceof BaseAgent ? "agent" : instance instanceof Tool ? "tool" : "llm",
framework_version: Version,
os_type: os.type(),
os_release: os.release(),
os_arch: os.arch(),
trace_id: traceId,
event_id: eventId,
});
}
Loading

0 comments on commit 89064aa

Please sign in to comment.