Updated docs

alexthomas93 · Sep 19, 2024 · 4de857c · 4de857c
1 parent f7f2451
commit 4de857c
Show file tree

Hide file tree

Showing 3 changed files with 21 additions and 3 deletions.
diff --git a/docs/source/api.rst b/docs/source/api.rst
@@ -174,11 +174,16 @@ LLMInterface
 
 
 OpenAILLM
----------
+=========
 
 .. autoclass:: neo4j_graphrag.llm.OpenAILLM
     :members:
 
+VertexAILLM
+===========
+
+.. autoclass:: neo4j_graphrag.llm.vertexai.VertexAILLM
+    :members:
 
 PromptTemplate
 ==============
@@ -389,4 +394,4 @@ PipelineStatusUpdateError
 =========================
 
 .. autoclass:: neo4j_graphrag.experimental.pipeline.exceptions.PipelineStatusUpdateError
-   :show-inheritance:
+   :show-inheritance:
diff --git a/src/neo4j_graphrag/llm/__init__.py b/src/neo4j_graphrag/llm/__init__.py
@@ -16,4 +16,4 @@
 from .openai_llm import OpenAILLM
 from .types import LLMResponse
 
-__all__ = ["LLMResponse", "LLMInterface", "OpenAILLM"]
+__all__ = ["LLMResponse", "LLMInterface", "OpenAILLM", "VertexAILLM"]
diff --git a/src/neo4j_graphrag/llm/vertexai.py b/src/neo4j_graphrag/llm/vertexai.py
@@ -36,6 +36,19 @@ class VertexAILLM(LLMInterface):
 
     Raises:
         LLMGenerationError: If there's an error generating the response from the model.
+
+    Example:
+
+    .. code-block:: 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("Who is the mother of Paul Atreides?")
     """
 
     def __init__(