From a305d43350014e621c67d568c84e588698898aca Mon Sep 17 00:00:00 2001 From: Alex Thomas Date: Wed, 18 Sep 2024 15:34:45 +0100 Subject: [PATCH] Updated docstrings --- src/neo4j_graphrag/llm/base.py | 8 +++++++- src/neo4j_graphrag/llm/openai_llm.py | 3 +-- src/neo4j_graphrag/llm/vertexai.py | 6 +++++- 3 files changed, 13 insertions(+), 4 deletions(-) diff --git a/src/neo4j_graphrag/llm/base.py b/src/neo4j_graphrag/llm/base.py index b9a639b7b..3d98423be 100644 --- a/src/neo4j_graphrag/llm/base.py +++ b/src/neo4j_graphrag/llm/base.py @@ -21,7 +21,13 @@ class LLMInterface(ABC): - """Interface for large language models.""" + """Interface for large language models. + + Args: + model_name (str): The name of the language model. + model_params (Optional[dict], optional): Additional parameters passed to the model when text is sent to it. Defaults to None. + **kwargs (Any): Arguments passed to the model when for the class is initialised. Defaults to None. + """ def __init__( self, diff --git a/src/neo4j_graphrag/llm/openai_llm.py b/src/neo4j_graphrag/llm/openai_llm.py index 334277b42..2c6225928 100644 --- a/src/neo4j_graphrag/llm/openai_llm.py +++ b/src/neo4j_graphrag/llm/openai_llm.py @@ -37,8 +37,7 @@ def __init__( Args: model_name (str): - model_params (str): Parameters like temperature and such that will be - passed to the model + model_params (str): Parameters like temperature that will be passed to the model when text is sent to it kwargs: All other parameters will be passed to the openai.OpenAI init. """ diff --git a/src/neo4j_graphrag/llm/vertexai.py b/src/neo4j_graphrag/llm/vertexai.py index 5e644ce05..694c7b538 100644 --- a/src/neo4j_graphrag/llm/vertexai.py +++ b/src/neo4j_graphrag/llm/vertexai.py @@ -31,7 +31,11 @@ class VertexAILLM(LLMInterface): Args: model_name (str, optional): Name of the LLM to use. Defaults to "gemini-1.5-flash-001". - model_params (Optional[Dict[str, Any]], optional): Parameters for passed to the LLM's invoke and ainvoke functions. + model_params (Optional[dict], optional): Additional parameters passed to the model when text is sent to it. Defaults to None. + **kwargs (Any): Arguments passed to the model when for the class is initialised. Defaults to None. + + Raises: + LLMGenerationError: If there's an error generating the response from the model. """ def __init__(