Skip to content

Commit 7fbab5f

Browse files
lalalunelalalune
committed
ensure connection to replace the ugly ensure participant, ensure room, etc
1 parent f1a5af8 commit 7fbab5f

File tree

14 files changed

+113
-219
lines changed

14 files changed

+113
-219
lines changed

‎core/src/clients/direct/index.ts

+7-20
Original file line numberDiff line numberDiff line change
@@ -144,26 +144,13 @@ class DirectClient {
144144
return;
145145
}
146146

147-
await Promise.all([
148-
runtime.ensureUserExists(
149-
runtime.agentId,
150-
runtime.character.name ?? "Agent",
151-
runtime.character.name ?? "Agent",
152-
"direct"
153-
),
154-
runtime.ensureUserExists(
155-
userId,
156-
req.body.userName ?? "User",
157-
req.body.name ?? "User",
158-
"direct"
159-
),
160-
runtime.ensureRoomExists(roomId),
161-
]);
162-
163-
await Promise.all([
164-
runtime.ensureParticipantInRoom(userId, roomId),
165-
runtime.ensureParticipantInRoom(runtime.agentId, roomId),
166-
]);
147+
await runtime.ensureConnection(
148+
userId,
149+
roomId,
150+
req.body.userName,
151+
req.body.name,
152+
"direct"
153+
);
167154

168155
const text = req.body.text;
169156
const messageId = stringToUuid(Date.now().toString());

‎core/src/clients/discord/index.ts

+17-51
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,7 @@ import { embeddingZeroVector } from "../../core/memory.ts";
1717
import { MessageManager } from "./messages.ts";
1818
import { VoiceManager } from "./voice.ts";
1919

20-
import { IAgentRuntime } from "../../core/types.ts";
20+
import { Character, IAgentRuntime } from "../../core/types.ts";
2121
import chat_with_attachments from "./actions/chat_with_attachments.ts";
2222
import joinvoice from "./actions/joinvoice.ts";
2323
import leavevoice from "./actions/leavevoice.ts";
@@ -31,7 +31,7 @@ export class DiscordClient extends EventEmitter {
3131
apiToken: string;
3232
private client: Client;
3333
private runtime: IAgentRuntime;
34-
character: any;
34+
character: Character;
3535
private messageManager: MessageManager;
3636
private voiceManager: VoiceManager;
3737

@@ -182,29 +182,16 @@ export class DiscordClient extends EventEmitter {
182182
console.error("Invalid user id or room id");
183183
return;
184184
}
185-
const agentId = this.runtime.agentId;
186185
const userName = reaction.message.author.username;
187186
const name = reaction.message.author.displayName;
188-
await Promise.all([
189-
this.runtime.ensureUserExists(
190-
agentId,
191-
this.client.user.username,
192-
this.runtime.character.name,
193-
"discord"
194-
),
195-
this.runtime.ensureUserExists(
196-
userIdUUID,
197-
userName,
198-
name,
199-
"discord"
200-
),
201-
this.runtime.ensureRoomExists(roomId),
202-
]);
203-
204-
await Promise.all([
205-
this.runtime.ensureParticipantInRoom(userIdUUID, roomId),
206-
this.runtime.ensureParticipantInRoom(agentId, roomId),
207-
]);
187+
188+
await this.runtime.ensureConnection(
189+
userIdUUID,
190+
roomId,
191+
userName,
192+
name,
193+
"discord"
194+
);
208195

209196
// Save the reaction as a message
210197
await this.runtime.messageManager.createMemory({
@@ -259,37 +246,16 @@ export class DiscordClient extends EventEmitter {
259246
`${reaction.message.id}-${user.id}-${emoji}-removed`
260247
);
261248

262-
const agentId = this.runtime.agentId;
263249
const userName = reaction.message.author.username;
264250
const name = reaction.message.author.displayName;
265251

266-
console.log("reactionUUID", reactionUUID);
267-
console.log("userIdUUID", userIdUUID);
268-
console.log("roomId", roomId);
269-
console.log("agentId", agentId);
270-
console.log("userName", userName);
271-
console.log("name", name);
272-
273-
await Promise.all([
274-
this.runtime.ensureUserExists(
275-
agentId,
276-
this.client.user.username,
277-
this.runtime.character.name,
278-
"discord"
279-
),
280-
this.runtime.ensureUserExists(
281-
userIdUUID,
282-
userName,
283-
name,
284-
"discord"
285-
),
286-
this.runtime.ensureRoomExists(roomId),
287-
]);
288-
289-
await Promise.all([
290-
this.runtime.ensureParticipantInRoom(userIdUUID, roomId),
291-
this.runtime.ensureParticipantInRoom(agentId, roomId),
292-
]);
252+
await this.runtime.ensureConnection(
253+
userIdUUID,
254+
roomId,
255+
userName,
256+
name,
257+
"discord"
258+
);
293259

294260
try {
295261
// Save the reaction removal as a message

‎core/src/clients/discord/messages.ts

+8-22
Original file line numberDiff line numberDiff line change
@@ -146,28 +146,14 @@ export class MessageManager {
146146

147147
const roomId = stringToUuid(channelId);
148148
const userIdUUID = stringToUuid(userId);
149-
const agentId = this.runtime.agentId;
150-
151-
await Promise.all([
152-
this.runtime.ensureUserExists(
153-
agentId,
154-
this.client.user.username,
155-
this.runtime.character.name,
156-
"discord"
157-
),
158-
this.runtime.ensureUserExists(
159-
userIdUUID,
160-
userName,
161-
name,
162-
"discord"
163-
),
164-
this.runtime.ensureRoomExists(roomId),
165-
]);
166-
167-
await Promise.all([
168-
this.runtime.ensureParticipantInRoom(userIdUUID, roomId),
169-
this.runtime.ensureParticipantInRoom(agentId, roomId),
170-
]);
149+
150+
await this.runtime.ensureConnection(
151+
userIdUUID,
152+
roomId,
153+
userName,
154+
name,
155+
"discord"
156+
);
171157

172158
const messageId = stringToUuid(message.id);
173159

‎core/src/clients/discord/voice.ts

+6-24
Original file line numberDiff line numberDiff line change
@@ -304,32 +304,14 @@ export class VoiceManager extends EventEmitter {
304304

305305
const roomId = stringToUuid(channelId);
306306
const userIdUUID = stringToUuid(userId);
307-
await this.runtime.ensureUserExists(
308-
this.runtime.agentId,
309-
this.client.user.username,
310-
this.runtime.character.name,
307+
308+
await this.runtime.ensureConnection(
309+
userIdUUID,
310+
roomId,
311+
userName,
312+
name,
311313
"discord"
312314
);
313-
await Promise.all([
314-
this.runtime.ensureUserExists(
315-
userIdUUID,
316-
userName,
317-
name,
318-
"discord"
319-
),
320-
this.runtime.ensureRoomExists(roomId),
321-
]);
322-
323-
await Promise.all([
324-
this.runtime.ensureParticipantInRoom(
325-
userIdUUID,
326-
roomId
327-
),
328-
this.runtime.ensureParticipantInRoom(
329-
this.runtime.agentId,
330-
roomId
331-
),
332-
]);
333315

334316
let state = await this.runtime.composeState(
335317
{

‎core/src/clients/telegram/src/messageManager.ts

+11-23
Original file line numberDiff line numberDiff line change
@@ -222,11 +222,10 @@ export class MessageManager {
222222
return; // Exit if no message or sender info
223223
}
224224

225-
//@ts-ignore
226-
if (ctx.message.text.startsWith("/")) {
227-
//Handle commands?
228-
return;
229-
}
225+
// TODO: Handle commands?
226+
// if (ctx.message.text?.startsWith("/")) {
227+
// return;
228+
// }
230229

231230
const message = ctx.message;
232231

@@ -239,24 +238,13 @@ export class MessageManager {
239238
const agentId = this.runtime.agentId;
240239
const roomId = chatId;
241240

242-
// Ensure user and room exist
243-
await Promise.all([
244-
this.runtime.ensureUserExists(
245-
agentId,
246-
this.bot.botInfo?.username || "Bot",
247-
this.runtime.character.name,
248-
"telegram"
249-
),
250-
this.runtime.ensureUserExists(
251-
userId,
252-
userName,
253-
userName,
254-
"telegram"
255-
),
256-
this.runtime.ensureRoomExists(roomId),
257-
this.runtime.ensureParticipantInRoom(userId, roomId),
258-
this.runtime.ensureParticipantInRoom(agentId, roomId),
259-
]);
241+
await this.runtime.ensureConnection(
242+
userId,
243+
roomId,
244+
userName,
245+
userName,
246+
"telegram"
247+
);
260248

261249
const messageId = stringToUuid(
262250
message.message_id.toString()

‎core/src/clients/twitter/base.ts

+4-21
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,3 @@
1-
// @ts-nocheck
2-
31
import {
42
QueryTweetsResponse,
53
Scraper,
@@ -410,22 +408,13 @@ export class ClientBase extends EventEmitter {
410408
? this.runtime.agentId
411409
: stringToUuid(tweet.userId);
412410

413-
await this.runtime.ensureRoomExists(roomId);
414-
await this.runtime.ensureParticipantExists(
415-
this.runtime.agentId,
416-
roomId
417-
);
418-
419-
await this.runtime.ensureUserExists(
411+
await this.runtime.ensureConnection(
420412
tweetuserId,
413+
roomId,
421414
tweet.username,
422415
tweet.name,
423416
"twitter"
424417
);
425-
await this.runtime.ensureParticipantExists(
426-
tweetuserId,
427-
roomId
428-
);
429418

430419
const content = {
431420
text: tweet.text,
@@ -521,19 +510,13 @@ export class ClientBase extends EventEmitter {
521510
? this.runtime.agentId
522511
: stringToUuid(tweet.userId);
523512

524-
await this.runtime.ensureRoomExists(roomId);
525-
await this.runtime.ensureParticipantExists(
526-
this.runtime.agentId,
527-
roomId
528-
);
529-
530-
await this.runtime.ensureUserExists(
513+
await this.runtime.ensureConnection(
531514
tweetuserId,
515+
roomId,
532516
tweet.username,
533517
tweet.name,
534518
"twitter"
535519
);
536-
await this.runtime.ensureParticipantExists(tweetuserId, roomId);
537520

538521
const content = {
539522
text: tweet.text,

‎core/src/clients/twitter/generate.ts

+2
Original file line numberDiff line numberDiff line change
@@ -162,11 +162,13 @@ export class TwitterGenerationClient extends ClientBase {
162162
const conversationId = tweet.conversationId;
163163
const roomId = stringToUuid(conversationId);
164164

165+
// make sure the agent is in the room
165166
await this.runtime.ensureRoomExists(roomId);
166167
await this.runtime.ensureParticipantInRoom(
167168
this.runtime.agentId,
168169
roomId
169170
);
171+
170172
await this.cacheTweet(tweet);
171173

172174
await this.runtime.messageManager.createMemory({

‎core/src/clients/twitter/interactions.ts

+7-25
Original file line numberDiff line numberDiff line change
@@ -117,37 +117,19 @@ export class TwitterInteractionClient extends ClientBase {
117117
!this.lastCheckedTweetId ||
118118
parseInt(tweet.id) > this.lastCheckedTweetId
119119
) {
120-
console.log("Processing tweet", tweet.id);
121120
const conversationId = tweet.conversationId;
122121

123122
const roomId = stringToUuid(conversationId);
124-
await this.runtime.ensureRoomExists(roomId);
125123

126124
const userIdUUID = stringToUuid(tweet.userId as string);
127-
const agentId = this.runtime.agentId;
128125

129-
await Promise.all([
130-
this.runtime.ensureUserExists(
131-
agentId,
132-
this.runtime.getSetting("TWITTER_USERNAME"),
133-
this.runtime.character.name,
134-
"twitter"
135-
),
136-
this.runtime.ensureUserExists(
137-
userIdUUID,
138-
tweet.username,
139-
tweet.name,
140-
"twitter"
141-
),
142-
]);
143-
144-
await Promise.all([
145-
this.runtime.ensureParticipantInRoom(
146-
userIdUUID,
147-
roomId
148-
),
149-
this.runtime.ensureParticipantInRoom(agentId, roomId),
150-
]);
126+
await this.runtime.ensureConnection(
127+
userIdUUID,
128+
roomId,
129+
tweet.username,
130+
tweet.name,
131+
"twitter"
132+
);
151133

152134
await buildConversationThread(tweet, this);
153135

0 commit comments

Comments
 (0)