From baaccb6cbfd8cc35fa853a3f6bc6c63b228f1e98 Mon Sep 17 00:00:00 2001
From: Galego <bmgalego@gmail.com>
Date: Mon, 25 Nov 2024 19:41:21 +0000
Subject: [PATCH] add knowledge to state

---
 packages/core/src/knowledge.ts |  7 +++----
 packages/core/src/runtime.ts   | 14 ++++++++++----
 packages/core/src/types.ts     |  5 +++++
 3 files changed, 18 insertions(+), 8 deletions(-)

diff --git a/packages/core/src/knowledge.ts b/packages/core/src/knowledge.ts
index e0d990e031..21f35f0f8f 100644
--- a/packages/core/src/knowledge.ts
+++ b/packages/core/src/knowledge.ts
@@ -6,7 +6,7 @@ import { embeddingZeroVector } from "./memory.ts";
 import { splitChunks } from "./generation.ts";
 import elizaLogger from "./logger.ts";
 
-async function get(runtime: AgentRuntime, message: Memory): Promise<string[]> {
+async function get(runtime: AgentRuntime, message: Memory): Promise<KnowledgeItem[]> {
     const processed = preprocess(message.content.text);
     elizaLogger.log(`Querying knowledge for: ${processed}`);
     const embedding = await embed(runtime, processed);
@@ -36,10 +36,9 @@ async function get(runtime: AgentRuntime, message: Memory): Promise<string[]> {
         )
     );
 
-    const knowledge = knowledgeDocuments
+    return knowledgeDocuments
         .filter((memory) => memory !== null)
-        .map((memory) => memory.content.text);
-    return knowledge;
+        .map((memory) => ({ id: memory.id, content: memory.content }));
 }
 
 async function set(
diff --git a/packages/core/src/runtime.ts b/packages/core/src/runtime.ts
index 53194f02fe..3a628d3dfe 100644
--- a/packages/core/src/runtime.ts
+++ b/packages/core/src/runtime.ts
@@ -29,6 +29,7 @@ import {
     ICacheManager,
     IDatabaseAdapter,
     IMemoryManager,
+    KnowledgeItem,
     ModelClass,
     ModelProviderName,
     Plugin,
@@ -952,10 +953,14 @@ Text: ${attachment.text}
                 .join(" ");
         }
 
+
+        const knowledegeData = await knowledge.get(this, message);
+
         const formattedKnowledge = formatKnowledge(
-            await knowledge.get(this, message)
+            knowledegeData
         );
 
+
         const initialState = {
             agentId: this.agentId,
             agentName,
@@ -971,6 +976,7 @@ Text: ${attachment.text}
                       ]
                     : "",
             knowledge: formattedKnowledge,
+            knowledgeData: knowledegeData,
             // Recent interactions between the sender and receiver, formatted as messages
             recentMessageInteractions: formattedMessageInteractions,
             // Recent interactions between the sender and receiver, formatted as posts
@@ -1097,7 +1103,7 @@ Text: ${attachment.text}
                     ? addHeader("# Attachments", formattedAttachments)
                     : "",
             ...additionalKeys,
-        };
+        } as State;
 
         const actionPromises = this.actions.map(async (action: Action) => {
             const result = await action.validate(this, message, initialState);
@@ -1235,6 +1241,6 @@ Text: ${attachment.text}
     }
 }
 
-const formatKnowledge = (knowledge: string[]) => {
-    return knowledge.map((knowledge) => `- ${knowledge}`).join("\n");
+const formatKnowledge = (knowledge: KnowledgeItem[]) => {
+    return knowledge.map((knowledge) => `- ${knowledge.content.text}`).join("\n");
 };
diff --git a/packages/core/src/types.ts b/packages/core/src/types.ts
index 733aa7f831..09c0607cd1 100644
--- a/packages/core/src/types.ts
+++ b/packages/core/src/types.ts
@@ -294,6 +294,11 @@ export interface State {
     /** Optional formatted conversation */
     formattedConversation?: string;
 
+    /** Optional formatted knowledge */
+    knowledge?: string,
+    /** Optional knowledge data */
+    knowledgeData?: KnowledgeItem[],
+
     /** Additional dynamic properties */
     [key: string]: unknown;
 }