Support for Azure OpenAI + Ollama + unit tests for the embeders and LLMs (#140)

* Add support for AzureOpenAI + remove OpenAI from llm imports + add tests for (Azure)OpenAI embedder and LLM + update doc

* Add UT for vertex AI

* Add google-cloud-aiplatform to dev dependencies

* Fix vertexai - other imports are trying to import from 'self' (the vertexai.py), instead of the vertexai package

* Ruff

* Update code after manual tests

* Ruff and co

* Fix typo :'(

* Update file names to prevent import errors

* Revert file renaming for embeddings

Author: stellasia

CHANGELOG.md

@@ -3,6 +3,7 @@
 ## Next
 
 ### Added
+- Added AzureOpenAILLM and AzureOpenAIEmbeddings to support Azure served OpenAI models
 - Add `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.

README.md

@@ -110,7 +110,7 @@ from neo4j import GraphDatabase
 from neo4j_graphrag.retrievers import VectorRetriever
 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
 
 URI = "neo4j://localhost:7687"
 AUTH = ("neo4j", "password")

docs/source/api.rst

@@ -176,7 +176,7 @@ LLMInterface
 OpenAILLM
 =========
 
-.. autoclass:: neo4j_graphrag.llm.OpenAILLM
+.. autoclass:: neo4j_graphrag.llm.openai.OpenAILLM
     :members:
 
 VertexAILLM

docs/source/user_guide_kg_builder.rst

@@ -246,7 +246,7 @@ It can be used in this way:
     from neo4j_graphrag.experimental.components.entity_relation_extractor import (
         LLMEntityRelationExtractor,
     )
-    from neo4j_graphrag.llm import OpenAILLM
+    from neo4j_graphrag.llm.openai import OpenAILLM
 
     extractor = LLMEntityRelationExtractor(
         llm=OpenAILLM(

docs/source/user_guide_rag.rst

@@ -23,7 +23,7 @@ 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 import OpenAILLM
+    from neo4j_graphrag.llm.openai import OpenAILLM
     from neo4j_graphrag.generation import GraphRAG
     from neo4j_graphrag.embeddings.openai import OpenAIEmbeddings
 
@@ -67,12 +67,47 @@ Each component can be configured individually: the LLM and the prompt.
 Using Another LLM Model
 ========================
 
-If OpenAI can not be used, there are two available alternatives:
+If OpenAI cannot be used directly, there are a few available alternatives:
 
-1. Utilize any LangChain chat model.
-2. Implement a custom interface.
+- Use Azure OpenAI.
+- Use a local Ollama model.
+- Implement a custom interface.
+- Utilize any LangChain chat model.
 
-Both options are illustrated below, using a local Ollama model as an example.
+All options are illustrated below, using a local Ollama model as an example.
+
+Using Azure Open AI LLM
+-----------------------
+
+It is possible to use Azure OpenAI switching to the `AzureOpenAILLM` class:
+
+.. code:: python
+
+    from neo4j_graphrag.llm.openai 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
+    )
+    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.
+
+
+Using a Local Model via Ollama
+-------------------------------
+
+Similarly to the official OpenAI Python client, the `OpenAILLM` can be
+used with Ollama. Assuming Ollama is running on the default address `127.0.0.1:11434`,
+it can be queried using the following:
+
+.. code:: python
+
+    from neo4j_graphrag.llm.openai import OpenAILLM
+    llm = OpenAILLM(api_key="ollama", base_url="http://127.0.0.1:11434/v1", model_name="orca-mini")
+    llm.invoke("say something")
 
 
 Using a Model from LangChain
@@ -99,10 +134,11 @@ It is however not mandatory to use LangChain. The alternative is to implement
 a custom model.
 
 Using a Custom Model
------------------------------
+--------------------
 
-To avoid LangChain, developers can create a custom LLM class by subclassing
-the `LLMInterface`. Here's an example using the Python Ollama client:
+If the provided implementations do not match their needs, developers can create a
+custom LLM class by subclassing the `LLMInterface`.
+Here's an example using the Python Ollama client:
 
 
 .. code:: python
@@ -123,6 +159,10 @@ the `LLMInterface`. Here's an example using the Python Ollama client:
                 content=response["message"]["content"]
             )
 
+        async def ainvoke(self, input: str) -> LLMResponse:
+            return self.invoke(input)  # TODO: implement async with ollama.AsyncClient
+
+
     # retriever = ...
 
     llm = OllamaLLM("llama3:8b")
@@ -607,7 +647,7 @@ LLMs can be different.
 
     from neo4j import GraphDatabase
     from neo4j_graphrag.retrievers import Text2CypherRetriever
-    from neo4j_graphrag.llm import OpenAILLM
+    from neo4j_graphrag.llm.openai import OpenAILLM
 
     URI = "neo4j://localhost:7687"
     AUTH = ("neo4j", "password")
@@ -772,4 +812,4 @@ Drop a Vector Index
 
     # Connect to Neo4j database
     driver = GraphDatabase.driver(URI, auth=AUTH)
-    drop_index_if_exists(driver, INDEX_NAME)
+    drop_index_if_exists(driver, INDEX_NAME)

examples/pipeline/kg_builder_from_pdf.py

@@ -16,7 +16,6 @@
 
 import asyncio
 import logging
-from typing import Any, Dict, List
 
 import neo4j
 from neo4j_graphrag.experimental.components.entity_relation_extractor import (
@@ -33,71 +32,12 @@
 from neo4j_graphrag.experimental.components.text_splitters.fixed_size_splitter import (
     FixedSizeSplitter,
 )
-from neo4j_graphrag.experimental.pipeline import Component, DataModel
 from neo4j_graphrag.experimental.pipeline.pipeline import PipelineResult
 from neo4j_graphrag.llm import OpenAILLM
-from pydantic import BaseModel, validate_call
 
 logging.basicConfig(level=logging.INFO)
 
 
-class DocumentChunkModel(DataModel):
-    chunks: list[str]
-
-
-class DocumentChunker(Component):
-    async def run(self, text: str) -> DocumentChunkModel:
-        chunks = [t.strip() for t in text.split(".") if t.strip()]
-        return DocumentChunkModel(chunks=chunks)
-
-
-class EntityModel(BaseModel):
-    label: str
-    properties: dict[str, str]
-
-
-class Neo4jGraph(DataModel):
-    entities: list[dict[str, Any]]
-    relations: list[dict[str, Any]]
-
-
-class ERExtractor(Component):
-    async def _process_chunk(self, chunk: str, schema: str) -> Dict[str, Any]:
-        return {
-            "entities": [{"label": "Person", "properties": {"name": "John Doe"}}],
-            "relations": [],
-        }
-
-    async def run(self, chunks: List[str], schema: str) -> Neo4jGraph:
-        tasks = [self._process_chunk(chunk, schema) for chunk in chunks]
-        result = await asyncio.gather(*tasks)
-        merged_result: dict[str, Any] = {"entities": [], "relations": []}
-        for res in result:
-            merged_result["entities"] += res["entities"]
-            merged_result["relations"] += res["relations"]
-        return Neo4jGraph(
-            entities=merged_result["entities"], relations=merged_result["relations"]
-        )
-
-
-class WriterModel(DataModel):
-    status: str
-    entities: list[EntityModel]
-    relations: list[EntityModel]
-
-
-class Writer(Component):
-    @validate_call
-    async def run(self, graph: Neo4jGraph) -> WriterModel:
-        entities = graph.entities
-        relations = graph.relations
-        return WriterModel(
-            status="OK",
-            entities=[EntityModel(**e) for e in entities],
-            relations=[EntityModel(**r) for r in relations],
-        )
-
-
 async def main(neo4j_driver: neo4j.Driver) -> PipelineResult:
     from neo4j_graphrag.experimental.pipeline import Pipeline
 

examples/text2cypher_search.py

@@ -1,6 +1,6 @@
 from neo4j import GraphDatabase
 from neo4j_graphrag.llm import OpenAILLM
-from neo4j_graphrag.retrievers.text2cypher import Text2CypherRetriever
+from neo4j_graphrag.retrievers import Text2CypherRetriever
 
 URI = "neo4j://localhost:7687"
 AUTH = ("neo4j", "password")