Refactored ingestion models by replacing DBNode with CodeModel and DevDocDBNode with DevDocModel, enhancing type safety and code organization. Updated related functions and types across the ingestion process to reflect these changes.

Author: yogesh-aggarwal

ingestion/src/lib/api/establish-relations.ts

@@ -16,4 +16,4 @@ export async function establishRelations(relations: DBNodeRelation[]): Promise<b
       console.log(e)
       return false
    }
-} 
+}

ingestion/src/lib/api/index.ts

@@ -1,3 +1,3 @@
-export { purgeDB } from './purge-db'
-export { insertBatch } from './insert-batch'
-export { establishRelations } from './establish-relations' 
+export { purgeDB } from "./purge-db"
+export { insertBatch } from "./insert-batch"
+export { establishRelations } from "./establish-relations"

ingestion/src/lib/api/insert-batch.ts

@@ -25,4 +25,4 @@ export async function insertBatch(batchID: string, nodes: DBNode[]): Promise<boo
       }
    }
    return false
-} 
+}

ingestion/src/lib/api/purge-db.ts

@@ -15,4 +15,4 @@ export async function purgeDB(): Promise<boolean> {
       console.log(e)
       return false
    }
-} 
+}

ingestion/src/lib/models/code.ts

@@ -0,0 +1,61 @@
+import { TreeNode } from "../../process/prepare/processor/core/treeNode"
+import { BaseNode, BaseRelation, EmbeddingsContainer, NodeDescriptor } from "./types"
+
+export type CodeModelRelation = BaseRelation
+
+export interface CodeModelProps extends BaseNode, EmbeddingsContainer {
+   name: string
+   type: string
+   code: string
+   filePath: string
+   isFile?: boolean
+   descriptor: NodeDescriptor
+}
+
+export class CodeModel implements CodeModelProps {
+   id: string
+   name: string
+   type: string
+   code: string
+   filePath: string
+   relations: CodeModelRelation[]
+   nameEmbeddings: number[]
+   codeEmbeddings: number[]
+   isFile: boolean
+   descriptor: NodeDescriptor
+
+   constructor(props: CodeModelProps) {
+      this.id = props.id
+      this.name = props.name
+      this.type = props.type
+      this.code = props.code
+      this.filePath = props.filePath
+      this.relations = props.relations
+      this.nameEmbeddings = props.nameEmbeddings || []
+      this.codeEmbeddings = props.codeEmbeddings || []
+      this.isFile = props.isFile || false
+      this.descriptor = props.descriptor
+   }
+
+   static fromTreeNode(node: TreeNode): CodeModel {
+      return new CodeModel({
+         id: node.getID(),
+         name: node.name,
+         type: node.type,
+         code: node.body,
+         filePath: node.sourceFileRelativePath,
+         relations: node.uses.map((use) => ({
+            target: use.name,
+            relation: "USES",
+         })),
+         nameEmbeddings: [],
+         codeEmbeddings: [],
+         isFile: false,
+         descriptor: "Node",
+      })
+   }
+
+   getNodeName(): string {
+      return this.name
+   }
+}

ingestion/src/lib/models/db-node.ts

@@ -0,0 +1,32 @@
+import { BaseNode, BaseRelation, EmbeddingsContainer } from "./types"
+
+export type DevDocModelRelation = BaseRelation
+
+export interface DevDocModelProps extends BaseNode, EmbeddingsContainer {
+   url: string
+   element: string
+   content: string
+}
+
+export class DevDocModel implements DevDocModelProps {
+   id: string
+   relations: DevDocModelRelation[]
+   nameEmbeddings: number[]
+   codeEmbeddings: number[]
+   contentEmbeddings: number[]
+
+   url: string
+   element: string
+   content: string
+
+   constructor(props: DevDocModelProps) {
+      this.id = props.id
+      this.relations = props.relations
+      this.nameEmbeddings = props.nameEmbeddings || []
+      this.codeEmbeddings = props.codeEmbeddings || []
+      this.contentEmbeddings = props.contentEmbeddings || []
+      this.url = props.url
+      this.element = props.element
+      this.content = props.content
+   }
+}

ingestion/src/lib/models/dev-doc-db-node.ts

@@ -0,0 +1,19 @@
+export type BaseRelationType = "CONTAINS" | "USES"
+
+export interface BaseRelation {
+   target: string
+   relation: BaseRelationType
+}
+
+export interface BaseNode {
+   id: string
+   relations: BaseRelation[]
+}
+
+export interface EmbeddingsContainer {
+   nameEmbeddings?: number[]
+   codeEmbeddings?: number[]
+   contentEmbeddings?: number[]
+}
+
+export type NodeDescriptor = "Node" | "Documentation" | "Code"

ingestion/src/lib/models/devdoc.ts

@@ -1,5 +1,5 @@
 import { writeFileSync } from "fs"
-import { DBNode } from "./models/db-node"
+import { CodeModel } from "./models/code"
 import { db } from "./neo4j"
 import { Styleguides } from "./styleguides.types"
 
@@ -54,7 +54,7 @@ export async function insertStyleguides() {
 
    const jobs = []
    for (const [filePath, data] of Object.entries(styleguides)) {
-      const node = new DBNode({
+      const node = new CodeModel({
          id: filePath,
          name: filePath,
          type: "",