Changes following review

Author: cbornet

libs/knowledge-store/ragstack_knowledge_store/__init__.py

@@ -1,3 +1,4 @@
-from .knowledge_store import KnowledgeStore
+from .embedding_model import EmbeddingModel
+from .knowledge_store import KnowledgeStore, Node, SetupMode, TextNode
 
-__all__ = ["KnowledgeStore"]
+__all__ = ["EmbeddingModel", "KnowledgeStore", "Node", "SetupMode", "TextNode"]

libs/knowledge-store/ragstack_knowledge_store/knowledge_store.py

@@ -29,7 +29,7 @@ class Node:
  """Node in the KnowledgeStore graph"""
 
  id: Optional[str] = None
- """Unique ID for the node. Shall be generated by the KnowledgeStore if not set"""
+ """Unique ID for the node. Will be generated by the KnowledgeStore if not set."""
  metadata: dict = field(default_factory=dict)
  """Metadata for the node. May contain information used to link this node 
  with other nodes."""
@@ -38,7 +38,7 @@ class Node:
 @dataclass
 class TextNode(Node):
  text: str = None
- """Text contained by the node"""
+ """Text contained by the node."""
 
 
 class SetupMode(Enum):

libs/knowledge-store/ragstack_knowledge_store/langchain/__init__.py

@@ -1,4 +1,4 @@
-from .base import KnowledgeStore
+from .base import KnowledgeStore, Node, TextNode
 from .cassandra import CassandraKnowledgeStore
 
-__all__ = ["CassandraKnowledgeStore", "KnowledgeStore"]
+__all__ = ["CassandraKnowledgeStore", "KnowledgeStore", "Node", "TextNode"]

libs/knowledge-store/ragstack_knowledge_store/langchain/base.py

@@ -31,18 +31,18 @@ def _has_next(iterator: Iterator) -> None:
 
 
 class Node(Serializable):
- """Node in the KnowledgeStore graph"""
+ """Node in the KnowledgeStore graph."""
 
  id: Optional[str]
- """Unique ID for the node. Shall be generated by the KnowledgeStore if not set"""
+ """Unique ID for the node. Will be generated by the KnowledgeStore if not set."""
  metadata: dict = Field(default_factory=dict)
  """Metadata for the node. May contain information used to link this node
  with other nodes."""
 
 
 class TextNode(Node):
  text: str
- """Text contained by the node"""
+ """Text contained by the node."""
 
 
 def _texts_to_nodes(

libs/knowledge-store/ragstack_knowledge_store/langchain/cassandra.py

@@ -12,12 +12,10 @@
 from langchain_core.embeddings import Embeddings
 
 from .base import KnowledgeStore, Node, TextNode
-from ragstack_knowledge_store import knowledge_store
-from ragstack_knowledge_store.edge_extractor import EdgeExtractor
-from ragstack_knowledge_store.embedding_model import EmbeddingModel
+from ragstack_knowledge_store import EmbeddingModel, knowledge_store
 
 
-class EmbeddingModelAdapter(EmbeddingModel):
+class _EmbeddingModelAdapter(EmbeddingModel):
  def __init__(self, embeddings: Embeddings):
  self.embeddings = embeddings
 
@@ -81,7 +79,7 @@ def __init__(
  _setup_mode = getattr(knowledge_store.SetupMode, setup_mode.name)
 
  self.store = knowledge_store.KnowledgeStore(
- embedding=EmbeddingModelAdapter(embedding),
+ embedding=_EmbeddingModelAdapter(embedding),
  node_table=node_table,
  edge_table=edge_table,
  session=session,

libs/knowledge-store/ragstack_knowledge_store/math.py

@@ -1,4 +1,6 @@
-"""Copied from langchain_community.utils.math"""
+"""Copied from langchain_community.utils.math
+See https://github.com/langchain-ai/langchain/blob/langchain-community%3D%3D0.0.38/libs/community/langchain_community/utils/math.py
+"""
 import logging
 from typing import List, Union