Updated component docs

Author: alexthomas93

docs/source/api.rst

@@ -3,6 +3,66 @@
 API Documentation
 #################
 
+.. _components-section:
+
+**********
+Components
+**********
+
+KGWriter
+========
+
+.. autoclass:: neo4j_genai.components.kg_writer.KGWriter
+    :members: run
+
+Neo4jWriter
+===========
+
+.. autoclass:: neo4j_genai.components.kg_writer.Neo4jWriter
+    :members: run
+
+TextSplitter
+============
+
+.. autoclass:: neo4j_genai.components.text_splitters.base.TextSplitter
+    :members: run
+
+LangChainTextSplitterAdapter
+============================
+
+.. autoclass:: neo4j_genai.components.text_splitters.langchain.LangChainTextSplitterAdapter
+    :members: run
+
+LlamaIndexTextSplitterAdapter
+=============================
+
+.. autoclass:: neo4j_genai.components.text_splitters.llamaindex.LlamaIndexTextSplitterAdapter
+    :members: run
+
+TextChunkEmbedder
+=================
+
+.. autoclass:: neo4j_genai.components.embedder.TextChunkEmbedder
+    :members: run
+
+SchemaBuilder
+=============
+
+.. autoclass:: neo4j_genai.components.schema.SchemaBuilder
+    :members: run
+
+EntityRelationExtractor
+=======================
+
+.. autoclass:: neo4j_genai.components.entity_relation_extractor.EntityRelationExtractor
+    :members: run
+
+LLMEntityRelationExtractor
+==========================
+
+.. autoclass:: neo4j_genai.components.entity_relation_extractor.LLMEntityRelationExtractor
+    :members: run
+
 .. _retrievers-section:
 
 **********

docs/source/types.rst

@@ -5,30 +5,80 @@ Types
 *****
 
 RawSearchResult
-==================
+===============
 
 .. autoclass:: neo4j_genai.types.RawSearchResult
 
 
 RetrieverResult
-==================
+===============
 
 .. autoclass:: neo4j_genai.types.RetrieverResult
 
 
 RetrieverResultItem
-====================
+===================
 
 .. autoclass:: neo4j_genai.types.RetrieverResultItem
 
 
 LLMResponse
-====================
+===========
 
 .. autoclass:: neo4j_genai.llm.types.LLMResponse
 
 
 RagResultModel
-====================
+==============
 
 .. autoclass:: neo4j_genai.generation.types.RagResultModel
+
+TextChunk
+=========
+
+.. autoclass:: neo4j_genai.components.types.TextChunk
+
+TextChunks
+==========
+
+.. autoclass:: neo4j_genai.components.types.TextChunks
+
+Neo4jNode
+=========
+
+.. autoclass:: neo4j_genai.components.types.Neo4jNode
+
+Neo4jRelationship
+=================
+
+.. autoclass:: neo4j_genai.components.types.Neo4jRelationship
+
+Neo4jGraph
+==========
+
+.. autoclass:: neo4j_genai.components.types.Neo4jGraph
+
+KGWriterModel
+=============
+
+.. autoclass:: neo4j_genai.components.kg_writer.KGWriterModel
+
+SchemaProperty
+==============
+
+.. autoclass:: neo4j_genai.components.schema.SchemaProperty
+
+SchemaEntity
+============
+
+.. autoclass:: neo4j_genai.components.schema.SchemaEntity
+
+SchemaRelation
+==============
+
+.. autoclass:: neo4j_genai.components.schema.SchemaEntity
+
+SchemaConfig
+==============
+
+.. autoclass:: neo4j_genai.components.schema.SchemaConfig

src/neo4j_genai/components/embedder.py

@@ -22,6 +22,20 @@ class TextChunkEmbedder(Component):
 
     Args:
         embedder (Embedder): The embedder to use to create the embeddings.
+
+    Example:
+
+    .. code-block:: python
+
+        from neo4j_genai.components.embedder import TextChunkEmbedder
+        from neo4j_genai.embeddings.openai import OpenAIEmbeddings
+        from neo4j_genai.pipeline import Pipeline
+
+        embedder = OpenAIEmbeddings(model="text-embedding-3-large")
+        chunk_embedder = TextChunkEmbedder(embedder)
+        pipeline = Pipeline()
+        pipeline.add_component("chunk_embedder", chunk_embedder)
+
     """
 
     def __init__(self, embedder: Embedder):

src/neo4j_genai/components/entity_relation_extractor.py

@@ -51,6 +51,13 @@ class OnError(enum.Enum):
 
 
 class EntityRelationExtractor(Component, abc.ABC):
+    """Abstract class for entity relation extraction components.
+
+    Args:
+        on_error (OnError): What to do when an error occurs during extraction. Defaults to raising an error.
+        create_lexical_graph (bool): Whether to include the text chunks in the graph in addition to the extracted entities and relations. Defaults to True.
+    """
+
     def __init__(
         self,
         *args: Any,
@@ -137,6 +144,31 @@ def build_lexical_graph(
 
 
 class LLMEntityRelationExtractor(EntityRelationExtractor):
+    """
+    Extracts a knowledge graph from a series of text chunks using a large language model.
+
+    Args:
+        llm (LLMInterface): The language model to use for extraction.
+        prompt_template (ERExtractionTemplate | str): A custom prompt template to use for extraction.
+        create_lexical_graph (bool): Whether to include the text chunks in the graph in addition to the extracted entities and relations. Defaults to True.
+        on_error (OnError): What to do when an error occurs during extraction. Defaults to raising an error.
+
+    Example:
+
+    .. code-block:: python
+
+        from neo4j_genai.components.entity_relation_extractor import LLMEntityRelationExtractor
+        from neo4j_genai.llm import OpenAILLM
+        from neo4j_genai.pipeline import Pipeline
+
+        llm = OpenAILLM(model_name="gpt-4o", model_params={"temperature": 0})
+
+        extractor = LLMEntityRelationExtractor(llm=llm)
+        pipe = Pipeline()
+        pipe.add_component("extractor", extractor)
+
+    """
+
     def __init__(
         self,
         llm: LLMInterface,

src/neo4j_genai/components/kg_writer.py

@@ -60,6 +60,25 @@ class Neo4jWriter(KGWriter):
     Args:
         driver (neo4j.driver): The Neo4j driver to connect to the database.
         neo4j_database (Optional[str]): The name of the Neo4j database to write to. Defaults to 'neo4j' if not provided.
+
+    Example:
+
+    .. code-block:: python
+
+        from neo4j import GraphDatabase
+        from neo4j_genai.components.kg_writer import Neo4jWriter
+        from neo4j_genai.pipeline import Pipeline
+
+        URI = "neo4j://localhost:7687"
+        AUTH = ("neo4j", "password")
+        DATABASE = "neo4j"
+
+        driver = GraphDatabase.driver(URI, auth=AUTH, database=DATABASE)
+        writer = Neo4jWriter(driver=driver, neo4j_database=DATABASE)
+
+        pipeline = Pipeline()
+        pipeline.add_component("writer", writer)
+
     """
 
     def __init__(

src/neo4j_genai/components/schema.py

@@ -12,20 +12,6 @@
 #  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
-#  Copyright (c) "Neo4j"
-#  Neo4j Sweden AB [https://neo4j.com]
-#  #
-#  Licensed under the Apache License, Version 2.0 (the "License");
-#  you may not use this file except in compliance with the License.
-#  You may obtain a copy of the License at
-#  #
-#      https://www.apache.org/licenses/LICENSE-2.0
-#  #
-#  Unless required by applicable law or agreed to in writing, software
-#  distributed under the License is distributed on an "AS IS" BASIS,
-#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-#  See the License for the specific language governing permissions and
-#  limitations under the License.
 from __future__ import annotations
 
 from typing import Any, Dict, List, Literal, Tuple
@@ -37,6 +23,10 @@
 
 
 class SchemaProperty(BaseModel):
+    """
+    Represents a property on a node or relationship in the graph.
+    """
+
     name: str
     # See https://neo4j.com/docs/cypher-manual/current/values-and-types/property-structural-constructed/#property-types
     type: Literal[
@@ -112,6 +102,59 @@ class SchemaBuilder(Component):
     """
     A builder class for constructing SchemaConfig objects from given entities,
     relations, and their interrelationships defined in a potential schema.
+
+    Example:
+
+    .. code-block:: python
+
+        from neo4j_genai.components.schema import (
+            SchemaBuilder,
+            SchemaEntity,
+            SchemaProperty,
+            SchemaRelation,
+        )
+        from neo4j_genai.pipeline import Pipeline
+
+        entities = [
+            SchemaEntity(
+                label="PERSON",
+                description="An individual human being.",
+                properties=[
+                    SchemaProperty(
+                        name="name", type="STRING", description="The name of the person"
+                    )
+                ],
+            ),
+            SchemaEntity(
+                label="ORGANIZATION",
+                description="A structured group of people with a common purpose.",
+                properties=[
+                    SchemaProperty(
+                        name="name", type="STRING", description="The name of the organization"
+                    )
+                ],
+            ),
+        ]
+        relations = [
+            SchemaRelation(
+                label="EMPLOYED_BY", description="Indicates employment relationship."
+            ),
+        ]
+        potential_schema = [
+            ("PERSON", "EMPLOYED_BY", "ORGANIZATION"),
+        ]
+        pipe = Pipeline()
+        schema_builder = SchemaBuilder()
+        pipe.add_component("schema_builder", schema_builder)
+        pipe_inputs = {
+            "schema": {
+                "entities": entities,
+                "relations": relations,
+                "potential_schema": potential_schema,
+            },
+            ...
+        }
+        pipe.run(pipe_inputs)
     """
 
     @staticmethod

src/neo4j_genai/components/text_splitters/langchain.py

@@ -26,6 +26,19 @@ class LangChainTextSplitterAdapter(TextSplitter):
 
     Args:
         text_splitter (LangChainTextSplitter): An instance of LangChain's TextSplitter class.
+
+    Example:
+
+    .. code-block:: python
+
+        from langchain_text_splitters import RecursiveCharacterTextSplitter
+        from neo4j_genai.components.text_splitters.langchain import LangChainTextSplitterAdapter
+        from neo4j_genai.pipeline import Pipeline
+
+        pipeline = Pipeline()
+        text_splitter = LangChainTextSplitterAdapter(RecursiveCharacterTextSplitter())
+        pipeline.add_component("text_splitter", text_splitter)
+
     """
 
     def __init__(self, text_splitter: LangChainTextSplitter) -> None:

src/neo4j_genai/components/text_splitters/llamaindex.py

@@ -26,6 +26,19 @@ class LlamaIndexTextSplitterAdapter(TextSplitter):
 
     Args:
         text_splitter (LlamaIndexTextSplitter): An instance of LlamaIndex's TextSplitter class.
+
+    Example:
+
+    .. code-block:: python
+
+        from llama_index.core.node_parser.text.sentence import SentenceSplitter
+        from neo4j_genai.components.text_splitters.langchain import LangChainTextSplitterAdapter
+        from neo4j_genai.pipeline import Pipeline
+
+        pipeline = Pipeline()
+        text_splitter = LlamaIndexTextSplitterAdapter(SentenceSplitter())
+        pipeline.add_component("text_splitter", text_splitter)
+
     """
 
     def __init__(self, text_splitter: LlamaIndexTextSplitter) -> None:

tests/e2e/weaviate_e2e/test_weaviate_e2e.py

@@ -131,6 +131,6 @@ def test_weaviate_neo4j_text_input_remote_embedder(
         r"labels=frozenset\({'Question'}\) properties={'question': 'In 1953 Watson \& "
         "Crick built a model of the molecular structure of this, the gene-carrying "
         "substance', 'id': 'question_c458c6f64d8d47429636bc5a94c97f51'}> "
-        r"score=0.5[0-9]+>"
+        r"score=0.6[0-9]+>"
     )
     assert re.match(pattern, results.items[0].content)