Mistral support (#142)

* Add MistralAIEmbeddings and MistralAILLM

* Add examples

* Add unit tests for Mistral embedder and LLM classes

* Update docs

* Ruff and modify example

* Fixed docstring

* Address comments

* Fix tests

* Fix get_message private method

* Fixed mypy errors

* Replace | operand with explicit typing.Union

* Update CHANGELOG

* Generic errors

* Check for api_key in kwargs first before using MISTRAL_API_KEY env var

* PR comments and update docs

* Removed types-requests

* Removed redundant items in CHANGELOG

* Add EmbeddingsGenerationError

* Resolve deps issues

* Upload poetry.lock

* Modified check for failure to retrieve embeddings

* EmbeddingsGenerationError for embed_query

Author: willtai

CHANGELOG.md

@@ -7,6 +7,7 @@
 - Added `template` validation in `PromptTemplate` class upon construction.
 - `custom_prompt` arg is now converted to `Text2CypherTemplate` class within the `Text2CypherRetriever.get_search_results` method.
 - `Text2CypherTemplate` and `RAGTemplate` prompt templates now require `query_text` arg and will error if it is not present. Previous `query_text` aliases may be used, but will warn of deprecation.
+- Examples demonstrating the use of Mistral embeddings and LLM in RAG pipelines.
 - Fixed bug in `Text2CypherRetriever` using `custom_prompt` arg where the `search` method would not inject the `query_text` content.
 - Added feature to include kwargs in `Text2CypherRetriever.search()` that will be injected into a custom prompt, if provided.
 - Added validation to `custom_prompt` parameter of `Text2CypherRetriever` to ensure that `query_text` placeholder exists in prompt.
@@ -15,6 +16,7 @@
 - Added unit tests for the Vertex AI LLM class.
 - Added support for Cohere LLM and embeddings - added optional dependency to `cohere`.
 - Added support for Anthropic LLM - added optional dependency to `anthropic`.
+- Added support for MistralAI LLM - added optional dependency to `mistralai`.
 
 ### Fixed
 - Resolved import issue with the Vertex AI Embeddings class.

docs/source/api.rst

@@ -156,12 +156,24 @@ OpenAIEmbeddings
 .. autoclass:: neo4j_graphrag.embeddings.openai.OpenAIEmbeddings
     :members:
 
+AzureOpenAIEmbeddings
+=====================
+
+.. autoclass:: neo4j_graphrag.embeddings.openai.AzureOpenAIEmbeddings
+    :members:
+
 VertexAIEmbeddings
 ==================
 
 .. autoclass:: neo4j_graphrag.embeddings.vertexai.VertexAIEmbeddings
     :members:
 
+MistralAIEmbeddings
+===================
+
+.. autoclass:: neo4j_graphrag.embeddings.mistral.MistralAIEmbeddings
+    :members:
+
 CohereEmbeddings
 ================
 
@@ -218,6 +230,13 @@ CohereLLM
     :members:
 
 
+MistralAILLM
+------------
+
+.. autoclass:: neo4j_graphrag.llm.mistralai_llm.MistralAILLM
+    :members:
+
+
 PromptTemplate
 ==============
 
@@ -282,6 +301,8 @@ Errors
 
   * :class:`neo4j_graphrag.exceptions.RetrieverInitializationError`
 
+  * :class:`neo4j_graphrag.exceptions.EmbeddingsGenerationError`
+
   * :class:`neo4j_graphrag.exceptions.SearchValidationError`
 
   * :class:`neo4j_graphrag.exceptions.FilterValidationError`

docs/source/user_guide_rag.rst

@@ -393,14 +393,16 @@ into a vector is required. Therefore, the retriever requires knowledge of an emb
 Embedders
 -----------------------------
 
-Currently, this package supports several embedders:
-- `OpenAIEmbeddings`
-- `AzureOpenAIEmbeddings`
-- `VertexAIEmbeddings`
-- `CohereEmbeddings`
-- `SentenceTransformerEmbeddings`.
-
-The `OpenAIEmbedder` was illustrated previously. Here is how to use the `SentenceTransformerEmbeddings`:
+Currently, this package supports the following embedders:
+
+- :ref:`openaiembeddings`
+- :ref:`sentencetransformerembeddings`
+- :ref:`vertexaiembeddings`
+- :ref:`mistralaiembeddings`
+- :ref:`cohereembeddings`
+- :ref:`azureopenaiembeddings`
+
+The `OpenAIEmbeddings` was illustrated previously. Here is how to use the `SentenceTransformerEmbeddings`:
 
 .. code:: python
 

examples/graphrag_with_mistral.py

@@ -0,0 +1,68 @@
+"""End to end example of building a RAG pipeline backed by a Neo4j database.
+Requires MISTRAL_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_graphrag.embeddings.mistral import MistralAIEmbeddings
+from neo4j_graphrag.generation import GraphRAG
+from neo4j_graphrag.indexes import create_vector_index
+from neo4j_graphrag.llm.mistralai_llm import MistralAILLM
+from neo4j_graphrag.retrievers import VectorCypherRetriever
+from neo4j_graphrag.types import RetrieverResultItem
+
+URI = "neo4j://localhost:7687"
+AUTH = ("neo4j", "password")
+DATABASE = "neo4j"
+INDEX_NAME = "moviePlotsEmbedding"
+
+
+# setup logger config
+logger = logging.getLogger("neo4j_graphrag")
+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,
+)
+
+create_vector_index(
+    driver,
+    INDEX_NAME,
+    label="Document",
+    embedding_property="vectorProperty",
+    dimensions=1024,
+    similarity_fn="cosine",
+)
+
+
+embedder = MistralAIEmbeddings()
+
+retriever = VectorCypherRetriever(
+    driver,
+    index_name=INDEX_NAME,
+    retrieval_query="with node, score return node.title as title, node.plot as plot",
+    result_formatter=formatter,
+    embedder=embedder,
+)
+
+llm = MistralAILLM(model_name="mistral-small-latest")
+
+rag = GraphRAG(retriever=retriever, llm=llm)
+
+result = rag.search("Tell me more about Avatar movies")
+print(result.answer)
+
+driver.close()

examples/similarity_search_for_text_mistral.py

@@ -0,0 +1,50 @@
+from __future__ import annotations
+
+from random import random
+
+from neo4j import GraphDatabase
+from neo4j_graphrag.embeddings.mistral import MistralAIEmbeddings
+from neo4j_graphrag.indexes import create_vector_index
+from neo4j_graphrag.retrievers import VectorRetriever
+
+URI = "neo4j://localhost:7687"
+AUTH = ("neo4j", "password")
+
+INDEX_NAME = "embedding-name"
+DIMENSION = 1024
+
+# Connect to Neo4j database
+driver = GraphDatabase.driver(URI, auth=AUTH)
+
+embedder = MistralAIEmbeddings()
+
+# Creating the index
+create_vector_index(
+    driver,
+    INDEX_NAME,
+    label="Document",
+    embedding_property="vectorProperty",
+    dimensions=DIMENSION,
+    similarity_fn="euclidean",
+)
+
+# Initialize the retriever
+retriever = VectorRetriever(driver, INDEX_NAME, embedder)
+
+# Upsert the query
+vector = [random() for _ in range(DIMENSION)]
+insert_query = (
+    "MERGE (n:Document {id: $id})"
+    "WITH n "
+    "CALL db.create.setNodeVectorProperty(n, 'vectorProperty', $vector)"
+    "RETURN n"
+)
+parameters = {
+    "id": 0,
+    "vector": vector,
+}
+driver.execute_query(insert_query, parameters)
+
+# Perform the similarity search for a text query
+query_text = "Find me a book about Fremen"
+print(retriever.search(query_text=query_text, top_k=5))