Merge pull request #1 from nodestream-proj/improvements

Improving code quality and adding tests

Author: zprobst

.github/workflows/ci.yaml

@@ -8,7 +8,7 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        python-version: ["3.10", "3.11"]
+        python-version: ["3.10", "3.11", "3.12", "3.13"]
 
     steps:
       - uses: actions/checkout@v3

.vscode/settings.json

@@ -0,0 +1,7 @@
+{
+    "python.testing.pytestArgs": [
+        "."
+    ],
+    "python.testing.unittestEnabled": false,
+    "python.testing.pytestEnabled": true
+}

nodestream_plugin_semantic/chunk.py

@@ -1,8 +1,8 @@
 from abc import ABC, abstractmethod
 from typing import Iterable
 
-from nodestream.subclass_registry import SubclassRegistry
 from nodestream.pluggable import Pluggable
+from nodestream.subclass_registry import SubclassRegistry
 
 from .model import Content
 
@@ -25,7 +25,8 @@ def from_file_data(type, **chunker_kwargs) -> "Chunker":
         return CHUNKER_SUBCLASS_REGISTRY.get(type)(**chunker_kwargs)
 
     @abstractmethod
-    def chunk(self, content: Content) -> Iterable[Content]: ...
+    def chunk(self, content: Content) -> Iterable[Content]:
+        ...
 
 
 class SplitOnDelimiterChunker(Chunker):

nodestream_plugin_semantic/content_types.py

@@ -1,10 +1,9 @@
 from abc import ABC, abstractmethod
-from typing import Iterable
 from pathlib import Path
+from typing import Iterable
 
 from nodestream.subclass_registry import SubclassRegistry
 
-
 CONTENT_TYPE_SUBCLASS_REGISTRY = SubclassRegistry()
 PLAIN_TEXT_ALIAS = "plain_text"
 PLAIN_TEXT_EXTENSIONS = {".txt", ".md"}

nodestream_plugin_semantic/embed.py

@@ -1,11 +1,10 @@
 from abc import ABC, abstractmethod
 
-from nodestream.subclass_registry import SubclassRegistry
 from nodestream.pluggable import Pluggable
+from nodestream.subclass_registry import SubclassRegistry
 
 from .model import Content, Embedding
 
-
 EMBEDDER_SUBCLASS_REGISTRY = SubclassRegistry()
 
 

nodestream_plugin_semantic/model.py

@@ -1,9 +1,8 @@
-from dataclasses import dataclass
 import hashlib
-from typing import List, Optional, Iterable
-
-from nodestream.model import DesiredIngestion, Node
+from dataclasses import dataclass
+from typing import Iterable, List, Optional
 
+from nodestream.model import DesiredIngestion, Node, Relationship
 
 Embedding = List[float | int]
 CONTENT_NODE_TYPE_ID_PROPERTY = "id"
@@ -66,8 +65,9 @@ def make_ingestible(
 
         if self.parent:
             self.parent.apply_to_node(node_type, related := Node())
+            relationship = Relationship(type=relationship_type)
             ingest.add_relationship(
-                related_node=related, type=relationship_type, outbound=False
+                related_node=related, relationship=relationship, outbound=False
             )
 
         return ingest

nodestream_plugin_semantic/pipeline.py

@@ -1,24 +1,24 @@
 from glob import glob
 from pathlib import Path
-from typing import Optional, List, Dict
+from typing import Dict, List, Optional
 
 from nodestream.model import DesiredIngestion
 from nodestream.pipeline import Extractor, Transformer
 from nodestream.pipeline.value_providers import (
-    ValueProvider,
     JmespathValueProvider,
     ProviderContext,
+    ValueProvider,
 )
 from nodestream.schema import (
+    Cardinality,
     ExpandsSchema,
-    SchemaExpansionCoordinator,
     GraphObjectSchema,
-    Cardinality,
+    SchemaExpansionCoordinator,
 )
 
 from .chunk import Chunker
+from .content_types import PLAIN_TEXT_ALIAS, ContentType
 from .embed import Embedder
-from .content_types import ContentType, PLAIN_TEXT_ALIAS
 from .model import Content
 
 DEFAULT_ID = JmespathValueProvider.from_string_expression("id")

poetry.lock

@@ -27,12 +27,16 @@ nodestream = "^0.13.0"
 semchunk = "^2.2.0"
 
 [tool.poetry.group.dev.dependencies]
-pytest = "^7.4.0"
+pytest = "^8.2.0"
 pytest-mock = "^3.11.1"
 ruff = "^0.0.284"
 isort = "^5.12.0"
 black = "^23.7.0"
 pytest-cov = "^4.1.0"
+pytest-asyncio = "^0.24.0"
+
+[tool.isort]
+profile = "black"
 
 [build-system]
 requires = ["poetry-core"]

pyproject.toml

@@ -0,0 +1,65 @@
+from unittest.mock import Mock
+
+from nodestream.model import DesiredIngestion, Node
+
+from nodestream_plugin_semantic.model import Content, hash
+
+
+def test_content_from_text():
+    content_text = "test content"
+    content = Content.from_text(content_text)
+    assert content.content == content_text
+    assert content.id == hash(content_text)
+    assert content.parent is None
+
+
+def test_content_add_metadata():
+    content = Content.from_text("test content")
+    content.add_metadata("key", "value")
+    assert content.metadata == {"key": "value"}
+
+
+def test_content_split_on_delimiter():
+    content_text = "line1\nline2\nline3"
+    content = Content.from_text(content_text)
+    lines = list(content.split_on_delimiter("\n"))
+    assert len(lines) == 3
+    assert lines[0].content == "line1"
+    assert lines[1].content == "line2"
+    assert lines[2].content == "line3"
+    assert all(line.parent == content for line in lines)
+
+
+def test_content_assign_embedding():
+    content = Content.from_text("test content")
+    embedding = [0.1, 0.2, 0.3]
+    content.assign_embedding(embedding)
+    assert content.embedding == embedding
+
+
+def test_content_apply_to_node():
+    content = Content.from_text("test content")
+    node = Mock(spec=Node)
+    content.apply_to_node("test_type", node)
+    node.type = "test_type"
+    node.key_values.set_property.assert_called_with("id", content.id)
+    node.properties.set_property.assert_any_call("content", content.content)
+
+
+def test_content_make_ingestible():
+    parent_content = Content.from_text("parent content")
+    child_content = Content.from_text("child content", parent=parent_content)
+    ingest = child_content.make_ingestible("test_type", "test_relationship")
+
+    assert isinstance(ingest, DesiredIngestion)
+    assert ingest.source.type == "test_type"
+    ingest.source.key_values == {"id": child_content.id}
+    ingest.source.properties == {"content": child_content.content}
+
+    assert len(ingest.relationships) == 1
+    relationship = ingest.relationships[0]
+    assert relationship.relationship.type == "test_relationship"
+    assert relationship.outbound is False
+    assert relationship.to_node.type == "test_type"
+    relationship.to_node.key_values == {"id": parent_content.id}
+    relationship.to_node.properties == {"content": parent_content.content}