Skip to content

Commit

Permalink
Merge pull request #3637 from elizaOS/shaw/v2-refactor-2
Browse files Browse the repository at this point in the history
feat: Add `agent` table and rename `user` to `entity` table, add multi-tenancy
  • Loading branch information
lalalune authored Feb 25, 2025
2 parents 092929f + 34f9fe7 commit 679e3a4
Show file tree
Hide file tree
Showing 47 changed files with 2,285 additions and 740 deletions.
Binary file removed bun.lockb
Binary file not shown.
272 changes: 199 additions & 73 deletions packages/agent/src/swarm/scenario.ts
Original file line number Diff line number Diff line change
@@ -1,81 +1,207 @@
import { ChannelType, type IAgentRuntime, Memory, stringToUuid, type UUID } from "@elizaos/core";
import { v4 as uuidv4 } from 'uuid';

async function sayMessage(participant: IAgentRuntime, roomMap, participants: IAgentRuntime[], message: string) {
const participantId = participant.agentId;

// for each participant, create the memory
for (const participant of [participants[0]]) {
console.log("participant");
console.log(participant.agentId, participantId);
if(participant.agentId === participantId) {
continue;
}

const roomId = roomMap.get(participant.agentId);
await participant.ensureConnection({
userId: participantId,
roomId,
userName: participant.character.name,
userScreenName: participant.character.name,
source: "scenario",
type: ChannelType.GROUP,
import {
ChannelType,
Client,
HandlerCallback,
IAgentRuntime,
Memory,
UUID,
stringToUuid,
} from "@elizaos/core";
import { v4 as uuidv4 } from "uuid";

export class ScenarioClient implements Client {
name = "scenario";
runtime: IAgentRuntime;
private messageHandlers: Map<UUID, HandlerCallback[]> = new Map();
private rooms: Map<string, { roomId: UUID }> = new Map();

async start(runtime: IAgentRuntime) {
this.runtime = runtime;
return this;
}

async stop() {
this.messageHandlers.clear();
this.rooms.clear();
}

// Create a room for an agent
async createRoom(agentId: string, name?: string) {
const roomId = uuidv4();

await this.runtime.ensureRoomExists({
id: roomId as UUID,
name: name || `Room for ${agentId}`,
source: "scenario",
type: ChannelType.GROUP,
channelId: roomId,
serverId: null,
});

this.rooms.set(agentId, { roomId: roomId as UUID });
}

// Save a message in all agents' memory without emitting events
async saveMessage(
sender: IAgentRuntime,
receivers: IAgentRuntime[],
text: string
) {

for (const receiver of receivers) {
const participantId = stringToUuid(sender.agentId + "-" + receiver.agentId);
const roomData = this.rooms.get(receiver.agentId);
if (!roomData) continue;

// Ensure connection exists
await receiver.ensureConnection({
userId: participantId,
roomId: roomData.roomId,
userName: sender.character.name,
userScreenName: sender.character.name,
source: "scenario",
type: ChannelType.GROUP,
});

const memory: Memory = {
userId: participantId,
agentId: receiver.agentId,
roomId: roomData.roomId,
content: {
text,
source: "scenario",
name: sender.character.name,
userName: sender.character.name,
},
};

await receiver.messageManager.createMemory(memory);
}
}

// Send a live message that triggers handlers
async sendMessage(
sender: IAgentRuntime,
receivers: IAgentRuntime[],
text: string
) {

for (const receiver of receivers) {
const participantId = stringToUuid(sender.agentId + "-" + receiver.agentId);
const roomData = this.rooms.get(receiver.agentId);
if (!roomData) continue;

if (receiver.agentId !== sender.agentId) {
// Ensure connection exists
await receiver.ensureConnection({
userId: participantId,
roomId: roomData.roomId,
userName: sender.character.name,
userScreenName: sender.character.name,
source: "scenario",
type: ChannelType.GROUP,
});
} else {
await receiver.ensureConnection({
userId: sender.agentId,
roomId: roomData.roomId,
userName: sender.character.name,
userScreenName: sender.character.name,
source: "scenario",
type: ChannelType.GROUP,
});
const memoryManager = participant.messageManager;
const memory: Memory = {
userId: participantId,
agentId: participant.agentId,
roomId,
content: {
text: message,
}
}
await memoryManager.createMemory(memory);

// participant.emitEvent("MESSAGE_RECEIVED", {
// runtime: participant,
// message: memory,
// roomId: roomId,
// userId: participantId,
// serverId: roomId,
// channelId: roomId,
// source: "scenario",
// type: ChannelType.GROUP,
// });
}

const memory: Memory = {
userId: receiver.agentId !== sender.agentId ? participantId : sender.agentId,
agentId: receiver.agentId,
roomId: roomData.roomId,
content: {
text,
source: "scenario",
name: sender.character.name,
userName: sender.character.name,
},
};

receiver.emitEvent("MESSAGE_RECEIVED", {
runtime: receiver,
message: memory,
roomId: roomData.roomId,
userId: receiver.agentId !== sender.agentId ? participantId : sender.agentId,
source: "scenario",
type: ChannelType.GROUP,
});
}
}

// Get conversation history for all participants
async getConversations(participants: IAgentRuntime[]) {
const conversations = await Promise.all(
participants.map(async (member) => {
const roomData = this.rooms.get(member.agentId);
if (!roomData) return [];
return member.messageManager.getMemories({
roomId: roomData.roomId,
});
})
);

console.log("\nConversation logs per agent:");
conversations.forEach((convo, i) => {
console.log(`\n${participants[i].character.name}'s perspective:`);
convo.forEach((msg) =>
console.log(`${msg.content.name}: ${msg.content.text}`)
);
});

return conversations;
}
}

// Updated scenario implementation using the new client
const scenarios = [
async function scenario1(members: IAgentRuntime[]) {
// create a map of member agentId to room UUID
const roomMap = new Map<string, UUID>();
for (const member of members) {
const roomId = uuidv4() as UUID;
roomMap.set(member.agentId, roomId);
}

await sayMessage(members[0], roomMap, members, "Hello bob!");
await sayMessage(members[1], roomMap, members, "Hello alice!");
await sayMessage(members[0], roomMap, members, "Hello bob again!");
await sayMessage(members[1], roomMap, members, "Hello alice again!");
await sayMessage(members[2], roomMap, members, "I'm charlie!");
await sayMessage(members[0], roomMap, members, "Hello bob, how are you?");
await sayMessage(members[1], roomMap, members, "Hello alice, I'm good!");
await sayMessage(members[2], roomMap, members, "I'm good too!");

const conversations = await Promise.all(members.map(async (member) => {
return member.messageManager.getMemories({
roomId: roomMap.get(member.agentId),
});
}));
console.log(conversations);
},
];
async function scenario1(members: IAgentRuntime[]) {
// Create and register test client
const client = new ScenarioClient();
await client.start(members[0]);
members[0].registerClient("scenario", client);

export async function startScenario(members: IAgentRuntime[]) {
// TODO: Connect to mock messaging protocol
for (const scenario of scenarios) {
await scenario(members);
// Create rooms for all members
for (const member of members) {
await client.createRoom(
member.agentId,
`Test Room for ${member.character.name}`
);
}

}
// Set up conversation history
await client.saveMessage(
members[0],
members,
"Earlier message from conversation..."
);
// await client.saveMessage(
// members[1],
// members,
// "Previous reply in history..."
// );

// // Send live message that triggers handlers
// await client.sendMessage(members[0], members, "Hello everyone!");

// // Get and display conversation logs
// // wait 5 seconds
// await new Promise((resolve) => setTimeout(resolve, 5000));
// await client.getConversations(members);

// Send a message to all members
// await client.sendMessage(members[0], members, "Hello everyone!");
},
];

export async function startScenario(members: IAgentRuntime[]) {
for (const scenario of scenarios) {
await scenario(members);
}
}
2 changes: 1 addition & 1 deletion packages/agent/src/swarm/shared/onboarding/initialize.ts
Original file line number Diff line number Diff line change
Expand Up @@ -279,7 +279,7 @@ async function startOnboardingDM(
serverId: guild.id
});

await runtime.ensureUserExists(
await runtime.getOrCreateUser(
runtime.agentId,
runtime.character.name,
runtime.character.name,
Expand Down
31 changes: 19 additions & 12 deletions packages/core/__tests__/database.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@ import { DatabaseAdapter } from "../src/database.ts";
import {
type Memory,
type Actor,
type Account,
type Entity,
type Goal,
GoalStatus,
type Participant,
Expand Down Expand Up @@ -120,6 +120,7 @@ class MockDatabaseAdapter extends DatabaseAdapter {
setParticipantUserState(
_roomId: UUID,
_userId: UUID,
_agentId: UUID,
_state: "FOLLOWED" | "MUTED" | null
): Promise<void> {
throw new Error("Method not implemented.");
Expand Down Expand Up @@ -199,16 +200,19 @@ class MockDatabaseAdapter extends DatabaseAdapter {
}

// Mock method for getting account by ID
async getAccountById(userId: UUID): Promise<Account | null> {
async getEntityById(userId: UUID): Promise<Entity | null> {
return {
id: userId,
username: "testuser",
name: "Test Account",
} as Account;
metadata: {
username: "testuser",
name: "Test Entity",
},
agentId: "agent-id" as UUID,
} as Entity;
}

// Other methods stay the same...
async createAccount(_account: Account): Promise<boolean> {
async createEntity(_account: Entity): Promise<boolean> {
return true;
}

Expand Down Expand Up @@ -307,18 +311,21 @@ describe("DatabaseAdapter Tests", () => {
});

it("should get an account by user ID", async () => {
const account = await adapter.getAccountById("test-user-id" as UUID);
const account = await adapter.getEntityById("test-user-id" as UUID);
expect(account).not.toBeNull();
expect(account?.username).toBe("testuser");
expect(account?.metadata?.username).toBe("testuser");
});

it("should create a new account", async () => {
const newAccount: Account = {
const newAccount: Entity = {
id: "new-user-id" as UUID,
username: "newuser",
name: "New Account",
metadata: {
username: "newuser",
name: "New Entity",
},
agentId: "agent-id" as UUID,
};
const result = await adapter.createAccount(newAccount);
const result = await adapter.createEntity(newAccount);
expect(result).toBe(true);
});

Expand Down
8 changes: 4 additions & 4 deletions packages/core/__tests__/messages.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,7 @@ describe("Messages Library", () => {
databaseAdapter: {
// Using vi.fn() instead of jest.fn()
getParticipantsForRoom: vi.fn(),
getAccountById: vi.fn(),
getEntityById: vi.fn(),
},
} as unknown as IAgentRuntime;

Expand All @@ -40,7 +40,7 @@ describe("Messages Library", () => {
vi.mocked(
runtime.databaseAdapter.getParticipantsForRoom
).mockResolvedValue([userId]);
vi.mocked(runtime.databaseAdapter.getAccountById).mockResolvedValue({
vi.mocked(runtime.databaseAdapter.getEntityById).mockResolvedValue({
id: userId,
name: "Test User",
username: "testuser",
Expand Down Expand Up @@ -191,7 +191,7 @@ describe("Messages", () => {
mockActors[0].id,
mockActors[1].id,
]),
getAccountById: vi.fn().mockImplementation((id) => {
getEntityById: vi.fn().mockImplementation((id) => {
const actor = mockActors.find((a) => a.id === id);
return Promise.resolve(actor);
}),
Expand All @@ -218,7 +218,7 @@ describe("Messages", () => {
getParticipantsForRoom: vi
.fn()
.mockResolvedValue([mockActors[0].id, invalidId]),
getAccountById: vi.fn().mockImplementation((id) => {
getEntityById: vi.fn().mockImplementation((id) => {
const actor = mockActors.find((a) => a.id === id);
return Promise.resolve(actor || null);
}),
Expand Down
Loading

0 comments on commit 679e3a4

Please sign in to comment.