Add knowledge tool to the assistant

Author: jan-osch

src/aggregator/knowledge-aggregator.ts

@@ -1,63 +1,25 @@
-import fs from "fs";
-import path from "path";
-import matter from "gray-matter";
-
-import { DynamicKnowledge } from "@prisma/client";
-import { prisma } from "../prisma";
 import { WebhookAlertDto } from "../checkly/alertDTO";
 import { CheckContext, ContextKey } from "./ContextAggregator";
+import { getAllDocuments, KnowledgeDocument, } from "../knowledge-base/knowledgeBase";
 
-const transformDocument = (fileContent: string, checkId: string): CheckContext => {
-  const { data } = matter(fileContent);
-
+const transformDocument = (document: KnowledgeDocument, checkId: string): CheckContext => {
   return {
     checkId,
-    value: fileContent,
+    value: document.fullContent,
     source: 'knowledge',
     key: ContextKey.Knowledge.replace(
       "$documentSlug",
-      data.slug
+      document.slug
     ),
-    analysis: data.summary,
+    analysis: document.summary,
   } as CheckContext;
 }
 
-const loadKnowledgeDocuments = async (directory: string): Promise<string[]> => {
-  const collectMarkdownFiles = (dir: string): string[] => {
-    const entries = fs.readdirSync(dir, { withFileTypes: true });
-    const files: string[] = [];
-
-    for (const entry of entries) {
-      const fullPath = path.join(dir, entry.name);
-      if (entry.isDirectory()) {
-        // Recursively collect files from subdirectories
-        files.push(...collectMarkdownFiles(fullPath));
-      } else if (entry.isFile() && fullPath.endsWith(".md")) {
-        // Add Markdown files
-        files.push(fullPath);
-      }
-    }
-
-    return files;
-  };
-
-  const markdownFiles = collectMarkdownFiles(directory);
-
-  return markdownFiles.map((filePath) => fs.readFileSync(filePath, "utf-8"));
-};
-
-const loadDynamicKnowledge = async (): Promise<string[]> => {
-  const documents = await prisma.dynamicKnowledge.findMany({}) as DynamicKnowledge[]
-
-  return documents.map((doc) => doc.content);
-};
-
-
 export const knowledgeAggregator = {
   name: "Knowledge",
   fetchContext: async (alert: WebhookAlertDto): Promise<CheckContext[]> => {
     console.log('Aggregating Knowledge Context...');
-    const documents = [...await loadKnowledgeDocuments("./data/knowledge"), ...await loadDynamicKnowledge()];
+    const documents = await getAllDocuments()
 
     return documents.map((doc) => transformDocument(doc, alert.CHECK_ID));
   },

src/knowledge-base/knowledgeBase.ts

@@ -0,0 +1,70 @@
+import fs from "fs";
+import path from "path";
+import matter from "gray-matter";
+
+import { DynamicKnowledge } from "@prisma/client";
+import { prisma } from "../prisma";
+
+// This path assumes that the working directory is the root directory of the project
+const STATIC_KNOWLEDGE_DIRECTORY = "./data/knowledge";
+
+export type KnowledgeDocument = {
+  content: string;
+  fullContent: string;
+  title: string
+  slug: string
+  summary: string
+  created: Date
+  updated: Date
+}
+
+const parseDocument = (rawDocument: string): KnowledgeDocument => {
+  const parsed = matter(rawDocument);
+
+  return {
+    slug: parsed.data.slug,
+    summary: parsed.data.summary,
+    created: new Date(parsed.data.created),
+    updated: new Date(parsed.data.updated),
+    fullContent: rawDocument,
+    content: parsed.content,
+    title: parsed.data.title
+  }
+}
+
+export const loadKnowledgeDocuments = async (directory: string = STATIC_KNOWLEDGE_DIRECTORY): Promise<KnowledgeDocument[]> => {
+  const collectMarkdownFiles = (dir: string): string[] => {
+    const entries = fs.readdirSync(dir, { withFileTypes: true });
+    const files: string[] = [];
+
+    for (const entry of entries) {
+      const fullPath = path.join(dir, entry.name);
+      if (entry.isDirectory()) {
+        // Recursively collect files from subdirectories
+        files.push(...collectMarkdownFiles(fullPath));
+      } else if (entry.isFile() && fullPath.endsWith(".md")) {
+        // Add Markdown files
+        files.push(fullPath);
+      }
+    }
+
+    return files;
+  };
+
+  const markdownFiles = collectMarkdownFiles(directory);
+
+  return markdownFiles.map((filePath) => parseDocument(fs.readFileSync(filePath, "utf-8")));
+};
+
+export const loadDynamicKnowledge = async (): Promise<KnowledgeDocument[]> => {
+  const documents = await prisma.dynamicKnowledge.findMany({}) as DynamicKnowledge[]
+
+  return documents.map((doc) => parseDocument(doc.content));
+};
+
+export const getAllDocuments = async (): Promise<KnowledgeDocument[]> => {
+  const staticKnowledge = await loadKnowledgeDocuments();
+  const dynamicKnowledge = await loadDynamicKnowledge();
+
+  return [...staticKnowledge, ...dynamicKnowledge];
+}

src/sre-assistant/SreAssistant.ts

@@ -6,6 +6,7 @@ import { ChecklyTool } from "./tools/ChecklyTool";
 import { GitHubTool } from "./tools/GitHubTool";
 import { prisma } from "../prisma";
 import { slackFormatInstructions } from "../slackbot/utils";
+import { KnowledgeTool } from "./tools/KnowledgeTool";
 
 export class SreAssistant extends BaseAssistant {
   alertId: string | undefined;
@@ -62,7 +63,7 @@ CONSTITUTION:
 7. Make active use of the tools (multiple times if needed) to get a holistic view of the situation
 8. Generate super short, concise and insightful messages. Users are experts, skip the fluff, no yapping.
 9. Context-Driven Analysis: Load the check to understand the failure context and its potential root cause. Cross-reference the alert with logs, metrics, and change histories to identify anomalies or patterns. Please remember that releases (changes) are the most common reason for alerting events, focus on analysing recent changes.
-10. Refer to the the knowledge context to build a better understanding of the product, systems and the organisation you are working for. Assume that the users have good knowledge of the systems, and do not proactively provide basic information unless explicitly asked.
+10. Refer to the the knowledge base to build a better understanding of the terminology, systems and the organisation you are working for. Assume that the users have good knowledge of the company, and do not proactively provide basic information unless explicitly asked.
 
 INTERACTION CONTEXT:
 Username: ${this.interactionContext["username"]}
@@ -76,11 +77,11 @@ ${alertSummary.length > 0 ? `Alert Summary:\n${alertSummary}` : ""}`;
 
   protected async getTools(): Promise<Tool[]> {
     if (!this.alertId) {
-      return [new ChecklyTool(this), new GitHubTool(this)];
+      return [new ChecklyTool(this), new GitHubTool(this), new KnowledgeTool(this)];
     }
 
     const searchContextTool = new SearchContextTool(this);
     await searchContextTool.init();
-    return [searchContextTool, new ChecklyTool(this), new GitHubTool(this)];
+    return [searchContextTool, new ChecklyTool(this), new GitHubTool(this), new KnowledgeTool(this)];
   }
 }

src/sre-assistant/tools/KnowledgeTool.ts

@@ -0,0 +1,70 @@
+import { z } from "zod";
+import { createToolOutput, createToolParameters, Tool } from "../../ai/Tool";
+import { SreAssistant } from "../SreAssistant";
+import { getAllDocuments } from "../../knowledge-base/knowledgeBase";
+
+const parameters = createToolParameters(
+  z.object({
+    action: z
+      .enum([
+        "listDocuments",
+        "getOneDocument",
+      ])
+      .describe("The action to perform on the Knowledge Base"),
+    documentSlug: z
+      .string()
+      .describe(
+        "The slug of the Document to get information about. Omit this field for the 'listDocuments' action. Required for the 'getOneDocument'"
+      )
+      .optional(),
+  })
+);
+
+const outputSchema = createToolOutput(
+  z.string().describe("The response from the Knowledge Base")
+);
+
+export class KnowledgeTool extends Tool<
+  typeof parameters,
+  typeof outputSchema,
+  SreAssistant
+> {
+  static parameters = parameters;
+  static outputSchema = outputSchema;
+
+  constructor(agent: SreAssistant) {
+    super({
+      name: "KnowledgeBase",
+      description:
+        "Interact with the Knowledge Base to retrieve relevant context about the organisation structure, projects and terminology.",
+      parameters,
+      agent,
+    });
+  }
+
+  async execute(input: z.infer<typeof parameters>) {
+    if (input.action === "listDocuments") {
+      const documents = await getAllDocuments();
+
+      return JSON.stringify(documents.map((doc) => ({
+        slug: doc.slug,
+        title: doc.title,
+        summary: doc.summary,
+      })));
+    } else if (input.action === "getOneDocument") {
+      if (!input.documentSlug) {
+        return "Document slug is required";
+      }
+
+      const document = await getAllDocuments().then(docs => docs.find(doc => doc.slug === input.documentSlug));
+
+      if (!document) {
+        return `Document for slug: ${input.documentSlug} not found`;
+      }
+
+      return document.fullContent;
+    }
+
+    return "Invalid action";
+  }
+}