Skip to content

Commit 9123996

Browse files
ponderingdemocritusponderingdemocritus
authoredNov 21, 2024··
Merge pull request elizaOS#472 from tarrencev/main
feat: Improve knowledge embeddings
2 parents 8450877 + 0d75334 commit 9123996

File tree

10 files changed

+201
-137
lines changed

10 files changed

+201
-137
lines changed
 

‎packages/client-discord/src/actions/summarize_conversation.ts

+6-1
Original file line numberDiff line numberDiff line change
@@ -251,7 +251,12 @@ const summarizeAction = {
251251
const model = models[runtime.character.settings.model];
252252
const chunkSize = model.settings.maxContextLength - 1000;
253253

254-
const chunks = await splitChunks(formattedMemories, chunkSize, 0);
254+
const chunks = await splitChunks(
255+
formattedMemories,
256+
chunkSize,
257+
"gpt-4o-mini",
258+
0
259+
);
255260

256261
const datestr = new Date().toUTCString().replace(/:/g, "-");
257262

‎packages/client-discord/src/messages.ts

+3-2
Original file line numberDiff line numberDiff line change
@@ -430,13 +430,13 @@ export class MessageManager {
430430
await this.runtime.messageManager.createMemory(memory);
431431
}
432432

433-
let state = (await this.runtime.composeState(userMessage, {
433+
let state = await this.runtime.composeState(userMessage, {
434434
discordClient: this.client,
435435
discordMessage: message,
436436
agentName:
437437
this.runtime.character.name ||
438438
this.client.user?.displayName,
439-
})) as State;
439+
});
440440

441441
if (!canSendMessage(message.channel).canSend) {
442442
return elizaLogger.warn(
@@ -649,6 +649,7 @@ export class MessageManager {
649649
message: DiscordMessage
650650
): Promise<{ processedContent: string; attachments: Media[] }> {
651651
let processedContent = message.content;
652+
652653
let attachments: Media[] = [];
653654

654655
// Process code blocks in the message content

‎packages/client-github/src/index.ts

+3-42
Original file line numberDiff line numberDiff line change
@@ -10,12 +10,8 @@ import {
1010
AgentRuntime,
1111
Client,
1212
IAgentRuntime,
13-
Content,
14-
Memory,
13+
knowledge,
1514
stringToUuid,
16-
embeddingZeroVector,
17-
splitChunks,
18-
embed,
1915
} from "@ai16z/eliza";
2016
import { validateGithubConfig } from "./enviroment";
2117

@@ -112,11 +108,8 @@ export class GitHubClient {
112108
relativePath
113109
);
114110

115-
const memory: Memory = {
111+
await knowledge.set(this.runtime, {
116112
id: knowledgeId,
117-
agentId: this.runtime.agentId,
118-
userId: this.runtime.agentId,
119-
roomId: this.runtime.agentId,
120113
content: {
121114
text: content,
122115
hash: contentHash,
@@ -128,39 +121,7 @@ export class GitHubClient {
128121
owner: this.config.owner,
129122
},
130123
},
131-
embedding: embeddingZeroVector,
132-
};
133-
134-
await this.runtime.documentsManager.createMemory(memory);
135-
136-
// Only split if content exceeds 4000 characters
137-
const fragments =
138-
content.length > 4000
139-
? await splitChunks(content, 2000, 200)
140-
: [content];
141-
142-
for (const fragment of fragments) {
143-
// Skip empty fragments
144-
if (!fragment.trim()) continue;
145-
146-
// Add file path context to the fragment before embedding
147-
const fragmentWithPath = `File: ${relativePath}\n\n${fragment}`;
148-
const embedding = await embed(this.runtime, fragmentWithPath);
149-
150-
await this.runtime.knowledgeManager.createMemory({
151-
// We namespace the knowledge base uuid to avoid id
152-
// collision with the document above.
153-
id: stringToUuid(knowledgeId + fragment),
154-
roomId: this.runtime.agentId,
155-
agentId: this.runtime.agentId,
156-
userId: this.runtime.agentId,
157-
content: {
158-
source: knowledgeId,
159-
text: fragment,
160-
},
161-
embedding,
162-
});
163-
}
124+
});
164125
}
165126
}
166127

‎packages/core/src/generation.ts

+15-11
Original file line numberDiff line numberDiff line change
@@ -463,34 +463,38 @@ export async function generateShouldRespond({
463463
* Splits content into chunks of specified size with optional overlapping bleed sections
464464
* @param content - The text content to split into chunks
465465
* @param chunkSize - The maximum size of each chunk in tokens
466-
* @param bleed - Number of characters to overlap between chunks (default: 100)
467466
* @param model - The model name to use for tokenization (default: runtime.model)
467+
* @param bleed - Number of characters to overlap between chunks (default: 100)
468468
* @returns Promise resolving to array of text chunks with bleed sections
469469
*/
470470
export async function splitChunks(
471471
content: string,
472472
chunkSize: number,
473+
model: string,
473474
bleed: number = 100
474475
): Promise<string[]> {
475-
const encoding = encoding_for_model("gpt-4o-mini");
476-
476+
const encoding = encoding_for_model(model as TiktokenModel);
477477
const tokens = encoding.encode(content);
478478
const chunks: string[] = [];
479479
const textDecoder = new TextDecoder();
480480

481481
for (let i = 0; i < tokens.length; i += chunkSize) {
482-
const chunk = tokens.slice(i, i + chunkSize);
483-
const decodedChunk = textDecoder.decode(encoding.decode(chunk));
482+
let chunk = tokens.slice(i, i + chunkSize);
484483

485484
// Append bleed characters from the previous chunk
486-
const startBleed = i > 0 ? content.slice(i - bleed, i) : "";
485+
if (i > 0) {
486+
chunk = new Uint32Array([...tokens.slice(i - bleed, i), ...chunk]);
487+
}
488+
487489
// Append bleed characters from the next chunk
488-
const endBleed =
489-
i + chunkSize < tokens.length
490-
? content.slice(i + chunkSize, i + chunkSize + bleed)
491-
: "";
490+
if (i + chunkSize < tokens.length) {
491+
chunk = new Uint32Array([
492+
...chunk,
493+
...tokens.slice(i + chunkSize, i + chunkSize + bleed),
494+
]);
495+
}
492496

493-
chunks.push(startBleed + decodedChunk + endBleed);
497+
chunks.push(textDecoder.decode(encoding.decode(chunk)));
494498
}
495499

496500
return chunks;

‎packages/core/src/index.ts

+1
Original file line numberDiff line numberDiff line change
@@ -20,3 +20,4 @@ export * from "./parsing.ts";
2020
export * from "./uuid.ts";
2121
export * from "./enviroment.ts";
2222
export * from "./cache.ts";
23+
export { default as knowledge } from "./knowledge.ts";

‎packages/core/src/knowledge.ts

+129
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,129 @@
1+
import { UUID } from "crypto";
2+
3+
import { AgentRuntime } from "./runtime.ts";
4+
import { embed } from "./embedding.ts";
5+
import { Content, ModelClass, type Memory } from "./types.ts";
6+
import { stringToUuid } from "./uuid.ts";
7+
import { embeddingZeroVector } from "./memory.ts";
8+
import { splitChunks } from "./generation.ts";
9+
import { models } from "./models.ts";
10+
import elizaLogger from "./logger.ts";
11+
12+
async function get(runtime: AgentRuntime, message: Memory): Promise<string[]> {
13+
const processed = preprocess(message.content.text);
14+
elizaLogger.log(`Querying knowledge for: ${processed}`);
15+
const embedding = await embed(runtime, processed);
16+
const fragments = await runtime.knowledgeManager.searchMemoriesByEmbedding(
17+
embedding,
18+
{
19+
roomId: message.agentId,
20+
agentId: message.agentId,
21+
count: 3,
22+
match_threshold: 0.1,
23+
}
24+
);
25+
26+
const uniqueSources = [
27+
...new Set(
28+
fragments.map((memory) => {
29+
elizaLogger.log(
30+
`Matched fragment: ${memory.content.text} with similarity: ${message.similarity}`
31+
);
32+
return memory.content.source;
33+
})
34+
),
35+
];
36+
37+
const knowledgeDocuments = await Promise.all(
38+
uniqueSources.map((source) =>
39+
runtime.documentsManager.getMemoryById(source as UUID)
40+
)
41+
);
42+
43+
const knowledge = knowledgeDocuments
44+
.filter((memory) => memory !== null)
45+
.map((memory) => memory.content.text);
46+
return knowledge;
47+
}
48+
49+
export type KnowledgeItem = {
50+
id: UUID;
51+
content: Content;
52+
};
53+
54+
async function set(runtime: AgentRuntime, item: KnowledgeItem) {
55+
await runtime.documentsManager.createMemory({
56+
embedding: embeddingZeroVector,
57+
id: item.id,
58+
agentId: runtime.agentId,
59+
roomId: runtime.agentId,
60+
userId: runtime.agentId,
61+
createdAt: Date.now(),
62+
content: item.content,
63+
});
64+
65+
const preprocessed = preprocess(item.content.text);
66+
const fragments = await splitChunks(
67+
preprocessed,
68+
10,
69+
models[runtime.character.modelProvider].model?.[ModelClass.EMBEDDING],
70+
5
71+
);
72+
73+
for (const fragment of fragments) {
74+
const embedding = await embed(runtime, fragment);
75+
await runtime.knowledgeManager.createMemory({
76+
// We namespace the knowledge base uuid to avoid id
77+
// collision with the document above.
78+
id: stringToUuid(item.id + fragment),
79+
roomId: runtime.agentId,
80+
agentId: runtime.agentId,
81+
userId: runtime.agentId,
82+
createdAt: Date.now(),
83+
content: {
84+
source: item.id,
85+
text: fragment,
86+
},
87+
embedding,
88+
});
89+
}
90+
}
91+
92+
export function preprocess(content: string): string {
93+
return (
94+
content
95+
// Remove code blocks and their content
96+
.replace(/```[\s\S]*?```/g, "")
97+
// Remove inline code
98+
.replace(/`.*?`/g, "")
99+
// Convert headers to plain text with emphasis
100+
.replace(/#{1,6}\s*(.*)/g, "$1")
101+
// Remove image links but keep alt text
102+
.replace(/!\[(.*?)\]\(.*?\)/g, "$1")
103+
// Remove links but keep text
104+
.replace(/\[(.*?)\]\(.*?\)/g, "$1")
105+
// Remove HTML tags
106+
.replace(/<[^>]*>/g, "")
107+
// Remove horizontal rules
108+
.replace(/^\s*[-*_]{3,}\s*$/gm, "")
109+
// Remove comments
110+
.replace(/\/\*[\s\S]*?\*\//g, "")
111+
.replace(/\/\/.*/g, "")
112+
// Normalize whitespace
113+
.replace(/\s+/g, " ")
114+
// Remove multiple newlines
115+
.replace(/\n{3,}/g, "\n\n")
116+
// strip all special characters
117+
.replace(/[^a-zA-Z0-9\s]/g, "")
118+
// Remove Discord mentions
119+
.replace(/<@!?\d+>/g, "")
120+
.trim()
121+
.toLowerCase()
122+
);
123+
}
124+
125+
export default {
126+
get,
127+
set,
128+
process,
129+
};

‎packages/core/src/runtime.ts

+38-80
Original file line numberDiff line numberDiff line change
@@ -14,8 +14,8 @@ import {
1414
} from "./evaluators.ts";
1515
import { generateText } from "./generation.ts";
1616
import { formatGoalsAsString, getGoals } from "./goals.ts";
17-
import { elizaLogger, embed, splitChunks } from "./index.ts";
18-
import { embeddingZeroVector, MemoryManager } from "./memory.ts";
17+
import { elizaLogger } from "./index.ts";
18+
import { MemoryManager } from "./memory.ts";
1919
import { formatActors, formatMessages, getActorDetails } from "./messages.ts";
2020
import { parseJsonArrayFromText } from "./parsing.ts";
2121
import { formatPosts } from "./posts.ts";
@@ -44,6 +44,7 @@ import {
4444
} from "./types.ts";
4545
import { stringToUuid } from "./uuid.ts";
4646
import { v4 as uuidv4 } from "uuid";
47+
import knowledge from "./knowledge.ts";
4748

4849
/**
4950
* Represents the runtime environment for an agent, handling message processing,
@@ -222,11 +223,21 @@ export class AgentRuntime implements IAgentRuntime {
222223
opts.character?.id ??
223224
opts?.agentId ??
224225
stringToUuid(opts.character?.name ?? uuidv4());
226+
this.character = opts.character || defaultCharacter;
227+
228+
// By convention, we create a user and room using the agent id.
229+
// Memories related to it are considered global context for the agent.
230+
this.ensureRoomExists(this.agentId);
231+
this.ensureUserExists(
232+
this.agentId,
233+
this.character.name,
234+
this.character.name
235+
);
236+
this.ensureParticipantExists(this.agentId, this.agentId);
225237

226238
elizaLogger.success("Agent ID", this.agentId);
227239

228240
this.fetch = (opts.fetch as typeof fetch) ?? this.fetch;
229-
this.character = opts.character || defaultCharacter;
230241
if (!opts.databaseAdapter) {
231242
throw new Error("No database adapter provided");
232243
}
@@ -348,60 +359,28 @@ export class AgentRuntime implements IAgentRuntime {
348359
* then chunks the content into fragments, embeds each fragment, and creates fragment memories.
349360
* @param knowledge An array of knowledge items containing id, path, and content.
350361
*/
351-
private async processCharacterKnowledge(knowledge: string[]) {
352-
// ensure the room exists and the agent exists in the room
353-
await this.ensureRoomExists(this.agentId);
354-
355-
await this.ensureUserExists(
356-
this.agentId,
357-
this.character.name,
358-
this.character.name
359-
);
360-
361-
await this.ensureParticipantExists(this.agentId, this.agentId);
362-
363-
for (const knowledgeItem of knowledge) {
364-
const knowledgeId = stringToUuid(knowledgeItem);
362+
private async processCharacterKnowledge(items: string[]) {
363+
for (const item of items) {
364+
const knowledgeId = stringToUuid(item);
365365
const existingDocument =
366366
await this.documentsManager.getMemoryById(knowledgeId);
367-
if (!existingDocument) {
368-
elizaLogger.success(
369-
"Processing knowledge for ",
370-
this.character.name,
371-
" - ",
372-
knowledgeItem.slice(0, 100)
373-
);
374-
await this.documentsManager.createMemory({
375-
embedding: embeddingZeroVector,
376-
id: knowledgeId,
377-
agentId: this.agentId,
378-
roomId: this.agentId,
379-
userId: this.agentId,
380-
createdAt: Date.now(),
381-
content: {
382-
text: knowledgeItem,
383-
},
384-
});
385-
386-
const fragments = await splitChunks(knowledgeItem, 1200, 200);
387-
for (const fragment of fragments) {
388-
const embedding = await embed(this, fragment);
389-
await this.knowledgeManager.createMemory({
390-
// We namespace the knowledge base uuid to avoid id
391-
// collision with the document above.
392-
id: stringToUuid(knowledgeId + fragment),
393-
roomId: this.agentId,
394-
agentId: this.agentId,
395-
userId: this.agentId,
396-
createdAt: Date.now(),
397-
content: {
398-
source: knowledgeId,
399-
text: fragment,
400-
},
401-
embedding,
402-
});
403-
}
367+
if (existingDocument) {
368+
return;
404369
}
370+
371+
console.log(
372+
"Processing knowledge for ",
373+
this.character.name,
374+
" - ",
375+
item.slice(0, 100)
376+
);
377+
378+
await knowledge.set(this, {
379+
id: knowledgeId,
380+
content: {
381+
text: item,
382+
},
383+
});
405384
}
406385
}
407386

@@ -935,33 +914,8 @@ Text: ${attachment.text}
935914
.join(" ");
936915
}
937916

938-
async function getKnowledge(
939-
runtime: AgentRuntime,
940-
message: Memory
941-
): Promise<string[]> {
942-
const embedding = await embed(runtime, message.content.text);
943-
944-
const memories =
945-
await runtime.knowledgeManager.searchMemoriesByEmbedding(
946-
embedding,
947-
{
948-
roomId: message.agentId,
949-
agentId: message.agentId,
950-
count: 3,
951-
}
952-
);
953-
954-
const knowledge = memories.map((memory) => memory.content.text);
955-
956-
return knowledge;
957-
}
958-
959-
const formatKnowledge = (knowledge: string[]) => {
960-
return knowledge.map((knowledge) => `- ${knowledge}`).join("\n");
961-
};
962-
963917
const formattedKnowledge = formatKnowledge(
964-
await getKnowledge(this, message)
918+
await knowledge.get(this, message)
965919
);
966920

967921
const initialState = {
@@ -1243,3 +1197,7 @@ Text: ${attachment.text}
12431197
} as State;
12441198
}
12451199
}
1200+
1201+
const formatKnowledge = (knowledge: string[]) => {
1202+
return knowledge.map((knowledge) => `- ${knowledge}`).join("\n");
1203+
};

‎packages/core/src/types.ts

+1
Original file line numberDiff line numberDiff line change
@@ -175,6 +175,7 @@ export interface Memory {
175175
embedding?: number[]; // An optional embedding vector representing the semantic content of the memory.
176176
roomId: UUID; // The room or conversation ID associated with the memory.
177177
unique?: boolean; // Whether the memory is unique or not
178+
similarity?: number; // embedding match similarity
178179
}
179180

180181
/**

‎packages/plugin-bootstrap/src/evaluators/fact.ts

+1-1
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,7 @@ import {
1313
export const formatFacts = (facts: Memory[]) => {
1414
const messageStrings = facts
1515
.reverse()
16-
.map((fact: Memory) => `${(fact.content as Content)?.content}`);
16+
.map((fact: Memory) => fact.content.text);
1717
const finalMessageStrings = messageStrings.join("\n");
1818
return finalMessageStrings;
1919
};

‎packages/plugin-bootstrap/src/index.ts

+4
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,10 @@ import { boredomProvider } from "./providers/boredom.ts";
1212
import { factsProvider } from "./providers/facts.ts";
1313
import { timeProvider } from "./providers/time.ts";
1414

15+
export * as actions from "./actions";
16+
export * as evaluators from "./evaluators";
17+
export * as providers from "./providers";
18+
1519
export const bootstrapPlugin: Plugin = {
1620
name: "bootstrap",
1721
description: "Agent bootstrap with basic actions and evaluators",

0 commit comments

Comments
 (0)
Please sign in to comment.