Skip to content

Commit b8b1f7a

Browse files
preeeshapreeesha
committed
Made Devdocs db node completely supported in ingestion and ai-assistant
1 parent 20e24c3 commit b8b1f7a

File tree

12 files changed

+167
-104
lines changed

12 files changed

+167
-104
lines changed

‎ai-assistant/src/core/services/db/devDocDBNode.ts‎

Lines changed: 71 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,71 @@
1+
import { IEmbeddingModel } from "../embeddings/embeddings.types";
2+
3+
export type DevDocDBNodeRelationType = "CONTAINS";
4+
export type DevDocDBNodeRelation = {
5+
target: string;
6+
relation: DevDocDBNodeRelationType;
7+
};
8+
9+
export class DevDocDBNode {
10+
id: string;
11+
relations: DevDocDBNodeRelation[];
12+
13+
url: string;
14+
element: string;
15+
16+
content: string;
17+
contentEmbeddings: number[];
18+
19+
constructor(node: {
20+
id: string;
21+
relations: DevDocDBNodeRelation[];
22+
23+
url: string;
24+
element: string;
25+
26+
content: string;
27+
contentEmbeddings: number[];
28+
}) {
29+
this.id = node.id;
30+
this.relations = node.relations;
31+
32+
this.url = node.url;
33+
this.element = node.element;
34+
35+
this.content = node.content;
36+
this.contentEmbeddings = node.contentEmbeddings;
37+
}
38+
39+
/**
40+
* Fills the embeddings for the given embedding model.
41+
*
42+
* @param {IEmbeddingModel} embeddingModel - The embedding model used to generate embeddings.
43+
* @returns {Promise<void>} - A promise that resolves when the embeddings are filled.
44+
*/
45+
async fillEmbeddings(embeddingModel: IEmbeddingModel): Promise<void> {
46+
this.contentEmbeddings =
47+
(await embeddingModel.generate(this.content)) ?? [];
48+
}
49+
50+
/**
51+
* Generates a database insert query for creating a new node with the specified properties.
52+
*
53+
* @returns The database insert query as a string.
54+
*/
55+
getDBInsertQuery(): string {
56+
let query = "";
57+
query += `
58+
CREATE (n:DevDocDBNode {
59+
id: $id,
60+
61+
url: $url,
62+
element: $element,
63+
64+
content: $content,
65+
contentEmbeddings: $contentEmbeddings
66+
})
67+
`;
68+
69+
return query;
70+
}
71+
}

‎ai-assistant/src/core/services/db/neo4j.ts‎

Lines changed: 14 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -33,12 +33,12 @@ export class Neo4j implements IDB {
3333
// password: string
3434
) {
3535
this.http = http;
36-
this.baseUrl = "http://neo4j:7474";
37-
this.username = "neo4j";
38-
this.password = "strongpasswordsafe123";
39-
// this.baseUrl = "http://44.192.104.170:7474";
36+
// this.baseUrl = "http://neo4j:7474";
4037
// this.username = "neo4j";
41-
// this.password = "individuals-societies-wools";
38+
// this.password = "strongpasswordsafe123";
39+
this.baseUrl = "http://44.192.104.170:7474";
40+
this.username = "neo4j";
41+
this.password = "individuals-societies-wools";
4242
}
4343

4444
/**
@@ -77,7 +77,7 @@ export class Neo4j implements IDB {
7777

7878
/**
7979
* Verifies the connectivity to the Neo4j database.
80-
*
80+
*
8181
* @returns A Promise that resolves to void.
8282
* @throws An Error if the connection to Neo4j fails.
8383
*/
@@ -92,7 +92,7 @@ export class Neo4j implements IDB {
9292

9393
/**
9494
* Closes the connection to the Neo4j database.
95-
*
95+
*
9696
* @returns A promise that resolves when the connection is closed.
9797
*/
9898
async closeDBConnection(): Promise<void> {
@@ -102,10 +102,10 @@ export class Neo4j implements IDB {
102102

103103
/**
104104
* Begins a new transaction in the Neo4j database.
105-
*
105+
*
106106
* @throws {Error} If a transaction already exists.
107107
* @throws {Error} If the transaction fails to begin.
108-
*
108+
*
109109
* @returns {Promise<void>} A promise that resolves when the transaction is successfully started.
110110
*/
111111
async beginTransaction(): Promise<void> {
@@ -124,10 +124,10 @@ export class Neo4j implements IDB {
124124

125125
/**
126126
* Commits the current transaction.
127-
*
127+
*
128128
* @throws {Error} If there is no transaction to commit.
129129
* @throws {Error} If the transaction commit fails.
130-
*
130+
*
131131
* @returns {Promise<void>} A promise that resolves when the transaction is successfully committed.
132132
*/
133133
async commitTransaction(): Promise<void> {
@@ -146,10 +146,10 @@ export class Neo4j implements IDB {
146146

147147
/**
148148
* Rolls back the current transaction.
149-
*
149+
*
150150
* @throws {Error} If there is no transaction to rollback.
151151
* @throws {Error} If the transaction rollback fails.
152-
*
152+
*
153153
* @returns {Promise<void>} A promise that resolves when the transaction is successfully rolled back.
154154
*/
155155
async rollbackTransaction(): Promise<void> {
@@ -200,6 +200,7 @@ export class Neo4j implements IDB {
200200
}
201201

202202
if (response.errors.length) {
203+
console.log(response.errors);
203204
throw new Error(
204205
response.errors.map((x) => JSON.stringify(x)).join("\n\n")
205206
);

‎ai-assistant/src/core/services/embeddings/minilml6.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -45,7 +45,7 @@ export class MiniLML6 implements IEmbeddingModel {
4545
* @returns A promise that resolves to an array of numbers representing the embeddings for the text, or null if the generation fails.
4646
*/
4747
async generate(text: string): Promise<number[] | null> {
48-
// return await this.fromHuggingFace(text);
48+
return await this.fromHuggingFace(text);
4949

5050
const res = await this.http.post(this.baseURL, {
5151
headers: {

‎ai-assistant/src/core/services/llm/llama3_70B.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -43,7 +43,7 @@ export class Llama3_70B implements ILLMModel {
4343
* @returns {Promise<string | null>} A promise that resolves with the response string or null if no response is available.
4444
*/
4545
async ask(prompt: Prompt): Promise<string | null> {
46-
// return await this.fromHuggingFace(prompt);
46+
return await this.fromHuggingFace(prompt);
4747

4848
const url = `${this.baseURL}/chat/completions`;
4949
const res = await this.http.post(url, {

‎ai-assistant/src/endpoints/ingest.ts‎

Lines changed: 13 additions & 34 deletions
Original file line numberDiff line numberDiff line change
@@ -11,45 +11,20 @@ import {
1111
IApiResponse,
1212
} from "@rocket.chat/apps-engine/definition/api";
1313
import { DBNode } from "../core/services/db/db";
14-
import { IDB } from "../core/services/db/db.types";
14+
import { DevDocDBNode } from "../core/services/db/devDocDBNode";
1515
import { Neo4j } from "../core/services/db/neo4j";
1616
import { MiniLML6 } from "../core/services/embeddings/minilml6";
1717
import {
1818
IngestEndpointRequestBody,
1919
IngestEndpointResponseBody,
2020
} from "./ingest.types";
2121

22-
namespace Helpers {
23-
/**
24-
* Inserts a node into the database.
25-
*
26-
* @param {IDB} db - The database instance.
27-
* @param {DBNode} node - The node to be inserted.
28-
* @returns {Promise<void>} - A promise that resolves when the node is successfully inserted.
29-
*/
30-
async function insertNode(db: IDB, node: DBNode) {
31-
const query = new DBNode(node).getDBInsertQuery();
32-
await db.run(query, node);
33-
}
34-
35-
/**
36-
* Inserts an array of nodes into the database.
37-
*
38-
* @param {IDB} db - The database object.
39-
* @param {DBNode[]} nodes - The array of nodes to be inserted.
40-
* @returns {Promise<void>} - A promise that resolves when all nodes have been inserted.
41-
*/
42-
export async function insertNodes(db: IDB, nodes: DBNode[]) {
43-
await Promise.all(nodes.map((node) => insertNode(db, node)));
44-
}
45-
}
46-
4722
export class IngestEndpoint extends ApiEndpoint {
4823
public path = "ingest";
4924

5025
/**
5126
* Generates the request and response bodies for the IngestEndpoint.
52-
*
27+
*
5328
* @param content - The content to be used for generating the request body.
5429
* @returns An array containing the generated request body and response body.
5530
*/
@@ -75,17 +50,21 @@ export class IngestEndpoint extends ApiEndpoint {
7550
): Promise<IApiResponse> {
7651
let [{ nodes }, responseBody] = this.makeBodies(request.content);
7752

78-
// -----------------------------------------------------------------------------------
79-
const embeddingModel = new MiniLML6(http);
80-
nodes = nodes.map((node) => new DBNode(node));
81-
await Promise.all(nodes.map((x) => x.fillEmbeddings(embeddingModel)));
8253
// -----------------------------------------------------------------------------------
8354
const db = new Neo4j(http);
8455
await db.verifyConnectivity();
56+
const embeddingModel = new MiniLML6(http);
57+
// -----------------------------------------------------------------------------------
58+
nodes = nodes.map((node) => {
59+
if ("element" in node) {
60+
return new DevDocDBNode(node);
61+
} else {
62+
return new DBNode(node);
63+
}
64+
});
65+
await Promise.all(nodes.map((x) => x.fillEmbeddings(embeddingModel)));
8566
// -----------------------------------------------------------------------------------
86-
const jobs = nodes.map((node) =>
87-
db.run(new DBNode(node).getDBInsertQuery(), node)
88-
);
67+
const jobs = nodes.map((node) => db.run(node.getDBInsertQuery(), node));
8968
await Promise.all(jobs);
9069
// -----------------------------------------------------------------------------------
9170

‎ai-assistant/src/endpoints/ingest.types.ts‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,9 @@
11
import { DBNode } from "../core/services/db/db";
2+
import { DevDocDBNode } from "../core/services/db/devDocDBNode";
23

34
export type IngestEndpointRequestBody = {
45
batchID: string;
5-
nodes: DBNode[];
6+
nodes: (DBNode | DevDocDBNode)[];
67
};
78

89
export type IngestEndpointResponseBody = {

‎ingestion/src/core/dbNode.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
import { TreeNode } from "../process/prepare/processor/core/treeNode";
1+
import { TreeNode } from "../process/prepare/processor/core/treeNode"
22

33
export type DBNodeRelationType = "CONTAINS" | "USES"
44
export type DBNodeRelation = { target: string; relation: DBNodeRelationType }

ingestion/src/core/devDocsDBNode.ts renamed to ingestion/src/core/devDocDBNode.ts

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@ export type DevDocDBNode = {
88
id: string
99
relations: DevDocDBNodeRelation[]
1010

11+
url: string
1112
element: string
1213

1314
content: string

‎ingestion/src/process/documentation/documentation.ts‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@ import puppeteer from "puppeteer"
22

33
import { writeFile } from "fs/promises"
44
import { DOCUMENTATION_URL } from "../../constants"
5-
import { DevDocDBNode } from "../../core/devDocsDBNode"
5+
import { DevDocDBNode } from "../../core/devDocDBNode"
66
import { IDocumentation } from "./documentation.types"
77
import { DocumentationPage } from "./documentationPage"
88

@@ -62,9 +62,10 @@ export class Documentation implements IDocumentation {
6262
jobs.push(
6363
writeFile(
6464
`${dataDirPath}/docs-${node.id}.json`,
65-
JSON.stringify(node, null, 2)
65+
JSON.stringify([node], null, 2)
6666
)
6767
)
6868
}
69+
await Promise.all(jobs)
6970
}
7071
}

‎ingestion/src/process/documentation/documentationPage.ts‎

Lines changed: 33 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
import { JSDOM } from "jsdom"
22

33
import { customAlphabet } from "nanoid"
4-
import { DevDocDBNode } from "../../core/devDocsDBNode"
4+
import { DevDocDBNode } from "../../core/devDocDBNode"
55
import {
66
DocumentPageElement_t,
77
IDocumentationPage,
@@ -79,39 +79,43 @@ export class DocumentationPage implements IDocumentationPage {
7979
return parseElements(elements)
8080
}
8181

82-
private convertHeirarchyToDevDocsDBNodes(
83-
hierarchy: DocumentPageElement_t[]
84-
): DevDocDBNode[] {
85-
const devDocsDBNodes: DevDocDBNode[] = []
82+
private traverseHierarchy(
83+
node: DocumentPageElement_t,
84+
devDocsDBNodes: DevDocDBNode[]
85+
): DevDocDBNode {
86+
const devDocDBNode: DevDocDBNode = {
87+
id: node.id,
88+
relations: [],
8689

87-
function traverseHierarchy(node: DocumentPageElement_t): DevDocDBNode {
88-
const devDocDBNode: DevDocDBNode = {
89-
id: node.id,
90-
element: node.element,
91-
relations: [],
92-
93-
content: node.content || "",
94-
contentEmbeddings: [],
95-
}
96-
97-
if (node.children) {
98-
node.children.forEach((child) => {
99-
const childNode = traverseHierarchy(child)
100-
devDocDBNode.relations.push({
101-
target: childNode.id,
102-
relation: "CONTAINS",
103-
})
104-
devDocsDBNodes.push(childNode)
105-
})
106-
}
90+
url: this._url,
91+
element: node.element,
92+
93+
content: node.content || "",
94+
contentEmbeddings: [],
95+
}
10796

108-
return devDocDBNode
97+
if (node.children) {
98+
node.children.forEach((child) => {
99+
const childNode = this.traverseHierarchy(child, devDocsDBNodes)
100+
devDocDBNode.relations.push({
101+
target: childNode.id,
102+
relation: "CONTAINS",
103+
})
104+
devDocsDBNodes.push(childNode)
105+
})
109106
}
110107

111-
hierarchy.forEach((node) => {
112-
const devDocDBNode = traverseHierarchy(node)
108+
return devDocDBNode
109+
}
110+
111+
private convertHeirarchyToDevDocsDBNodes(
112+
hierarchy: DocumentPageElement_t[]
113+
): DevDocDBNode[] {
114+
const devDocsDBNodes: DevDocDBNode[] = []
115+
for (const node of hierarchy) {
116+
const devDocDBNode = this.traverseHierarchy(node, devDocsDBNodes)
113117
devDocsDBNodes.push(devDocDBNode)
114-
})
118+
}
115119

116120
return devDocsDBNodes
117121
}

0 commit comments

Comments
 (0)