Added message history examples

Author: alexthomas93

examples/README.md

@@ -52,7 +52,7 @@ are listed in [the last section of this file](#customize).
 
 - [End to end GraphRAG](./answer/graphrag.py)
 - [GraphRAG with message history](./question_answering/graphrag_with_message_history.py)
-
+- [GraphRAG with Neo4j message history](./question_answering/graphrag_with_neo4j_message_history.py)
 
 ## Customize
 
@@ -75,6 +75,7 @@ are listed in [the last section of this file](#customize).
 - [Custom LLM](./customize/llms/custom_llm.py)
 
 - [Message history](./customize/llms/llm_with_message_history.py)
+- [Message history with Neo4j](./customize/llms/llm_with_neo4j_message_history.py)
 - [System Instruction](./customize/llms/llm_with_system_instructions.py)
 
 

examples/customize/llms/llm_with_neo4j_message_history.py

@@ -0,0 +1,61 @@
+"""This example illustrates the message_history feature
+of the LLMInterface by mocking a conversation between a user
+and an LLM about Tom Hanks.
+
+Neo4j is used as the database for storing the message history.
+
+OpenAILLM can be replaced by any supported LLM from this package.
+"""
+
+import neo4j
+from neo4j_graphrag.llm import LLMResponse, OpenAILLM
+from neo4j_graphrag.message_history import Neo4jMessageHistory
+
+# Define database credentials
+URI = "neo4j+s://demo.neo4jlabs.com"
+AUTH = ("recommendations", "recommendations")
+DATABASE = "recommendations"
+INDEX = "moviePlotsEmbedding"
+
+# set api key here on in the OPENAI_API_KEY env var
+api_key = None
+
+llm = OpenAILLM(model_name="gpt-4o", api_key=api_key)
+
+questions = [
+    "What are some movies Tom Hanks starred in?",
+    "Is he also a director?",
+    "Wow, that's impressive. And what about his personal life, does he have children?",
+]
+
+driver = neo4j.GraphDatabase.driver(
+    URI,
+    auth=AUTH,
+    database=DATABASE,
+)
+
+history = Neo4jMessageHistory(
+    session_id="123", driver=driver, node_label="Message", window=10
+)
+
+for question in questions:
+    res: LLMResponse = llm.invoke(
+        question,
+        message_history=history,
+    )
+    history.add_message(
+        {
+            "role": "user",
+            "content": question,
+        }
+    )
+    history.add_message(
+        {
+            "role": "assistant",
+            "content": res.content,
+        }
+    )
+
+    print("#" * 50, question)
+    print(res.content)
+    print("#" * 50)

examples/question_answering/graphrag_with_neo4j_message_history.py

@@ -0,0 +1,89 @@
+"""End to end example of building a RAG pipeline backed by a Neo4j database,
+simulating a chat with message history which is also stored in Neo4j.
+
+Requires OPENAI_API_KEY to be in the env var.
+"""
+
+import neo4j
+from neo4j_graphrag.embeddings.openai import OpenAIEmbeddings
+from neo4j_graphrag.generation import GraphRAG
+from neo4j_graphrag.llm import OpenAILLM
+from neo4j_graphrag.message_history import Neo4jMessageHistory
+from neo4j_graphrag.retrievers import VectorCypherRetriever
+
+# Define database credentials
+URI = "neo4j+s://demo.neo4jlabs.com"
+AUTH = ("recommendations", "recommendations")
+DATABASE = "recommendations"
+INDEX = "moviePlotsEmbedding"
+
+
+driver = neo4j.GraphDatabase.driver(
+    URI,
+    auth=AUTH,
+)
+
+embedder = OpenAIEmbeddings()
+
+retriever = VectorCypherRetriever(
+    driver,
+    index_name=INDEX,
+    retrieval_query="""
+        WITH node as movie, score
+        CALL(movie) {
+            MATCH (movie)<-[:ACTED_IN]-(p:Person)
+            RETURN collect(p.name) as actors
+        }
+        CALL(movie) {
+            MATCH (movie)<-[:DIRECTED]-(p:Person)
+            RETURN collect(p.name) as directors
+        }
+        RETURN movie.title as title, movie.plot as plot, movie.year as year, actors, directors
+    """,
+    embedder=embedder,
+    neo4j_database=DATABASE,
+)
+
+llm = OpenAILLM(model_name="gpt-4o", model_params={"temperature": 0})
+
+rag = GraphRAG(
+    retriever=retriever,
+    llm=llm,
+)
+
+history = Neo4jMessageHistory(
+    session_id="123", driver=driver, node_label="Message", window=10
+)
+
+questions = [
+    "Who starred in the Apollo 13 movies?",
+    "Who was its director?",
+    "In which year was this movie released?",
+]
+
+for question in questions:
+    result = rag.search(
+        question,
+        return_context=False,
+        message_history=history,
+    )
+
+    answer = result.answer
+    print("#" * 50, question)
+    print(answer)
+    print("#" * 50)
+
+    history.add_message(
+        {
+            "role": "user",
+            "content": question,
+        }
+    )
+    history.add_message(
+        {
+            "role": "assistant",
+            "content": answer,
+        }
+    )
+
+driver.close()