Skip to content

Commit b5c0af9

Browse files
wtfsayowtfsayo
committed
update mongodb implementation and tests
1 parent 91a8e61 commit b5c0af9

File tree

9 files changed

+374
-332
lines changed

9 files changed

+374
-332
lines changed

‎.env.example

+2-2
Original file line numberDiff line numberDiff line change
@@ -16,8 +16,8 @@ SUPABASE_URL=
1616
SUPABASE_ANON_KEY=
1717

1818
# MongoDB
19-
MONGODB_CONNECTION_STRING= #mongodb connection string
20-
MONGODB_DATABASE= #name of the database in mongoDB atlas
19+
MONGODB_CONNECTION_STRING= #mongodb connection string
20+
MONGODB_DATABASE= #name of the database in mongoDB atlas #default: 'elizaAgent'
2121

2222
# Comma separated list of remote character urls (optional)
2323
REMOTE_CHARACTER_URLS=

‎.turbo/config.json

-3
This file was deleted.

‎agent/src/index.ts

+2-1
Original file line numberDiff line numberDiff line change
@@ -139,6 +139,7 @@ import { zerionPlugin } from "@elizaos/plugin-zerion"
139139
import { minaPlugin } from "@elizaos/plugin-mina"
140140
import { ankrPlugin } from "@elizaos/plugin-ankr";
141141
import { formPlugin } from "@elizaos/plugin-form";
142+
import { MongoClient } from "mongodb";
142143

143144
const __filename = fileURLToPath(import.meta.url) // get the resolved path to the file
144145
const __dirname = path.dirname(__filename) // get the name of the directory
@@ -526,7 +527,7 @@ function initializeDatabase(dataDir: string) {
526527
retryReads: true
527528
});
528529

529-
const dbName = process.env.MONGODB_DATABASE || 'CumulusAiAgent'; // Default database name
530+
const dbName = process.env.MONGODB_DATABASE || 'elizaAgent';
530531
const db = new MongoDBDatabaseAdapter(client, dbName);
531532

532533
// Test the connection

‎packages/adapter-mongodb/package.json

+18-8
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
{
22
"name": "@elizaos/adapter-mongodb",
3-
"version": "0.1.0",
3+
"version": "0.0.1",
4+
"description": "MongoDB adapter for ElizaOS",
45
"main": "dist/index.js",
56
"type": "module",
67
"types": "dist/index.d.ts",
@@ -10,18 +11,27 @@
1011
"uuid": "^9.0.1"
1112
},
1213
"devDependencies": {
13-
"@types/node": "^20.10.0",
14-
"@types/uuid": "^10.0.0",
15-
"tsup": "8.3.5",
16-
"typescript": "^5.0.0",
17-
"mongodb": "6.3.0"
14+
"@types/jest": "^29.5.11",
15+
"@types/node": "^20.11.5",
16+
"@types/uuid": "^9.0.7",
17+
"jest": "^29.7.0",
18+
"ts-jest": "^29.1.1",
19+
"tsup": "^8.0.1",
20+
"typescript": "^5.3.3"
1821
},
1922
"scripts": {
20-
"build": "tsup --format esm --dts",
23+
"build": "tsup",
2124
"dev": "tsup --format esm --dts --watch",
22-
"lint": "eslint --fix --cache ."
25+
"lint": "eslint --fix --cache .",
26+
"test": "cd src/__tests__ && ./run_tests.sh",
27+
"test:watch": "jest --watch"
2328
},
2429
"engines": {
2530
"node": ">=16.0.0"
31+
},
32+
"jest": {
33+
"preset": "ts-jest",
34+
"testEnvironment": "node",
35+
"testMatch": ["<rootDir>/src/__tests__/**/*.test.ts"]
2636
}
2737
}

‎packages/adapter-mongodb/src/__tests__/docker-compose.test.yml

+16
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,16 @@
1+
# yaml-language-server: $schema=https://raw.githubusercontent.com/compose-spec/compose-spec/master/schema/compose-spec.json
2+
version: '3.8'
3+
services:
4+
mongodb-test:
5+
image: mongo:latest
6+
environment:
7+
MONGO_INITDB_ROOT_USERNAME: mongodb
8+
MONGO_INITDB_ROOT_PASSWORD: mongodb
9+
MONGO_INITDB_DATABASE: eliza_test
10+
ports:
11+
- "27018:27017"
12+
healthcheck:
13+
test: ["CMD", "mongosh", "--eval", "db.adminCommand('ping')"]
14+
interval: 5s
15+
timeout: 5s
16+
retries: 5

‎packages/adapter-mongodb/src/__tests__/mongodb-adapter.test.ts

+134
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,134 @@
1+
import { MongoClient } from 'mongodb';
2+
import { MongoDBDatabaseAdapter } from '../index';
3+
import { v4 as uuidv4 } from 'uuid';
4+
import { RAGKnowledgeItem, type UUID } from '@elizaos/core';
5+
6+
describe('MongoDBAdapter', () => {
7+
let adapter: MongoDBDatabaseAdapter;
8+
const testUrl = 'mongodb://mongodb:mongodb@localhost:27018';
9+
const dbName = 'eliza_test';
10+
11+
beforeAll(async () => {
12+
adapter = new MongoDBDatabaseAdapter(new MongoClient(testUrl), dbName);
13+
await adapter.init();
14+
});
15+
16+
afterAll(async () => {
17+
await adapter.close();
18+
});
19+
20+
beforeEach(async () => {
21+
// Clear the collections before each test
22+
const client = new MongoClient(testUrl);
23+
await client.connect();
24+
const db = client.db(dbName);
25+
await db.collection('knowledge').deleteMany({});
26+
await client.close();
27+
});
28+
29+
describe('Knowledge Management', () => {
30+
it('should create and retrieve knowledge', async () => {
31+
const testKnowledge: RAGKnowledgeItem = {
32+
id: uuidv4() as UUID,
33+
agentId: uuidv4() as UUID,
34+
content: {
35+
text: 'Test knowledge content',
36+
metadata: {
37+
isShared: false,
38+
isMain: true
39+
}
40+
},
41+
embedding: new Float32Array([0.1, 0.2, 0.3]),
42+
createdAt: Date.now()
43+
};
44+
45+
await adapter.createKnowledge(testKnowledge);
46+
47+
const retrieved = await adapter.getKnowledge({ id: testKnowledge.id, agentId: testKnowledge.agentId });
48+
expect(retrieved).toHaveLength(1);
49+
expect(retrieved[0].id).toBe(testKnowledge.id);
50+
expect(retrieved[0].content.text).toBe(testKnowledge.content.text);
51+
});
52+
53+
it('should search knowledge by embedding', async () => {
54+
const testKnowledge1: RAGKnowledgeItem = {
55+
id: uuidv4() as UUID,
56+
agentId: uuidv4() as UUID,
57+
content: {
58+
text: 'First test knowledge',
59+
metadata: { isShared: false }
60+
},
61+
embedding: new Float32Array([0.1, 0.2, 0.3]),
62+
createdAt: Date.now()
63+
};
64+
65+
const testKnowledge2: RAGKnowledgeItem = {
66+
id: uuidv4() as UUID,
67+
agentId: uuidv4() as UUID,
68+
content: {
69+
text: 'Second test knowledge',
70+
metadata: { isShared: false }
71+
},
72+
embedding: new Float32Array([0.4, 0.5, 0.6]),
73+
createdAt: Date.now()
74+
};
75+
76+
await adapter.createKnowledge(testKnowledge1);
77+
await adapter.createKnowledge(testKnowledge2);
78+
79+
const searchResults = await adapter.searchKnowledge({
80+
agentId: testKnowledge1.agentId,
81+
embedding: new Float32Array([0.1, 0.2, 0.3]),
82+
match_count: 1,
83+
match_threshold: 0.8
84+
});
85+
86+
expect(searchResults).toHaveLength(1);
87+
expect(searchResults[0].id).toBe(testKnowledge1.id);
88+
});
89+
90+
it('should remove knowledge', async () => {
91+
const testKnowledge: RAGKnowledgeItem = {
92+
id: uuidv4() as UUID,
93+
agentId: uuidv4() as UUID,
94+
content: {
95+
text: 'Test knowledge to remove',
96+
metadata: { isShared: false }
97+
},
98+
embedding: new Float32Array([0.1, 0.2, 0.3]),
99+
createdAt: Date.now()
100+
};
101+
102+
await adapter.createKnowledge(testKnowledge);
103+
await adapter.removeKnowledge(testKnowledge.id);
104+
105+
const retrieved = await adapter.getKnowledge({ id: testKnowledge.id, agentId: testKnowledge.agentId });
106+
expect(retrieved).toHaveLength(0);
107+
});
108+
});
109+
110+
describe('Cache Operations', () => {
111+
it('should set and get cache', async () => {
112+
const agentId = uuidv4();
113+
const key = 'test-key';
114+
const value = 'test-value';
115+
116+
await adapter.setCache({ key, agentId: agentId as UUID, value });
117+
const retrieved = await adapter.getCache({ key, agentId: agentId as UUID });
118+
119+
expect(retrieved).toBe(value);
120+
});
121+
122+
it('should delete cache', async () => {
123+
const agentId = uuidv4();
124+
const key = 'test-key';
125+
const value = 'test-value';
126+
127+
await adapter.setCache({ key, agentId: agentId as UUID, value });
128+
await adapter.deleteCache({ key, agentId: agentId as UUID });
129+
130+
const retrieved = await adapter.getCache({ key, agentId: agentId as UUID });
131+
expect(retrieved).toBeUndefined();
132+
});
133+
});
134+
});

‎packages/adapter-mongodb/src/__tests__/run_tests.sh

+31
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,31 @@
1+
#!/bin/bash
2+
3+
# Function to cleanup resources
4+
cleanup() {
5+
echo "Cleaning up resources..."
6+
docker-compose -f docker-compose.test.yml down
7+
exit 0
8+
}
9+
10+
# Trap SIGINT and SIGTERM signals and cleanup
11+
trap cleanup SIGINT SIGTERM
12+
13+
# Start MongoDB container
14+
echo "Starting MongoDB container..."
15+
docker-compose -f docker-compose.test.yml up -d
16+
17+
# Wait for MongoDB to be ready
18+
echo "Waiting for MongoDB to be ready..."
19+
until docker-compose -f docker-compose.test.yml exec -T mongodb-test mongosh --eval "db.adminCommand('ping')" > /dev/null 2>&1; do
20+
echo "MongoDB is not ready yet..."
21+
sleep 1
22+
done
23+
24+
echo "MongoDB is ready!"
25+
26+
# Run tests
27+
echo "Running tests..."
28+
jest --runInBand --forceExit
29+
30+
# Cleanup after tests
31+
cleanup

‎packages/adapter-mongodb/src/index.ts

+1
Original file line numberDiff line numberDiff line change
@@ -44,6 +44,7 @@ export class MongoDBDatabaseAdapter
4444
private hasVectorSearch: boolean;
4545
private isConnected: boolean = false;
4646
private isVectorSearchIndexComputable: boolean;
47+
public db: MongoClient;
4748

4849
constructor(client: MongoClient, databaseName: string) {
4950
super();

0 commit comments

Comments
 (0)