vellum-ai 0.9.16__py3-none-any.whl → 0.9.16rc4__py3-none-any.whl

vellum_ee/workflows/display/workflows/vellum_workflow_display.py

@@ -0,0 +1,386 @@
+import logging
+from uuid import UUID
+from typing import Dict, List, Optional, Type, cast
+
+from vellum_ee.workflows.display.nodes.base_node_vellum_display import BaseNodeVellumDisplay
+from vellum_ee.workflows.display.nodes.types import PortDisplay
+from vellum_ee.workflows.display.nodes.vellum.utils import create_node_input
+from vellum_ee.workflows.display.utils.uuids import uuid4_from_hash
+from vellum_ee.workflows.display.utils.vellum import infer_vellum_variable_type
+from vellum_ee.workflows.display.vellum import (
+    EdgeVellumDisplay,
+    EdgeVellumDisplayOverrides,
+    EntrypointVellumDisplay,
+    EntrypointVellumDisplayOverrides,
+    NodeDisplayData,
+    WorkflowInputsVellumDisplay,
+    WorkflowInputsVellumDisplayOverrides,
+    WorkflowMetaVellumDisplay,
+    WorkflowMetaVellumDisplayOverrides,
+    WorkflowOutputVellumDisplay,
+    WorkflowOutputVellumDisplayOverrides,
+)
+from vellum_ee.workflows.display.workflows.base_workflow_display import BaseWorkflowDisplay
+from vellum.workflows.descriptors.base import BaseDescriptor
+from vellum.workflows.edges import Edge
+from vellum.workflows.nodes.bases import BaseNode
+from vellum.workflows.nodes.displayable.final_output_node import FinalOutputNode
+from vellum.workflows.nodes.utils import get_wrapped_node, has_wrapped_node
+from vellum.workflows.ports import Port
+from vellum.workflows.references import WorkflowInputReference
+from vellum.workflows.references.output import OutputReference
+from vellum.workflows.types.core import JsonArray, JsonObject
+from vellum.workflows.types.generics import WorkflowType
+
+logger = logging.getLogger(__name__)
+
+
+class VellumWorkflowDisplay(
+    BaseWorkflowDisplay[
+        WorkflowType,
+        WorkflowMetaVellumDisplay,
+        WorkflowMetaVellumDisplayOverrides,
+        WorkflowInputsVellumDisplay,
+        WorkflowInputsVellumDisplayOverrides,
+        BaseNodeVellumDisplay,
+        EntrypointVellumDisplay,
+        EntrypointVellumDisplayOverrides,
+        EdgeVellumDisplay,
+        EdgeVellumDisplayOverrides,
+        WorkflowOutputVellumDisplay,
+        WorkflowOutputVellumDisplayOverrides,
+    ]
+):
+    node_display_base_class = BaseNodeVellumDisplay
+
+    def serialize(self, raise_errors: bool = True) -> JsonObject:
+        input_variables: JsonArray = []
+        for workflow_input, workflow_input_display in self.display_context.workflow_input_displays.items():
+            input_variables.append(
+                {
+                    "id": str(workflow_input_display.id),
+                    "key": workflow_input.name,
+                    "type": infer_vellum_variable_type(workflow_input),
+                    # TODO: Add support for serializing default, required, and extensions
+                    #   https://app.shortcut.com/vellum/story/5429
+                    "default": None,
+                    "required": None,
+                    "extensions": None,
+                }
+            )
+
+        nodes: JsonArray = []
+        edges: JsonArray = []
+
+        # Add a single synthetic node for the workflow entrypoint
+        base_node_definition: JsonObject = {
+            "name": BaseNode.__name__,
+            "module": cast(JsonArray, BaseNode.__module__.split(".")),
+            "bases": [],
+        }
+        nodes.append(
+            {
+                "id": str(self.display_context.workflow_display.entrypoint_node_id),
+                "type": "ENTRYPOINT",
+                "inputs": [],
+                "data": {
+                    "label": "Entrypoint Node",
+                    "source_handle_id": str(self.display_context.workflow_display.entrypoint_node_source_handle_id),
+                },
+                "display_data": self.display_context.workflow_display.entrypoint_node_display.dict(),
+                "definition": base_node_definition,
+            },
+        )
+
+        # Add all the nodes in the workflow
+        for node, node_display in self.display_context.node_displays.items():
+            if getattr(node, "_is_wrapped_node") is True:
+                # Nodes that are wrapped or decorated by other nodes are not serialized here
+                # They are instead serialized by the wrapper node
+                continue
+
+            try:
+                serialized_node = node_display.serialize(self.display_context)
+            except NotImplementedError:
+                logger.warning("Unable to serialize node", extra={"node": node.__name__})
+                if raise_errors:
+                    raise
+                else:
+                    continue
+
+            nodes.append(serialized_node)
+
+        synthetic_output_edges: JsonArray = []
+        output_variables: JsonArray = []
+        final_output_nodes = [
+            node for node in self.display_context.node_displays.keys() if issubclass(node, FinalOutputNode)
+        ]
+        final_output_node_outputs = {node.Outputs.value for node in final_output_nodes}
+        unreferenced_final_output_node_outputs = final_output_node_outputs.copy()
+        final_output_node_definition: JsonObject = {
+            "name": FinalOutputNode.__name__,
+            "module": cast(JsonArray, FinalOutputNode.__module__.split(".")),
+            "bases": [base_node_definition],
+        }
+
+        # Add a synthetic Terminal Node and track the Workflow's output variables for each Workflow output
+        for workflow_output, workflow_output_display in self.display_context.workflow_output_displays.items():
+            final_output_node_id = workflow_output_display.node_id
+            inferred_type = infer_vellum_variable_type(workflow_output)
+
+            # Remove the terminal node output from the unreferenced set
+            unreferenced_final_output_node_outputs.discard(cast(OutputReference, workflow_output.instance))
+
+            if workflow_output.instance not in final_output_node_outputs:
+                # Create a synthetic terminal node only if there is no terminal node for this output
+                node_input = create_node_input(
+                    final_output_node_id,
+                    "node_input",
+                    # This is currently the wrapper node's output, but we want the wrapped node
+                    workflow_output.instance,
+                    self.display_context,
+                    workflow_output_display.node_input_id,
+                )
+
+                source_node_display: Optional[BaseNodeVellumDisplay]
+                first_rule = node_input.value.rules[0]
+                if first_rule.type == "NODE_OUTPUT":
+                    source_node_id = UUID(first_rule.data.node_id)
+                    try:
+                        source_node_display = [
+                            node_display
+                            for node_display in self.display_context.node_displays.values()
+                            if node_display.node_id == source_node_id
+                        ][0]
+                    except IndexError:
+                        source_node_display = None
+
+                nodes.append(
+                    {
+                        "id": str(final_output_node_id),
+                        "type": "TERMINAL",
+                        "data": {
+                            "label": workflow_output_display.label,
+                            "name": workflow_output_display.name,
+                            "target_handle_id": str(workflow_output_display.target_handle_id),
+                            "output_id": str(workflow_output_display.id),
+                            "output_type": inferred_type,
+                            "node_input_id": str(node_input.id),
+                        },
+                        "inputs": [node_input.dict()],
+                        "display_data": workflow_output_display.display_data.dict(),
+                        "definition": final_output_node_definition,
+                    }
+                )
+
+                if source_node_display:
+                    synthetic_output_edges.append(
+                        {
+                            "id": str(workflow_output_display.edge_id),
+                            "source_node_id": str(source_node_display.node_id),
+                            "source_handle_id": str(
+                                source_node_display.get_source_handle_id(
+                                    port_displays=self.display_context.port_displays
+                                )
+                            ),
+                            "target_node_id": str(workflow_output_display.node_id),
+                            "target_handle_id": str(workflow_output_display.target_handle_id),
+                            "type": "DEFAULT",
+                        }
+                    )
+
+            output_variables.append(
+                {
+                    "id": str(workflow_output_display.id),
+                    "key": workflow_output_display.name,
+                    "type": inferred_type,
+                }
+            )
+
+        # If there are terminal nodes with no workflow output reference,
+        # raise a serialization error
+        if len(unreferenced_final_output_node_outputs) > 0:
+            raise ValueError("Unable to serialize terminal nodes that are not referenced by workflow outputs.")
+
+        # Add an edge for each edge in the workflow
+        all_edge_displays: List[EdgeVellumDisplay] = [
+            # Create a synthetic edge from the synthetic entrypoint node to each actual entrypoint
+            *[
+                entrypoint_display.edge_display
+                for entrypoint_display in self.display_context.entrypoint_displays.values()
+            ],
+            # Include the concrete edges in the workflow
+            *self.display_context.edge_displays.values(),
+        ]
+
+        for edge_display in all_edge_displays:
+            edges.append(
+                {
+                    "id": str(edge_display.id),
+                    "source_node_id": str(edge_display.source_node_id),
+                    "source_handle_id": str(edge_display.source_handle_id),
+                    "target_node_id": str(edge_display.target_node_id),
+                    "target_handle_id": str(edge_display.target_handle_id),
+                    "type": edge_display.type,
+                }
+            )
+
+        edges.extend(synthetic_output_edges)
+
+        return {
+            "workflow_raw_data": {
+                "nodes": nodes,
+                "edges": edges,
+                "display_data": self.display_context.workflow_display.display_data.dict(),
+                "definition": {
+                    "name": self._workflow.__name__,
+                    "module": cast(JsonArray, self._workflow.__module__.split(".")),
+                },
+            },
+            "input_variables": input_variables,
+            "output_variables": output_variables,
+        }
+
+    def _generate_workflow_meta_display(self) -> WorkflowMetaVellumDisplay:
+        overrides = self.workflow_display
+        if overrides:
+            return WorkflowMetaVellumDisplay(
+                entrypoint_node_id=overrides.entrypoint_node_id,
+                entrypoint_node_source_handle_id=overrides.entrypoint_node_source_handle_id,
+                entrypoint_node_display=overrides.entrypoint_node_display,
+                display_data=overrides.display_data,
+            )
+
+        entrypoint_node_id = uuid4_from_hash(f"{self.workflow_id}|entrypoint_node_id")
+        entrypoint_node_source_handle_id = uuid4_from_hash(f"{self.workflow_id}|entrypoint_node_source_handle_id")
+
+        return WorkflowMetaVellumDisplay(
+            entrypoint_node_id=entrypoint_node_id,
+            entrypoint_node_source_handle_id=entrypoint_node_source_handle_id,
+            entrypoint_node_display=NodeDisplayData(),
+        )
+
+    def _generate_workflow_input_display(
+        self, workflow_input: WorkflowInputReference, overrides: Optional[WorkflowInputsVellumDisplayOverrides] = None
+    ) -> WorkflowInputsVellumDisplay:
+        workflow_input_id: UUID
+        if overrides:
+            workflow_input_id = overrides.id
+        else:
+            workflow_input_id = uuid4_from_hash(f"{self.workflow_id}|inputs|id|{workflow_input.name}")
+
+        return WorkflowInputsVellumDisplay(id=workflow_input_id)
+
+    def _generate_entrypoint_display(
+        self,
+        entrypoint: Type[BaseNode],
+        workflow_display: WorkflowMetaVellumDisplay,
+        node_displays: Dict[Type[BaseNode], BaseNodeVellumDisplay],
+        overrides: Optional[EntrypointVellumDisplayOverrides] = None,
+    ) -> EntrypointVellumDisplay:
+        entrypoint_node_id = workflow_display.entrypoint_node_id
+        source_handle_id = workflow_display.entrypoint_node_source_handle_id
+
+        edge_display_overrides = overrides.edge_display if overrides else None
+        entrypoint_id = (
+            edge_display_overrides.id
+            if edge_display_overrides
+            else uuid4_from_hash(f"{self.workflow_id}|id|{entrypoint_node_id}")
+        )
+
+        if has_wrapped_node(entrypoint):
+            wrapped_node = get_wrapped_node(entrypoint)
+            if wrapped_node._is_wrapped_node:
+                entrypoint = wrapped_node
+
+        target_node_id = node_displays[entrypoint].node_id
+        target_handle_id = node_displays[entrypoint].get_target_handle_id()
+
+        edge_display = self._generate_edge_display_from_source(
+            entrypoint_node_id, source_handle_id, target_node_id, target_handle_id, overrides=edge_display_overrides
+        )
+
+        return EntrypointVellumDisplay(id=entrypoint_id, edge_display=edge_display)
+
+    def _generate_workflow_output_display(
+        self,
+        output: BaseDescriptor,
+        overrides: Optional[WorkflowOutputVellumDisplayOverrides] = None,
+    ) -> WorkflowOutputVellumDisplay:
+        if overrides:
+            return WorkflowOutputVellumDisplay(
+                id=overrides.id,
+                name=overrides.name,
+                label=overrides.label,
+                node_id=overrides.node_id,
+                node_input_id=overrides.node_input_id,
+                target_handle_id=overrides.target_handle_id,
+                edge_id=overrides.edge_id,
+                display_data=overrides.display_data,
+            )
+
+        output_id = uuid4_from_hash(f"{self.workflow_id}|id|{output.name}")
+        edge_id = uuid4_from_hash(f"{self.workflow_id}|edge_id|{output.name}")
+        node_id = uuid4_from_hash(f"{self.workflow_id}|node_id|{output.name}")
+        node_input_id = uuid4_from_hash(f"{self.workflow_id}|node_input_id|{output.name}")
+        target_handle_id = uuid4_from_hash(f"{self.workflow_id}|target_handle_id|{output.name}")
+
+        return WorkflowOutputVellumDisplay(
+            id=output_id,
+            node_id=node_id,
+            node_input_id=node_input_id,
+            name=output.name,
+            label="Final Output",
+            target_handle_id=target_handle_id,
+            edge_id=edge_id,
+            display_data=NodeDisplayData(),
+        )
+
+    def _generate_edge_display(
+        self,
+        edge: Edge,
+        node_displays: Dict[Type[BaseNode], BaseNodeVellumDisplay],
+        port_displays: Dict[Port, PortDisplay],
+        overrides: Optional[EdgeVellumDisplayOverrides] = None,
+    ) -> EdgeVellumDisplay:
+        source_node = edge.from_port.node_class
+        target_node = edge.to_node
+
+        if has_wrapped_node(source_node):
+            source_node = get_wrapped_node(source_node)
+
+        if has_wrapped_node(target_node):
+            target_node = get_wrapped_node(target_node)
+
+        source_node_id = node_displays[source_node].node_id
+        source_handle_id = port_displays[edge.from_port].id
+
+        target_node_display = node_displays[target_node]
+        target_node_id = target_node_display.node_id
+        target_handle_id = target_node_display.get_target_handle_id()
+
+        return self._generate_edge_display_from_source(
+            source_node_id, source_handle_id, target_node_id, target_handle_id, overrides
+        )
+
+    def _generate_edge_display_from_source(
+        self,
+        source_node_id: UUID,
+        source_handle_id: UUID,
+        target_node_id: UUID,
+        target_handle_id: UUID,
+        overrides: Optional[EdgeVellumDisplayOverrides] = None,
+    ) -> EdgeVellumDisplay:
+        edge_id: UUID
+        if overrides:
+            edge_id = overrides.id
+        else:
+            edge_id = uuid4_from_hash(f"{self.workflow_id}|id|{source_node_id}|{target_node_id}")
+
+        return EdgeVellumDisplay(
+            id=edge_id,
+            source_node_id=source_node_id,
+            target_node_id=target_node_id,
+            source_handle_id=source_handle_id,
+            target_handle_id=target_handle_id,
+        )

vellum/lib/__init__.py

@@ -1,5 +0,0 @@
-from .test_suites import VellumTestSuite
-
-__all__ = [
-    "VellumTestSuite",
-]

vellum_ai-0.9.16.dist-info/METADATA

@@ -1,148 +0,0 @@
-Metadata-Version: 2.1
-Name: vellum-ai
-Version: 0.9.16
-Summary: 
-License: MIT
-Requires-Python: >=3.9,<4.0
-Classifier: Intended Audience :: Developers
-Classifier: License :: OSI Approved :: MIT License
-Classifier: Operating System :: MacOS
-Classifier: Operating System :: Microsoft :: Windows
-Classifier: Operating System :: OS Independent
-Classifier: Operating System :: POSIX
-Classifier: Operating System :: POSIX :: Linux
-Classifier: Programming Language :: Python
-Classifier: Programming Language :: Python :: 3
-Classifier: Programming Language :: Python :: 3.9
-Classifier: Programming Language :: Python :: 3.10
-Classifier: Programming Language :: Python :: 3.11
-Classifier: Programming Language :: Python :: 3.12
-Classifier: Programming Language :: Python :: 3.8
-Classifier: Topic :: Software Development :: Libraries :: Python Modules
-Classifier: Typing :: Typed
-Requires-Dist: Jinja2 (==3.1.2)
-Requires-Dist: cdktf (>=0.20.5,<0.21.0)
-Requires-Dist: click (==8.1.7)
-Requires-Dist: docker (==7.1.0)
-Requires-Dist: httpx (>=0.21.2)
-Requires-Dist: orderly-set (>=5.2.2,<6.0.0)
-Requires-Dist: publication (==0.0.3)
-Requires-Dist: pydantic (>=1.9.2)
-Requires-Dist: pydantic-core (>=2.18.2,<3.0.0)
-Requires-Dist: pydash (==7.0.6)
-Requires-Dist: python-dotenv (==1.0.1)
-Requires-Dist: pytz (==2022.6)
-Requires-Dist: pyyaml (==6.0.1)
-Requires-Dist: requests (==2.32.3)
-Requires-Dist: typing_extensions (>=4.0.0)
-Description-Content-Type: text/markdown
-
-# Vellum Python Library
-
-[![pypi](https://img.shields.io/pypi/v/vellum-ai.svg)](https://pypi.python.org/pypi/vellum-ai)
-![license badge](https://img.shields.io/github/license/vellum-ai/vellum-client-python)
-[![fern shield](https://img.shields.io/badge/%F0%9F%8C%BF-SDK%20generated%20by%20Fern-brightgreen)](https://buildwithfern.com/?utm_source=vellum-ai/vellum-client-python/readme)
-
-The Vellum Python SDK provides access to the Vellum API from python.
-
-
-## API Docs
-You can find Vellum's complete API docs at [docs.vellum.ai](https://docs.vellum.ai/api-reference/introduction/getting-started).
-
-## Installation
-
-```sh
-pip install --upgrade vellum-ai
-```
-
-## Usage
-Below is how you would invoke a deployed Prompt from the Vellum API. For a complete list of all APIs
-that Vellum supports, check out our [API Reference](https://docs.vellum.ai/api-reference/introduction/getting-started).
-
-```python
-from vellum import (
-    StringInputRequest,
-)
-from vellum.client import Vellum
-
-client = Vellum(
-    api_key="YOUR_API_KEY",
-)
-
-def execute() -> str:
-    result = client.execute_prompt(
-        prompt_deployment_name="<example-deployment-name>>",
-        release_tag="LATEST",
-        inputs=[
-            StringInputRequest(
-                name="input_a",
-                type="STRING",
-                value="Hello, world!",
-            )
-        ],
-    )
-    
-    if result.state == "REJECTED":
-        raise Exception(result.error.message)
-
-    return result.outputs[0].value
-
-if __name__ == "__main__":
-    print(execute())
-```
-
-> [!TIP]
-> You can set a system environment variable `VELLUM_API_KEY` to avoid writing your api key within your code. To do so, add `export VELLUM_API_KEY=<your-api-token>`
-> to your ~/.zshrc or ~/.bashrc, open a new terminal, and then any code calling `vellum.Vellum()` will read this key.
-
-## Async Client
-This SDK has an async version. Here's how to use it:
-
-
-
-```python
-import asyncio
-
-import vellum
-from vellum.client import AsyncVellum
-
-client = AsyncVellum(api_key="YOUR_API_KEY")
-
-async def execute() -> str:
-    result = await client.execute_prompt(
-        prompt_deployment_name="<example-deployment-name>>",
-        release_tag="LATEST",
-        inputs=[
-            vellum.StringInputRequest(
-                name="input_a",
-                value="Hello, world!",
-            )
-        ],
-    )
-
-    if result.state == "REJECTED":
-        raise Exception(result.error.message)
-    
-    return result.outputs[0].value
-
-if __name__ == "__main__":
-    print(asyncio.run(execute()))
-```
-
-## Contributing
-
-While we value open-source contributions to this SDK, most of this library is generated programmatically.
-
-Please feel free to make contributions to any of the directories or files below:
-```plaintext
-examples/*
-src/vellum/lib/*
-tests/*
-README.md
-```
-
-Any additions made to files beyond those directories and files above would have to be moved over to our generation code
-(found in the separate [vellum-client-generator](https://github.com/vellum-ai/vellum-client-generator) repo),
-otherwise they would be overwritten upon the next generated release. Feel free to open a PR as a proof of concept,
-but know that we will not be able to merge it as-is. We suggest opening an issue first to discuss with us!
-