feat(observe): observe and mlflow implementation (#2)

.gitignore

@@ -104,6 +104,7 @@ dist
 # vuepress v2.x temp and cache directory
 .temp
 .cache
+tmp
 
 # Docusaurus cache and generated files
 .docusaurus

README.md

@@ -6,3 +6,23 @@
 2. Install dependencies `npm ci`
 3. Create `.env` (from `.env.template`) and fill in missing values (if any).
 4. Start the example `npm run start` (it runs the ./src/agent.ts file).
+
+## Using observability
+
+> [!NOTE]
+>
+> Docker distribution with support for compose is required, the following are supported:
+>
+> - [Docker](https://www.docker.com/)
+> - [Rancher](https://www.rancher.com/) - macOS users may want to use VZ instead of QEMU
+
+Get full visibility of the agent's inner working via our observability stack.
+
+- The [MLFlow](https://mlflow.org/) is used as UI for observability.
+- The [Bee Observe](https://github.com/i-am-bee/bee-observe) is the main Open-source observability service for Bee Agent Framework.
+- The [Bee Observe Connector](https://github.com/i-am-bee/bee-observe-connector) is the observability connector for Bee Agent Framework
+
+### Steps
+
+1. Start all services related to Observe `npm run infra:start-observe`
+2. Start the agent using the observe and MLFlow `npm run start:observe` (it runs ./src/agent_observe.ts file). The output of the `curl` command is saved in the **./tmp/observe/trace.json** file

docker-compose.yml

@@ -0,0 +1,59 @@
+services:
+  mongo:
+    image: mongo:7.0.14
+    environment:
+      MONGO_INITDB_ROOT_USERNAME: mongo
+      MONGO_INITDB_ROOT_PASSWORD: mongo
+    healthcheck:
+      test: |
+        mongosh --quiet --eval 'db.getSiblingDB("bee-observe").getCollection("span") ? quit(0) : quit(1)' || exit 1
+      interval: 10s
+      timeout: 5s
+      retries: 5
+  redis:
+    image: redis:7
+    command: redis-server --save 20 1 --loglevel warning
+    healthcheck:
+      test: ["CMD-SHELL", "redis-cli ping | grep PONG"]
+      interval: 10s
+      timeout: 5s
+      retries: 5
+  mlflow:
+    image: bitnami/mlflow:2.14.1
+    ports:
+      - "8080:8080"
+    entrypoint:
+      [
+        "/bin/bash",
+        "-c",
+        "/entrypoint.sh && mlflow server --app-name basic-auth --host 0.0.0.0 --port 8080",
+      ]
+    security_opt:
+      - "label=disable"
+    volumes:
+      - ./infra/observe/entrypoint.sh:/entrypoint.sh:ro
+  observe_api_migration:
+    image: iambeeagent/bee-observe:0.0.3
+    entrypoint: "npx mikro-orm migration:up --config ./dist/mikro-orm.config.js"
+    env_file:
+      - ./infra/observe/.env.docker
+    environment:
+      - NODE_ENV=production
+    depends_on:
+      mongo:
+        condition: service_healthy
+      redis:
+        condition: service_started
+  observe_api:
+    image: iambeeagent/bee-observe:0.0.3
+    ports:
+      - "4002:3000"
+    env_file:
+      - ./infra/observe/.env.docker
+    healthcheck:
+      test: wget --no-verbose --tries=1 --spider http://0.0.0.0:3000/health || exit 1
+      interval: 10s
+      timeout: 5s
+      retries: 5
+    depends_on:
+      - observe_api_migration

infra/observe/.env.docker

@@ -0,0 +1,16 @@
+NODE_ENV=production
+PORT=3000
+AUTH_KEY=testing-api-key
+FASTIFY_BODY_LIMIT=10485760
+
+REDIS_URL=redis://redis:6379/0
+MONGODB_URL=mongodb://mongo:mongo@mongo:27017
+DATA_EXPIRATION_IN_DAYS=7
+
+MLFLOW_API_URL=http://mlflow:8080/
+MLFLOW_AUTHORIZATION=BASE_AUTH
+MLFLOW_USERNAME=admin
+MLFLOW_PASSWORD=password
+MLFLOW_DEFAULT_EXPERIMENT_ID=0
+MLFLOW_TRACE_DELETE_IN_BATCHES_CRON_PATTERN=0 */1 * * * *
+MLFLOW_TRACE_DELETE_IN_BATCHES_BATCH_SIZE=100

infra/observe/entrypoint.sh

@@ -0,0 +1,19 @@
+#!/bin/bash
+
+set -e
+
+# Create experiment directory and ensure it has the correct permissions
+mkdir -p /app/mlruns/0
+chmod 755 /app/mlruns/0
+
+# The same operation for the important meta.yaml file, where the experiment id is defined
+echo "artifact_location: mlflow-artifacts:/0
+creation_time: 1720092866890
+experiment_id: '0'
+last_update_time: 1720092866890
+lifecycle_stage: active
+name: Default
+" > /app/mlruns/0/meta.yaml
+chmod 755 /app/mlruns/0/meta.yaml
+
+exec "$@"

package.json

@@ -21,14 +21,20 @@
   },
   "scripts": {
     "start": "tsx ./src/agent.ts -p tsconfig.json",
+    "start:observe": "tsx ./src/agent_observe.ts -p tsconfig.json",
     "build": "rimraf dist && tsc",
     "lint": "eslint",
     "lint:fix": "eslint --fix",
     "format": "prettier --check .",
-    "format:fix": "prettier --write ."
+    "format:fix": "prettier --write .",
+    "infra:start-all": "docker compose up -d",
+    "infra:start-observe": "docker compose up -d redis mongo mlflow observe_api_migration observe_api ",
+    "infra:stop-all": "docker compose down",
+    "infra:clean-all": "docker compose down --volumes"
   },
   "dependencies": {
     "bee-agent-framework": "^0.0.27",
+    "bee-observe-connector": "^0.0.3",
     "dotenv": "^16.4.5",
     "ollama": "^0.5.9",
     "pino": "^9.4.0"

src/agent.ts

@@ -13,7 +13,7 @@ const llm = new OllamaChatLLM({
   parameters: {
     temperature: 0,
     repeat_penalty: 1,
-    num_predict: 2000
+    num_predict: 2000,
   },
 });
 
@@ -24,12 +24,12 @@ const agent = new BeeAgent({
 });
 
 const getPrompt = () => {
-  const fallback = `What is the current weather in Las Vegas?`
+  const fallback = `What is the current weather in Las Vegas?`;
   if (process.stdin.isTTY) {
-    return fallback
+    return fallback;
   }
-  return fs.readFileSync(process.stdin.fd).toString().trim() || fallback
-}
+  return fs.readFileSync(process.stdin.fd).toString().trim() || fallback;
+};
 
 try {
   const prompt = getPrompt();

src/agent_observe.ts

@@ -0,0 +1,79 @@
+import "dotenv/config.js";
+import { BeeAgent } from "bee-agent-framework/agents/bee/agent";
+import { FrameworkError } from "bee-agent-framework/errors";
+import { TokenMemory } from "bee-agent-framework/memory/tokenMemory";
+import { DuckDuckGoSearchTool } from "bee-agent-framework/tools/search/duckDuckGoSearch";
+import { OpenMeteoTool } from "bee-agent-framework/tools/weather/openMeteo";
+import { OllamaChatLLM } from "bee-agent-framework/adapters/ollama/chat";
+import * as fs from "node:fs";
+import * as process from "node:process";
+import { createObserveConnector, ObserveError } from "bee-observe-connector";
+import { beeObserveApiSetting } from "./helpers/observe.js";
+
+const llm = new OllamaChatLLM({
+  modelId: "llama3.1",
+  parameters: {
+    temperature: 0,
+    repeat_penalty: 1,
+    num_predict: 2000,
+  },
+});
+
+const agent = new BeeAgent({
+  llm,
+  memory: new TokenMemory({ llm }),
+  tools: [new DuckDuckGoSearchTool(), new OpenMeteoTool()],
+});
+
+const getPrompt = () => {
+  const fallback = `What is the current weather in Las Vegas?`;
+  if (process.stdin.isTTY) {
+    return fallback;
+  }
+  return fs.readFileSync(process.stdin.fd).toString().trim() || fallback;
+};
+
+try {
+  const prompt = getPrompt();
+  console.info(`User 👤 : ${prompt}`);
+
+  const response = await agent
+    .run(
+      { prompt },
+      {
+        execution: {
+          maxIterations: 8,
+          maxRetriesPerStep: 3,
+          totalMaxRetries: 10,
+        },
+      },
+    )
+    .middleware(
+      createObserveConnector({
+        api: beeObserveApiSetting,
+        cb: async (err, data) => {
+          if (err) {
+            console.error(`Agent 🤖 : `, ObserveError.ensure(err).explain());
+          } else {
+            const { id, response } = data?.result || {};
+            console.log(`Agent 🤖 : `, response?.text || "Invalid output");
+
+            // you can use `&include_mlflow_tree=true` as well to return all sent data to mlflow
+            console.log(
+              `Agent 🤖 : Call the Observe API via this curl command outside of this Interactive session and see the trace data in the "trace.json" file: \n\n`,
+              `curl -X GET "${beeObserveApiSetting.baseUrl}/trace/${id}?include_tree=true&include_mlflow=true" \\
+\t-H "x-bee-authorization: ${beeObserveApiSetting.apiAuthKey}" \\
+\t-H "Content-Type: application/json" \\
+\t-o tmp/observe/trace.json`,
+            );
+          }
+        },
+      }),
+    );
+
+  console.info(`Agent 🤖 : ${response.result.text}`);
+} catch (error) {
+  console.error(FrameworkError.ensure(error).dump());
+} finally {
+  process.exit(0);
+}

src/helpers/observe.ts

@@ -0,0 +1,4 @@
+export const beeObserveApiSetting = {
+  baseUrl: "http://127.0.0.1:4002",
+  apiAuthKey: "testing-api-key",
+};

tsconfig.json

@@ -1,7 +1,7 @@
 {
   "compilerOptions": {
     "target": "ES2022",
-    "module": "ES2022",
+    "module": "NodeNext",
     "rootDir": "src",
     "baseUrl": ".",
     "moduleResolution": "NodeNext",
@@ -18,5 +18,6 @@
     "strict": true,
     "skipLibCheck": true,
     "strictNullChecks": true
-  }
+  },
+  "exclude": ["node_modules", "dist", "eslint.config.js", "prettier.config.js"]
 }