Skip to content

Commit

Permalink
feat: support short type names and reorg vector index
Browse files Browse the repository at this point in the history 
Signed-off-by: Dan Selman <danscode@selman.org>
  • Loading branch information
@dselman
dselman committed Dec 12, 2023
1 parent 7e543c7 commit ea49e0c
Show file tree
Hide file tree
Showing 3 changed files with 66 additions and 38 deletions.
48 changes: 24 additions & 24 deletions src/demo/index.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@ concept Address {
o String country
}
// show how maps get flattened
scalar PersonName extends String
scalar PersonEmail extends String
map AddressBook {
Expand Down Expand Up @@ -47,9 +48,8 @@ concept Genre extends GraphNode {
}
concept Movie extends GraphNode {
@vector_index("summary", 1536, "COSINE")
o Double[] embedding optional
@embedding
@vector_index("embedding", 1536, "COSINE")
o String summary optional
@label("IN_GENRE")
--> Genre[] genres optional
Expand All @@ -65,7 +65,7 @@ async function run() {
logQueries: false,
embeddingFunction: process.env.OPENAI_API_KEY ? getOpenAiEmbedding : undefined
}
const graphModel = new GraphModel(MODEL, options);
const graphModel = new GraphModel([MODEL], options);
await graphModel.connect();
await graphModel.dropIndexes();
await graphModel.createConstraints();
Expand All @@ -84,37 +84,37 @@ async function run() {
const addressBook = {
'Dan' : 'dan@example.com'
};
await graphModel.mergeNode(transaction, `${NS}.Movie`, {identifier: 'Brazil', summary: 'The film centres on Sam Lowry, a low-ranking bureaucrat trying to find a woman who appears in his dreams while he is working in a mind-numbing job and living in a small apartment, set in a dystopian world in which there is an over-reliance on poorly maintained (and rather whimsical) machines'} );
await graphModel.mergeNode(transaction, `${NS}.Movie`, {identifier: 'The Man Who Killed Don Quixote', summary: 'Instead of a literal adaptation, Gilliam\'s film was about "an old, retired, and slightly kooky nobleman named Alonso Quixano".'} );
await graphModel.mergeNode(transaction, `${NS}.Movie`, {identifier: 'Fear and Loathing in Las Vegas', summary: 'Duke, under the influence of mescaline, complains of a swarm of giant bats, and inventories their drug stash. They pick up a young hitchhiker and explain their mission: Duke has been assigned by a magazine to cover the Mint 400 motorcycle race in Las Vegas. They bought excessive drugs for the trip, and rented a red Chevrolet Impala convertible.'} );
await graphModel.mergeNode(transaction, 'Movie', {identifier: 'Brazil', summary: 'The film centres on Sam Lowry, a low-ranking bureaucrat trying to find a woman who appears in his dreams while he is working in a mind-numbing job and living in a small apartment, set in a dystopian world in which there is an over-reliance on poorly maintained (and rather whimsical) machines'} );
await graphModel.mergeNode(transaction, 'Movie', {identifier: 'The Man Who Killed Don Quixote', summary: 'Instead of a literal adaptation, Gilliam\'s film was about "an old, retired, and slightly kooky nobleman named Alonso Quixano".'} );
await graphModel.mergeNode(transaction, 'Movie', {identifier: 'Fear and Loathing in Las Vegas', summary: 'Duke, under the influence of mescaline, complains of a swarm of giant bats, and inventories their drug stash. They pick up a young hitchhiker and explain their mission: Duke has been assigned by a magazine to cover the Mint 400 motorcycle race in Las Vegas. They bought excessive drugs for the trip, and rented a red Chevrolet Impala convertible.'} );

await graphModel.mergeNode(transaction, `${NS}.Genre`, {identifier: 'Comedy'} );
await graphModel.mergeNode(transaction, `${NS}.Genre`, {identifier: 'Science Fiction'} );
await graphModel.mergeNode(transaction, 'Genre', {identifier: 'Comedy'} );
await graphModel.mergeNode(transaction, 'Genre', {identifier: 'Science Fiction'} );

await graphModel.mergeRelationship(transaction, `${NS}.Movie`, 'Brazil', `${NS}.Genre`, 'Comedy', 'genres' );
await graphModel.mergeRelationship(transaction, `${NS}.Movie`, 'Brazil', `${NS}.Genre`, 'Science Fiction', 'genres' );
await graphModel.mergeRelationship(transaction, `${NS}.Movie`, 'The Man Who Killed Don Quixote', `${NS}.Genre`, 'Comedy', 'genres' );
await graphModel.mergeRelationship(transaction, `${NS}.Movie`, 'Fear and Loathing in Las Vegas', `${NS}.Genre`, 'Comedy', 'genres' );
await graphModel.mergeRelationship(transaction, 'Movie', 'Brazil', 'Genre', 'Comedy', 'genres' );
await graphModel.mergeRelationship(transaction, 'Movie', 'Brazil', 'Genre', 'Science Fiction', 'genres' );
await graphModel.mergeRelationship(transaction, 'Movie', 'The Man Who Killed Don Quixote', 'Genre', 'Comedy', 'genres' );
await graphModel.mergeRelationship(transaction, 'Movie', 'Fear and Loathing in Las Vegas', 'Genre', 'Comedy', 'genres' );

await graphModel.mergeNode(transaction, `${NS}.Director`, {identifier: 'Terry Gilliam'} );
await graphModel.mergeRelationship(transaction, `${NS}.Director`, 'Terry Gilliam', `${NS}.Movie`, 'Brazil', 'directed' );
await graphModel.mergeRelationship(transaction, `${NS}.Director`, 'Terry Gilliam', `${NS}.Movie`, 'The Man Who Killed Don Quixote', 'directed' );
await graphModel.mergeRelationship(transaction, `${NS}.Director`, 'Terry Gilliam', `${NS}.Movie`, 'Fear and Loathing in Las Vegas', 'directed' );
await graphModel.mergeNode(transaction, 'Director', {identifier: 'Terry Gilliam'} );
await graphModel.mergeRelationship(transaction, 'Director', 'Terry Gilliam', 'Movie', 'Brazil', 'directed' );
await graphModel.mergeRelationship(transaction, 'Director', 'Terry Gilliam', 'Movie', 'The Man Who Killed Don Quixote', 'directed' );
await graphModel.mergeRelationship(transaction, 'Director', 'Terry Gilliam', 'Movie', 'Fear and Loathing in Las Vegas', 'directed' );

await graphModel.mergeNode(transaction, `${NS}.User`, {identifier: 'Dan', address, addressBook} );
await graphModel.mergeRelationship(transaction, `${NS}.User`, 'Dan', `${NS}.Movie`, 'Brazil', 'ratedMovies' );
await graphModel.mergeNode(transaction, 'User', {identifier: 'Dan', address, addressBook} );
await graphModel.mergeRelationship(transaction, 'User', 'Dan', 'Movie', 'Brazil', 'ratedMovies' );

await graphModel.mergeNode(transaction, `${NS}.Actor`, {identifier: 'Jonathan Pryce'} );
await graphModel.mergeRelationship(transaction, `${NS}.Actor`, 'Jonathan Pryce', `${NS}.Movie`, 'Brazil', 'actedIn' );
await graphModel.mergeRelationship(transaction, `${NS}.Actor`, 'Jonathan Pryce', `${NS}.Movie`, 'The Man Who Killed Don Quixote', 'actedIn' );
await graphModel.mergeNode(transaction, 'Actor', {identifier: 'Jonathan Pryce'} );
await graphModel.mergeRelationship(transaction, 'Actor', 'Jonathan Pryce', 'Movie', 'Brazil', 'actedIn' );
await graphModel.mergeRelationship(transaction, 'Actor', 'Jonathan Pryce', 'Movie', 'The Man Who Killed Don Quixote', 'actedIn' );

await graphModel.mergeNode(transaction, `${NS}.Actor`, {identifier: 'Johnny Depp'} );
await graphModel.mergeRelationship(transaction, `${NS}.Actor`, 'Johnny Depp', `${NS}.Movie`, 'Fear and Loathing in Las Vegas', 'actedIn' );
await graphModel.mergeNode(transaction, 'Actor', {identifier: 'Johnny Depp'} );
await graphModel.mergeRelationship(transaction, 'Actor', 'Johnny Depp', 'Movie', 'Fear and Loathing in Las Vegas', 'actedIn' );
});
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}'`);
const results = await graphModel.similarityQuery(`${NS}.Movie`, 'embedding', search, 3);
const results = await graphModel.similarityQuery('Movie', 'summary', search, 3);
console.log(results);
}
await graphModel.closeSession(context);
Expand Down
2 changes: 1 addition & 1 deletion src/graphmodel.test.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,7 @@ describe('GraphModel', () => {
NEO4J_URL: 'Pasta',
NEO4J_PASS: 'pasta',
}
const graphModel = new GraphModel(cto, options);
const graphModel = new GraphModel([cto], options);
expect(graphModel).not.toBeNull()
})
})
Expand Down
54 changes: 41 additions & 13 deletions src/graphmodel.ts
View file
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
import { ClassDeclaration, Factory, Introspector, ModelManager, RelationshipDeclaration, Serializer } from "@accordproject/concerto-core";
import { ClassDeclaration, Factory, Introspector, ModelManager, ModelUtil, RelationshipDeclaration, Serializer } from "@accordproject/concerto-core";
import neo4j, { DateTime, Driver, ManagedTransaction, Session } from 'neo4j-driver';
import * as crypto from 'crypto'
import OpenAI from "openai";
Expand Down Expand Up @@ -68,8 +68,8 @@ concept GraphNode identified by identifier {
o String identifier
}
concept EmbeddingCacheNode extends GraphNode {
@vector_index("content", 1536, "COSINE")
o Double[] embedding
@vector_index("embedding", 1536, "COSINE")
o String content
}
`;
Expand All @@ -83,12 +83,17 @@ export class GraphModel {
modelManager: ModelManager;
driver: Driver | undefined = undefined;
options: GraphModelOptions;
defaultNamespace: string|undefined;

constructor(graphModel: string, options: GraphModelOptions) {
constructor(graphModels: Array<string>, options: GraphModelOptions) {
this.options = options;
this.modelManager = new ModelManager({ strict: true, enableMapType: true });
this.modelManager.addCTOModel(ROOT_MODEL, 'root.cto');
this.modelManager.addCTOModel(graphModel, 'model.cto');
graphModels.forEach( (model, index) => {
const mf = this.modelManager.addCTOModel(model, `model-${index}.cto`, true);
this.defaultNamespace = mf.getNamespace();
})
this.modelManager.validateModelFiles();
}

async connect() {
Expand All @@ -113,7 +118,15 @@ export class GraphModel {
context.session.close();
}

private getGraphNodeDeclaration(fqn: string) {
private getFullyQualifiedType(type:string) {
const typeNs = ModelUtil.getNamespace(type);
const typeShortName = ModelUtil.getShortName(type);
const ns = typeNs.length > 0 ? typeNs : this.defaultNamespace;
return ModelUtil.getFullyQualifiedName(ns ? ns : '',typeShortName);
}

private getGraphNodeDeclaration(typeName: string) {
const fqn = this.getFullyQualifiedType(typeName);
const decl = this.modelManager.getType(fqn);
const superTypesNames = decl.getAllSuperTypeDeclarations().map(s => s.getFullyQualifiedName());
if (!superTypesNames.includes('org.accordproject.graph@1.0.0.GraphNode')) {
Expand Down Expand Up @@ -150,8 +163,21 @@ export class GraphModel {
if (typeof args[2] !== 'string') {
throw new Error(`@vector_index decorator on property ${property.getFullyQualifiedName()} does not have a second argument that is a string`);
}

if(property.getType() !== 'String') {
throw new Error(`@vector_index decorator on property ${property.getFullyQualifiedName()} is invalid. Can only be added to String properties.`);
}
const propertyName = property.getDecorator('vector_index').getArguments()[0] as unknown as string;
const embeddingProperty = property.getParent().getProperty(propertyName);
if(!embeddingProperty) {
throw new Error(`@vector_index decorator on property ${property.getFullyQualifiedName()} is invalid. References the property ${propertyName} which does not exist.`);
}

if(embeddingProperty.getType() !== 'Double' && !embeddingProperty.isArray() ) {
throw new Error(`@vector_index decorator on property ${property.getFullyQualifiedName()} is invalid. It references the property ${propertyName} but the property is not Double[].`);
}
return {
property: property.getDecorator('vector_index').getArguments()[0] as unknown as string,
property: propertyName,
size: property.getDecorator('vector_index').getArguments()[1] as unknown as number,
type: property.getDecorator('vector_index').getArguments()[2] as unknown as string,
}
Expand Down Expand Up @@ -207,7 +233,7 @@ export class GraphModel {
const vectorProperty = vectorProperties[i];
const vectorIndex = this.getPropertyVectorIndex(vectorProperty);
const indexName = this.getPropertyVectorIndexName(graphNode, vectorProperty);
await tx.run(`CALL db.index.vector.createNodeIndex("${indexName}", "${graphNode.getName()}", "${vectorProperty.getName()}", ${vectorIndex.size}, "${vectorIndex.type}")`);
await tx.run(`CALL db.index.vector.createNodeIndex("${indexName}", "${graphNode.getName()}", "${vectorIndex.property}", ${vectorIndex.size}, "${vectorIndex.type}")`);
}
}
})
Expand Down Expand Up @@ -238,14 +264,14 @@ export class GraphModel {
throw new Error(`${typeName} does not have a property ${propertyName}`);
}
// check this is a vector index property
const vectorIndex = this.getPropertyVectorIndex(vectorProperty);
this.getPropertyVectorIndex(vectorProperty);
const index = this.getPropertyVectorIndexName(decl, vectorProperty);
const q = `MATCH (l:${decl.getName()})
CALL db.index.vector.queryNodes('${index}', ${count}, ${JSON.stringify(embedding)} )
YIELD node AS similar, score
MATCH (similar)
RETURN similar.identifier as identifier, similar.${vectorIndex.property} as content, score limit ${count}`;
const queryResult = await this.query(q);
RETURN similar.identifier as identifier, similar.${propertyName} as content, score limit ${count}`;
const queryResult = await this.query(q);
return queryResult ? queryResult.records.map(v => {
return {
identifier: v.get('identifier'),
Expand Down Expand Up @@ -384,13 +410,15 @@ export class GraphModel {
if(key !== '$class') {
const value = properties[key];
const property = decl.getProperty(key);
const embeddingDecorator = property.getDecorator('embedding');
if(embeddingDecorator && this.options.embeddingFunction) {
const embeddingDecorator = property.getDecorator('vector_index');
if(decl.getFullyQualifiedName() !== `${ROOT_NAMESPACE}.EmbeddingCacheNode` &&
embeddingDecorator && this.options.embeddingFunction) {
const vectorIndex = this.getPropertyVectorIndex(property);
if(typeof value !== 'string') {
throw new Error(`Can only calculate embedding for string properties`);
}
const cacheNode = await this.mergeEmbeddingCacheNode(transaction, value as string);
newProperties['embedding'] = cacheNode.embedding;
newProperties[vectorIndex.property] = cacheNode.embedding;
newProperties[key] = value;
}
else if (property.getType() === 'DateTime') {
Expand Down

0 comments on commit ea49e0c

Please sign in to comment.