Support for Cohere embeddings and LLM (#147)

* Support Cohere embeddings and LLM

* Fix tests

* Ruff

* Fix dependencies

* Fix mock

* Add test for cohere LLM failure

* Rename file for consistency

* Ruff

* Mypy

* Fix tests

* Documentation & CHANGELOG

* Recreate lock file after merge

Author: stellasia

CHANGELOG.md

@@ -4,15 +4,16 @@
 
 ### Added
 - Added AzureOpenAILLM and AzureOpenAIEmbeddings to support Azure served OpenAI models
-- Add `template` validation in `PromptTemplate` class upon construction.
+- 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.
-- Fix bug in `Text2CypherRetriever` using `custom_prompt` arg where the `search` method would not inject the `query_text` content.
-- Add feature to include kwargs in `Text2CypherRetriever.search()` that will be injected into a custom prompt, if provided.
-- Add validation to `custom_prompt` parameter of `Text2CypherRetriever` to ensure that `query_text` placeholder exists in prompt.
+- 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.
 - Introduced a fixed size text splitter component for splitting text into specified fixed size chunks with overlap. Updated examples and tests to utilize this new component.
 - Introduced Vertex AI LLM class for integrating Vertex AI models.
 - Added unit tests for the Vertex AI LLM class.
+- Added support for Cohere LLM and embeddings - added optional dependency to `cohere`.
 
 ### Fixed
 - Resolved import issue with the Vertex AI Embeddings class.
@@ -31,6 +32,9 @@
 - Updated documentation to include OpenAI and Vertex AI embeddings classes.
 - Added google-cloud-aiplatform as an optional dependency for Vertex AI embeddings.
 
+### Fixed
+- Make `pygraphviz` an optional dependency - it is now only required when calling `pipeline.draw`.
+
 ## 0.6.2
 
 ### Fixed

docs/source/api.rst

@@ -162,6 +162,12 @@ VertexAIEmbeddings
 .. autoclass:: neo4j_graphrag.embeddings.vertexai.VertexAIEmbeddings
     :members:
 
+CohereEmbeddings
+================
+
+.. autoclass:: neo4j_graphrag.embeddings.cohere.CohereEmbeddings
+    :members:
+
 **********
 Generation
 **********
@@ -174,17 +180,26 @@ LLMInterface
 
 
 OpenAILLM
-=========
+---------
 
-.. autoclass:: neo4j_graphrag.llm.openai.OpenAILLM
+.. autoclass:: neo4j_graphrag.llm.openai_llm.OpenAILLM
     :members:
+    :undoc-members: get_messages, client_class, async_client_class
+
 
 VertexAILLM
-===========
+-----------
 
-.. autoclass:: neo4j_graphrag.llm.vertexai.VertexAILLM
+.. autoclass:: neo4j_graphrag.llm.vertexai_llm.VertexAILLM
     :members:
 
+CohereLLM
+---------
+
+.. autoclass:: neo4j_graphrag.llm.cohere_llm.CohereLLM
+    :members:
+
+
 PromptTemplate
 ==============
 

docs/source/user_guide_rag.rst

@@ -23,9 +23,9 @@ In practice, it's done with only a few lines of code:
 
     from neo4j import GraphDatabase
     from neo4j_graphrag.retrievers import VectorRetriever
-    from neo4j_graphrag.llm.openai import OpenAILLM
+    from neo4j_graphrag.llm import OpenAILLM
     from neo4j_graphrag.generation import GraphRAG
-    from neo4j_graphrag.embeddings.openai import OpenAIEmbeddings
+    from neo4j_graphrag.embeddings import OpenAIEmbeddings
 
     # 1. Neo4j driver
     URI = "neo4j://localhost:7687"
@@ -56,6 +56,12 @@ In practice, it's done with only a few lines of code:
     print(response.answer)
 
 
+.. note::
+
+    In order to run this code, the `openai` Python package needs to be installed:
+    `pip install openai`
+
+
 The following sections provide more details about how to customize this code.
 
 ******************************
@@ -70,6 +76,8 @@ Using Another LLM Model
 If OpenAI cannot be used directly, there are a few available alternatives:
 
 - Use Azure OpenAI.
+- Use Google VertexAI.
+- Use Cohere.
 - Use a local Ollama model.
 - Implement a custom interface.
 - Utilize any LangChain chat model.
@@ -83,18 +91,46 @@ It is possible to use Azure OpenAI switching to the `AzureOpenAILLM` class:
 
 .. code:: python
 
-    from neo4j_graphrag.llm.openai import AzureOpenAILLM
+    from neo4j_graphrag.llm import AzureOpenAILLM
     llm = AzureOpenAILLM(
         model_name="gpt-4o",
         azure_endpoint="https://example-endpoint.openai.azure.com/",  # update with your endpoint
         api_version="2024-06-01",  # update appropriate version
-        api_key="ba3a46d86259405385f73f08078f588b",  # api_key is optional and can also be set with OPENAI_API_KEY env var
+        api_key="...",  # api_key is optional and can also be set with OPENAI_API_KEY env var
     )
     llm.invoke("say something")
 
 Check the OpenAI Python client [documentation](https://github.com/openai/openai-python?tab=readme-ov-file#microsoft-azure-openai)
 to learn more about the configuration.
 
+.. note::
+
+    In order to run this code, the `openai` Python package needs to be installed:
+    `pip install openai`
+
+
+Using VertexAI LLM
+------------------
+
+To use VertexAI, instantiate the `VertexAILLM` class:
+
+.. code:: python
+
+    from neo4j_graphrag.llm import VertexAILLM
+    from vertexai.generative_models import GenerationConfig
+
+    generation_config = GenerationConfig(temperature=0.0)
+    llm = VertexAILLM(
+        model_name="gemini-1.5-flash-001", generation_config=generation_config
+    )
+    llm.invoke("say something")
+
+
+.. note::
+
+    In order to run this code, the `google-cloud-aiplatform` Python package needs to be installed:
+    `pip install google-cloud-aiplatform`
+
 
 Using a Local Model via Ollama
 -------------------------------
@@ -105,7 +141,7 @@ it can be queried using the following:
 
 .. code:: python
 
-    from neo4j_graphrag.llm.openai import OpenAILLM
+    from neo4j_graphrag.llm import OpenAILLM
     llm = OpenAILLM(api_key="ollama", base_url="http://127.0.0.1:11434/v1", model_name="orca-mini")
     llm.invoke("say something")
 
@@ -300,13 +336,18 @@ into a vector is required. Therefore, the retriever requires knowledge of an emb
 Embedders
 -----------------------------
 
-Currently, this package supports two embedders: `OpenAIEmbeddings` and `SentenceTransformerEmbeddings`.
+Currently, this package supports several embedders:
+- `OpenAIEmbeddings`
+- `AzureOpenAIEmbeddings`
+- `VertexAIEmbeddings`
+- `CohereEmbeddings`
+- `SentenceTransformerEmbeddings`.
 
 The `OpenAIEmbedder` was illustrated previously. Here is how to use the `SentenceTransformerEmbeddings`:
 
 .. code:: python
 
-    from neo4j_graphrag.embeddings.sentence_transformers import SentenceTransformerEmbeddings
+    from neo4j_graphrag.embeddings import SentenceTransformerEmbeddings
 
     embedder = SentenceTransformerEmbeddings(model="all-MiniLM-L6-v2")  # Note: this is the default model
 
@@ -330,10 +371,10 @@ the following implementation of an embedder that wraps the `OllamaEmbedding` mod
             return embedding[0]
 
     ollama_embedding = OllamaEmbedding(
-                model_name="llama3",
-                base_url="http://localhost:11434",
-                ollama_additional_kwargs={"mirostat": 0},
-            )
+        model_name="llama3",
+        base_url="http://localhost:11434",
+        ollama_additional_kwargs={"mirostat": 0},
+    )
     embedder = OllamaEmbedder(ollama_embedding)
     vector = embedder.embed_query("some text")