feat: add full text index

Signed-off-by: Dan Selman <danscode@selman.org>
accordproject · Dec 20, 2023 · e2d7657 · e2d7657
1 parent 6e1d423
commit e2d7657
Show file tree

Hide file tree

Showing 2 changed files with 73 additions and 0 deletions.
diff --git a/src/demo/index.ts b/src/demo/index.ts
@@ -55,6 +55,7 @@ concept Genre extends GraphNode {
 concept Movie extends GraphNode {
   o Double[] embedding optional
   @vector_index("embedding", 1536, "COSINE")
+  @fulltext_index
   o String summary optional
   @label("IN_GENRE")
   --> Genre[] genres optional
@@ -75,6 +76,7 @@ async function run() {
   await graphModel.dropIndexes();
   await graphModel.createConstraints();
   await graphModel.createVectorIndexes();
+  await graphModel.createFullTextIndexes();
   const context = await graphModel.openSession();
 
   const { session } = context;
@@ -122,6 +124,10 @@ async function run() {
     await graphModel.mergeNode(transaction, 'Actor', {identifier: 'Johnny Depp'} );
     await graphModel.mergeRelationship(transaction, 'Actor', 'Johnny Depp', 'Movie', 'Fear and Loathing in Las Vegas', 'actedIn' );
   });
+  const fullTextSearch = 'film';
+  console.log(`Full text search for movies with: '${fullTextSearch}'`);
+  const fullTextResults = await graphModel.fullTextQuery('Movie', fullTextSearch, 2);
+  console.log(fullTextResults);
   if(process.env.OPENAI_API_KEY) {
     const search = 'Working in a boring job and looking for love.';
     console.log(`Searching for movies related to: '${search}'`);

diff --git a/src/graphmodel.ts b/src/graphmodel.ts
@@ -9,6 +9,10 @@ type VectorIndex = {
     type: string;
 }
 
+type FullTextIndex = {
+    properties: Array<string>;
+}
+
 export type Context = {
     session: Session;
 }
@@ -142,6 +146,13 @@ export class GraphModel {
                 .find(s => s.getFullyQualifiedName() === `${ROOT_NAMESPACE}.GraphNode`));
     }
 
+    private getFullTextIndex(decl) : FullTextIndex|undefined {
+        const properties = decl.getProperties()
+        .filter(p => p.getDecorator('fulltext_index'))
+        .map( p => p.getName());
+        return properties.length > 0 ? { properties } : undefined;
+    }
+
     private getPropertyVectorIndex(property): VectorIndex {
         const decorator = property.getDecorator('vector_index');
         if (!decorator) {
@@ -187,6 +198,10 @@ export class GraphModel {
         return `${decl.getName()}_${vectorProperty.getName()}`.toLowerCase();
     }
 
+    private getFullTextIndexName(decl) {
+        return `${decl.getName()}_fulltext`.toLowerCase();
+    }
+
     async dropIndexes() {
         this.options.logger?.log('Dropping indexes...');
         const { session } = await this.openSession();
@@ -201,6 +216,11 @@ export class GraphModel {
                     const indexName = this.getPropertyVectorIndexName(graphNode, vectorProperty);
                     await tx.run(`DROP INDEX ${indexName} IF EXISTS`);
                 }
+                const fullTextIndex = this.getFullTextIndex(graphNode);
+                if(fullTextIndex) {
+                    const indexName = this.getFullTextIndexName(graphNode);
+                    await tx.run(`DROP INDEX ${indexName} IF EXISTS`);
+                }
             }
         })
         await session.close();
@@ -241,6 +261,25 @@ export class GraphModel {
         this.options.logger?.log('Create vector indexes completed');
     }
 
+    async createFullTextIndexes() {
+        this.options.logger?.log('Creating full text indexes...');
+        const { session } = await this.openSession();
+        await session.executeWrite(async tx => {
+            const graphNodes = this.getGraphNodeDeclarations();
+            for (let n = 0; n < graphNodes.length; n++) {
+                const graphNode = graphNodes[n];
+                const fullTextIndex = this.getFullTextIndex(graphNode);
+                if(fullTextIndex) {
+                    const indexName = this.getFullTextIndexName(graphNode);
+                    const props = fullTextIndex.properties.map( p => `n.${p}`);
+                    await tx.run(`CREATE FULLTEXT INDEX ${indexName} FOR (n:${graphNode.getName()}) ON EACH [${props.join(',')}];`);
+                }
+            }
+        })
+        await session.close();
+        this.options.logger?.log('Create full text indexes completed');
+    }
+
     async deleteGraph() {
         const { session } = await this.openSession();
         await session.executeWrite(async tx => {
@@ -392,6 +431,34 @@ export class GraphModel {
         }
     }
 
+    async fullTextQuery(typeName: string, searchText:string, count: number) {
+        try {
+            const graphNode = this.getGraphNodeDeclaration(typeName);
+            const fullTextIndex = this.getFullTextIndex(graphNode);
+            if(!fullTextIndex) {
+                throw new Error(`No full text index for properties of ${typeName}`);
+            }
+            const indexName = this.getFullTextIndexName(graphNode);
+            const props = fullTextIndex.properties.map( p => `node.${p}`);
+            props.push('node.identifier');
+            const q = `CALL db.index.fulltext.queryNodes("${indexName}", "${searchText}") YIELD node, score RETURN ${props.join(',')}, score limit ${count}`;
+            const queryResult = await this.query(q);
+            return queryResult ? queryResult.records.map(v => {
+                const result = {};
+                fullTextIndex.properties.forEach( p => {
+                    result[p] = v.get(`node.${p}`)
+                });
+                result['score'] = v.get('score');
+                result['identifier'] = v.get('node.identifier');
+                return result;
+            }) : [];
+        }
+        catch (err) {
+            this.options.logger?.log((err as object).toString());
+            throw err;
+        }
+    }
+
     private async validateAndTransformProperties(transaction, decl: ClassDeclaration, properties: PropertyBag): Promise<PropertyBag> {
         const factory = new Factory(this.modelManager);
         const serializer = new Serializer(factory, this.modelManager);