RAG Pipeline (neo4j#50)

* RAG Pipeline

* Ruff

* Copyright header

* Add doc

* Examples

* Add examples to template

* RAG Pipeline

* Ruff

* Copyright header

* Add doc

* Examples

* Add examples to template

* Fix tests

* Observer interface + types for RAG

* RAG Pipeline

* Ruff

* Copyright header

* Add doc

* Examples

* Add examples to template

* RAG Pipeline

* Fix tests

* Observer interface + types for RAG

* Mypy-related stuffs

* Fix e2e test

* More e2e tests

* Fix merge

* Turn prompt package into a module since it contains a single file

* Remove observer

* Update examples/rag_custom_prompt.py

Co-authored-by: Oskar Hane <oh@oskarhane.com>

* Use Literal type

* Update error handling (required small refactoring to avoid circular imports)

* Merge with main

* Rename RAG to GraphRAG

* Improved doc and naming

* Ruff

* Fix class name in docs

* Rename files to match class name "GraphRAG"

* Rename test cases for consistency with name "GraphRAG"

* Introduce LLMResponse model

* Update changelog

* Fix examples

* Ignore type so that mypy does not complain about None VS module types

* Changes to accept langchain (or other) models as input to the RAG pipeline + working example

* Update Text2CypherRetriever with new LLM interface

* Update poetry.lock as required by poetry during CI

* Fix tests

* Move example in README (and fix the example...)

* Update imports

---------

Co-authored-by: Oskar Hane <oh@oskarhane.com>

Author: stellasia

CHANGELOG.md

@@ -7,6 +7,9 @@
 -   Added `upsert_vector` utility function for attaching vectors to node properties.
 -   Introduced `Neo4jInsertionError` for handling insertion failures in Neo4j.
 -   Included Pinecone and Weaviate retrievers in neo4j_genai.retrievers.
+-   Introduced the GraphRAG object, enabling a full RAG (Retrieval-Augmented Generation) pipeline with context retrieval, prompt formatting, and answer generation.
+-   Added PromptTemplate and RagTemplate for customizable prompt generation.
+-   Added LLMInterface with implementation for OpenAI LLM.
 
 ### Changed
 -   Refactored import paths for retrievers to neo4j_genai.retrievers.

README.md

@@ -28,15 +28,21 @@ While the library has more retrievers than shown here, the following examples sh
 
 Assumption: Neo4j running with populated vector index in place.
 
+In the following example, we use a simple vector search as retriever,
+that will perform a similarity search over the `index-name` vector index
+in Neo4j.
+
 ```python
 from neo4j import GraphDatabase
 from neo4j_genai.retrievers import VectorRetriever
-from langchain_openai import OpenAIEmbeddings
+from neo4j_genai.llm import OpenAILLM
+from neo4j_genai.generation import GraphRAG
+from neo4j_genai.embeddings.openai import OpenAIEmbeddings
 
 URI = "neo4j://localhost:7687"
 AUTH = ("neo4j", "password")
 
-INDEX_NAME = "embedding-name"
+INDEX_NAME = "index-name"
 
 # Connect to Neo4j database
 driver = GraphDatabase.driver(URI, auth=AUTH)
@@ -47,9 +53,17 @@ embedder = OpenAIEmbeddings(model="text-embedding-3-large")
 # Initialize the retriever
 retriever = VectorRetriever(driver, INDEX_NAME, embedder)
 
-# Run the similarity search
+# Initialize the LLM
+# Note: the OPENAI_API_KEY must be in the env vars
+llm = OpenAILLM(model_name="gpt-4o", model_params={"temperature": 0})
+
+# Initialize the RAG pipeline
+rag = GraphRAG(retriever=retriever, llm=llm)
+
+# Query the graph
 query_text = "How do I do similarity search in Neo4j?"
-response = retriever.search(query_text=query_text, top_k=5)
+response = rag.search(query_text=query_text, retriever_config={"top_k": 5})
+print(response.answer)
 ```
 
 ### Creating a vector index

docs/source/api.rst

@@ -3,18 +3,17 @@
 API Documentation
 #################
 
-************************************
-Retrieval-Augmented Generation (RAG)
-************************************
-RAG is a technique that enhances Large Language Model (LLM) responses by retrieving
-source information from external data stores to augment generated responses.
-
-This package enables Python developers to perform RAG using Neo4j.
-
 **********
 Retrievers
 **********
 
+RetrieverInterface
+===================
+
+.. autoclass:: neo4j_genai.retrievers.base.Retriever
+    :members:
+
+
 VectorRetriever
 ===============
 
@@ -89,6 +88,12 @@ Errors
 
   * :class:`neo4j_genai.exceptions.SchemaFetchError`
 
+  * :class:`neo4j_genai.exceptions.RagInitializationError`
+
+  * :class:`neo4j_genai.exceptions.PromptMissingInputError`
+
+  * :class:`neo4j_genai.exceptions.LLMGenerationError`
+
 
 Neo4jGenAiError
 ===============
@@ -165,3 +170,24 @@ SchemaFetchError
 
 .. autoclass:: neo4j_genai.exceptions.SchemaFetchError
    :show-inheritance:
+
+
+RagInitializationError
+==========================
+
+.. autoclass:: neo4j_genai.exceptions.RagInitializationError
+   :show-inheritance:
+
+
+PromptMissingInputError
+==========================
+
+.. autoclass:: neo4j_genai.exceptions.PromptMissingInputError
+   :show-inheritance:
+
+
+LLMGenerationError
+==========================
+
+.. autoclass:: neo4j_genai.exceptions.LLMGenerationError
+   :show-inheritance:

docs/source/conf.py

@@ -44,6 +44,7 @@
     "sphinx.ext.intersphinx",
     "sphinx.ext.napoleon",
     "sphinx.ext.viewcode",
+    "sphinx.ext.autosectionlabel",
 ]
 
 # The suffix(es) of source filenames.

docs/source/index.rst

@@ -16,6 +16,7 @@ fast to ship new features and high performing patterns and methods.
 Topics
 ******
 
++ :ref:`rag-documentation`
 + :ref:`api-documentation`
 + :ref:`types-documentation`
 
@@ -24,6 +25,7 @@ Topics
     :caption: Contents:
     :hidden:
 
+    rag.rst
     api.rst
     types.rst
 

docs/source/rag.rst

@@ -0,0 +1,103 @@
+.. _rag-documentation:
+
+RAG Documentation
+#################
+
+************************************
+Retrieval-Augmented Generation (RAG)
+************************************
+RAG is a technique that enhances Large Language Model (LLM) responses by retrieving
+source information from external data stores to augment generated responses.
+
+This package enables Python developers to perform RAG using Neo4j.
+
+
+************************************
+Overview
+************************************
+
+.. code:: python
+
+    from neo4j import GraphDatabase
+    from neo4j_genai.retrievers import VectorRetriever
+    from neo4j_genai.llm import OpenAILLM
+    from neo4j_genai.generation import GraphRAG
+    from neo4j_genai.embeddings.openai import OpenAIEmbeddings
+
+    URI = "neo4j://localhost:7687"
+    AUTH = ("neo4j", "password")
+
+    INDEX_NAME = "index-name"
+
+    # Connect to Neo4j database
+    driver = GraphDatabase.driver(URI, auth=AUTH)
+
+    # Create Embedder object
+    embedder = OpenAIEmbeddings(model="text-embedding-3-large")
+
+    # Initialize the retriever
+    retriever = VectorRetriever(driver, INDEX_NAME, embedder)
+
+    # Initialize the LLM
+    # Note: the OPENAI_API_KEY must be in the env vars
+    llm = OpenAILLM(model_name="gpt-4o", model_params={"temperature": 0})
+
+    # Initialize the RAG pipeline
+    rag = GraphRAG(retriever=retriever, llm=llm)
+
+    # Query the graph
+    query_text = "How do I do similarity search in Neo4j?"
+    response = rag.search(query_text=query_text, retriever_config={"top_k": 5})
+    print(response.answer)
+
+
+The retriever can be any of the :ref:`supported retrievers<retrievers>`, or any class
+inheriting from the `Retriever` interface.
+
+***************
+Advanced usage
+***************
+
+
+Using another LLM
+==================
+
+This package only provide support for OpenAI LLM. If you need to use another LLM,
+you need to subclass the `LLMInterface`:
+
+.. autoclass:: neo4j_genai.llm.LLMInterface
+    :members:
+    :show-inheritance:
+
+Configuring the prompt
+=======================
+
+Prompt are managed through `PromptTemplate` classes. More
+specifically, the `RAG` pipeline uses a `RagTemplate` with
+a default prompt. You can use another prompt by subclassing
+the `RagTemplate` class and passing it to the `RAG` pipeline
+object during initialization:
+
+.. code:: python
+
+    from neo4j_genai.generation import RagTemplate, GraphRAG
+
+    # ...
+
+    prompt_template = RagTemplate(
+        prompt="Answer the question {question} using context {context} and examples {examples}",
+        expected_inputs=["context", "question", "examples"]
+    )
+    rag = GraphRAG(retriever=vector_retriever, llm=llm, prompt_template=prompt_template)
+
+    # ...
+
+For more details, see:
+
+.. autoclass:: neo4j_genai.generation.prompts.PromptTemplate
+    :members:
+
+and
+
+.. autoclass:: neo4j_genai.generation.prompts.RagTemplate
+    :members:

examples/graphrag.py

@@ -0,0 +1,58 @@
+"""End to end example of building a RAG pipeline backed by a Neo4j database.
+Requires OPENAI_API_KEY to be in the env var.
+
+This example illustrates:
+- VectorCypherRetriever with a custom formatter function to extract relevant
+    context from neo4j result
+- Logging configuration
+"""
+
+import logging
+import neo4j
+
+from neo4j_genai.embeddings.openai import OpenAIEmbeddings
+from neo4j_genai.types import RetrieverResultItem
+from neo4j_genai.retrievers import VectorCypherRetriever
+from neo4j_genai.llm import OpenAILLM
+from neo4j_genai.generation import GraphRAG
+
+URI = "neo4j://localhost:7687"
+AUTH = ("neo4j", "password")
+DATABASE = "neo4j"
+INDEX = "moviePlotsEmbedding"
+
+
+# setup logger config
+logger = logging.getLogger("neo4j_genai")
+logging.basicConfig(format="%(asctime)s - %(message)s")
+logger.setLevel(logging.DEBUG)
+
+
+def formatter(record: neo4j.Record) -> RetrieverResultItem:
+    return RetrieverResultItem(content=f'{record.get("title")}: {record.get("plot")}')
+
+
+driver = neo4j.GraphDatabase.driver(
+    URI,
+    auth=AUTH,
+    database=DATABASE,
+)
+
+embedder = OpenAIEmbeddings()
+
+retriever = VectorCypherRetriever(
+    driver,
+    index_name=INDEX,
+    retrieval_query="with node, score return node.title as title, node.plot as plot",
+    format_record_function=formatter,
+    embedder=embedder,
+)
+
+llm = OpenAILLM(model_name="gpt-4o", model_params={"temperature": 0})
+
+rag = GraphRAG(retriever=retriever, llm=llm)
+
+result = rag.search("Tell me more about Avatar movies")
+print(result.answer)
+
+driver.close()

examples/graphrag_custom_prompt.py

@@ -0,0 +1,74 @@
+"""End to end example of building a RAG pipeline backed by a Neo4j database.
+Requires OPENAI_API_KEY to be in the env var.
+
+This example illustrates:
+- VectorCypherRetriever with a custom formatter function to extract relevant
+    context from neo4j result
+- Use of a custom prompt for RAG
+- Logging configuration
+"""
+
+import logging
+import neo4j
+
+from neo4j_genai.types import RetrieverResultItem
+from neo4j_genai.embeddings.openai import OpenAIEmbeddings
+from neo4j_genai.retrievers import VectorCypherRetriever
+from neo4j_genai.generation import GraphRAG, RagTemplate
+from neo4j_genai.llm import OpenAILLM
+
+URI = "neo4j://localhost:7687"
+AUTH = ("neo4j", "password")
+DATABASE = "neo4j"
+INDEX = "moviePlotsEmbedding"
+
+
+# setup logger config
+logger = logging.getLogger("neo4j_genai")
+logging.basicConfig(format="%(asctime)s - %(message)s")
+logger.setLevel(logging.DEBUG)
+
+
+def formatter(record: neo4j.Record) -> RetrieverResultItem:
+    return RetrieverResultItem(content=f'{record.get("title")}: {record.get("plot")}')
+
+
+driver = neo4j.GraphDatabase.driver(
+    URI,
+    auth=AUTH,
+    database=DATABASE,
+)
+
+embedder = OpenAIEmbeddings()
+
+retriever = VectorCypherRetriever(
+    driver,
+    index_name=INDEX,
+    retrieval_query="with node, score return node.title as title, node.plot as plot",
+    format_record_function=formatter,
+    embedder=embedder,
+)
+
+llm = OpenAILLM(model_name="gpt-4o", model_params={"temperature": 0})
+
+template = RagTemplate(
+    template="""You are an expert at movies and actors. Your task is to
+    answer the user's question based on the provided context. Use only the
+    information within that context.
+
+    Context:
+    {context}
+
+    Question:
+    {query}
+
+    Answer:
+    """
+)
+
+rag = GraphRAG(retriever=retriever, llm=llm, prompt_template=template)
+
+result = rag.search("Tell me more about Avatar movies")
+print(result.answer)
+
+driver.close()