feat: tracing framework

agent/src/index.ts

@@ -34,6 +34,7 @@ import {
     stringToUuid,
     validateCharacterConfig,
     parseBooleanFromText,
+    Instrumentation,
 } from "@elizaos/core";
 import { zgPlugin } from "@elizaos/plugin-0g";
 
@@ -50,7 +51,7 @@ import { artheraPlugin } from "@elizaos/plugin-arthera";
 import { autonomePlugin } from "@elizaos/plugin-autonome";
 import { availPlugin } from "@elizaos/plugin-avail";
 import { avalanchePlugin } from "@elizaos/plugin-avalanche";
-import { b2Plugin } from "@elizaos/plugin-b2";
+// import { b2Plugin } from "@elizaos/plugin-b2";
 import { binancePlugin } from "@elizaos/plugin-binance";
 import { birdeyePlugin } from "@elizaos/plugin-birdeye";
 import {
@@ -133,6 +134,7 @@ const logFetch = async (url: string, options: any) => {
     elizaLogger.debug(`Fetching ${url}`);
     // Disabled to avoid disclosure of sensitive information such as API keys
     // elizaLogger.debug(JSON.stringify(options, null, 2));
+    Instrumentation.trace("fetch", {url, options});
     return fetch(url, options);
 };
 
@@ -1047,6 +1049,7 @@ async function startAgent(
     directClient: DirectClient
 ): Promise<AgentRuntime> {
     let db: IDatabaseAdapter & IDatabaseCacheAdapter;
+    Instrumentation.trace("startAgent", {character});
     try {
         character.id ??= stringToUuid(character.name);
         character.username ??= character.name;
@@ -1168,9 +1171,15 @@ const startAgents = async () => {
     );
 };
 
-startAgents().catch((error) => {
-    elizaLogger.error("Unhandled error in startAgents:", error);
-    process.exit(1);
+const db = new PostgresDatabaseAdapter({connectionString: process.env.POSTGRES_URL, parseInputs: true});
+await db.init();
+Instrumentation.init((run, time, name, data) => {db.createTrace(run, time, name, data)});
+Instrumentation.run(() => {
+    Instrumentation.trace("TEST", {a: 123});
+    startAgents().catch((error) => {
+        elizaLogger.error("Unhandled error in startAgents:", error);
+        process.exit(1);
+    });
 });
 
 // Prevent unhandled exceptions from crashing the process if desired

packages/adapter-postgres/schema.sql

@@ -131,6 +131,15 @@ CREATE TABLE IF NOT EXISTS  cache (
     PRIMARY KEY ("key", "agentId")
 );
 
+CREATE TABLE IF NOT EXISTS traces (
+    "id" BIGSERIAL NOT NULL,
+    "run" UUID,
+    "time" TIMESTAMP NOT NULL,
+    "name" VARCHAR(80) NOT NULL,
+    "data" JSON,
+    PRIMARY KEY ("id")
+);
+
 DO $$
 DECLARE
     vector_dim INTEGER;
@@ -164,5 +173,8 @@ CREATE INDEX IF NOT EXISTS idx_knowledge_original ON knowledge("originalId");
 CREATE INDEX IF NOT EXISTS idx_knowledge_created ON knowledge("agentId", "createdAt");
 CREATE INDEX IF NOT EXISTS idx_knowledge_shared ON knowledge("isShared");
 CREATE INDEX IF NOT EXISTS idx_knowledge_embedding ON knowledge USING ivfflat (embedding vector_cosine_ops);
+CREATE INDEX IF NOT EXISTS idx_traces_1 ON traces ("run", "name", "time");
+CREATE INDEX IF NOT EXISTS idx_traces_2 ON traces ("name", "run", "time");
+CREATE INDEX IF NOT EXISTS idx_traces_3 ON traces ("time", "name");
 
 COMMIT;

packages/adapter-postgres/src/index.ts

@@ -1810,6 +1810,23 @@ export class PostgresDatabaseAdapter
             ]
         );
     }
+
+    async createTrace(run: string, time: Date, name: string, data: any): Promise<void> {
+        return this.withDatabase(async () => {
+            const client = await this.pool.connect();
+            try {
+                await client.query("BEGIN");
+                await client.query(`INSERT INTO traces ("run", "time", "name", "data") VALUES ($1, $2, $3, $4)`,
+                    [run, time.toISOString(), name, data]);
+                await client.query("COMMIT");
+            } catch (error) {
+                await client.query("ROLLBACK");
+                throw error;
+            } finally {
+                client.release();
+            }
+        }, "createTrace");
+    }
 }
 
 export default PostgresDatabaseAdapter;

packages/core/src/generation.ts

@@ -18,7 +18,7 @@ import { encodingForModel, TiktokenModel } from "js-tiktoken";
 import { AutoTokenizer } from "@huggingface/transformers";
 import Together from "together-ai";
 import { ZodSchema } from "zod";
-import { elizaLogger } from "./index.ts";
+import { elizaLogger, Instrumentation } from "./index.ts";
 import {
     models,
     getModelSettings,
@@ -258,6 +258,7 @@ export async function generateText({
         verifiableInference,
     });
     elizaLogger.log("Using provider:", runtime.modelProvider);
+    Instrumentation.trace("generateText", {modelClass, context, runtime});
     // If verifiable inference is requested and adapter is provided, use it
     if (verifiableInference && runtime.verifiableInferenceAdapter) {
         elizaLogger.log(
@@ -1003,6 +1004,8 @@ export async function generateText({
             }
         }
 
+        Instrumentation.trace("response", {value: response});
+
         return response;
     } catch (error) {
         elizaLogger.error("Error in generateText:", error);
@@ -1298,6 +1301,8 @@ export async function generateMessageResponse({
         try {
             elizaLogger.log("Generating message response..");
 
+            Instrumentation.trace("generateMessageResponse", {runtime, context});
+
             const response = await generateText({
                 runtime,
                 context,

packages/core/src/index.ts

@@ -25,3 +25,4 @@ export * from "./cache.ts";
 export { default as knowledge } from "./knowledge.ts";
 export * from "./ragknowledge.ts";
 export * from "./utils.ts";
+export * from "./instrumentation.ts";

packages/core/src/instrumentation.ts

@@ -0,0 +1,46 @@
+import { AsyncLocalStorage } from "async_hooks";
+import { inspect } from "util";
+import { v4 as uuidv4 } from "uuid";
+
+export type ITrace = (run: string | null, time: Date, name: string, data: any) => void;
+
+export function generateRunUUID(): string {
+    return uuidv4();
+}
+
+export class Instrumentation {
+    private static store = new AsyncLocalStorage<Instrumentation>();
+    private static tracer: ITrace = null;
+    private run: string = null;
+
+    static init(tracer: ITrace) {
+        this.tracer = tracer;
+    }
+
+    static trace(name: string, data: any) {
+        if (!this.tracer) return;
+        var jsonPayload;
+        try {
+            jsonPayload = JSON.stringify(data);
+        } catch (error) {
+            name = `ERROR: ${name}`;
+            jsonPayload = {ok: false, name: name, error: `${error}, inspect: ${inspect(data)}`};
+            console.log(`Could not convert object to JSON for ${name}. Please select individual fields that are serializable.`);
+            console.log(data);
+        }
+        const instance = this.store.getStore();
+        const run = instance.run;
+        const time = new Date();
+        (async () => {
+            await this.tracer(run, time, name, jsonPayload);
+        })();
+    }
+
+    static run(f: () => any, run?: string | null): any {
+        const i = new Instrumentation();
+        if (run === undefined) run = null;
+        if (!run) run = uuidv4();
+        i.run = run;
+        return this.store.run(i, f);
+    }
+}

packages/core/src/memory.ts

@@ -1,4 +1,5 @@
 import { embed, getEmbeddingZeroVector } from "./embedding.ts";
+import { Instrumentation } from "./instrumentation.ts";
 import elizaLogger from "./logger.ts";
 import {
     IAgentRuntime,
@@ -97,7 +98,7 @@ export class MemoryManager implements IMemoryManager {
         start?: number;
         end?: number;
     }): Promise<Memory[]> {
-        return await this.runtime.databaseAdapter.getMemories({
+        const result = await this.runtime.databaseAdapter.getMemories({
             roomId,
             count,
             unique,
@@ -106,6 +107,8 @@ export class MemoryManager implements IMemoryManager {
             start,
             end,
         });
+        Instrumentation.trace("MemoryManager.getMemories", {roomId: roomId, count: count, result: result});
+        return result;
     }
 
     async getCachedEmbeddings(content: string): Promise<
@@ -187,6 +190,7 @@ export class MemoryManager implements IMemoryManager {
             this.tableName,
             unique
         );
+        Instrumentation.trace("MemoryManager.createMemory", {memoryId: memory.id, text: memory.content.text})
     }
 
     async getMemoriesByRoomIds(params: { roomIds: UUID[], limit?: number; }): Promise<Memory[]> {

packages/core/src/runtime.ts

@@ -17,7 +17,7 @@ import {
 } from "./evaluators.ts";
 import { generateText } from "./generation.ts";
 import { formatGoalsAsString, getGoals } from "./goals.ts";
-import { elizaLogger } from "./index.ts";
+import { elizaLogger, Instrumentation } from "./index.ts";
 import knowledge from "./knowledge.ts";
 import { MemoryManager } from "./memory.ts";
 import { formatActors, formatMessages, getActorDetails } from "./messages.ts";
@@ -410,6 +410,8 @@ export class AgentRuntime implements IAgentRuntime {
     }
 
     async initialize() {
+        Instrumentation.trace("AgentRuntime.initialize()", {});
+
         for (const [serviceType, service] of this.services.entries()) {
             try {
                 await service.initialize(this);
@@ -858,6 +860,8 @@ export class AgentRuntime implements IAgentRuntime {
             verifiableInferenceAdapter: this.verifiableInferenceAdapter,
         });
 
+        Instrumentation.trace("evaluate", {context: context, result: result});
+
         const evaluators = parseJsonArrayFromText(
             result
         ) as unknown as string[];