From 5d30e3cf29c8b84c4d23ee33ec8546f97f22bfdd Mon Sep 17 00:00:00 2001
From: Dan Selman <danscode@selman.org>
Date: Fri, 28 Jun 2024 18:54:39 +0100
Subject: [PATCH] feat: max content size option for TextToGraph

Signed-off-by: Dan Selman <danscode@selman.org>
---
 src/TextToGraph.ts |  3 ++-
 src/demo/index.ts  |  2 +-
 src/prompt.ts      | 50 ++++++++++++++++++++++++++++++++++++++++++++++
 src/types.ts       |  7 +++++++
 4 files changed, 60 insertions(+), 2 deletions(-)

diff --git a/src/TextToGraph.ts b/src/TextToGraph.ts
index b046f21..27fe375 100644
--- a/src/TextToGraph.ts
+++ b/src/TextToGraph.ts
@@ -25,6 +25,7 @@ export class TextToGraph {
             logger: options.logger,
             openAiOptions: options.openAiOptions,
             systemPrompt: options.textToGraphPrompt ? options.textToGraphPrompt : TEXT_TO_GRAPH_PROMPT,
+            maxContextSize: options.maxContextSize
         });
     }
 
@@ -34,7 +35,7 @@ export class TextToGraph {
      * @returns an object that describes which nodes and relationships were added
      */
     async mergeText(text: string) {
-        const messages = await this.conversation.runMessages([this.conversation.getSystemMessage()], `Add the nodes in this text to the knowledge graph: ${text}`);
+        const messages = await this.conversation.runMessages([this.conversation.getSystemMessage()], text);
         const relationships:Array<string> = [];
         const nodes:Array<string> = [];
         for(let n=0; n < messages.length; n++) {
diff --git a/src/demo/index.ts b/src/demo/index.ts
index de1f031..b18f5c0 100644
--- a/src/demo/index.ts
+++ b/src/demo/index.ts
@@ -147,7 +147,7 @@ War for the Planet of the Apes. Pitched by Ball as Apocalypto but with apes and
 the story follows a new protagonist Noa (Owen Teague) as he tries to steer the apes away from the totalitarian future 
 they are headed towards. Freya Allan, Kevin Durand, Dichen Lachman, and William H. Macy star.`;
 
-    const textToGraph = new TextToGraph(graphModel, {logger});
+    const textToGraph = new TextToGraph(graphModel, {logger, maxContextSize: 64000});
     const results = await textToGraph.mergeText(text);
     logger.success('added nodes and relationships to graph', results);  
   }
diff --git a/src/prompt.ts b/src/prompt.ts
index 3ba8024..f9b10e9 100644
--- a/src/prompt.ts
+++ b/src/prompt.ts
@@ -52,6 +52,56 @@ Do not enclose the result in a markdown code block.
 `}
 }
 
+export function getTextToGraphPrompt(ctoModel:string, text:string) : ChatCompletionSystemMessageParam {
+    return {
+        role: 'system',
+        content: `Convert the natural language query delimited by triple quotes to a Neo4J Cypher query.
+Just return the Cypher query, without an explanation for how it works.
+
+The nodes and edges in Neo4j are described by the following Accord Project Concerto model:
+\`\`\`
+${ctoModel}
+\`\`\`
+
+Concerto properties with the @vector_index decorator have a Neo4J vector index. The name
+of the vector index is the lowercase name of the declaration with '_' and the lowercase 
+name of the property appended. 
+For example: 'movie_summary' is the name of the vector index for the 'Movie.summary' property.
+
+Concerto declarations with any properties with the @fulltext_index decorator have a 
+full text index. The name of the full text index is the lowercase name of the declaration
+with '_fulltext' appended.
+For example: movie_fulltext is the name of the full text index for the 'Movie' declaration.
+
+Use the token ${EMBEDDINGS_MAGIC} to denote the embedding vector for input text.
+
+Here is an example Neo4J Cypher query that matches 3 movies summaries by conceptual similarity 
+(using vector cosine similarity):
+\`\`\`
+MATCH (l:Movie)
+    CALL db.index.vector.queryNodes('movie_summary', 10, ${EMBEDDINGS_MAGIC} )
+    YIELD node AS similar1, score
+    MATCH (similar)
+    RETURN similar.identifier as identifier, similar.summary as content, score limit 3
+\`\`\`
+
+Use 10 as the second argument to db.index.vector.queryNodes.
+
+Here is an example  Neo4J Cypher query that finds the shortest path between two people 'Dan Selman' and 'Ann Selman':
+\`\`\`
+MATCH
+  (a:Person {identifier: 'Dan Selman'}),
+  (b:Person {identifier: 'Ann Selman'}),
+  p = shortestPath((a)-[:RELATED_TO*]-(b))
+RETURN p
+\`\`\`
+
+Convert the following natural language query to Neo4J Cypher: """${text}
+"""
+Do not enclose the result in a markdown code block.
+`}
+}
+
 export const CONVERSATION_PROMPT = `Please use our knowledge graph, which you can access using tools to answer the following questions. 
 Query the knowledge graph using the 'chat_with_data' tool before trying other tools.`;
 
diff --git a/src/types.ts b/src/types.ts
index eba05dd..ba5af19 100644
--- a/src/types.ts
+++ b/src/types.ts
@@ -131,6 +131,13 @@ export type TextToGraphOptions = {
      * An optional prompt to further guide conversion of text to graph
      */
     textToGraphPrompt?: string
+    
+    /**
+     * The maximum context size for the conversation. Old messages
+     * will be automatically removed once the context size limit is
+     * reached
+     */
+    maxContextSize?: number;
 }
 
 /**