Minor changes

RocketChat · Sep 12, 2024 · b9a805f · b9a805f
2 parents 8ead1fd + 4d2915c
commit b9a805f
Show file tree

Hide file tree

Showing 16 changed files with 195 additions and 210 deletions.
diff --git a/ai-assistant/src/commands/AskCodeCommand.ts b/ai-assistant/src/commands/AskCodeCommand.ts
@@ -21,7 +21,7 @@ export class AskCodeCommand implements ISlashCommand {
     * @param {string} query - The user's query.
     * @returns {Promise<string | null>} A promise that resolves to the response to be given to the user or `null` if no answer or no reference is found.
     */
-   private async process(http: IHttp, query: string): Promise<string | null> {
+   private async process(http: IHttp, query: string): Promise<string> {
       const db = new Neo4j(http)
       const llm = new Llama3_70B(http)
       const embeddingModel = new MiniLML6(http)
@@ -33,7 +33,7 @@ export class AskCodeCommand implements ISlashCommand {
        * ---------------------------------------------------------------------------------------------
        */
       const keywords = await Query.getDBKeywordsFromQuery(llm, query)
-      if (!keywords.length) return null
+      if (!keywords.length) return "I'm sorry, I couldn't understand your query. Please try again."
 
       /**
        * ---------------------------------------------------------------------------------------------
@@ -42,7 +42,7 @@ export class AskCodeCommand implements ISlashCommand {
        * ---------------------------------------------------------------------------------------------
        */
       const results = await Query.getCodeNodesFromKeywords(db, embeddingModel, keywords)
-      if (!results.length) return null
+      if (!results.length) return "I'm sorry, I couldn't find any code related to your query."
 
       /**
        * ---------------------------------------------------------------------------------------------
@@ -53,7 +53,7 @@ export class AskCodeCommand implements ISlashCommand {
       const answer = await llm.ask(
          PromptFactory.makeAskCodePrompt(results.map((x) => x.code).join("\n\n"), query)
       )
-      if (!answer) return null
+      if (!answer) return "I'm sorry, I'm having trouble connecting to the server. Please try again later."
 
       return answer
    }
@@ -85,10 +85,6 @@ export class AskCodeCommand implements ISlashCommand {
       )
 
       const res = await this.process(http, query)
-      if (res) {
-         await sendEditedMessage(res)
-      } else {
-         await sendEditedMessage("❌ Unable to process your query")
-      }
+      await sendEditedMessage(res)
    }
 }
diff --git a/ai-assistant/src/commands/AskDocsCommand.ts b/ai-assistant/src/commands/AskDocsCommand.ts
@@ -21,7 +21,7 @@ export class AskDocsCommand implements ISlashCommand {
     * @param {string} query - The user's query.
     * @returns {Promise<string | null>} A promise that resolves to the response to be given to the user or `null` if no answer or no reference is found.
     */
-   private async process(http: IHttp, query: string): Promise<string | null> {
+   private async process(http: IHttp, query: string): Promise<string> {
       const db = new Neo4j(http)
       const llm = new Llama3_70B(http)
       const embeddingModel = new MiniLML6(http)
@@ -33,7 +33,7 @@ export class AskDocsCommand implements ISlashCommand {
        * ---------------------------------------------------------------------------------------------
        */
       const results = await Query.getDocsNodesFromQuery(db, embeddingModel, query)
-      if (!results.length) return null
+      if (!results.length) return "I'm sorry, I couldn't find any documentation related to your query."
 
       /**
        * ---------------------------------------------------------------------------------------------
@@ -45,7 +45,7 @@ export class AskDocsCommand implements ISlashCommand {
       const answer = await llm.ask(
          PromptFactory.makeAskDocsPrompt(results.map((x) => x.content).join("\n\n"), uniqueSources, query)
       )
-      if (!answer) return null
+      if (!answer) return "I'm sorry, I'm having trouble connecting to the server. Please try again later."
 
       return answer
    }
@@ -68,10 +68,6 @@ export class AskDocsCommand implements ISlashCommand {
       )
 
       const res = await this.process(http, query)
-      if (res) {
-         await sendEditedMessage(res)
-      } else {
-         await sendEditedMessage("❌ Unable to process your query")
-      }
+      await sendEditedMessage(res)
    }
 }
diff --git a/ai-assistant/src/commands/DiagramCommand.ts b/ai-assistant/src/commands/DiagramCommand.ts
@@ -47,9 +47,9 @@ export class DiagramCommand implements ISlashCommand {
       const diagram = await llm.ask(
          PromptFactory.makeDiagramPrompt(results.map((x) => x.code).join("\n\n"), query)
       )
+      console.log(diagram)
       if (!diagram) return null
 
-      // @ts-ignore
       const diagramContent = diagram
          .replace("```mermaid", "")
          .replace("```", "")

diff --git a/ai-assistant/src/commands/DocumentCommand.ts b/ai-assistant/src/commands/DocumentCommand.ts
@@ -23,7 +23,7 @@ export class DocumentCommand implements ISlashCommand {
    private async process(
       http: IHttp,
       query: string
-   ): Promise<{ jsDoc: string; explanation: string | null } | null> {
+   ): Promise<{ jsDoc: string; explanation: string | null } | string> {
       const db = new Neo4j(http)
       const llm = new Llama3_70B(http)
       const embeddingModel = new MiniLML6(http)
@@ -35,7 +35,7 @@ export class DocumentCommand implements ISlashCommand {
        * ---------------------------------------------------------------------------------------------
        */
       const keywords = await Query.getDBKeywordsFromQuery(llm, query)
-      if (!keywords.length) return null
+      if (!keywords.length) return "I'm sorry, I couldn't understand your query. Please try again."
 
       /**
        * ---------------------------------------------------------------------------------------------
@@ -44,7 +44,7 @@ export class DocumentCommand implements ISlashCommand {
        * ---------------------------------------------------------------------------------------------
        */
       const codeNodes = await Query.getCodeNodesFromKeywords(db, embeddingModel, keywords)
-      if (!codeNodes.length) return null
+      if (!codeNodes.length) return "I'm sorry, I couldn't find any code related to your query."
 
       /**
        * ---------------------------------------------------------------------------------------------
@@ -53,7 +53,7 @@ export class DocumentCommand implements ISlashCommand {
        * ---------------------------------------------------------------------------------------------
        */
       const result = await llm.ask(PromptFactory.makeDocumentPrompt(JSON.stringify(codeNodes), query))
-      if (!result) return null
+      if (!result) return "I'm sorry, I couldn't generate documentation for your query."
 
       //@ts-ignore
       const jsDoc = result.split("<JSDOC_START>")[1].split("<JSDOC_END>")[0].trim()
@@ -93,10 +93,10 @@ export class DocumentCommand implements ISlashCommand {
       )
 
       let res = await this.process(http, query)
-      if (res) {
-         await sendEditedMessage(`${res.jsDoc}\n\n${res.explanation}`)
+      if (typeof res === "string") {
+         await sendEditedMessage(res)
       } else {
-         await sendEditedMessage("❌ No references found!")
+         await sendEditedMessage(`${res.jsDoc}\n\n${res.explanation}`)
       }
    }
 }
diff --git a/ai-assistant/src/commands/TranslateCommand.ts b/ai-assistant/src/commands/TranslateCommand.ts
@@ -21,7 +21,7 @@ export class TranslateCommand implements ISlashCommand {
     * @param {string} targetEntity - The target entity for translation.
     * @returns {Promise<string | null>} A promise that resolves to the translated result or null if no translation is found.
     */
-   private async process(http: IHttp, targetLanguage: string, targetEntity: string): Promise<string | null> {
+   private async process(http: IHttp, targetLanguage: string, targetEntity: string): Promise<string> {
       const db = new Neo4j(http)
       const llm = new Llama3_70B(http)
       const embeddingModel = new MiniLML6(http)
@@ -33,7 +33,7 @@ export class TranslateCommand implements ISlashCommand {
        * ---------------------------------------------------------------------------------------------
        */
       const codeNodes = await Query.getCodeNodesFromKeywords(db, embeddingModel, [targetEntity])
-      if (!codeNodes.length) return null
+      if (!codeNodes.length) return "I'm sorry, I couldn't find any code related to your query."
 
       /**
        * ---------------------------------------------------------------------------------------------
@@ -48,7 +48,7 @@ export class TranslateCommand implements ISlashCommand {
             targetLanguage
          )
       )
-      if (!res) return null
+      if (!res) return "I'm sorry, I'm having trouble connecting to the server. Please try again later."
 
       return res
    }
@@ -85,10 +85,6 @@ export class TranslateCommand implements ISlashCommand {
       )
 
       const res = await this.process(http, targetEntity, targetLanguage)
-      if (res) {
-         await sendEditedMessage(res)
-      } else {
-         await sendEditedMessage("❌ Translation failed")
-      }
+      await sendEditedMessage(res)
    }
 }
diff --git a/ai-assistant/src/commands/WhyUsedCommand.ts b/ai-assistant/src/commands/WhyUsedCommand.ts
@@ -7,6 +7,7 @@ import { Query } from "../core/query"
 import { MiniLML6 } from "../core/services/embeddings/minilml6"
 import { Llama3_70B } from "../core/services/llm/llama3_70B"
 import { handleCommandResponse } from "../utils/handleCommandResponse"
+import { renderDiagramToBase64URI } from "../core/diagram"
 
 export class WhyUsedCommand implements ISlashCommand {
    public command = "rcc-whyused"
@@ -24,10 +25,13 @@ export class WhyUsedCommand implements ISlashCommand {
    private async process(
       http: IHttp,
       query: string
-   ): Promise<{
-      explanation: string
-      diagram: string
-   } | null> {
+   ): Promise<
+      | {
+           explanation: string
+           diagram: string
+        }
+      | string
+   > {
       const db = new Neo4j(http)
       const llm = new Llama3_70B(http)
       const embeddingModel = new MiniLML6(http)
@@ -39,7 +43,7 @@ export class WhyUsedCommand implements ISlashCommand {
        * ---------------------------------------------------------------------------------------------
        */
       const keywords = await Query.getDBKeywordsFromQuery(llm, query)
-      if (!keywords.length) return null
+      if (!keywords.length) return "I'm sorry, I couldn't understand your query. Please try again."
 
       /**
        * ---------------------------------------------------------------------------------------------
@@ -48,7 +52,7 @@ export class WhyUsedCommand implements ISlashCommand {
        * ---------------------------------------------------------------------------------------------
        */
       const codeNodes = await Query.getCodeNodesFromKeywords(db, embeddingModel, keywords)
-      if (!codeNodes.length) return null
+      if (!codeNodes.length) return "I'm sorry, I couldn't find any code related to your query."
 
       /**
        * ---------------------------------------------------------------------------------------------
@@ -59,7 +63,7 @@ export class WhyUsedCommand implements ISlashCommand {
       const result = await llm.ask(
          PromptFactory.makeWhyUsedPrompt(codeNodes.map((x) => x.code).join("\n\n"), query)
       )
-      if (!result) return null
+      if (!result) return "I'm sorry, I couldn't find any references for your query."
 
       const explanation = result.split("<EXPLANATION>")[1].split("</EXPLANATION>")[0].trim()
       const diagram = result.split("<DIAGRAM>")[1].split("</DIAGRAM>")[0].trim()
@@ -71,17 +75,12 @@ export class WhyUsedCommand implements ISlashCommand {
        * ---------------------------------------------------------------------------------------------
        */
       const data = { explanation, diagram: "" }
-      // TODO:
-      // if (diagram) {
-      //     const parsedDiagram = diagram
-      //         .replace("```mermaid", "")
-      //         .replace("```", "")
-      //         .trim();
-      //     writeFileSync("output.txt", parsedDiagram);
-      //     try {
-      //         // data.diagram = await renderDiagramToBase64URI(parsedDiagram);
-      //     } catch {}
-      // }
+      if (diagram) {
+         const parsedDiagram = diagram.replace("```mermaid", "").replace("```", "").trim()
+         try {
+            data.diagram = await renderDiagramToBase64URI(http, parsedDiagram)
+         } catch {}
+      }
 
       return data
    }
@@ -113,11 +112,10 @@ export class WhyUsedCommand implements ISlashCommand {
       )
 
       const res = await this.process(http, query)
-      if (!res) {
-         await sendEditedMessage("❌ No references found!")
+      if (typeof res === "string") {
+         await sendEditedMessage(res)
          return
       }
-
       await sendEditedMessage(res.explanation, [res.diagram!])
    }
 }
diff --git a/ai-assistant/src/core/query.ts b/ai-assistant/src/core/query.ts
@@ -27,12 +27,12 @@ export namespace Query {
       const result = await db.run(
          `
 				CALL db.index.vector.queryNodes("${indexName}", 2, $vector)
-				YIELD node, score 
-                WHERE score >= ${threshold}
-                WITH node, score
-                OPTIONAL MATCH (node)-[r]->(relatedNode)
-                RETURN node, COLLECT(relatedNode) AS relatedNodes, score
-                ORDER BY score DESC
+				YIELD node, score
+            WHERE score >= ${threshold}
+            WITH node, score
+            OPTIONAL MATCH (node)-[r]->(relatedNode)
+            RETURN node, COLLECT(relatedNode) AS relatedNodes, score
+            ORDER BY score DESC
 			`,
          { vector }
       )
@@ -67,13 +67,16 @@ export namespace Query {
       keywords: string[]
    ): Promise<DBNode[]> {
       const results: DBNode[] = []
-      for (const keyword of keywords) {
-         const queryVector = await embeddingModel.generate(keyword)
-         if (!queryVector) continue
 
-         const result = await getDBNodesFromVectorQuery(db, "nameEmbeddings", queryVector, 0.85)
-         results.push(...result)
-      }
+      try {
+         for (const keyword of keywords) {
+            const queryVector = await embeddingModel.generate(keyword)
+            if (!queryVector) continue
+
+            const result = await getDBNodesFromVectorQuery(db, "nameEmbeddings", queryVector, 0.85)
+            results.push(...result)
+         }
+      } catch {}
 
       return results
    }

diff --git a/ai-assistant/src/core/services/db/neo4j.ts b/ai-assistant/src/core/services/db/neo4j.ts
@@ -33,12 +33,14 @@ export class Neo4j implements IDB {
       // password: string
    ) {
       this.http = http
+
       this.baseUrl = "http://neo4j:7474"
       this.username = "neo4j"
       this.password = "strongpasswordsafe123"
-      // this.baseUrl = "http://44.192.104.170:7474";
-      // this.username = "neo4j";
-      // this.password = "individuals-societies-wools";
+
+      // this.baseUrl = "http://44.192.104.170:7474"
+      // this.username = "neo4j"
+      // this.password = "individuals-societies-wools"
    }
 
    /**