vs-graph 0.1.0__py3-none-any.whl

vs_graph/__init__.py

@@ -0,0 +1,35 @@
+from vs_graph.vs_base_graph import VsBaseGraph
+from vs_graph.node.vs_base_node import VsBaseNode
+from vs_graph.edge.vs_base_edge import VsBaseEdge
+from vs_graph.edge.vs_base_conditional_edge import VsBaseConditionalEdge
+from vs_graph.edge.vs_direct_edge import VsDirectEdge
+from vs_graph.edge.vs_fan_out_edge import VsFanOutEdge
+from vs_graph.edge.vs_dynamic_fan_out_edge import VsDynamicFanOutEdge
+from vs_graph.schema.vs_base_graph_state import VsBaseGraphState
+from vs_graph.schema.vs_graph_error import VsGraphError
+from vs_graph.decorator.vs_node_decorator import node
+from vs_graph.decorator.vs_edge_decorator import edge
+from vs_graph.exception.vs_graph_exceptions import (
+    VsUserException,
+    VsGraphBuildException,
+    VsNodeExecutionException,
+    VsEdgeEvaluationException,
+)
+
+__all__ = [
+    "VsBaseGraph",
+    "VsBaseNode",
+    "VsBaseEdge",
+    "VsBaseConditionalEdge",
+    "VsDirectEdge",
+    "VsFanOutEdge",
+    "VsDynamicFanOutEdge",
+    "VsBaseGraphState",
+    "VsGraphError",
+    "node",
+    "edge",
+    "VsUserException",
+    "VsGraphBuildException",
+    "VsNodeExecutionException",
+    "VsEdgeEvaluationException",
+]

vs_graph/checkpointer/vs_base_checkpointer.py

@@ -0,0 +1,33 @@
+from abc import ABC, abstractmethod
+from typing import Any, Dict, Optional
+
+from vs_common.log.vs_log_manager import VsLogManager
+
+
+class VsBaseCheckpointer(ABC):
+
+    def __init__(self, config: Optional[Dict[str, Any]] = None):
+        self.config = config or {}
+        self.logger = VsLogManager.get_instance(self.__class__.__name__)
+        self._checkpointer: Optional[Any] = None
+
+    @abstractmethod
+    def initialize(self) -> None:
+        pass
+
+    @abstractmethod
+    def health_check(self) -> bool:
+        pass
+
+    def cleanup(self) -> None:
+        pass
+
+    async def ensure_ready(self) -> None:
+        pass
+
+    def get_checkpointer(self) -> Any:
+        if self._checkpointer is None:
+            raise RuntimeError(
+                f"{self.__class__.__name__} has not been initialized. Call initialize() first."
+            )
+        return self._checkpointer

vs_graph/checkpointer/vs_checkpointer_factory.py

@@ -0,0 +1,44 @@
+from typing import Any, Callable, Dict
+
+from vs_common.log.vs_log_manager import VsLogManager
+
+
+class VsCheckpointerFactory:
+
+    _registry: Dict[str, Callable] = {}
+    _logger = None
+
+    @classmethod
+    def _get_logger(cls):
+        if cls._logger is None:
+            cls._logger = VsLogManager.get_instance("VsCheckpointerFactory")
+        return cls._logger
+
+    @classmethod
+    def register(cls, name: str, factory: Callable) -> None:
+        cls._registry[name] = factory
+        cls._get_logger().info(f"Registered checkpointer factory: '{name}'")
+
+    @classmethod
+    def create(cls, checkpointer_type: str, config: Dict[str, Any]):
+        if checkpointer_type not in cls._registry:
+            raise ValueError(
+                f"Unknown checkpointer type: '{checkpointer_type}'. "
+                f"Registered: {sorted(cls._registry.keys())}"
+            )
+        return cls._registry[checkpointer_type](config)
+
+
+VsCheckpointerFactory.register(
+    "memory",
+    lambda cfg: __import__(
+        "vs_graph.checkpointer.vs_memory_checkpointer", fromlist=["VsMemoryCheckpointer"]
+    ).VsMemoryCheckpointer(config=cfg),
+)
+
+VsCheckpointerFactory.register(
+    "redis",
+    lambda cfg: __import__(
+        "vs_graph.checkpointer.vs_redis_checkpointer", fromlist=["VsRedisCheckpointer"]
+    ).VsRedisCheckpointer(config=cfg),
+)

vs_graph/checkpointer/vs_memory_checkpointer.py

@@ -0,0 +1,22 @@
+from typing import Any, Dict, Optional
+
+from vs_graph.checkpointer.vs_base_checkpointer import VsBaseCheckpointer
+
+
+class VsMemoryCheckpointer(VsBaseCheckpointer):
+
+    def __init__(self, config: Optional[Dict[str, Any]] = None):
+        super().__init__(config)
+
+    def initialize(self) -> None:
+        try:
+            from langgraph.checkpoint.memory import MemorySaver
+            self._checkpointer = MemorySaver()
+            self.logger.info("Memory checkpointer initialized")
+        except ImportError as e:
+            raise ImportError("langgraph not installed. Run: pip install langgraph") from e
+        except Exception as e:
+            raise RuntimeError(f"Failed to initialize memory checkpointer: {e}") from e
+
+    def health_check(self) -> bool:
+        return self._checkpointer is not None

vs_graph/checkpointer/vs_redis_checkpointer.py

@@ -0,0 +1,46 @@
+from typing import Any, Dict, Optional
+
+from vs_graph.checkpointer.vs_base_checkpointer import VsBaseCheckpointer
+
+
+class VsRedisCheckpointer(VsBaseCheckpointer):
+
+    def __init__(self, config: Optional[Dict[str, Any]] = None):
+        super().__init__(config)
+        if not self.config.get("redis_url"):
+            raise ValueError("redis_url is required in config for VsRedisCheckpointer.")
+
+    def initialize(self) -> None:
+        try:
+            from langgraph.checkpoint.redis.aio import AsyncRedisSaver
+
+            redis_url = self.config["redis_url"]
+            key_prefix = self.config.get("key_prefix", "vs:checkpoint:")
+
+            self._checkpointer = AsyncRedisSaver(redis_url=redis_url, checkpoint_prefix=key_prefix)
+            self._setup_done = False
+            self.logger.info(f"Redis checkpointer initialized: {redis_url}")
+        except ImportError as e:
+            raise ImportError("Install with: pip install redis langgraph[redis]") from e
+        except Exception as e:
+            raise RuntimeError(f"Failed to initialize Redis checkpointer: {e}") from e
+
+    async def ensure_ready(self) -> None:
+        if not getattr(self, "_setup_done", False):
+            await self._checkpointer.setup()
+            self._setup_done = True
+
+    def cleanup(self) -> None:
+        self._checkpointer = None
+
+    def health_check(self) -> bool:
+        if not self._checkpointer:
+            return False
+        try:
+            import redis
+            r = redis.from_url(self.config["redis_url"])
+            r.ping()
+            r.close()
+            return True
+        except Exception:
+            return False

vs_graph/decorator/vs_edge_decorator.py

@@ -0,0 +1,14 @@
+from typing import Optional, Type
+
+
+def edge(_cls=None, *, name: Optional[str] = None):
+    def decorator(cls: Type) -> Type:
+        from vs_graph.registry.vs_edge_registry import VsEdgeRegistry
+        registered_name = name or cls.__name__
+        VsEdgeRegistry.register(registered_name, cls)
+        return cls
+
+    if _cls is not None:
+        return decorator(_cls)
+
+    return decorator

vs_graph/decorator/vs_graph_decorator.py

@@ -0,0 +1,23 @@
+import importlib
+import pkgutil
+from typing import List, Optional, Type, Union
+
+
+def graph(
+    nodes: Optional[Union[str, List[str]]] = None,
+    edges: Optional[Union[str, List[str]]] = None,
+):
+    def decorator(cls: Type) -> Type:
+        cls._vs_node_packages = [nodes] if isinstance(nodes, str) else (nodes or [])
+        cls._vs_edge_packages = [edges] if isinstance(edges, str) else (edges or [])
+        return cls
+
+    return decorator
+
+
+def _import_package(package_path: str) -> None:
+    pkg = importlib.import_module(package_path)
+    if not hasattr(pkg, "__path__"):
+        return
+    for _, modname, _ in pkgutil.walk_packages(pkg.__path__, prefix=pkg.__name__ + "."):
+        importlib.import_module(modname)

vs_graph/decorator/vs_node_decorator.py

@@ -0,0 +1,15 @@
+from typing import Optional, Type
+
+
+def node(_cls=None, *, name: Optional[str] = None):
+    def decorator(cls: Type) -> Type:
+        from vs_graph.registry.vs_node_registry import VsNodeRegistry
+        registered_name = name or cls.__name__
+        VsNodeRegistry.register(registered_name, cls)
+        cls._node_name = registered_name
+        return cls
+
+    if _cls is not None:
+        return decorator(_cls)
+
+    return decorator

vs_graph/edge/vs_base_conditional_edge.py

@@ -0,0 +1,57 @@
+import time
+from abc import abstractmethod
+from typing import Any, Dict, Optional, TypeVar
+
+from vs_common.publisher.vs_base_trace_publisher import VsBaseTracePublisher
+from vs_graph.edge.vs_base_edge import VsBaseEdge
+from vs_graph.schema.vs_base_graph_state import VsBaseGraphState
+
+S = TypeVar("S", bound=VsBaseGraphState)
+
+
+class VsBaseConditionalEdge(VsBaseEdge[S]):
+
+    def __init__(
+        self,
+        source: str,
+        route: Dict[str, str],
+        config: Optional[Dict[str, Any]] = None,
+        trace_publisher: Optional[VsBaseTracePublisher] = None,
+    ):
+        super().__init__(source=source, trace_publisher=trace_publisher)
+
+        if not route or "default" not in route:
+            raise ValueError(f"Conditional edge from '{source}' must have a 'default' route.")
+        if not route["default"]:
+            raise ValueError(f"Conditional edge from '{source}' has an empty 'default' target.")
+
+        self.route_map = route
+        self.config = config or {}
+
+    async def route(self, state: S) -> str:
+        trace_id = state.get("trace_id", "")
+        self.logger.info(f"[{self.source}] routing started")
+        start_time = time.time()
+        try:
+            route_key = self._evaluate(state)
+            next_node = self.route_map.get(route_key) or self.route_map.get("default")
+            if next_node is None:
+                raise RuntimeError(f"Conditional edge from '{self.source}' has no default route.")
+            duration = round(time.time() - start_time, 3)
+            self.logger.info(
+                f"[{self.source}] routed to '{next_node}' (key='{route_key}') in {duration}s"
+            )
+            await self.trace(
+                trace_id,
+                f"edge '{self.source}' routing to '{next_node}'",
+                data={"route_key": route_key, "next_node": next_node, "duration_seconds": duration},
+            )
+            return next_node
+        except Exception as e:
+            duration = round(time.time() - start_time, 3)
+            self.logger.error(f"[{self.source}] routing failed in {duration}s: {e}", exc_info=True)
+            raise
+
+    @abstractmethod
+    def _evaluate(self, state: S) -> str:
+        pass

vs_graph/edge/vs_base_edge.py

@@ -0,0 +1,48 @@
+from abc import ABC
+from typing import Any, Dict, Generic, Optional, TypeVar
+
+from vs_common.log.vs_log_manager import VsLogManager
+from vs_common.publisher.vs_base_trace_publisher import VsBaseTracePublisher
+from vs_common.schema.vs_trace_level import VsTraceLevel
+from vs_graph.schema.vs_base_graph_state import VsBaseGraphState
+
+S = TypeVar("S", bound=VsBaseGraphState)
+
+
+class VsBaseEdge(ABC, Generic[S]):
+
+    def __init__(
+        self,
+        source: str,
+        trace_publisher: Optional[VsBaseTracePublisher] = None,
+    ):
+        if not source:
+            raise ValueError("source cannot be empty or None")
+
+        self.source = source
+        self.trace_publisher = trace_publisher
+        self.logger = VsLogManager.get_instance(self.__class__.__name__)
+
+    async def route(self, state: S):
+        """Async routing function registered with LangGraph for edges that make a
+        decision at runtime — conditional and dynamic fan-out. Direct and static
+        fan-out edges carry no decision and are wired as plain LangGraph edges,
+        so they never reach this method."""
+        raise NotImplementedError(
+            f"{type(self).__name__} does not implement route(). "
+            "Only conditional and dynamic fan-out edges are registered as routing functions."
+        )
+
+    async def trace(
+        self,
+        trace_id: str,
+        message: str,
+        data: Optional[Dict[str, Any]] = None,
+        level: VsTraceLevel = VsTraceLevel.INFO,
+    ) -> None:
+        if not self.trace_publisher:
+            return
+        try:
+            await self.trace_publisher.publish(trace_id=trace_id, message=message, level=level, data=data)
+        except Exception as e:
+            self.logger.warning(f"[{self.source}] trace publish failed: {e}")

vs_graph/edge/vs_direct_edge.py

@@ -0,0 +1,23 @@
+from typing import Optional, TypeVar
+
+from vs_common.publisher.vs_base_trace_publisher import VsBaseTracePublisher
+from vs_graph.edge.vs_base_edge import VsBaseEdge
+from vs_graph.schema.vs_base_graph_state import VsBaseGraphState
+
+S = TypeVar("S", bound=VsBaseGraphState)
+
+
+class VsDirectEdge(VsBaseEdge[S]):
+
+    def __init__(
+        self,
+        source: str,
+        target: str,
+        trace_publisher: Optional[VsBaseTracePublisher] = None,
+    ):
+        super().__init__(source=source, trace_publisher=trace_publisher)
+
+        if not target:
+            raise ValueError(f"DirectEdge from '{source}' requires a target.")
+
+        self.target = target

vs_graph/edge/vs_dynamic_fan_out_edge.py

@@ -0,0 +1,49 @@
+import time
+from abc import abstractmethod
+from typing import Any, Dict, List, Optional, TypeVar
+
+from vs_common.publisher.vs_base_trace_publisher import VsBaseTracePublisher
+from vs_graph.edge.vs_base_edge import VsBaseEdge
+from vs_graph.schema.vs_base_graph_state import VsBaseGraphState
+
+S = TypeVar("S", bound=VsBaseGraphState)
+
+
+class VsDynamicFanOutEdge(VsBaseEdge[S]):
+
+    def __init__(
+        self,
+        source: str,
+        config: Optional[Dict[str, Any]] = None,
+        trace_publisher: Optional[VsBaseTracePublisher] = None,
+    ):
+        super().__init__(source=source, trace_publisher=trace_publisher)
+        self.config = config or {}
+
+    async def route(self, state: S) -> List:
+        trace_id = state.get("trace_id", "")
+        self.logger.info(f"[{self.source}] dynamic fan-out started")
+        start_time = time.time()
+        try:
+            sends = self.get_sends(state)
+            targets = [getattr(s, "node", None) for s in sends]
+            duration = round(time.time() - start_time, 3)
+            self.logger.info(
+                f"[{self.source}] fanning out to {len(sends)} branch(es) {targets} in {duration}s"
+            )
+            await self.trace(
+                trace_id,
+                f"edge '{self.source}' fanning out to {len(sends)} branch(es)",
+                data={"branch_count": len(sends), "targets": targets, "duration_seconds": duration},
+            )
+            return sends
+        except Exception as e:
+            duration = round(time.time() - start_time, 3)
+            self.logger.error(
+                f"[{self.source}] dynamic fan-out failed in {duration}s: {e}", exc_info=True
+            )
+            raise
+
+    @abstractmethod
+    def get_sends(self, state: S) -> List:
+        pass

vs_graph/edge/vs_fan_out_edge.py

@@ -0,0 +1,30 @@
+from typing import List, Optional, TypeVar
+
+from vs_common.publisher.vs_base_trace_publisher import VsBaseTracePublisher
+from vs_graph.edge.vs_base_edge import VsBaseEdge
+from vs_graph.schema.vs_base_graph_state import VsBaseGraphState
+
+S = TypeVar("S", bound=VsBaseGraphState)
+
+
+class VsFanOutEdge(VsBaseEdge[S]):
+
+    def __init__(
+        self,
+        source: str,
+        targets: List[str],
+        trace_publisher: Optional[VsBaseTracePublisher] = None,
+    ):
+        super().__init__(source=source, trace_publisher=trace_publisher)
+
+        if not targets or len(targets) < 2:
+            raise ValueError(f"FanOutEdge from '{source}' requires at least 2 targets.")
+
+        duplicates = [t for t in targets if targets.count(t) > 1]
+        if duplicates:
+            raise ValueError(f"FanOutEdge from '{source}' has duplicate targets: {sorted(set(duplicates))}.")
+
+        self.targets = targets
+
+    def get_targets(self) -> List[str]:
+        return self.targets

vs_graph/exception/vs_graph_exceptions.py

@@ -0,0 +1,17 @@
+from vs_common.exception.vs_base_exception import VsBaseException
+
+
+class VsUserException(VsBaseException):
+    pass
+
+
+class VsGraphBuildException(VsBaseException):
+    pass
+
+
+class VsNodeExecutionException(VsBaseException):
+    pass
+
+
+class VsEdgeEvaluationException(VsBaseException):
+    pass

vs_graph/node/vs_base_node.py

@@ -0,0 +1,82 @@
+import time
+import traceback
+from abc import ABC, abstractmethod
+from typing import Any, Dict, Generic, Optional, TypeVar
+
+from vs_common.log.vs_log_manager import VsLogManager
+from vs_common.publisher.vs_base_trace_publisher import VsBaseTracePublisher
+from vs_common.schema.vs_trace_level import VsTraceLevel
+from vs_graph.exception.vs_graph_exceptions import VsUserException, VsNodeExecutionException
+from vs_graph.schema.vs_base_graph_state import VsBaseGraphState
+
+S = TypeVar("S", bound=VsBaseGraphState)
+
+
+class VsBaseNode(ABC, Generic[S]):
+
+    def __init__(
+        self,
+        node_id: str,
+        config: Optional[Dict[str, Any]] = None,
+        trace_publisher: Optional[VsBaseTracePublisher] = None,
+    ):
+        if not node_id:
+            raise ValueError("node_id cannot be empty or None")
+
+        self.node_id = node_id
+        self.config = config or {}
+        self.trace_publisher = trace_publisher
+        self.logger = VsLogManager.get_instance(self.__class__.__name__)
+
+    async def invoke(self, state: S) -> dict:
+        trace_id = state.get("trace_id", "")
+        self.logger.info(f"[{self.node_id}] execution started")
+        await self.trace(trace_id, f"node '{self.node_id}' started")
+        start_time = time.time()
+        try:
+            result = await self.execute(state)
+            duration = round(time.time() - start_time, 3)
+            self.logger.info(f"[{self.node_id}] execution completed in {duration}s")
+            await self.trace(trace_id, f"node '{self.node_id}' completed", data={"duration_seconds": duration})
+            return result
+        except VsUserException as e:
+            duration = round(time.time() - start_time, 3)
+            self.logger.warning(f"[{self.node_id}] user-facing error in {duration}s: {e.message}")
+            await self.trace(trace_id, f"node '{self.node_id}' user error", data={"error": e.message}, level=VsTraceLevel.WARN)
+            state["error"] = {
+                "message": e.message,
+                "cause": None,
+                "stack_trace": None,
+            }
+            return state
+        except VsNodeExecutionException as e:
+            duration = round(time.time() - start_time, 3)
+            self.logger.error(f"[{self.node_id}] node execution error in {duration}s: {e.message}", exc_info=True)
+            await self.trace(trace_id, f"node '{self.node_id}' failed", data={"error": e.message}, level=VsTraceLevel.ERROR)
+            raise
+        except Exception as e:
+            duration = round(time.time() - start_time, 3)
+            self.logger.error(f"[{self.node_id}] unexpected error in {duration}s: {e}", exc_info=True)
+            await self.trace(trace_id, f"node '{self.node_id}' failed", data={"error": str(e)}, level=VsTraceLevel.ERROR)
+            raise
+
+    async def trace(
+        self,
+        trace_id: str,
+        message: str,
+        data: Optional[Dict[str, Any]] = None,
+        level: VsTraceLevel = VsTraceLevel.INFO,
+    ) -> None:
+        if not self.trace_publisher:
+            return
+        try:
+            await self.trace_publisher.publish(trace_id=trace_id,
+                                               message=message,
+                                               level=level,
+                                               data=data)
+        except Exception as e:
+            self.logger.warning(f"[{self.node_id}] trace publish failed: {e}")
+
+    @abstractmethod
+    async def execute(self, state: S) -> dict:
+        pass

vs_graph/registry/vs_edge_registry.py

@@ -0,0 +1,65 @@
+import threading
+from typing import Any, Dict, List, Type
+
+from vs_common.log.vs_log_manager import VsLogManager
+
+
+class VsEdgeRegistry:
+
+    def __new__(cls, *args, **kwargs):
+        raise TypeError("VsEdgeRegistry cannot be instantiated.")
+
+    _edges: Dict[str, Type] = {}
+    _lock = threading.Lock()
+    _logger = None
+
+    @classmethod
+    def _get_logger(cls):
+        if cls._logger is None:
+            cls._logger = VsLogManager.get_instance("VsEdgeRegistry")
+        return cls._logger
+
+    @classmethod
+    def register(cls, name: str, edge_class: Type) -> None:
+        from vs_graph.edge.vs_base_edge import VsBaseEdge
+        with cls._lock:
+            if name in cls._edges:
+                existing = cls._edges[name]
+                raise ValueError(
+                    f"Edge '{name}' is already registered to {existing.__module__}.{existing.__name__}."
+                )
+            if not issubclass(edge_class, VsBaseEdge):
+                raise TypeError(f"Edge class {edge_class.__name__} must extend VsBaseEdge.")
+            cls._edges[name] = edge_class
+            cls._get_logger().info(f"Registered edge: '{name}' -> {edge_class.__module__}.{edge_class.__name__}")
+
+    @classmethod
+    def get(cls, name: str) -> Type:
+        if name not in cls._edges:
+            raise ValueError(
+                f"Edge '{name}' not found in registry. Available: {sorted(cls._edges.keys())}"
+            )
+        return cls._edges[name]
+
+    @classmethod
+    def create(cls, name: str, source: str, **kwargs: Any) -> Any:
+        edge_class = cls.get(name)
+        try:
+            instance = edge_class(source=source, **kwargs)
+            cls._get_logger().info(f"Created edge: '{name}' source='{source}'")
+            return instance
+        except Exception as e:
+            raise RuntimeError(f"Failed to instantiate edge '{name}' (source='{source}'): {e}") from e
+
+    @classmethod
+    def is_registered(cls, name: str) -> bool:
+        return name in cls._edges
+
+    @classmethod
+    def list_edges(cls) -> List[str]:
+        return sorted(cls._edges.keys())
+
+    @classmethod
+    def clear(cls) -> None:
+        with cls._lock:
+            cls._edges.clear()