diff --git a/src/dynavec/__init__.py b/src/dynavec/__init__.py index 8aee485..bfd5499 100644 --- a/src/dynavec/__init__.py +++ b/src/dynavec/__init__.py @@ -63,6 +63,25 @@ UpsertResult, ) from .namespace import NamespaceView +from .nodes import ( + ApprovalNode, + CodeNode, + Edge, + LLMNode, + LoopNode, + Node, + NodePort, + NodeRegistry, + NodeSchema, + PortType, + RetrieverNode, + RouterNode, + SimpleGraph, + SubGraphNode, + ToolNode, + WebhookNode, + node_registry, +) from .quantization import ( OPQRotation, OptimizedProductQuantizer, @@ -153,4 +172,22 @@ "ItemTooLargeError", "ConflictError", "MissingDependencyError", + # Nodes & Graph + "Node", + "NodePort", + "NodeSchema", + "PortType", + "NodeRegistry", + "node_registry", + "Edge", + "SimpleGraph", + "RetrieverNode", + "LLMNode", + "ToolNode", + "RouterNode", + "CodeNode", + "ApprovalNode", + "LoopNode", + "WebhookNode", + "SubGraphNode", ] diff --git a/src/dynavec/nodes/__init__.py b/src/dynavec/nodes/__init__.py new file mode 100644 index 0000000..57feb80 --- /dev/null +++ b/src/dynavec/nodes/__init__.py @@ -0,0 +1,38 @@ +"""Typed node library and execution graph for dynaflow.""" + +from __future__ import annotations + +from .base import Node, NodePort, NodeSchema, PortType +from .builtin import ( + ApprovalNode, + CodeNode, + LLMNode, + LoopNode, + RetrieverNode, + RouterNode, + SubGraphNode, + ToolNode, + WebhookNode, +) +from .graph import Edge, SimpleGraph +from .registry import NodeRegistry, node_registry + +__all__ = [ + "Node", + "NodePort", + "NodeSchema", + "PortType", + "NodeRegistry", + "node_registry", + "RetrieverNode", + "LLMNode", + "ToolNode", + "RouterNode", + "CodeNode", + "ApprovalNode", + "LoopNode", + "WebhookNode", + "SubGraphNode", + "Edge", + "SimpleGraph", +] diff --git a/src/dynavec/nodes/base.py b/src/dynavec/nodes/base.py new file mode 100644 index 0000000..77661ec --- /dev/null +++ b/src/dynavec/nodes/base.py @@ -0,0 +1,100 @@ +"""Base abstractions and schemas for the dynaflow node system.""" + +from __future__ import annotations + +import uuid +from abc import ABC, abstractmethod +from dataclasses import dataclass, field +from typing import Any, Literal + +PortType = Literal["string", "number", "boolean", "array", "object", "any"] + + +@dataclass +class NodePort: + """A typed input or output connection port on a node.""" + + name: str + port_type: PortType = "any" + description: str = "" + required: bool = True + default: Any = None + + def to_dict(self) -> dict[str, Any]: + """Convert port definition to a dictionary.""" + return { + "name": self.name, + "port_type": self.port_type, + "description": self.description, + "required": self.required, + "default": self.default, + } + + +@dataclass +class NodeSchema: + """Metadata describing a node type's ports, config, and canvas properties.""" + + node_type: str + label: str + description: str + category: str = "general" + inputs: list[NodePort] = field(default_factory=list) + outputs: list[NodePort] = field(default_factory=list) + config_schema: dict[str, Any] = field(default_factory=dict) + + def to_dict(self) -> dict[str, Any]: + """Serialize schema to dictionary.""" + return { + "node_type": self.node_type, + "label": self.label, + "description": self.description, + "category": self.category, + "inputs": [p.to_dict() for p in self.inputs], + "outputs": [p.to_dict() for p in self.outputs], + "config_schema": self.config_schema, + } + + def to_reactflow_dict(self) -> dict[str, Any]: + """Export ReactFlow-compatible component metadata for visual canvas builders.""" + return { + "type": self.node_type, + "label": self.label, + "category": self.category, + "description": self.description, + "inputPorts": [p.name for p in self.inputs], + "outputPorts": [p.name for p in self.outputs], + "config": self.config_schema, + } + + +class Node(ABC): + """Abstract base class for all execution nodes in dynaflow.""" + + node_type: str + schema: NodeSchema + + def __init__( + self, + node_id: str | None = None, + config: dict[str, Any] | None = None, + ) -> None: + self.id = node_id or f"{self.node_type}_{uuid.uuid4().hex[:8]}" + self.config = dict(config or {}) + + @abstractmethod + def execute(self, inputs: dict[str, Any]) -> dict[str, Any]: + """Execute the node logic given an inputs dictionary, returning outputs.""" + + def to_dict(self) -> dict[str, Any]: + """Serialize node to a ReactFlow-compatible node object.""" + return { + "id": self.id, + "type": self.node_type, + "data": { + "label": self.schema.label, + "config": self.config, + "category": self.schema.category, + }, + "position": {"x": 0, "y": 0}, + } diff --git a/src/dynavec/nodes/builtin.py b/src/dynavec/nodes/builtin.py new file mode 100644 index 0000000..283afc6 --- /dev/null +++ b/src/dynavec/nodes/builtin.py @@ -0,0 +1,628 @@ +"""Built-in standard node library for dynaflow agent orchestrations.""" + +from __future__ import annotations + +import json +from collections.abc import Callable +from typing import Any + +from ..chat.base import ChatModel, Message +from .base import Node, NodePort, NodeSchema +from .registry import node_registry + +# --- 1. RetrieverNode --- + + +@node_registry.register +class RetrieverNode(Node): + """Executes dense/hybrid vector retrieval against dynavec vector store.""" + + node_type = "retriever" + schema = NodeSchema( + node_type="retriever", + label="Vector Retriever", + description="Searches dynavec vector index and formats context for LLMs.", + category="retrieval", + inputs=[ + NodePort( + name="query", + port_type="string", + description="Natural language query string.", + required=True, + ), + NodePort( + name="filter", + port_type="object", + description="Metadata filter dictionary.", + required=False, + default=None, + ), + NodePort( + name="top_k", + port_type="number", + description="Maximum number of documents to return.", + required=False, + default=5, + ), + ], + outputs=[ + NodePort( + name="documents", + port_type="array", + description="Retrieved document records.", + ), + NodePort( + name="context", + port_type="string", + description="Concatenated document text context block.", + ), + ], + config_schema={ + "top_k": {"type": "integer", "default": 5}, + "namespace": {"type": "string", "default": "default"}, + }, + ) + + def __init__( + self, + retriever_fn: Callable[..., Any] | None = None, + node_id: str | None = None, + config: dict[str, Any] | None = None, + ) -> None: + super().__init__(node_id=node_id, config=config) + self.retriever_fn = retriever_fn + + def execute(self, inputs: dict[str, Any]) -> dict[str, Any]: + query = str(inputs.get("query", "")) + top_k = inputs.get("top_k", self.config.get("top_k", 5)) + filter_dict = inputs.get("filter", self.config.get("filter")) + + documents: list[dict[str, Any]] = [] + if self.retriever_fn is not None: + raw_docs = self.retriever_fn(query, top_k=top_k, filter=filter_dict) + for d in raw_docs: + if hasattr(d, "text") and hasattr(d, "id"): + documents.append( + { + "id": d.id, + "text": getattr(d, "text", ""), + "score": getattr(d, "score", 0.0), + "metadata": getattr(d, "metadata", {}), + } + ) + elif isinstance(d, dict): + documents.append(d) + else: + documents.append({"text": str(d)}) + + context_parts = [d.get("text", "") for d in documents if d.get("text")] + context_str = "\n\n".join(context_parts) + + return { + "documents": documents, + "context": context_str, + } + + +# --- 2. LLMNode --- + + +@node_registry.register +class LLMNode(Node): + """Executes a language model prompt invocation.""" + + node_type = "llm" + schema = NodeSchema( + node_type="llm", + label="LLM Generation", + description="Invokes a ChatModel with prompt and contextual grounding.", + category="model", + inputs=[ + NodePort( + name="prompt", + port_type="string", + description="The prompt or user question.", + required=True, + ), + NodePort( + name="context", + port_type="string", + description="Optional retrieved reference context.", + required=False, + default="", + ), + NodePort( + name="system_prompt", + port_type="string", + description="System instruction prompt.", + required=False, + default=None, + ), + ], + outputs=[ + NodePort( + name="response", + port_type="string", + description="Generated text output.", + ), + NodePort( + name="message", + port_type="object", + description="Full ChatModel output message object.", + ), + ], + config_schema={ + "model_name": {"type": "string"}, + "system_prompt": {"type": "string"}, + "temperature": {"type": "number", "default": 0.0}, + }, + ) + + def __init__( + self, + model: ChatModel | Callable[[list[Message]], Any], + node_id: str | None = None, + config: dict[str, Any] | None = None, + ) -> None: + super().__init__(node_id=node_id, config=config) + self.model = model + + def execute(self, inputs: dict[str, Any]) -> dict[str, Any]: + prompt = str(inputs.get("prompt", "")) + context = str(inputs.get("context", "")) + sys_prompt = inputs.get("system_prompt", self.config.get("system_prompt")) + + full_prompt = prompt + if context: + full_prompt = f"Context:\n{context}\n\nUser Question:\n{prompt}" + + messages: list[Message] = [] + if sys_prompt: + messages.append(Message(role="system", content=sys_prompt)) + messages.append(Message(role="user", content=full_prompt)) + + if hasattr(self.model, "invoke"): + chat_res = self.model.invoke(messages) + out_text = chat_res.message.content or "" + out_msg = chat_res.message + elif callable(self.model): + res = self.model(messages) + out_text = str(res) + out_msg = Message(role="assistant", content=out_text) + else: + raise TypeError("Model must be a ChatModel or callable.") + + return { + "response": out_text, + "message": out_msg, + } + + +# --- 3. ToolNode --- + + +@node_registry.register +class ToolNode(Node): + """Invokes a callable tool or registered AgentTool with input arguments.""" + + node_type = "tool" + schema = NodeSchema( + node_type="tool", + label="Tool Execution", + description="Invokes external APIs, code functions, or MCP connectors.", + category="tools", + inputs=[ + NodePort( + name="arguments", + port_type="object", + description="JSON parameters dictionary for tool invocation.", + required=False, + default={}, + ), + NodePort( + name="input", + port_type="any", + description="Alternative single parameter input.", + required=False, + default=None, + ), + ], + outputs=[ + NodePort( + name="result", + port_type="any", + description="Output returned from tool execution.", + ), + NodePort( + name="status", + port_type="string", + description="Execution status ('success' or 'error').", + ), + ], + config_schema={ + "tool_name": {"type": "string"}, + }, + ) + + def __init__( + self, + tool: Callable[..., Any], + node_id: str | None = None, + config: dict[str, Any] | None = None, + ) -> None: + super().__init__(node_id=node_id, config=config) + self.tool = tool + + def execute(self, inputs: dict[str, Any]) -> dict[str, Any]: + args = inputs.get("arguments") + if args is None and "input" in inputs: + args = inputs["input"] + + try: + if hasattr(self.tool, "execute"): + res = self.tool.execute(args) + elif isinstance(args, dict): + res = self.tool(**args) + elif isinstance(args, str): + try: + parsed = json.loads(args) + if isinstance(parsed, dict): + res = self.tool(**parsed) + else: + res = self.tool(args) + except Exception: + res = self.tool(args) + elif args is not None: + res = self.tool(args) + else: + res = self.tool() + return {"result": res, "status": "success"} + except Exception as exc: + return {"result": str(exc), "status": "error"} + + +# --- 4. RouterNode --- + + +@node_registry.register +class RouterNode(Node): + """Routes execution along one of multiple named branches based on condition.""" + + node_type = "router" + schema = NodeSchema( + node_type="router", + label="Router / Branch", + description="Selects an execution path based on predicate function or condition.", + category="control_flow", + inputs=[ + NodePort( + name="data", + port_type="any", + description="Input data payload passing through router.", + required=True, + ), + NodePort( + name="route", + port_type="string", + description="Explicit route key override.", + required=False, + default=None, + ), + ], + outputs=[ + NodePort( + name="selected_route", + port_type="string", + description="The route key chosen for downstream branching.", + ), + NodePort( + name="data", + port_type="any", + description="Forwarded payload unchanged.", + ), + ], + config_schema={ + "routes": { + "type": "array", + "items": {"type": "string"}, + "default": ["default"], + }, + }, + ) + + def __init__( + self, + route_fn: Callable[[Any], str] | None = None, + node_id: str | None = None, + config: dict[str, Any] | None = None, + ) -> None: + super().__init__(node_id=node_id, config=config) + self.route_fn = route_fn + + def execute(self, inputs: dict[str, Any]) -> dict[str, Any]: + data = inputs.get("data") + explicit_route = inputs.get("route") + + if explicit_route is not None: + chosen = str(explicit_route) + elif self.route_fn is not None: + chosen = str(self.route_fn(data)) + else: + chosen = self.config.get("default_route", "default") + + return { + "selected_route": chosen, + "data": data, + } + + +# --- 5. CodeNode --- + + +@node_registry.register +class CodeNode(Node): + """Executes arbitrary user-defined Python logic.""" + + node_type = "code" + schema = NodeSchema( + node_type="code", + label="Custom Python Code", + description="Runs Python transform function over inputs.", + category="custom", + inputs=[ + NodePort( + name="inputs", + port_type="object", + description="Dictionary of inputs passed to the code function.", + required=True, + ), + ], + outputs=[ + NodePort( + name="outputs", + port_type="object", + description="Dictionary returned by the code function.", + ), + ], + ) + + def __init__( + self, + code_fn: Callable[[dict[str, Any]], dict[str, Any]], + node_id: str | None = None, + config: dict[str, Any] | None = None, + ) -> None: + super().__init__(node_id=node_id, config=config) + self.code_fn = code_fn + + def execute(self, inputs: dict[str, Any]) -> dict[str, Any]: + in_data = inputs.get("inputs", inputs) + result = self.code_fn(in_data) + if isinstance(result, dict): + return result + return {"outputs": result} + + +# --- 6. ApprovalNode --- + + +@node_registry.register +class ApprovalNode(Node): + """Human-in-the-loop approval or interrupt checkpoint node.""" + + node_type = "approval" + schema = NodeSchema( + node_type="approval", + label="Human Approval", + description="Gates workflow progression on human review or decision.", + category="control_flow", + inputs=[ + NodePort( + name="request", + port_type="any", + description="Payload submitted for human review.", + required=True, + ), + NodePort( + name="approved", + port_type="boolean", + description="Review decision (true to continue, false to reject).", + required=False, + default=True, + ), + NodePort( + name="notes", + port_type="string", + description="Reviewer comments or revision notes.", + required=False, + default=None, + ), + ], + outputs=[ + NodePort( + name="approved", + port_type="boolean", + description="Final approval verdict.", + ), + NodePort( + name="data", + port_type="any", + description="Payload forwarded upon approval.", + ), + NodePort( + name="notes", + port_type="string", + description="Feedback notes from reviewer.", + ), + ], + config_schema={ + "role": {"type": "string", "default": "admin"}, + }, + ) + + def execute(self, inputs: dict[str, Any]) -> dict[str, Any]: + return { + "approved": bool(inputs.get("approved", True)), + "data": inputs.get("request"), + "notes": inputs.get("notes"), + } + + +# --- 7. LoopNode --- + + +@node_registry.register +class LoopNode(Node): + """Iterates or maps an array of elements through execution.""" + + node_type = "loop" + schema = NodeSchema( + node_type="loop", + label="Loop / Map", + description="Splits an array of items for parallel or sequential processing.", + category="control_flow", + inputs=[ + NodePort( + name="items", + port_type="array", + description="List of items to iterate over.", + required=True, + ), + ], + outputs=[ + NodePort( + name="items", + port_type="array", + description="Processed output elements.", + ), + NodePort( + name="count", + port_type="number", + description="Total number of elements processed.", + ), + ], + config_schema={ + "max_iterations": {"type": "integer", "default": 100}, + }, + ) + + def __init__( + self, + map_fn: Callable[[Any], Any] | None = None, + node_id: str | None = None, + config: dict[str, Any] | None = None, + ) -> None: + super().__init__(node_id=node_id, config=config) + self.map_fn = map_fn + + def execute(self, inputs: dict[str, Any]) -> dict[str, Any]: + raw_items = list(inputs.get("items", [])) + max_iters = self.config.get("max_iterations", 100) + items = raw_items[:max_iters] + + if self.map_fn is not None: + processed = [self.map_fn(item) for item in items] + else: + processed = items + + return { + "items": processed, + "count": len(processed), + } + + +# --- 8. WebhookNode --- + + +@node_registry.register +class WebhookNode(Node): + """Trigger / entrypoint node ingesting external webhook payloads.""" + + node_type = "webhook" + schema = NodeSchema( + node_type="webhook", + label="Webhook Trigger", + description="Entrypoint receiving external HTTP webhooks and event triggers.", + category="triggers", + inputs=[ + NodePort( + name="payload", + port_type="object", + description="Incoming JSON request body.", + required=True, + ), + NodePort( + name="headers", + port_type="object", + description="HTTP request headers.", + required=False, + default={}, + ), + ], + outputs=[ + NodePort( + name="payload", + port_type="object", + description="Sanitized event payload to start workflow.", + ), + ], + config_schema={ + "path": {"type": "string", "default": "/webhook"}, + }, + ) + + def execute(self, inputs: dict[str, Any]) -> dict[str, Any]: + return {"payload": inputs.get("payload", {})} + + +# --- 9. SubGraphNode --- + + +@node_registry.register +class SubGraphNode(Node): + """Encapsulates a nested graph or composite workflow as a single node.""" + + node_type = "subgraph" + schema = NodeSchema( + node_type="subgraph", + label="Sub-Graph", + description="Executes a modular reusable sub-workflow.", + category="composite", + inputs=[ + NodePort( + name="inputs", + port_type="object", + description="Input dictionary for sub-graph execution.", + required=True, + ), + ], + outputs=[ + NodePort( + name="outputs", + port_type="object", + description="Outputs resulting from sub-graph execution.", + ), + ], + config_schema={ + "subgraph_id": {"type": "string"}, + }, + ) + + def __init__( + self, + subgraph: Any = None, + node_id: str | None = None, + config: dict[str, Any] | None = None, + ) -> None: + super().__init__(node_id=node_id, config=config) + self.subgraph = subgraph + + def execute(self, inputs: dict[str, Any]) -> dict[str, Any]: + in_data = inputs.get("inputs", inputs) + if hasattr(self.subgraph, "run"): + res = self.subgraph.run(in_data) + elif callable(self.subgraph): + res = self.subgraph(in_data) + else: + res = in_data + return {"outputs": res} diff --git a/src/dynavec/nodes/graph.py b/src/dynavec/nodes/graph.py new file mode 100644 index 0000000..778c7e4 --- /dev/null +++ b/src/dynavec/nodes/graph.py @@ -0,0 +1,89 @@ +"""Pipeline execution graph for wiring and running connected nodes.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any + +from .base import Node + + +@dataclass +class Edge: + """Directed connection passing outputs from source node to target node.""" + + source_node: str + target_node: str + source_port: str = "output" + target_port: str = "input" + mapping: dict[str, str] = field(default_factory=dict) + + +class SimpleGraph: + """Directed execution pipeline to compose and run node graphs.""" + + def __init__(self) -> None: + self.nodes: dict[str, Node] = {} + self.edges: list[Edge] = [] + self._execution_order: list[str] = [] + + def add_node(self, node: Node) -> SimpleGraph: + """Add a node instance to the graph.""" + self.nodes[node.id] = node + if node.id not in self._execution_order: + self._execution_order.append(node.id) + return self + + def add_edge( + self, + source_id: str, + target_id: str, + source_port: str = "output", + target_port: str = "input", + mapping: dict[str, str] | None = None, + ) -> SimpleGraph: + """Connect source node output to target node input.""" + if source_id not in self.nodes: + raise ValueError(f"Source node {source_id!r} not in graph.") + if target_id not in self.nodes: + raise ValueError(f"Target node {target_id!r} not in graph.") + + edge = Edge( + source_node=source_id, + target_node=target_id, + source_port=source_port, + target_port=target_port, + mapping=dict(mapping or {}), + ) + self.edges.append(edge) + return self + + def run(self, initial_inputs: dict[str, Any]) -> dict[str, Any]: + """Execute the graph sequentially, threading state along connected edges.""" + node_outputs: dict[str, dict[str, Any]] = {} + current_state: dict[str, Any] = dict(initial_inputs) + + for node_id in self._execution_order: + node = self.nodes[node_id] + + # Collect inputs for this node from incoming edges or current state + node_inputs: dict[str, Any] = dict(current_state) + + incoming_edges = [e for e in self.edges if e.target_node == node_id] + for edge in incoming_edges: + prev_outputs = node_outputs.get(edge.source_node, {}) + if edge.mapping: + for src_k, tgt_k in edge.mapping.items(): + if src_k in prev_outputs: + node_inputs[tgt_k] = prev_outputs[src_k] + else: + if edge.source_port in prev_outputs: + node_inputs[edge.target_port] = prev_outputs[edge.source_port] + else: + node_inputs.update(prev_outputs) + + outputs = node.execute(node_inputs) + node_outputs[node_id] = outputs + current_state.update(outputs) + + return current_state diff --git a/src/dynavec/nodes/registry.py b/src/dynavec/nodes/registry.py new file mode 100644 index 0000000..a86cf95 --- /dev/null +++ b/src/dynavec/nodes/registry.py @@ -0,0 +1,45 @@ +"""Node registry for node discovery and canvas palette serialization.""" + +from __future__ import annotations + +from typing import Any + +from .base import Node, NodeSchema + + +class NodeRegistry: + """Registry maintaining available node types and their schemas.""" + + def __init__(self) -> None: + self._registry: dict[str, type[Node]] = {} + + def register(self, node_cls: type[Node]) -> type[Node]: + """Register a node class under its node_type.""" + if not hasattr(node_cls, "node_type"): + raise ValueError(f"Class {node_cls.__name__} must define a 'node_type' attribute.") + self._registry[node_cls.node_type] = node_cls + return node_cls + + def get(self, node_type: str) -> type[Node] | None: + """Lookup a registered node class by type.""" + return self._registry.get(node_type) + + def list_types(self) -> list[str]: + """Return all registered node type identifiers.""" + return sorted(self._registry.keys()) + + def get_schemas(self) -> list[NodeSchema]: + """Return schemas for all registered nodes.""" + schemas: list[NodeSchema] = [] + for n_type in self.list_types(): + cls = self._registry[n_type] + if hasattr(cls, "schema"): + schemas.append(cls.schema) + return schemas + + def export_canvas_palette(self) -> list[dict[str, Any]]: + """Export ReactFlow-compatible component palette for canvas composition.""" + return [s.to_reactflow_dict() for s in self.get_schemas()] + + +node_registry = NodeRegistry() diff --git a/tests/test_nodes.py b/tests/test_nodes.py new file mode 100644 index 0000000..4936dea --- /dev/null +++ b/tests/test_nodes.py @@ -0,0 +1,348 @@ +"""Unit tests for dynaflow built-in node library and graph execution pipeline.""" + +from __future__ import annotations + +from typing import Any + +import pytest + +from dynavec import ( + ApprovalNode, + CodeNode, + LLMNode, + LoopNode, + NodePort, + NodeSchema, + RetrieverNode, + RouterNode, + SimpleGraph, + SubGraphNode, + ToolNode, + WebhookNode, + node_registry, +) +from dynavec.chat.base import ChatModel, ChatResult, Message + + +class MockChatModel(ChatModel): + """Simple mock chat model for testing LLMNode.""" + + def invoke( + self, + messages: list[Message], + tools: list[Any] | None = None, + **kwargs: Any, + ) -> ChatResult: + last = messages[-1].content + return ChatResult( + message=Message( + role="assistant", + content=f"Processed: {last}", + ), + finish_reason="stop", + ) + + def stream( + self, + messages: list[Message], + tools: list[Any] | None = None, + **kwargs: Any, + ): + last = messages[-1].content + from dynavec.chat.base import ChatChunk + + yield ChatChunk(content=f"Processed: {last}") + + +class MockDoc: + def __init__(self, doc_id: str, text: str, score: float = 0.9): + self.id = doc_id + self.text = text + self.score = score + self.metadata = {"source": "unit-test"} + + +# ===================================================================== +# Registry & Schema Tests +# ===================================================================== + + +def test_registry_registration_and_palette(): + expected_types = { + "approval", + "code", + "llm", + "loop", + "retriever", + "router", + "subgraph", + "tool", + "webhook", + } + registered = set(node_registry.list_types()) + assert expected_types.issubset(registered) + + palette = node_registry.export_canvas_palette() + palette_types = {item["type"] for item in palette} + assert expected_types.issubset(palette_types) + + # Check palette item structure + for item in palette: + assert "type" in item + assert "label" in item + assert "category" in item + assert "inputPorts" in item + assert "outputPorts" in item + + +def test_port_and_schema_serialization(): + port = NodePort( + name="test_in", + port_type="string", + description="test input port", + required=True, + default="abc", + ) + port_dict = port.to_dict() + assert port_dict["name"] == "test_in" + assert port_dict["port_type"] == "string" + assert port_dict["required"] is True + assert port_dict["default"] == "abc" + + schema = NodeSchema( + node_type="custom_test", + label="Custom Test", + description="A test node schema", + inputs=[port], + outputs=[], + ) + s_dict = schema.to_dict() + assert s_dict["node_type"] == "custom_test" + assert len(s_dict["inputs"]) == 1 + + rf_dict = schema.to_reactflow_dict() + assert rf_dict["type"] == "custom_test" + assert rf_dict["inputPorts"] == ["test_in"] + + +# ===================================================================== +# Individual Built-in Node Tests +# ===================================================================== + + +def test_retriever_node(): + def dummy_retriever(query: str, top_k: int = 5, filter: dict[str, Any] | None = None): + return [ + MockDoc("d1", "Dynavec is a serverless vector database."), + MockDoc("d2", "DynamoDB provides scalable key-value storage."), + ] + + node = RetrieverNode(retriever_fn=dummy_retriever, config={"top_k": 2}) + assert node.node_type == "retriever" + + rf_node = node.to_dict() + assert rf_node["type"] == "retriever" + assert rf_node["data"]["category"] == "retrieval" + + outputs = node.execute({"query": "What is Dynavec?"}) + assert len(outputs["documents"]) == 2 + assert outputs["documents"][0]["id"] == "d1" + assert "Dynavec is a serverless" in outputs["context"] + assert "DynamoDB provides" in outputs["context"] + + +def test_llm_node_chat_model(): + model = MockChatModel() + node = LLMNode(model=model, config={"system_prompt": "You are a helpful assistant."}) + + outputs = node.execute( + { + "prompt": "Explain vector indexing", + "context": "Context snippet here.", + } + ) + assert "response" in outputs + assert ( + "Processed: Context:\nContext snippet here.\n\nUser Question:\nExplain vector indexing" + in outputs["response"] + ) + assert outputs["message"].role == "assistant" + + +def test_llm_node_callable(): + def simple_callable(messages: list[Message]) -> str: + return f"Echo: {messages[-1].content}" + + node = LLMNode(model=simple_callable) + outputs = node.execute({"prompt": "Hello world"}) + assert outputs["response"] == "Echo: Hello world" + + +def test_tool_node_success(): + def add(a: int, b: int) -> int: + return a + b + + node = ToolNode(tool=add) + outputs = node.execute({"arguments": {"a": 10, "b": 25}}) + assert outputs["status"] == "success" + assert outputs["result"] == 35 + + +def test_tool_node_json_string_and_error(): + def multiply(x: int, y: int) -> int: + return x * y + + node = ToolNode(tool=multiply) + outputs = node.execute({"arguments": '{"x": 6, "y": 7}'}) + assert outputs["status"] == "success" + assert outputs["result"] == 42 + + # Test error handling + fail_node = ToolNode(tool=lambda: 1 / 0) + err_out = fail_node.execute({}) + assert err_out["status"] == "error" + assert "division by zero" in err_out["result"] + + +def test_router_node(): + node = RouterNode(route_fn=lambda x: "admin_flow" if x.get("is_admin") else "user_flow") + + out1 = node.execute({"data": {"is_admin": True}}) + assert out1["selected_route"] == "admin_flow" + assert out1["data"] == {"is_admin": True} + + out2 = node.execute({"data": {"is_admin": False}}) + assert out2["selected_route"] == "user_flow" + + # Explicit route override + out3 = node.execute({"data": {}, "route": "manual_override"}) + assert out3["selected_route"] == "manual_override" + + +def test_code_node(): + def custom_transform(inputs: dict[str, Any]) -> dict[str, Any]: + return {"uppercased": inputs["text"].upper(), "len": len(inputs["text"])} + + node = CodeNode(code_fn=custom_transform) + res = node.execute({"inputs": {"text": "dynaflow"}}) + assert res["uppercased"] == "DYNAFLOW" + assert res["len"] == 8 + + # When code returns non-dict + scalar_node = CodeNode(code_fn=lambda inp: len(inp["text"])) + res2 = scalar_node.execute({"inputs": {"text": "test"}}) + assert res2["outputs"] == 4 + + +def test_approval_node(): + node = ApprovalNode() + out = node.execute( + { + "request": {"action": "delete_database", "target": "prod"}, + "approved": True, + "notes": "Approved by SecOps", + } + ) + assert out["approved"] is True + assert out["data"]["action"] == "delete_database" + assert out["notes"] == "Approved by SecOps" + + out_rejected = node.execute( + { + "request": {"action": "drop_table"}, + "approved": False, + } + ) + assert out_rejected["approved"] is False + + +def test_loop_node(): + node = LoopNode(map_fn=lambda x: x * 2, config={"max_iterations": 3}) + out = node.execute({"items": [1, 2, 3, 4, 5]}) + assert out["items"] == [2, 4, 6] # clipped to max 3 + assert out["count"] == 3 + + +def test_webhook_node(): + node = WebhookNode() + payload = {"event": "user.signup", "user_id": "usr_123"} + out = node.execute({"payload": payload, "headers": {"Authorization": "Bearer tok"}}) + assert out["payload"] == payload + + +def test_subgraph_node(): + def nested_workflow(inputs: dict[str, Any]) -> dict[str, Any]: + return {"processed_by_subgraph": True, "value": inputs.get("val", 0) + 10} + + node = SubGraphNode(subgraph=nested_workflow) + out = node.execute({"inputs": {"val": 5}}) + assert out["outputs"]["processed_by_subgraph"] is True + assert out["outputs"]["value"] == 15 + + +# ===================================================================== +# SimpleGraph Pipeline Integration: Retriever -> LLM -> Tool +# ===================================================================== + + +def test_simple_graph_pipeline_retriever_llm_tool(): + # 1. RetrieverNode: searches knowledge base + def mock_retriever(query: str, top_k: int = 5, filter: dict[str, Any] | None = None): + return [MockDoc("doc-1", "Dynavec provides DynamoDB vectors.")] + + retriever = RetrieverNode(retriever_fn=mock_retriever, node_id="retriever_1") + + # 2. LLMNode: formats tool call request based on context + def mock_llm(messages: list[Message]) -> str: + # LLM decides to call calculator tool with argument 42 + return '{"operation": "double", "value": 42}' + + llm = LLMNode(model=mock_llm, node_id="llm_1") + + # 3. ToolNode: executes calculation + def mock_tool(operation: str, value: int) -> int: + if operation == "double": + return value * 2 + return value + + tool = ToolNode(tool=mock_tool, node_id="tool_1") + + # Assemble graph + graph = SimpleGraph() + graph.add_node(retriever) + graph.add_node(llm) + graph.add_node(tool) + + # retriever.context -> llm.context + graph.add_edge( + source_id="retriever_1", + target_id="llm_1", + source_port="context", + target_port="context", + ) + + # llm.response -> tool.arguments + graph.add_edge( + source_id="llm_1", + target_id="tool_1", + source_port="response", + target_port="arguments", + ) + + result = graph.run({"query": "How many vectors?", "prompt": "Calculate vector capacity"}) + + assert "documents" in result + assert "DynamoDB vectors" in result["context"] + assert result["status"] == "success" + assert result["result"] == 84 + + +def test_simple_graph_validation_errors(): + graph = SimpleGraph() + with pytest.raises(ValueError, match="Source node 'nonexistent' not in graph"): + graph.add_edge("nonexistent", "target") + + node = CodeNode(code_fn=lambda x: x, node_id="valid_node") + graph.add_node(node) + with pytest.raises(ValueError, match="Target node 'missing_target' not in graph"): + graph.add_edge("valid_node", "missing_target")