From dd3dcfed64a4ead034a684b54e40dd3367057bcc Mon Sep 17 00:00:00 2001 From: Alex Thomas Date: Tue, 13 Aug 2024 14:58:23 +0100 Subject: [PATCH] Moves the KG builder code to the experimental folder (#101) --- docs/source/api.rst | 30 +++++++++---------- docs/source/types.rst | 20 ++++++------- docs/source/user_guide_pipeline.rst | 8 ++--- examples/pipeline/kg_builder.py | 16 +++++----- examples/pipeline/kg_builder_with_schema.py | 6 ++-- examples/pipeline/rag.py | 10 +++++-- .../__init__.py | 0 .../{ => experimental}/components/__init__.py | 0 .../{ => experimental}/components/embedder.py | 8 ++--- .../components/entity_relation_extractor.py | 12 ++++---- .../components/kg_writer.py | 12 +++++--- .../{ => experimental}/components/schema.py | 6 ++-- .../components/text_splitters}/__init__.py | 0 .../components/text_splitters/base.py | 4 +-- .../components/text_splitters/langchain.py | 8 ++--- .../components/text_splitters/llamaindex.py | 8 ++--- .../{ => experimental}/components/types.py | 2 +- .../{ => experimental}/pipeline/__init__.py | 0 .../{ => experimental}/pipeline/component.py | 2 +- .../{ => experimental}/pipeline/exceptions.py | 0 .../{ => experimental}/pipeline/pipeline.py | 14 +++++---- .../pipeline/pipeline_graph.py | 0 .../{ => experimental}/pipeline/stores.py | 0 .../{ => experimental}/pipeline/types.py | 2 +- tests/e2e/test_kb_builder_pipeline_e2e.py | 18 ++++++----- tests/e2e/test_kg_construction_e2e.py | 4 +-- tests/unit/experimental/__init__.py | 0 .../{ => experimental}/components/__init__.py | 0 .../components/test_embedder.py | 4 +-- .../test_entity_relation_extractor.py | 11 +++++-- .../components/test_kg_writer.py | 4 +-- .../components/test_schema.py | 4 +-- .../components/test_types.py | 2 +- .../text_splitters/test_langchain.py | 6 ++-- .../text_splitters/test_llamaindex.py | 4 +-- tests/unit/experimental/pipeline/__init__.py | 0 .../{ => experimental}/pipeline/components.py | 2 +- .../pipeline/test_component.py | 0 .../pipeline/test_orchestrator.py | 4 +-- .../pipeline/test_pipeline.py | 8 ++--- .../pipeline/test_pipeline_graph.py | 2 +- .../{ => experimental}/pipeline/test_store.py | 2 +- 42 files changed, 133 insertions(+), 110 deletions(-) rename src/neo4j_genai/{components/text_splitters => experimental}/__init__.py (100%) rename src/neo4j_genai/{ => experimental}/components/__init__.py (100%) rename src/neo4j_genai/{ => experimental}/components/embedder.py (89%) rename src/neo4j_genai/{ => experimental}/components/entity_relation_extractor.py (96%) rename src/neo4j_genai/{ => experimental}/components/kg_writer.py (94%) rename src/neo4j_genai/{ => experimental}/components/schema.py (97%) rename {tests/unit/pipeline => src/neo4j_genai/experimental/components/text_splitters}/__init__.py (100%) rename src/neo4j_genai/{ => experimental}/components/text_splitters/base.py (86%) rename src/neo4j_genai/{ => experimental}/components/text_splitters/langchain.py (84%) rename src/neo4j_genai/{ => experimental}/components/text_splitters/llamaindex.py (84%) rename src/neo4j_genai/{ => experimental}/components/types.py (97%) rename src/neo4j_genai/{ => experimental}/pipeline/__init__.py (100%) rename src/neo4j_genai/{ => experimental}/pipeline/component.py (97%) rename src/neo4j_genai/{ => experimental}/pipeline/exceptions.py (100%) rename src/neo4j_genai/{ => experimental}/pipeline/pipeline.py (97%) rename src/neo4j_genai/{ => experimental}/pipeline/pipeline_graph.py (100%) rename src/neo4j_genai/{ => experimental}/pipeline/stores.py (100%) rename src/neo4j_genai/{ => experimental}/pipeline/types.py (94%) create mode 100644 tests/unit/experimental/__init__.py rename tests/unit/{ => experimental}/components/__init__.py (100%) rename tests/unit/{ => experimental}/components/test_embedder.py (90%) rename tests/unit/{ => experimental}/components/test_entity_relation_extractor.py (97%) rename tests/unit/{ => experimental}/components/test_kg_writer.py (97%) rename tests/unit/{ => experimental}/components/test_schema.py (99%) rename tests/unit/{ => experimental}/components/test_types.py (94%) rename tests/unit/{ => experimental}/components/text_splitters/test_langchain.py (90%) rename tests/unit/{ => experimental}/components/text_splitters/test_llamaindex.py (92%) create mode 100644 tests/unit/experimental/pipeline/__init__.py rename tests/unit/{ => experimental}/pipeline/components.py (95%) rename tests/unit/{ => experimental}/pipeline/test_component.py (100%) rename tests/unit/{ => experimental}/pipeline/test_orchestrator.py (94%) rename tests/unit/{ => experimental}/pipeline/test_pipeline.py (95%) rename tests/unit/{ => experimental}/pipeline/test_pipeline_graph.py (98%) rename tests/unit/{ => experimental}/pipeline/test_store.py (80%) diff --git a/docs/source/api.rst b/docs/source/api.rst index 0e309b7a..c6bc67ab 100644 --- a/docs/source/api.rst +++ b/docs/source/api.rst @@ -12,55 +12,55 @@ Components KGWriter ======== -.. autoclass:: neo4j_genai.components.kg_writer.KGWriter +.. autoclass:: neo4j_genai.experimental.components.kg_writer.KGWriter :members: run Neo4jWriter =========== -.. autoclass:: neo4j_genai.components.kg_writer.Neo4jWriter +.. autoclass:: neo4j_genai.experimental.components.kg_writer.Neo4jWriter :members: run TextSplitter ============ -.. autoclass:: neo4j_genai.components.text_splitters.base.TextSplitter +.. autoclass:: neo4j_genai.experimental.components.text_splitters.base.TextSplitter :members: run LangChainTextSplitterAdapter ============================ -.. autoclass:: neo4j_genai.components.text_splitters.langchain.LangChainTextSplitterAdapter +.. autoclass:: neo4j_genai.experimental.components.text_splitters.langchain.LangChainTextSplitterAdapter :members: run LlamaIndexTextSplitterAdapter ============================= -.. autoclass:: neo4j_genai.components.text_splitters.llamaindex.LlamaIndexTextSplitterAdapter +.. autoclass:: neo4j_genai.experimental.components.text_splitters.llamaindex.LlamaIndexTextSplitterAdapter :members: run TextChunkEmbedder ================= -.. autoclass:: neo4j_genai.components.embedder.TextChunkEmbedder +.. autoclass:: neo4j_genai.experimental.components.embedder.TextChunkEmbedder :members: run SchemaBuilder ============= -.. autoclass:: neo4j_genai.components.schema.SchemaBuilder +.. autoclass:: neo4j_genai.experimental.components.schema.SchemaBuilder :members: run EntityRelationExtractor ======================= -.. autoclass:: neo4j_genai.components.entity_relation_extractor.EntityRelationExtractor +.. autoclass:: neo4j_genai.experimental.components.entity_relation_extractor.EntityRelationExtractor :members: run LLMEntityRelationExtractor ========================== -.. autoclass:: neo4j_genai.components.entity_relation_extractor.LLMEntityRelationExtractor +.. autoclass:: neo4j_genai.experimental.components.entity_relation_extractor.LLMEntityRelationExtractor :members: run .. _retrievers-section: @@ -219,11 +219,11 @@ Errors * :class:`neo4j_genai.exceptions.LLMGenerationError` - * :class:`neo4j_genai.pipeline.exceptions.PipelineDefinitionError` + * :class:`neo4j_genai.experimental.pipeline.exceptions.PipelineDefinitionError` - * :class:`neo4j_genai.pipeline.exceptions.PipelineMissingDependencyError` + * :class:`neo4j_genai.experimental.pipeline.exceptions.PipelineMissingDependencyError` - * :class:`neo4j_genai.pipeline.exceptions.PipelineStatusUpdateError` + * :class:`neo4j_genai.experimental.pipeline.exceptions.PipelineStatusUpdateError` Neo4jGenAiError @@ -327,19 +327,19 @@ LLMGenerationError PipelineDefinitionError ======================= -.. autoclass:: neo4j_genai.pipeline.exceptions.PipelineDefinitionError +.. autoclass:: neo4j_genai.experimental.pipeline.exceptions.PipelineDefinitionError :show-inheritance: PipelineMissingDependencyError ============================== -.. autoclass:: neo4j_genai.pipeline.exceptions.PipelineMissingDependencyError +.. autoclass:: neo4j_genai.experimental.pipeline.exceptions.PipelineMissingDependencyError :show-inheritance: PipelineStatusUpdateError ========================= -.. autoclass:: neo4j_genai.pipeline.exceptions.PipelineStatusUpdateError +.. autoclass:: neo4j_genai.experimental.pipeline.exceptions.PipelineStatusUpdateError :show-inheritance: diff --git a/docs/source/types.rst b/docs/source/types.rst index ec763a6e..2bc87ffb 100644 --- a/docs/source/types.rst +++ b/docs/source/types.rst @@ -36,49 +36,49 @@ RagResultModel TextChunk ========= -.. autoclass:: neo4j_genai.components.types.TextChunk +.. autoclass:: neo4j_genai.experimental.components.types.TextChunk TextChunks ========== -.. autoclass:: neo4j_genai.components.types.TextChunks +.. autoclass:: neo4j_genai.experimental.components.types.TextChunks Neo4jNode ========= -.. autoclass:: neo4j_genai.components.types.Neo4jNode +.. autoclass:: neo4j_genai.experimental.components.types.Neo4jNode Neo4jRelationship ================= -.. autoclass:: neo4j_genai.components.types.Neo4jRelationship +.. autoclass:: neo4j_genai.experimental.components.types.Neo4jRelationship Neo4jGraph ========== -.. autoclass:: neo4j_genai.components.types.Neo4jGraph +.. autoclass:: neo4j_genai.experimental.components.types.Neo4jGraph KGWriterModel ============= -.. autoclass:: neo4j_genai.components.kg_writer.KGWriterModel +.. autoclass:: neo4j_genai.experimental.components.kg_writer.KGWriterModel SchemaProperty ============== -.. autoclass:: neo4j_genai.components.schema.SchemaProperty +.. autoclass:: neo4j_genai.experimental.components.schema.SchemaProperty SchemaEntity ============ -.. autoclass:: neo4j_genai.components.schema.SchemaEntity +.. autoclass:: neo4j_genai.experimental.components.schema.SchemaEntity SchemaRelation ============== -.. autoclass:: neo4j_genai.components.schema.SchemaEntity +.. autoclass:: neo4j_genai.experimental.components.schema.SchemaEntity SchemaConfig ============ -.. autoclass:: neo4j_genai.components.schema.SchemaConfig +.. autoclass:: neo4j_genai.experimental.components.schema.SchemaConfig diff --git a/docs/source/user_guide_pipeline.rst b/docs/source/user_guide_pipeline.rst index e6745f1a..2f2c629d 100644 --- a/docs/source/user_guide_pipeline.rst +++ b/docs/source/user_guide_pipeline.rst @@ -19,8 +19,8 @@ such as chunking documents or saving results to Neo4j. This package includes a few default components, but developers can create their own by following these steps: -1. Create a subclass of the Pydantic `neo4j_genai.pipeline.DataModel` to represent the data being returned by the component -2. Create a subclass of `neo4j_genai.pipeline.Component` +1. Create a subclass of the Pydantic `neo4j_genai.experimental.pipeline.DataModel` to represent the data being returned by the component +2. Create a subclass of `neo4j_genai.experimental.pipeline.Component` 3. Create a run method in this new class and specify the required inputs and output model using the just created `DataModel` 4. Implement the run method: it's an `async` method, allowing tasks to be parallelized and awaited within this method. @@ -29,7 +29,7 @@ the resulting sum: .. code:: python - from neo4j_genai.pipeline import Component, DataModel + from neo4j_genai.experimental.pipeline import Component, DataModel class IntResultModel(DataModel): result: int @@ -53,7 +53,7 @@ Here's how to create a simple pipeline and propagate results from one component .. code:: python import asyncio - from neo4j_genai.pipeline import Pipeline + from neo4j_genai.experimental.pipeline import Pipeline pipe = Pipeline() pipe.add_component("a", ComponentAdd()) diff --git a/examples/pipeline/kg_builder.py b/examples/pipeline/kg_builder.py index 754d3a0f..dcaa102f 100644 --- a/examples/pipeline/kg_builder.py +++ b/examples/pipeline/kg_builder.py @@ -20,22 +20,24 @@ import neo4j from langchain_text_splitters import CharacterTextSplitter -from neo4j_genai.components.embedder import TextChunkEmbedder -from neo4j_genai.components.entity_relation_extractor import ( +from neo4j_genai.embeddings.openai import OpenAIEmbeddings +from neo4j_genai.experimental.components.embedder import TextChunkEmbedder +from neo4j_genai.experimental.components.entity_relation_extractor import ( LLMEntityRelationExtractor, OnError, ) -from neo4j_genai.components.kg_writer import Neo4jWriter -from neo4j_genai.components.schema import ( +from neo4j_genai.experimental.components.kg_writer import Neo4jWriter +from neo4j_genai.experimental.components.schema import ( SchemaBuilder, SchemaEntity, SchemaProperty, SchemaRelation, ) -from neo4j_genai.components.text_splitters.langchain import LangChainTextSplitterAdapter -from neo4j_genai.embeddings.openai import OpenAIEmbeddings +from neo4j_genai.experimental.components.text_splitters.langchain import ( + LangChainTextSplitterAdapter, +) +from neo4j_genai.experimental.pipeline import Pipeline from neo4j_genai.llm import OpenAILLM -from neo4j_genai.pipeline import Pipeline # set log level to DEBUG for all neo4j_genai.* loggers logging.config.dictConfig( diff --git a/examples/pipeline/kg_builder_with_schema.py b/examples/pipeline/kg_builder_with_schema.py index b2470d01..17ee1a9e 100644 --- a/examples/pipeline/kg_builder_with_schema.py +++ b/examples/pipeline/kg_builder_with_schema.py @@ -18,12 +18,12 @@ import logging from typing import Any -from neo4j_genai.components.schema import ( +from neo4j_genai.experimental.components.schema import ( SchemaBuilder, SchemaEntity, SchemaRelation, ) -from neo4j_genai.pipeline import Component, DataModel +from neo4j_genai.experimental.pipeline import Component, DataModel from pydantic import BaseModel, validate_call logging.basicConfig(level=logging.DEBUG) @@ -87,7 +87,7 @@ async def run(self, graph: Neo4jGraph) -> WriterModel: if __name__ == "__main__": - from neo4j_genai.pipeline import Pipeline + from neo4j_genai.experimental.pipeline import Pipeline # Instantiate Entity and Relation objects entities = [ diff --git a/examples/pipeline/rag.py b/examples/pipeline/rag.py index f645f45c..1afb3920 100644 --- a/examples/pipeline/rag.py +++ b/examples/pipeline/rag.py @@ -24,11 +24,15 @@ import neo4j from neo4j_genai.embeddings.openai import OpenAIEmbeddings +from neo4j_genai.experimental.pipeline import Component, Pipeline +from neo4j_genai.experimental.pipeline.component import DataModel +from neo4j_genai.experimental.pipeline.types import ( + ComponentConfig, + ConnectionConfig, + PipelineConfig, +) from neo4j_genai.generation import PromptTemplate, RagTemplate from neo4j_genai.llm import LLMInterface, OpenAILLM -from neo4j_genai.pipeline import Component, Pipeline -from neo4j_genai.pipeline.component import DataModel -from neo4j_genai.pipeline.types import ComponentConfig, ConnectionConfig, PipelineConfig from neo4j_genai.retrievers import VectorRetriever from neo4j_genai.retrievers.base import Retriever diff --git a/src/neo4j_genai/components/text_splitters/__init__.py b/src/neo4j_genai/experimental/__init__.py similarity index 100% rename from src/neo4j_genai/components/text_splitters/__init__.py rename to src/neo4j_genai/experimental/__init__.py diff --git a/src/neo4j_genai/components/__init__.py b/src/neo4j_genai/experimental/components/__init__.py similarity index 100% rename from src/neo4j_genai/components/__init__.py rename to src/neo4j_genai/experimental/components/__init__.py diff --git a/src/neo4j_genai/components/embedder.py b/src/neo4j_genai/experimental/components/embedder.py similarity index 89% rename from src/neo4j_genai/components/embedder.py rename to src/neo4j_genai/experimental/components/embedder.py index cf38a8bd..0683bef1 100644 --- a/src/neo4j_genai/components/embedder.py +++ b/src/neo4j_genai/experimental/components/embedder.py @@ -14,9 +14,9 @@ # limitations under the License. from pydantic import validate_call -from neo4j_genai.components.types import TextChunk, TextChunks from neo4j_genai.embedder import Embedder -from neo4j_genai.pipeline.component import Component +from neo4j_genai.experimental.components.types import TextChunk, TextChunks +from neo4j_genai.experimental.pipeline.component import Component class TextChunkEmbedder(Component): @@ -29,9 +29,9 @@ class TextChunkEmbedder(Component): .. code-block:: python - from neo4j_genai.components.embedder import TextChunkEmbedder + from neo4j_genai.experimental.components.embedder import TextChunkEmbedder from neo4j_genai.embeddings.openai import OpenAIEmbeddings - from neo4j_genai.pipeline import Pipeline + from neo4j_genai.experimental.pipeline import Pipeline embedder = OpenAIEmbeddings(model="text-embedding-3-large") chunk_embedder = TextChunkEmbedder(embedder) diff --git a/src/neo4j_genai/components/entity_relation_extractor.py b/src/neo4j_genai/experimental/components/entity_relation_extractor.py similarity index 96% rename from src/neo4j_genai/components/entity_relation_extractor.py rename to src/neo4j_genai/experimental/components/entity_relation_extractor.py index 19be51bf..cb46e17f 100644 --- a/src/neo4j_genai/components/entity_relation_extractor.py +++ b/src/neo4j_genai/experimental/components/entity_relation_extractor.py @@ -24,18 +24,18 @@ from pydantic import ValidationError, validate_call -from neo4j_genai.components.schema import SchemaConfig -from neo4j_genai.components.types import ( +from neo4j_genai.exceptions import LLMGenerationError +from neo4j_genai.experimental.components.schema import SchemaConfig +from neo4j_genai.experimental.components.types import ( Neo4jGraph, Neo4jNode, Neo4jRelationship, TextChunk, TextChunks, ) -from neo4j_genai.exceptions import LLMGenerationError +from neo4j_genai.experimental.pipeline.component import Component from neo4j_genai.generation.prompts import ERExtractionTemplate, PromptTemplate from neo4j_genai.llm import LLMInterface -from neo4j_genai.pipeline.component import Component logger = logging.getLogger(__name__) @@ -163,9 +163,9 @@ class LLMEntityRelationExtractor(EntityRelationExtractor): .. code-block:: python - from neo4j_genai.components.entity_relation_extractor import LLMEntityRelationExtractor + from neo4j_genai.experimental.components.entity_relation_extractor import LLMEntityRelationExtractor from neo4j_genai.llm import OpenAILLM - from neo4j_genai.pipeline import Pipeline + from neo4j_genai.experimental.pipeline import Pipeline llm = OpenAILLM(model_name="gpt-4o", model_params={"temperature": 0, "response_format": {"type": "object"}}) diff --git a/src/neo4j_genai/components/kg_writer.py b/src/neo4j_genai/experimental/components/kg_writer.py similarity index 94% rename from src/neo4j_genai/components/kg_writer.py rename to src/neo4j_genai/experimental/components/kg_writer.py index d32b90ab..6f561e25 100644 --- a/src/neo4j_genai/components/kg_writer.py +++ b/src/neo4j_genai/experimental/components/kg_writer.py @@ -21,10 +21,14 @@ import neo4j from pydantic import validate_call -from neo4j_genai.components.types import Neo4jGraph, Neo4jNode, Neo4jRelationship +from neo4j_genai.experimental.components.types import ( + Neo4jGraph, + Neo4jNode, + Neo4jRelationship, +) +from neo4j_genai.experimental.pipeline.component import Component, DataModel from neo4j_genai.indexes import upsert_vector, upsert_vector_on_relationship from neo4j_genai.neo4j_queries import UPSERT_NODE_QUERY, UPSERT_RELATIONSHIP_QUERY -from neo4j_genai.pipeline.component import Component, DataModel logger = logging.getLogger(__name__) @@ -66,8 +70,8 @@ class Neo4jWriter(KGWriter): .. code-block:: python from neo4j import GraphDatabase - from neo4j_genai.components.kg_writer import Neo4jWriter - from neo4j_genai.pipeline import Pipeline + from neo4j_genai.experimental.components.kg_writer import Neo4jWriter + from neo4j_genai.experimental.pipeline import Pipeline URI = "neo4j://localhost:7687" AUTH = ("neo4j", "password") diff --git a/src/neo4j_genai/components/schema.py b/src/neo4j_genai/experimental/components/schema.py similarity index 97% rename from src/neo4j_genai/components/schema.py rename to src/neo4j_genai/experimental/components/schema.py index 72c96828..e05e7506 100644 --- a/src/neo4j_genai/components/schema.py +++ b/src/neo4j_genai/experimental/components/schema.py @@ -19,7 +19,7 @@ from pydantic import BaseModel, ValidationError, model_validator, validate_call from neo4j_genai.exceptions import SchemaValidationError -from neo4j_genai.pipeline import Component, DataModel +from neo4j_genai.experimental.pipeline import Component, DataModel class SchemaProperty(BaseModel): @@ -107,13 +107,13 @@ class SchemaBuilder(Component): .. code-block:: python - from neo4j_genai.components.schema import ( + from neo4j_genai.experimental.components.schema import ( SchemaBuilder, SchemaEntity, SchemaProperty, SchemaRelation, ) - from neo4j_genai.pipeline import Pipeline + from neo4j_genai.experimental.pipeline import Pipeline entities = [ SchemaEntity( diff --git a/tests/unit/pipeline/__init__.py b/src/neo4j_genai/experimental/components/text_splitters/__init__.py similarity index 100% rename from tests/unit/pipeline/__init__.py rename to src/neo4j_genai/experimental/components/text_splitters/__init__.py diff --git a/src/neo4j_genai/components/text_splitters/base.py b/src/neo4j_genai/experimental/components/text_splitters/base.py similarity index 86% rename from src/neo4j_genai/components/text_splitters/base.py rename to src/neo4j_genai/experimental/components/text_splitters/base.py index f0bc9fa5..5708d400 100644 --- a/src/neo4j_genai/components/text_splitters/base.py +++ b/src/neo4j_genai/experimental/components/text_splitters/base.py @@ -16,8 +16,8 @@ from abc import abstractmethod -from neo4j_genai.components.types import TextChunks -from neo4j_genai.pipeline.component import Component +from neo4j_genai.experimental.components.types import TextChunks +from neo4j_genai.experimental.pipeline.component import Component class TextSplitter(Component): diff --git a/src/neo4j_genai/components/text_splitters/langchain.py b/src/neo4j_genai/experimental/components/text_splitters/langchain.py similarity index 84% rename from src/neo4j_genai/components/text_splitters/langchain.py rename to src/neo4j_genai/experimental/components/text_splitters/langchain.py index 6ef28823..077748f0 100644 --- a/src/neo4j_genai/components/text_splitters/langchain.py +++ b/src/neo4j_genai/experimental/components/text_splitters/langchain.py @@ -16,8 +16,8 @@ from langchain_text_splitters import TextSplitter as LangChainTextSplitter -from neo4j_genai.components.text_splitters.base import TextSplitter -from neo4j_genai.components.types import TextChunk, TextChunks +from neo4j_genai.experimental.components.text_splitters.base import TextSplitter +from neo4j_genai.experimental.components.types import TextChunk, TextChunks class LangChainTextSplitterAdapter(TextSplitter): @@ -32,8 +32,8 @@ class LangChainTextSplitterAdapter(TextSplitter): .. code-block:: python from langchain_text_splitters import RecursiveCharacterTextSplitter - from neo4j_genai.components.text_splitters.langchain import LangChainTextSplitterAdapter - from neo4j_genai.pipeline import Pipeline + from neo4j_genai.experimental.components.text_splitters.langchain import LangChainTextSplitterAdapter + from neo4j_genai.experimental.pipeline import Pipeline pipeline = Pipeline() text_splitter = LangChainTextSplitterAdapter(RecursiveCharacterTextSplitter()) diff --git a/src/neo4j_genai/components/text_splitters/llamaindex.py b/src/neo4j_genai/experimental/components/text_splitters/llamaindex.py similarity index 84% rename from src/neo4j_genai/components/text_splitters/llamaindex.py rename to src/neo4j_genai/experimental/components/text_splitters/llamaindex.py index 9b489d86..24289593 100644 --- a/src/neo4j_genai/components/text_splitters/llamaindex.py +++ b/src/neo4j_genai/experimental/components/text_splitters/llamaindex.py @@ -16,8 +16,8 @@ from llama_index.core.node_parser import TextSplitter as LlamaIndexTextSplitter -from neo4j_genai.components.text_splitters.base import TextSplitter -from neo4j_genai.components.types import TextChunk, TextChunks +from neo4j_genai.experimental.components.text_splitters.base import TextSplitter +from neo4j_genai.experimental.components.types import TextChunk, TextChunks class LlamaIndexTextSplitterAdapter(TextSplitter): @@ -32,8 +32,8 @@ class LlamaIndexTextSplitterAdapter(TextSplitter): .. 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 + from neo4j_genai.experimental.components.text_splitters.langchain import LangChainTextSplitterAdapter + from neo4j_genai.experimental.pipeline import Pipeline pipeline = Pipeline() text_splitter = LlamaIndexTextSplitterAdapter(SentenceSplitter()) diff --git a/src/neo4j_genai/components/types.py b/src/neo4j_genai/experimental/components/types.py similarity index 97% rename from src/neo4j_genai/components/types.py rename to src/neo4j_genai/experimental/components/types.py index ff7fdee5..a22443a1 100644 --- a/src/neo4j_genai/components/types.py +++ b/src/neo4j_genai/experimental/components/types.py @@ -18,7 +18,7 @@ from pydantic import BaseModel, field_validator -from neo4j_genai.pipeline.component import DataModel +from neo4j_genai.experimental.pipeline.component import DataModel class TextChunk(BaseModel): diff --git a/src/neo4j_genai/pipeline/__init__.py b/src/neo4j_genai/experimental/pipeline/__init__.py similarity index 100% rename from src/neo4j_genai/pipeline/__init__.py rename to src/neo4j_genai/experimental/pipeline/__init__.py diff --git a/src/neo4j_genai/pipeline/component.py b/src/neo4j_genai/experimental/pipeline/component.py similarity index 97% rename from src/neo4j_genai/pipeline/component.py rename to src/neo4j_genai/experimental/pipeline/component.py index 4446d251..6625d11a 100644 --- a/src/neo4j_genai/pipeline/component.py +++ b/src/neo4j_genai/experimental/pipeline/component.py @@ -20,7 +20,7 @@ from pydantic import BaseModel -from neo4j_genai.pipeline.exceptions import PipelineDefinitionError +from neo4j_genai.experimental.pipeline.exceptions import PipelineDefinitionError class DataModel(BaseModel): diff --git a/src/neo4j_genai/pipeline/exceptions.py b/src/neo4j_genai/experimental/pipeline/exceptions.py similarity index 100% rename from src/neo4j_genai/pipeline/exceptions.py rename to src/neo4j_genai/experimental/pipeline/exceptions.py diff --git a/src/neo4j_genai/pipeline/pipeline.py b/src/neo4j_genai/experimental/pipeline/pipeline.py similarity index 97% rename from src/neo4j_genai/pipeline/pipeline.py rename to src/neo4j_genai/experimental/pipeline/pipeline.py index 7c8e35d5..c036b780 100644 --- a/src/neo4j_genai/pipeline/pipeline.py +++ b/src/neo4j_genai/experimental/pipeline/pipeline.py @@ -22,19 +22,23 @@ from pydantic import BaseModel, Field -from neo4j_genai.pipeline.component import Component, DataModel -from neo4j_genai.pipeline.exceptions import ( +from neo4j_genai.experimental.pipeline.component import Component, DataModel +from neo4j_genai.experimental.pipeline.exceptions import ( PipelineDefinitionError, PipelineMissingDependencyError, PipelineStatusUpdateError, ) -from neo4j_genai.pipeline.pipeline_graph import ( +from neo4j_genai.experimental.pipeline.pipeline_graph import ( PipelineEdge, PipelineGraph, PipelineNode, ) -from neo4j_genai.pipeline.stores import InMemoryStore, Store -from neo4j_genai.pipeline.types import ComponentConfig, ConnectionConfig, PipelineConfig +from neo4j_genai.experimental.pipeline.stores import InMemoryStore, Store +from neo4j_genai.experimental.pipeline.types import ( + ComponentConfig, + ConnectionConfig, + PipelineConfig, +) logger = logging.getLogger(__name__) diff --git a/src/neo4j_genai/pipeline/pipeline_graph.py b/src/neo4j_genai/experimental/pipeline/pipeline_graph.py similarity index 100% rename from src/neo4j_genai/pipeline/pipeline_graph.py rename to src/neo4j_genai/experimental/pipeline/pipeline_graph.py diff --git a/src/neo4j_genai/pipeline/stores.py b/src/neo4j_genai/experimental/pipeline/stores.py similarity index 100% rename from src/neo4j_genai/pipeline/stores.py rename to src/neo4j_genai/experimental/pipeline/stores.py diff --git a/src/neo4j_genai/pipeline/types.py b/src/neo4j_genai/experimental/pipeline/types.py similarity index 94% rename from src/neo4j_genai/pipeline/types.py rename to src/neo4j_genai/experimental/pipeline/types.py index 89ad169d..9f03b8db 100644 --- a/src/neo4j_genai/pipeline/types.py +++ b/src/neo4j_genai/experimental/pipeline/types.py @@ -16,7 +16,7 @@ from pydantic import BaseModel, ConfigDict -from neo4j_genai.pipeline.component import Component +from neo4j_genai.experimental.pipeline.component import Component class ComponentConfig(BaseModel): diff --git a/tests/e2e/test_kb_builder_pipeline_e2e.py b/tests/e2e/test_kb_builder_pipeline_e2e.py index 28d8ebcd..6d9a5248 100644 --- a/tests/e2e/test_kb_builder_pipeline_e2e.py +++ b/tests/e2e/test_kb_builder_pipeline_e2e.py @@ -21,23 +21,25 @@ import neo4j import pytest from langchain_text_splitters import CharacterTextSplitter -from neo4j_genai.components.embedder import TextChunkEmbedder -from neo4j_genai.components.entity_relation_extractor import ( +from neo4j_genai.embedder import Embedder +from neo4j_genai.exceptions import LLMGenerationError +from neo4j_genai.experimental.components.embedder import TextChunkEmbedder +from neo4j_genai.experimental.components.entity_relation_extractor import ( LLMEntityRelationExtractor, OnError, ) -from neo4j_genai.components.kg_writer import Neo4jWriter -from neo4j_genai.components.schema import ( +from neo4j_genai.experimental.components.kg_writer import Neo4jWriter +from neo4j_genai.experimental.components.schema import ( SchemaBuilder, SchemaEntity, SchemaProperty, SchemaRelation, ) -from neo4j_genai.components.text_splitters.langchain import LangChainTextSplitterAdapter -from neo4j_genai.embedder import Embedder -from neo4j_genai.exceptions import LLMGenerationError +from neo4j_genai.experimental.components.text_splitters.langchain import ( + LangChainTextSplitterAdapter, +) +from neo4j_genai.experimental.pipeline import Pipeline from neo4j_genai.llm import LLMInterface, LLMResponse -from neo4j_genai.pipeline import Pipeline BASE_DIR = os.path.dirname(os.path.abspath(__file__)) diff --git a/tests/e2e/test_kg_construction_e2e.py b/tests/e2e/test_kg_construction_e2e.py index 548ba7cd..5c9cf42a 100644 --- a/tests/e2e/test_kg_construction_e2e.py +++ b/tests/e2e/test_kg_construction_e2e.py @@ -15,8 +15,8 @@ import neo4j import pytest -from neo4j_genai.components.kg_writer import Neo4jWriter -from neo4j_genai.components.types import ( +from neo4j_genai.experimental.components.kg_writer import Neo4jWriter +from neo4j_genai.experimental.components.types import ( Neo4jGraph, Neo4jNode, Neo4jRelationship, diff --git a/tests/unit/experimental/__init__.py b/tests/unit/experimental/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/unit/components/__init__.py b/tests/unit/experimental/components/__init__.py similarity index 100% rename from tests/unit/components/__init__.py rename to tests/unit/experimental/components/__init__.py diff --git a/tests/unit/components/test_embedder.py b/tests/unit/experimental/components/test_embedder.py similarity index 90% rename from tests/unit/components/test_embedder.py rename to tests/unit/experimental/components/test_embedder.py index ff87f8e1..98064184 100644 --- a/tests/unit/components/test_embedder.py +++ b/tests/unit/experimental/components/test_embedder.py @@ -15,8 +15,8 @@ from unittest.mock import MagicMock import pytest -from neo4j_genai.components.embedder import TextChunkEmbedder -from neo4j_genai.components.types import TextChunk, TextChunks +from neo4j_genai.experimental.components.embedder import TextChunkEmbedder +from neo4j_genai.experimental.components.types import TextChunk, TextChunks @pytest.mark.asyncio diff --git a/tests/unit/components/test_entity_relation_extractor.py b/tests/unit/experimental/components/test_entity_relation_extractor.py similarity index 97% rename from tests/unit/components/test_entity_relation_extractor.py rename to tests/unit/experimental/components/test_entity_relation_extractor.py index 8e1e6200..e502a21f 100644 --- a/tests/unit/components/test_entity_relation_extractor.py +++ b/tests/unit/experimental/components/test_entity_relation_extractor.py @@ -17,13 +17,18 @@ from unittest.mock import MagicMock import pytest -from neo4j_genai.components.entity_relation_extractor import ( +from neo4j_genai.exceptions import LLMGenerationError +from neo4j_genai.experimental.components.entity_relation_extractor import ( EntityRelationExtractor, LLMEntityRelationExtractor, OnError, ) -from neo4j_genai.components.types import Neo4jGraph, Neo4jNode, TextChunk, TextChunks -from neo4j_genai.exceptions import LLMGenerationError +from neo4j_genai.experimental.components.types import ( + Neo4jGraph, + Neo4jNode, + TextChunk, + TextChunks, +) from neo4j_genai.llm import LLMInterface, LLMResponse diff --git a/tests/unit/components/test_kg_writer.py b/tests/unit/experimental/components/test_kg_writer.py similarity index 97% rename from tests/unit/components/test_kg_writer.py rename to tests/unit/experimental/components/test_kg_writer.py index f2b90ba0..23a0c6f8 100644 --- a/tests/unit/components/test_kg_writer.py +++ b/tests/unit/experimental/components/test_kg_writer.py @@ -17,8 +17,8 @@ from unittest.mock import MagicMock import pytest -from neo4j_genai.components.kg_writer import Neo4jWriter -from neo4j_genai.components.types import ( +from neo4j_genai.experimental.components.kg_writer import Neo4jWriter +from neo4j_genai.experimental.components.types import ( Neo4jGraph, Neo4jNode, Neo4jRelationship, diff --git a/tests/unit/components/test_schema.py b/tests/unit/experimental/components/test_schema.py similarity index 99% rename from tests/unit/components/test_schema.py rename to tests/unit/experimental/components/test_schema.py index c4e9074f..6be43a19 100644 --- a/tests/unit/components/test_schema.py +++ b/tests/unit/experimental/components/test_schema.py @@ -15,13 +15,13 @@ from __future__ import annotations import pytest -from neo4j_genai.components.schema import ( +from neo4j_genai.exceptions import SchemaValidationError +from neo4j_genai.experimental.components.schema import ( SchemaBuilder, SchemaEntity, SchemaProperty, SchemaRelation, ) -from neo4j_genai.exceptions import SchemaValidationError from pydantic import ValidationError diff --git a/tests/unit/components/test_types.py b/tests/unit/experimental/components/test_types.py similarity index 94% rename from tests/unit/components/test_types.py rename to tests/unit/experimental/components/test_types.py index 840fda8c..c812f599 100644 --- a/tests/unit/components/test_types.py +++ b/tests/unit/experimental/components/test_types.py @@ -13,7 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. import pytest -from neo4j_genai.components.types import Neo4jNode +from neo4j_genai.experimental.components.types import Neo4jNode def test_neo4j_node_invalid_property() -> None: diff --git a/tests/unit/components/text_splitters/test_langchain.py b/tests/unit/experimental/components/text_splitters/test_langchain.py similarity index 90% rename from tests/unit/components/text_splitters/test_langchain.py rename to tests/unit/experimental/components/text_splitters/test_langchain.py index 6b523547..af95dc29 100644 --- a/tests/unit/components/text_splitters/test_langchain.py +++ b/tests/unit/experimental/components/text_splitters/test_langchain.py @@ -14,8 +14,10 @@ # limitations under the License. import pytest from langchain_text_splitters import RecursiveCharacterTextSplitter -from neo4j_genai.components.text_splitters.langchain import LangChainTextSplitterAdapter -from neo4j_genai.components.types import TextChunk, TextChunks +from neo4j_genai.experimental.components.text_splitters.langchain import ( + LangChainTextSplitterAdapter, +) +from neo4j_genai.experimental.components.types import TextChunk, TextChunks text = """ Lorem ipsum dolor sit amet, consectetur adipiscing elit. diff --git a/tests/unit/components/text_splitters/test_llamaindex.py b/tests/unit/experimental/components/text_splitters/test_llamaindex.py similarity index 92% rename from tests/unit/components/text_splitters/test_llamaindex.py rename to tests/unit/experimental/components/text_splitters/test_llamaindex.py index 4e9b784b..a9e0832d 100644 --- a/tests/unit/components/text_splitters/test_llamaindex.py +++ b/tests/unit/experimental/components/text_splitters/test_llamaindex.py @@ -14,10 +14,10 @@ # limitations under the License. import pytest from llama_index.core.node_parser.text.sentence import SentenceSplitter -from neo4j_genai.components.text_splitters.llamaindex import ( +from neo4j_genai.experimental.components.text_splitters.llamaindex import ( LlamaIndexTextSplitterAdapter, ) -from neo4j_genai.components.types import TextChunk, TextChunks +from neo4j_genai.experimental.components.types import TextChunk, TextChunks text = """ Lorem ipsum dolor sit amet, consectetur adipiscing elit. diff --git a/tests/unit/experimental/pipeline/__init__.py b/tests/unit/experimental/pipeline/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/unit/pipeline/components.py b/tests/unit/experimental/pipeline/components.py similarity index 95% rename from tests/unit/pipeline/components.py rename to tests/unit/experimental/pipeline/components.py index 25fac635..93e74370 100644 --- a/tests/unit/pipeline/components.py +++ b/tests/unit/experimental/pipeline/components.py @@ -12,7 +12,7 @@ # 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 neo4j_genai.pipeline import Component, DataModel +from neo4j_genai.experimental.pipeline import Component, DataModel class StringResultModel(DataModel): diff --git a/tests/unit/pipeline/test_component.py b/tests/unit/experimental/pipeline/test_component.py similarity index 100% rename from tests/unit/pipeline/test_component.py rename to tests/unit/experimental/pipeline/test_component.py diff --git a/tests/unit/pipeline/test_orchestrator.py b/tests/unit/experimental/pipeline/test_orchestrator.py similarity index 94% rename from tests/unit/pipeline/test_orchestrator.py rename to tests/unit/experimental/pipeline/test_orchestrator.py index e311549a..8e319df9 100644 --- a/tests/unit/pipeline/test_orchestrator.py +++ b/tests/unit/experimental/pipeline/test_orchestrator.py @@ -13,8 +13,8 @@ # See the License for the specific language governing permissions and # limitations under the License. import pytest -from neo4j_genai.pipeline import Component -from neo4j_genai.pipeline.pipeline import Orchestrator, Pipeline, RunStatus +from neo4j_genai.experimental.pipeline import Component +from neo4j_genai.experimental.pipeline.pipeline import Orchestrator, Pipeline, RunStatus @pytest.fixture(scope="function") diff --git a/tests/unit/pipeline/test_pipeline.py b/tests/unit/experimental/pipeline/test_pipeline.py similarity index 95% rename from tests/unit/pipeline/test_pipeline.py rename to tests/unit/experimental/pipeline/test_pipeline.py index 1839874b..f226dc34 100644 --- a/tests/unit/pipeline/test_pipeline.py +++ b/tests/unit/experimental/pipeline/test_pipeline.py @@ -18,8 +18,8 @@ from unittest.mock import AsyncMock import pytest -from neo4j_genai.pipeline import Component, Pipeline -from neo4j_genai.pipeline.exceptions import PipelineDefinitionError +from neo4j_genai.experimental.pipeline import Component, Pipeline +from neo4j_genai.experimental.pipeline.exceptions import PipelineDefinitionError from .components import ( ComponentAdd, @@ -39,7 +39,7 @@ async def test_simple_pipeline_two_components() -> None: pipe.add_component("b", component_b) pipe.connect("a", "b", {}) with mock.patch( - "tests.unit.pipeline.test_pipeline.ComponentNoParam.run" + "tests.unit.experimental.pipeline.test_pipeline.ComponentNoParam.run" ) as mock_run: mock_run.side_effect = [ StringResultModel(result=""), @@ -68,7 +68,7 @@ async def test_pipeline_parameter_propagation() -> None: }, ) with mock.patch( - "tests.unit.pipeline.test_pipeline.ComponentPassThrough.run" + "tests.unit.experimental.pipeline.test_pipeline.ComponentPassThrough.run" ) as mock_run: mock_run.side_effect = [ StringResultModel(result="text"), diff --git a/tests/unit/pipeline/test_pipeline_graph.py b/tests/unit/experimental/pipeline/test_pipeline_graph.py similarity index 98% rename from tests/unit/pipeline/test_pipeline_graph.py rename to tests/unit/experimental/pipeline/test_pipeline_graph.py index 8a2b1576..85b18867 100644 --- a/tests/unit/pipeline/test_pipeline_graph.py +++ b/tests/unit/experimental/pipeline/test_pipeline_graph.py @@ -13,7 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. import pytest -from neo4j_genai.pipeline.pipeline_graph import ( +from neo4j_genai.experimental.pipeline.pipeline_graph import ( PipelineEdge, PipelineGraph, PipelineNode, diff --git a/tests/unit/pipeline/test_store.py b/tests/unit/experimental/pipeline/test_store.py similarity index 80% rename from tests/unit/pipeline/test_store.py rename to tests/unit/experimental/pipeline/test_store.py index 6864e3f8..5de7c1f8 100644 --- a/tests/unit/pipeline/test_store.py +++ b/tests/unit/experimental/pipeline/test_store.py @@ -1,5 +1,5 @@ import pytest -from neo4j_genai.pipeline.stores import InMemoryStore +from neo4j_genai.experimental.pipeline.stores import InMemoryStore def test_memory_store() -> None: