lmnr 0.2.5__py3-none-any.whl → 0.2.7__py3-none-any.whl

lmnr/cli/parser/nodes/code.py

@@ -0,0 +1,27 @@
+from dataclasses import dataclass
+import uuid
+
+from lmnr.cli.parser.nodes import Handle, NodeFunctions
+from lmnr.cli.parser.utils import map_handles
+
+
+@dataclass
+class CodeNode(NodeFunctions):
+    id: uuid.UUID
+    name: str
+    inputs: list[Handle]
+    outputs: list[Handle]
+    inputs_mappings: dict[uuid.UUID, uuid.UUID]
+
+    def handles_mapping(
+        self, output_handle_id_to_node_name: dict[str, str]
+    ) -> list[tuple[str, str]]:
+        return map_handles(
+            self.inputs, self.inputs_mappings, output_handle_id_to_node_name
+        )
+
+    def node_type(self) -> str:
+        return "Code"
+
+    def config(self) -> dict:
+        return {}

lmnr/cli/parser/nodes/condition.py

@@ -0,0 +1,30 @@
+from dataclasses import dataclass
+import uuid
+
+from lmnr.cli.parser.nodes import Handle, NodeFunctions
+from lmnr.cli.parser.utils import map_handles
+
+
+@dataclass
+class ConditionNode(NodeFunctions):
+    id: uuid.UUID
+    name: str
+    inputs: list[Handle]
+    outputs: list[Handle]
+    inputs_mappings: dict[uuid.UUID, uuid.UUID]
+    condition: str
+
+    def handles_mapping(
+        self, output_handle_id_to_node_name: dict[str, str]
+    ) -> list[tuple[str, str]]:
+        return map_handles(
+            self.inputs, self.inputs_mappings, output_handle_id_to_node_name
+        )
+
+    def node_type(self) -> str:
+        return "Condition"
+
+    def config(self) -> dict:
+        return {
+            "condition": self.condition,
+        }

lmnr/cli/parser/nodes/input.py

@@ -0,0 +1,26 @@
+from dataclasses import dataclass
+from typing import Optional
+import uuid
+
+from lmnr.cli.parser.nodes import Handle, HandleType, NodeFunctions
+from lmnr.types import NodeInput
+
+
+@dataclass
+class InputNode(NodeFunctions):
+    id: uuid.UUID
+    name: str
+    outputs: list[Handle]
+    input: Optional[NodeInput]
+    input_type: HandleType
+
+    def handles_mapping(
+        self, output_handle_id_to_node_name: dict[str, str]
+    ) -> list[tuple[str, str]]:
+        return []
+
+    def node_type(self) -> str:
+        return "Input"
+
+    def config(self) -> dict:
+        return {}

lmnr/cli/parser/nodes/llm.py

@@ -0,0 +1,51 @@
+from dataclasses import dataclass
+from typing import Optional
+import uuid
+
+from lmnr.cli.parser.nodes import Handle, NodeFunctions
+from lmnr.cli.parser.utils import map_handles
+
+
+@dataclass
+class LLMNode(NodeFunctions):
+    id: uuid.UUID
+    name: str
+    inputs: list[Handle]
+    dynamic_inputs: list[Handle]
+    outputs: list[Handle]
+    inputs_mappings: dict[uuid.UUID, uuid.UUID]
+    prompt: str
+    model: str
+    model_params: Optional[str]
+    stream: bool
+    structured_output_enabled: bool
+    structured_output_max_retries: int
+    structured_output_schema: Optional[str]
+    structured_output_schema_target: Optional[str]
+
+    def handles_mapping(
+        self, output_handle_id_to_node_name: dict[str, str]
+    ) -> list[tuple[str, str]]:
+        combined_inputs = self.inputs + self.dynamic_inputs
+        return map_handles(
+            combined_inputs, self.inputs_mappings, output_handle_id_to_node_name
+        )
+
+    def node_type(self) -> str:
+        return "LLM"
+
+    def config(self) -> dict:
+        # For easier access in the template separate the provider and model here
+        provider, model = self.model.split(":", maxsplit=1)
+
+        return {
+            "prompt": self.prompt,
+            "provider": provider,
+            "model": model,
+            "model_params": self.model_params,
+            "stream": self.stream,
+            "structured_output_enabled": self.structured_output_enabled,
+            "structured_output_max_retries": self.structured_output_max_retries,
+            "structured_output_schema": self.structured_output_schema,
+            "structured_output_schema_target": self.structured_output_schema_target,
+        }

lmnr/cli/parser/nodes/output.py

@@ -0,0 +1,27 @@
+from dataclasses import dataclass
+import uuid
+
+from lmnr.cli.parser.nodes import Handle, NodeFunctions
+from lmnr.cli.parser.utils import map_handles
+
+
+@dataclass
+class OutputNode(NodeFunctions):
+    id: uuid.UUID
+    name: str
+    inputs: list[Handle]
+    outputs: list[Handle]
+    inputs_mappings: dict[uuid.UUID, uuid.UUID]
+
+    def handles_mapping(
+        self, output_handle_id_to_node_name: dict[str, str]
+    ) -> list[tuple[str, str]]:
+        return map_handles(
+            self.inputs, self.inputs_mappings, output_handle_id_to_node_name
+        )
+
+    def node_type(self) -> str:
+        return "Output"
+
+    def config(self) -> dict:
+        return {}

lmnr/cli/parser/nodes/router.py

@@ -0,0 +1,37 @@
+from dataclasses import dataclass
+import uuid
+
+from lmnr.cli.parser.nodes import Handle, NodeFunctions
+from lmnr.cli.parser.utils import map_handles
+
+
+@dataclass
+class Route:
+    name: str
+
+
+@dataclass
+class RouterNode(NodeFunctions):
+    id: uuid.UUID
+    name: str
+    inputs: list[Handle]
+    outputs: list[Handle]
+    inputs_mappings: dict[uuid.UUID, uuid.UUID]
+    routes: list[Route]
+    has_default_route: bool
+
+    def handles_mapping(
+        self, output_handle_id_to_node_name: dict[str, str]
+    ) -> list[tuple[str, str]]:
+        return map_handles(
+            self.inputs, self.inputs_mappings, output_handle_id_to_node_name
+        )
+
+    def node_type(self) -> str:
+        return "Router"
+
+    def config(self) -> dict:
+        return {
+            "routes": str([route.name for route in self.routes]),
+            "has_default_route": str(self.has_default_route),
+        }

lmnr/cli/parser/nodes/semantic_search.py

@@ -0,0 +1,81 @@
+from dataclasses import dataclass
+from datetime import datetime
+
+import uuid
+
+from lmnr.cli.parser.nodes import Handle, NodeFunctions
+from lmnr.cli.parser.utils import map_handles
+
+
+@dataclass
+class FileMetadata:
+    id: uuid.UUID
+    created_at: datetime
+    project_id: uuid.UUID
+    filename: str
+
+
+@dataclass
+class Dataset:
+    id: uuid.UUID
+    created_at: datetime
+    project_id: uuid.UUID
+    name: str
+
+
+@dataclass
+class SemanticSearchDatasource:
+    type: str
+    id: uuid.UUID
+    # TODO: Paste other fields here, use Union[FileMetadata, Dataset]
+
+    @classmethod
+    def from_dict(cls, datasource_dict: dict) -> "SemanticSearchDatasource":
+        if datasource_dict["type"] == "File":
+            return cls(
+                type="File",
+                id=uuid.UUID(datasource_dict["id"]),
+            )
+        elif datasource_dict["type"] == "Dataset":
+            return cls(
+                type="Dataset",
+                id=uuid.UUID(datasource_dict["id"]),
+            )
+        else:
+            raise ValueError(
+                f"Invalid SemanticSearchDatasource type: {datasource_dict['type']}"
+            )
+
+
+@dataclass
+class SemanticSearchNode(NodeFunctions):
+    id: uuid.UUID
+    name: str
+    inputs: list[Handle]
+    outputs: list[Handle]
+    inputs_mappings: dict[uuid.UUID, uuid.UUID]
+    limit: int
+    threshold: float
+    template: str
+    datasources: list[SemanticSearchDatasource]
+
+    def handles_mapping(
+        self, output_handle_id_to_node_name: dict[str, str]
+    ) -> list[tuple[str, str]]:
+        return map_handles(
+            self.inputs, self.inputs_mappings, output_handle_id_to_node_name
+        )
+
+    def node_type(self) -> str:
+        return "SemanticSearch"
+
+    def config(self) -> dict:
+        return {
+            "limit": self.limit,
+            "threshold": self.threshold,
+            "template": self.template,
+            "datasource_ids": [str(datasource.id) for datasource in self.datasources],
+            "datasource_ids_list": str(
+                [str(datasource.id) for datasource in self.datasources]
+            ),
+        }

lmnr/cli/parser/nodes/types.py

@@ -1,8 +1,17 @@
-from dataclasses import dataclass
-from typing import Any, Optional, Union
+from typing import Any, Union
 import uuid
-from lmnr.cli.parser.nodes import Handle, HandleType, NodeFunctions
-from lmnr.cli.parser.utils import map_handles
+
+from lmnr.cli.parser.nodes import Handle
+from lmnr.cli.parser.nodes.code import CodeNode
+from lmnr.cli.parser.nodes.condition import ConditionNode
+from lmnr.cli.parser.nodes.input import InputNode
+from lmnr.cli.parser.nodes.llm import LLMNode
+from lmnr.cli.parser.nodes.output import OutputNode
+from lmnr.cli.parser.nodes.router import Route, RouterNode
+from lmnr.cli.parser.nodes.semantic_search import (
+    SemanticSearchDatasource,
+    SemanticSearchNode,
+)
 from lmnr.types import NodeInput, ChatMessage
 
 
@@ -15,97 +24,15 @@ def node_input_from_json(json_val: Any) -> NodeInput:
         raise ValueError(f"Invalid NodeInput value: {json_val}")
 
 
-# TODO: Convert to Pydantic
-@dataclass
-class InputNode(NodeFunctions):
-    id: uuid.UUID
-    name: str
-    outputs: list[Handle]
-    input: Optional[NodeInput]
-    input_type: HandleType
-
-    def handles_mapping(
-        self, output_handle_id_to_node_name: dict[str, str]
-    ) -> list[tuple[str, str]]:
-        return []
-
-    def node_type(self) -> str:
-        return "Input"
-
-    def config(self) -> dict:
-        return {}
-
-
-# TODO: Convert to Pydantic
-@dataclass
-class LLMNode(NodeFunctions):
-    id: uuid.UUID
-    name: str
-    inputs: list[Handle]
-    dynamic_inputs: list[Handle]
-    outputs: list[Handle]
-    inputs_mappings: dict[uuid.UUID, uuid.UUID]
-    prompt: str
-    model: str
-    model_params: Optional[str]
-    stream: bool
-    structured_output_enabled: bool
-    structured_output_max_retries: int
-    structured_output_schema: Optional[str]
-    structured_output_schema_target: Optional[str]
-
-    def handles_mapping(
-        self, output_handle_id_to_node_name: dict[str, str]
-    ) -> list[tuple[str, str]]:
-        combined_inputs = self.inputs + self.dynamic_inputs
-        return map_handles(
-            combined_inputs, self.inputs_mappings, output_handle_id_to_node_name
-        )
-
-    def node_type(self) -> str:
-        return "LLM"
-
-    def config(self) -> dict:
-        # For easier access in the template separate the provider and model here
-        provider, model = self.model.split(":", maxsplit=1)
-
-        return {
-            "prompt": self.prompt,
-            "provider": provider,
-            "model": model,
-            "model_params": self.model_params,
-            "stream": self.stream,
-            "structured_output_enabled": self.structured_output_enabled,
-            "structured_output_max_retries": self.structured_output_max_retries,
-            "structured_output_schema": self.structured_output_schema,
-            "structured_output_schema_target": self.structured_output_schema_target,
-        }
-
-
-# TODO: Convert to Pydantic
-@dataclass
-class OutputNode(NodeFunctions):
-    id: uuid.UUID
-    name: str
-    inputs: list[Handle]
-    outputs: list[Handle]
-    inputs_mappings: dict[uuid.UUID, uuid.UUID]
-
-    def handles_mapping(
-        self, output_handle_id_to_node_name: dict[str, str]
-    ) -> list[tuple[str, str]]:
-        return map_handles(
-            self.inputs, self.inputs_mappings, output_handle_id_to_node_name
-        )
-
-    def node_type(self) -> str:
-        return "Output"
-
-    def config(self) -> dict:
-        return {}
-
-
-Node = Union[InputNode, OutputNode, LLMNode]
+Node = Union[
+    InputNode,
+    OutputNode,
+    ConditionNode,
+    LLMNode,
+    RouterNode,
+    SemanticSearchNode,
+    CodeNode,
+]
 
 
 def node_from_dict(node_dict: dict) -> Node:
@@ -128,6 +55,18 @@ def node_from_dict(node_dict: dict) -> Node:
                 for k, v in node_dict["inputsMappings"].items()
             },
         )
+    elif node_dict["type"] == "Condition":
+        return ConditionNode(
+            id=uuid.UUID(node_dict["id"]),
+            name=node_dict["name"],
+            inputs=[Handle.from_dict(handle) for handle in node_dict["inputs"]],
+            outputs=[Handle.from_dict(handle) for handle in node_dict["outputs"]],
+            inputs_mappings={
+                uuid.UUID(k): uuid.UUID(v)
+                for k, v in node_dict["inputsMappings"].items()
+            },
+            condition=node_dict["condition"],
+        )
     elif node_dict["type"] == "LLM":
         return LLMNode(
             id=uuid.UUID(node_dict["id"]),
@@ -153,5 +92,47 @@ def node_from_dict(node_dict: dict) -> Node:
             structured_output_schema=None,
             structured_output_schema_target=None,
         )
+    elif node_dict["type"] == "Router":
+        return RouterNode(
+            id=uuid.UUID(node_dict["id"]),
+            name=node_dict["name"],
+            inputs=[Handle.from_dict(handle) for handle in node_dict["inputs"]],
+            outputs=[Handle.from_dict(handle) for handle in node_dict["outputs"]],
+            inputs_mappings={
+                uuid.UUID(k): uuid.UUID(v)
+                for k, v in node_dict["inputsMappings"].items()
+            },
+            routes=[Route(name=route["name"]) for route in node_dict["routes"]],
+            has_default_route=node_dict["hasDefaultRoute"],
+        )
+    elif node_dict["type"] == "SemanticSearch":
+        return SemanticSearchNode(
+            id=uuid.UUID(node_dict["id"]),
+            name=node_dict["name"],
+            inputs=[Handle.from_dict(handle) for handle in node_dict["inputs"]],
+            outputs=[Handle.from_dict(handle) for handle in node_dict["outputs"]],
+            inputs_mappings={
+                uuid.UUID(k): uuid.UUID(v)
+                for k, v in node_dict["inputsMappings"].items()
+            },
+            limit=node_dict["limit"],
+            threshold=node_dict["threshold"],
+            template=node_dict["template"],
+            datasources=[
+                SemanticSearchDatasource.from_dict(ds)
+                for ds in node_dict["datasources"]
+            ],
+        )
+    elif node_dict["type"] == "Code":
+        return CodeNode(
+            id=uuid.UUID(node_dict["id"]),
+            name=node_dict["name"],
+            inputs=[Handle.from_dict(handle) for handle in node_dict["inputs"]],
+            outputs=[Handle.from_dict(handle) for handle in node_dict["outputs"]],
+            inputs_mappings={
+                uuid.UUID(k): uuid.UUID(v)
+                for k, v in node_dict["inputsMappings"].items()
+            },
+        )
     else:
         raise ValueError(f"Node type {node_dict['type']} not supported")

lmnr/cli/parser/parser.py

@@ -17,6 +17,9 @@ def runnable_graph_to_template_vars(graph: dict) -> dict:
         node = node_from_dict(node_obj)
         handles_mapping = node.handles_mapping(output_handle_id_to_node_name)
         node_type = node.node_type()
+
+        unique_handles = set([handle_name for (handle_name, _) in handles_mapping])
+
         tasks.append(
             {
                 "name": node.name,
@@ -28,10 +31,7 @@ def runnable_graph_to_template_vars(graph: dict) -> dict:
                     handle_name for (handle_name, _) in handles_mapping
                 ],
                 "handle_args": ", ".join(
-                    [
-                        f"{handle_name}: NodeInput"
-                        for (handle_name, _) in handles_mapping
-                    ]
+                    [f"{handle_name}: NodeInput" for handle_name in unique_handles]
                 ),
                 "prev": [],
                 "next": [],

lmnr/cli/{{cookiecutter.lmnr_pipelines_dir_name}}/engine/engine.py

@@ -9,7 +9,8 @@ import queue
 from .task import Task
 from .action import NodeRunError, RunOutput
 from .state import State
-from lmnr_engine.types import Message, NodeInput
+from lmnr.types import NodeInput
+from lmnr_engine.types import Message
 
 
 logger = logging.getLogger(__name__)
@@ -169,13 +170,17 @@ class Engine:
                 active_tasks.remove(task_id)
 
             if depth > 0:
-                # propagate reset once we enter the loop
-                # TODO: Implement this for cycles
-                raise NotImplementedError()
+                self.propagate_reset(task_id, task_id, tasks)
 
+            # terminate graph on recursion depth exceeding 10
             if depth == 10:
-                # TODO: Implement this for cycles
-                raise NotImplementedError()
+                logging.error("Max recursion depth exceeded, terminating graph")
+                error = Message(
+                    id=uuid.uuid4(),
+                    value="Max recursion depth exceeded",
+                    start_time=start_time,
+                    end_time=datetime.datetime.now(),
+                )
 
             if not next:
                 # if there are no next tasks, we can terminate the graph
@@ -259,3 +264,30 @@ class Engine:
             task,
             queue,
         )
+
+    def propagate_reset(
+        self, current_task_name: str, start_task_name: str, tasks: dict[str, Task]
+    ) -> None:
+        task = tasks[current_task_name]
+
+        for next_task_name in task.next:
+            if next_task_name == start_task_name:
+                return
+
+            next_task = tasks[next_task_name]
+
+            # in majority of cases there will be only one handle name
+            # however we need to handle the case when single output is mapped
+            # to multiple inputs on the next node
+            handle_names = []
+            for k, v in next_task.handles_mapping:
+                if v == task.name:
+                    handle_names.append(k)
+
+            for handle_name in handle_names:
+                next_state = next_task.input_states[handle_name]
+
+                if next_state.get_state().is_success():
+                    next_state.set_state(State.empty())
+                    next_state.semaphore.release()
+                    self.propagate_reset(next_task_name, start_task_name, tasks)

lmnr/cli/{{cookiecutter.lmnr_pipelines_dir_name}}/pipelines/{{cookiecutter.pipeline_dir_name}}/nodes/functions.py

@@ -1,6 +1,7 @@
 import requests
 import json
 
+from lmnr.types import ConditionedValue
 from lmnr_engine.engine.action import NodeRunError, RunOutput
 from lmnr_engine.types import ChatMessage, NodeInput
 
@@ -24,8 +25,8 @@ def {{task.function_name}}({{ task.handle_args }}, _env: dict[str, str]) -> RunO
     rendered_prompt = """{{task.config.prompt}}"""
     {% set prompt_variables = task.input_handle_names|reject("equalto", "chat_messages") %}
     {% for prompt_variable in prompt_variables %}
-    # TODO: Fix this. Using double curly braces in quotes because normal double curly braces
-    # get replaced during rendering by Cookiecutter. This is a hacky solution.
+    {# TODO: Fix this. Using double curly braces in quotes because normal double curly braces
+    # get replaced during rendering by Cookiecutter. This is a hacky solution.#}
     rendered_prompt = rendered_prompt.replace("{{'{{'}}{{prompt_variable}}{{'}}'}}", {{prompt_variable}})  # type: ignore
     {% endfor %}
 
@@ -68,7 +69,8 @@ def {{task.function_name}}({{ task.handle_args }}, _env: dict[str, str]) -> RunO
     completion_message = chat_completion["choices"][0]["message"]["content"]
 
     meta_log = {}
-    meta_log["node_chunk_id"] = None  # TODO: Add node chunk id
+    {# TODO: Add node chunk id #}
+    meta_log["node_chunk_id"] = None
     meta_log["model"] = "{{task.config.model}}"
     meta_log["prompt"] = rendered_prompt
     meta_log["input_message_count"] = len(messages)
@@ -77,7 +79,8 @@ def {{task.function_name}}({{ task.handle_args }}, _env: dict[str, str]) -> RunO
     meta_log["total_token_count"] = (
         chat_completion["usage"]["prompt_tokens"] + chat_completion["usage"]["completion_tokens"]
     )
-    meta_log["approximate_cost"] = None  # TODO: Add approximate cost
+    {# TODO: Add approximate cost #}
+    meta_log["approximate_cost"] = None
     {% elif task.config.provider == "anthropic" %}
     data = {
         "model": "{{task.config.model}}",
@@ -85,7 +88,7 @@ def {{task.function_name}}({{ task.handle_args }}, _env: dict[str, str]) -> RunO
     }
     data.update(params)
 
-    # TODO: Generate appropriate code based on this if-else block
+    {# TODO: Generate appropriate code based on this if-else block #}
     if len(messages) == 1 and messages[0].role == "system":
         messages[0].role = "user"
         message_jsons = [
@@ -116,7 +119,8 @@ def {{task.function_name}}({{ task.handle_args }}, _env: dict[str, str]) -> RunO
     completion_message = chat_completion["content"][0]["text"]
 
     meta_log = {}
-    meta_log["node_chunk_id"] = None  # TODO: Add node chunk id
+    {# TODO: Add node chunk id#}
+    meta_log["node_chunk_id"] = None
     meta_log["model"] = "{{task.config.model}}"
     meta_log["prompt"] = rendered_prompt
     meta_log["input_message_count"] = len(messages)
@@ -125,13 +129,81 @@ def {{task.function_name}}({{ task.handle_args }}, _env: dict[str, str]) -> RunO
     meta_log["total_token_count"] = (
         chat_completion["usage"]["input_tokens"] + chat_completion["usage"]["output_tokens"]
     )
-    meta_log["approximate_cost"] = None  # TODO: Add approximate cost
+    {# TODO: Add approximate cost#}
+    meta_log["approximate_cost"] = None
     {% else %}
     {% endif %}
 
     return RunOutput(status="Success", output=completion_message)
 
 
+{% elif task.node_type == "SemanticSearch" %}
+def {{task.function_name}}(query: NodeInput, _env: dict[str, str]) -> RunOutput:
+    {% set datasources_length=task.config.datasource_ids|length %}
+    {% if datasources_length == 0 %}
+    raise NodeRunError("No datasources provided")
+    {% endif %}
+
+    headers = {
+        "Authorization": f"Bearer {_env['LMNR_PROJECT_API_KEY']}",
+    }
+    data = {
+        "query": query,
+        "limit": {{ task.config.limit }},
+        "threshold": {{ task.config.threshold }},
+        "datasourceIds": {{ task.config.datasource_ids_list }},
+    }
+    query_res = requests.post("https://api.lmnr.ai/v2/semantic-search", headers=headers, json=data)
+    if query_res.status_code != 200:
+        raise NodeRunError(f"Vector search request failed:{query_res.status_code}\n{query_res.text}")
+
+    results = query_res.json()
+
+    def render_query_res_point(template: str, point: dict, relevance_index: int) -> str:
+        data = point["data"]
+        data["relevance_index"] = relevance_index
+        res = template
+        for key, value in data.items():
+            res = res.replace("{{'{{'}}" + key + "{{'}}'}}", str(value))
+        return res
+
+    rendered_res_points = [render_query_res_point("""{{task.config.template}}""", res_point, index + 1) for (index, res_point) in enumerate(results)]
+    output = "\n".join(rendered_res_points)
+
+    return RunOutput(status="Success", output=output)
+
+
+{% elif task.node_type == "Router" %}
+def {{task.function_name}}(condition: NodeInput, input: NodeInput, _env: dict[str, str]) -> RunOutput:
+    routes = {{ task.config.routes }}
+    has_default_route = {{ task.config.has_default_route }}
+
+    for route in routes:
+        if route == condition:
+            return RunOutput(status="Success", output=ConditionedValue(condition=route, value=input))
+        
+    if has_default_route:
+        return RunOutput(status="Success", output=ConditionedValue(condition=routes[-1], value=input))
+
+    raise NodeRunError(f"No route found for condition {condition}")
+
+
+{% elif task.node_type == "Condition" %}
+def {{task.function_name}}(input: NodeInput, _env: dict[str, str]) -> RunOutput:
+    condition = "{{task.config.condition}}"
+
+    if input.condition == condition:
+        return RunOutput(status="Success", output=input.value)
+    else:
+        return RunOutput(status="Termination", output=None)
+
+
+{% elif task.node_type == "Code" %}
+def {{task.function_name}}({{ task.handle_args }}, _env: dict[str, str]) -> RunOutput:
+    # Implement any functionality you want here
+    raise NodeRunError("Implement your code here")
+
+
 {% elif task.node_type == "Output" %}
 def {{task.function_name}}(output: NodeInput, _env: dict[str, str]) -> RunOutput:
     return RunOutput(status="Success", output=output)

lmnr/sdk/endpoint.py

@@ -4,7 +4,7 @@ import pydantic
 import requests
 from lmnr.types import (
     EndpointRunError, EndpointRunResponse, NodeInput, EndpointRunRequest,
-    ToolCall, SDKError
+    ToolCallError, ToolCallRequest, ToolCallResponse, SDKError
 )
 from typing import Callable, Optional
 from websockets.sync.client import connect
@@ -126,31 +126,52 @@ class Laminar:
             }
         ) as websocket:
             websocket.send(request.model_dump_json())
+            req_id = None
 
             while True:
                 message = websocket.recv()
                 try:
-                    tool_call = ToolCall.model_validate_json(message)
+                    tool_call = ToolCallRequest.model_validate_json(message)
+                    req_id = tool_call.req_id
                     matching_tools = [
                         tool for tool in tools
-                        if tool.__name__ == tool_call.function.name
+                        if tool.__name__ == tool_call.toolCall.function.name
                     ]
                     if not matching_tools:
                         raise SDKError(
-                            f'Tool {tool_call.function.name} not found.'
+                            f'Tool {tool_call.toolCall.function.name} not found.'
                             ' Registered tools: '
                             f'{", ".join([tool.__name__ for tool in tools])}'
                         )
                     tool = matching_tools[0]
-                    if tool.__name__ == tool_call.function.name:
                         # default the arguments to an empty dictionary
+                    if tool.__name__ == tool_call.toolCall.function.name:
                         arguments = {}
                         try:
-                            arguments = json.loads(tool_call.function.arguments)
+                            arguments = json.loads(tool_call.toolCall.function.arguments)
                         except:
                             pass
-                        response = tool(**arguments)
-                        websocket.send(json.dumps(response))
+                        try:
+                            response = tool(**arguments)
+                        except Exception as e:
+                            error_message = 'Error occurred while running tool' +\
+                             f'{tool.__name__}: {e}'
+                            e = ToolCallError(error=error_message, reqId=req_id)
+                            websocket.send(e.model_dump_json())
+                        formatted_response = None
+                        try:
+                            formatted_response = ToolCallResponse(
+                                reqId=tool_call.reqId,
+                                response=response
+                            )
+                        except pydantic.ValidationError as e:
+                            formatted_response = ToolCallResponse(
+                                reqId=tool_call.reqId,
+                                response=str(response)
+                            )
+                        websocket.send(
+                            formatted_response.model_dump_json()
+                        )
                 except pydantic.ValidationError as e:
                     message_json = json.loads(message)
                     keys = list(message_json.keys())
@@ -161,6 +182,5 @@ class Laminar:
                     result = EndpointRunResponse.model_validate(message_json)
                     websocket.close()
                     return result
-                except Exception:
-                    websocket.close()
-                    raise SDKError('Error communicating to backend through websocket')
+                except Exception as e:
+                    raise SDKError(f'Error communicating to backend through websocket {e}')

lmnr/sdk/remote_debugger.py

@@ -1,13 +1,21 @@
 from typing import Callable, Optional
 from websockets.sync.client import connect
+import pydantic
 import websockets
-from lmnr.types import DeregisterDebuggerRequest, NodeInput, RegisterDebuggerRequest, SDKError, ToolCall, ToolCallError
+from lmnr.types import (
+    DeregisterDebuggerRequest, NodeInput, RegisterDebuggerRequest,
+    SDKError, ToolCallError, ToolCallRequest, ToolCallResponse
+)
 import uuid
 import json
 from threading import Thread
 
 class RemoteDebugger:
-    def __init__(self, project_api_key: str, tools: list[Callable[..., NodeInput]] = []):
+    def __init__(
+        self,
+        project_api_key: str,
+        tools: list[Callable[..., NodeInput]] = []
+    ):
         self.project_api_key = project_api_key
         self.url = 'wss://api.lmnr.ai/v2/endpoint/ws'
         self.tools = tools
@@ -25,7 +33,8 @@ class RemoteDebugger:
         self.stop_flag = True
         self.thread.join()
         self.session = None
-        # python allows running threads only once, so we need to create a new thread
+        # python allows running threads only once, so we need to create
+        # a new thread
         # in case the user wants to start the debugger again
         self.thread = Thread(target=self._run)
     
@@ -41,10 +50,13 @@ class RemoteDebugger:
             }
         ) as websocket:
             websocket.send(request.model_dump_json())
-            print(self._format_session_id())
+            print(self._format_session_id_and_registerd_functions())
+            req_id = None
+
             while not self.stop_flag:
                 try:
-                    # blocks the thread until a message is received or a timeout (3 seconds) occurs
+                    # blocks the thread until a message 
+                    # is received or a timeout (3 seconds) occurs
                     message = websocket.recv(3)
                 except TimeoutError:
                     continue
@@ -52,45 +64,76 @@ class RemoteDebugger:
                     print("Connection closed. Please restart the debugger.")
                     return
                 try:
-                    tool_call = ToolCall.model_validate_json(message)
+                    tool_call = ToolCallRequest.model_validate_json(message)
+                    req_id = tool_call.reqId
                 except:
                     raise SDKError(f'Invalid message received:\n{message}')
                 matching_tools = [
                     tool for tool in self.tools
-                    if tool.__name__ == tool_call.function.name
+                    if tool.__name__ == tool_call.toolCall.function.name
                 ]
                 if not matching_tools:
-                    error_message = f'Tool {tool_call.function.name} not found.' +\
-                    f' Registered tools: {", ".join([tool.__name__ for tool in self.tools])}'
-                    e = ToolCallError(error=error_message)
+                    error_message = \
+                        f'Tool {tool_call.toolCall.function.name} not found' +\
+                        '. Registered tools: ' +\
+                        {", ".join([tool.__name__ for tool in self.tools])}
+                    e = ToolCallError(error=error_message, reqId=req_id)
                     websocket.send(e.model_dump_json())
                     continue
                 tool = matching_tools[0]
-                if tool.__name__ == tool_call.function.name:
+                if tool.__name__ == tool_call.toolCall.function.name:
                     # default the arguments to an empty dictionary
                     arguments = {}
                     try:
-                        arguments = json.loads(tool_call.function.arguments)
+                        arguments = json.loads(
+                            tool_call.toolCall.function.arguments)
                     except:
                         pass
                     try:
-                        response = tool(**arguments)  # of type NodeInput
-                        websocket.send(json.dumps(response))
+                        response = tool(**arguments)
                     except Exception as e:
-                        error_message = f'Error occurred while running tool {tool.__name__}: {e}'
-                        e = ToolCallError(error=error_message)
+                        error_message = 'Error occurred while running tool' +\
+                            f'{tool.__name__}: {e}'
+                        e = ToolCallError(error=error_message, reqId=req_id)
                         websocket.send(e.model_dump_json())
-            websocket.send(DeregisterDebuggerRequest(debuggerSessionId=self.session, deregister=True).model_dump_json())
+                    formatted_response = None
+                    try:
+                        formatted_response = ToolCallResponse(
+                            reqId=tool_call.reqId,
+                            response=response
+                        )
+                    except pydantic.ValidationError as e:
+                        formatted_response = ToolCallResponse(
+                            reqId=tool_call.reqId,
+                            response=str(response)
+                        )
+                    websocket.send(
+                        formatted_response.model_dump_json()
+                    )
+            websocket.send(
+                DeregisterDebuggerRequest(
+                    debuggerSessionId=self.session,
+                    deregister=True
+                ).model_dump_json()
+            )
 
     def _generate_session_id(self) -> str:
         return uuid.uuid4().urn[9:]
 
-    def _format_session_id(self) -> str:
+    def _format_session_id_and_registerd_functions(self) -> str:
+        registered_functions = \
+            ',\n'.join(['- ' + tool.__name__ for tool in self.tools])
         return \
 f"""
 ========================================
 Debugger Session ID:
 {self.session}
+========================================
+
+Registered functions:
+{registered_functions}
+
 ========================================
 """
+
     

lmnr/types.py

@@ -1,13 +1,21 @@
-
 import requests
 import pydantic
+import uuid
 from typing import Union, Optional
 
+
 class ChatMessage(pydantic.BaseModel):
     role: str
     content: str
 
-NodeInput = Union[str, list[ChatMessage]] # TypeAlias
+
+class ConditionedValue(pydantic.BaseModel):
+    condition: str
+    value: "NodeInput"
+
+
+NodeInput = Union[str, list[ChatMessage], ConditionedValue]  # TypeAlias
+
 
 class EndpointRunRequest(pydantic.BaseModel):
     inputs: dict[str, NodeInput]
@@ -15,10 +23,12 @@ class EndpointRunRequest(pydantic.BaseModel):
     env: dict[str, str] = pydantic.Field(default_factory=dict)
     metadata: dict[str, str] = pydantic.Field(default_factory=dict)
 
+
 class EndpointRunResponse(pydantic.BaseModel):
     outputs: dict[str, dict[str, NodeInput]]
     run_id: str
 
+
 class EndpointRunError(Exception):
     error_code: str
     error_message: str
@@ -26,39 +36,57 @@ class EndpointRunError(Exception):
     def __init__(self, response: requests.Response):
         try:
             resp_json = response.json()
-            self.error_code = resp_json['error_code']
-            self.error_message = resp_json['error_message']
+            self.error_code = resp_json["error_code"]
+            self.error_message = resp_json["error_message"]
             super().__init__(self.error_message)
-        except:
+        except Exception:
             super().__init__(response.text)
-    
+
     def __str__(self) -> str:
         try:
-            return str({'error_code': self.error_code, 'error_message': self.error_message})
-        except:
+            return str(
+                {"error_code": self.error_code, "error_message": self.error_message}
+            )
+        except Exception:
             return super().__str__()
-        
+
+
 class SDKError(Exception):
     def __init__(self, error_message: str):
         super().__init__(error_message)
 
-class ToolCallRequest(pydantic.BaseModel):
+
+class ToolCallFunction(pydantic.BaseModel):
     name: str
     arguments: str
 
+
 class ToolCall(pydantic.BaseModel):
     id: Optional[str]
     type: Optional[str]
-    function: ToolCallRequest
+    function: ToolCallFunction
+
+
+# TODO: allow snake_case and manually convert to camelCase
+class ToolCallRequest(pydantic.BaseModel):
+    reqId: uuid.UUID
+    toolCall: ToolCall
+
+
+class ToolCallResponse(pydantic.BaseModel):
+    reqId: uuid.UUID
+    response: NodeInput
+
 
-ToolCallResponse = NodeInput
 class ToolCallError(pydantic.BaseModel):
+    reqId: uuid.UUID
     error: str
 
-# TODO: allow snake_case and manually convert to camelCase
+
 class RegisterDebuggerRequest(pydantic.BaseModel):
     debuggerSessionId: str
 
+
 class DeregisterDebuggerRequest(pydantic.BaseModel):
     debuggerSessionId: str
     deregister: bool

{lmnr-0.2.5.dist-info → lmnr-0.2.7.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: lmnr
-Version: 0.2.5
+Version: 0.2.7
 Summary: Python SDK for Laminar AI
 License: Apache-2.0
 Author: lmnr.ai
@@ -17,11 +17,12 @@ Requires-Dist: cookiecutter (>=2.6.0,<3.0.0)
 Requires-Dist: pydantic (>=2.7.4,<3.0.0)
 Requires-Dist: python-dotenv (>=1.0.1,<2.0.0)
 Requires-Dist: requests (>=2.32.3,<3.0.0)
-Requires-Dist: urllib3 (==1.26.6)
 Requires-Dist: websockets (>=12.0,<13.0)
 Description-Content-Type: text/markdown
 
-# Python SDK for Laminar AI
+# Laminar AI
+
+This repo provides core for code generation, Laminar CLI, and Laminar SDK.
 
 ## Quickstart
 ```sh
@@ -139,14 +140,6 @@ Set up `DEBUGGER_SESSION_ID` environment variable in your pipeline.
 
 You can run as many sessions as you need, experimenting with your flows.
 
-#### 5. Stop the debugger
-
-In order to stop the session, do
-
-```python
-debugger.stop()
-```
-
 ## CLI for code generation
 
 ### Basic usage

{lmnr-0.2.5.dist-info → lmnr-0.2.7.dist-info}/RECORD

@@ -5,24 +5,31 @@ lmnr/cli/cli.py,sha256=pzr5LUILi7TcaJIkC-CzmT7RG7-HWApQmUpgK9bc7mI,2847
 lmnr/cli/cookiecutter.json,sha256=PeiMMzCPzDhsapqYoAceYGPI5lOUNimvFzh5KeQv5QE,359
 lmnr/cli/parser/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 lmnr/cli/parser/nodes/__init__.py,sha256=BNbbfn0WwbFDA6TNhLOaT_Ji69rCL5voUibqMD7Knng,1163
-lmnr/cli/parser/nodes/types.py,sha256=fqw_0-szpcZ8G3bADjgG97gCng2DRIkL72nLV11gTP4,5081
-lmnr/cli/parser/parser.py,sha256=UtuupKj2Dh-47HVcYU4PI2s4Erh1gTjKiuGCcc-nkbM,2163
+lmnr/cli/parser/nodes/code.py,sha256=GXqOxN6tdiStZGWLbN3WZCmDfzwYIgSRmZ5t72AOIXc,661
+lmnr/cli/parser/nodes/condition.py,sha256=AJny0ILXbSy1hTwsRvZvDUqts9INNx63yQSkD7Dp7KU,740
+lmnr/cli/parser/nodes/input.py,sha256=Xwktcih7Mezqv4cEejfPkpG8uJxDsbqRytBvKmlJDYE,578
+lmnr/cli/parser/nodes/llm.py,sha256=iQWYFnQi5PcQD9WJpTSHbSzClM6s0wBOoEqyN5c4yQo,1674
+lmnr/cli/parser/nodes/output.py,sha256=1XBppSscxM01kfZhE9oOh2GgdCVzyPVe2RAxLI5HmUc,665
+lmnr/cli/parser/nodes/router.py,sha256=dmCx4ho8_GdFJXQa8UevMf_uEP7AKBv_MJ2zpLC6Vck,894
+lmnr/cli/parser/nodes/semantic_search.py,sha256=o_XCR7BShAq8VGeKjPTwL6MxLdB07XHSd5CE71sFFiY,2105
+lmnr/cli/parser/nodes/types.py,sha256=NRhlgI3WGd86AToo-tU974DEZzbLaH4iDdP-hEEQiIo,5343
+lmnr/cli/parser/parser.py,sha256=kAZEeg358lyj_Q1IIhQB_bA7LW3Aw6RduShIfUSmLqQ,2173
 lmnr/cli/parser/utils.py,sha256=wVaqHVOR9VXl8Og9nkVyCVgHIcgbtYGkDOEGPtmjZ8g,715
 lmnr/cli/{{cookiecutter.lmnr_pipelines_dir_name}}/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 lmnr/cli/{{cookiecutter.lmnr_pipelines_dir_name}}/engine/__init__.py,sha256=pLVZqvDnNf9foGR-HXnX2F7WC2TWmyCTNcUctG8SXAI,27
 lmnr/cli/{{cookiecutter.lmnr_pipelines_dir_name}}/engine/action.py,sha256=mZMQwwPV5LtSfwdwQ7HefI3ttvwuokp4mhVI_Xn1Zck,274
-lmnr/cli/{{cookiecutter.lmnr_pipelines_dir_name}}/engine/engine.py,sha256=8DPh3zLhs4oDepuT5g6Z8mxuy2561v8P7Dhd4hUOn6g,8375
+lmnr/cli/{{cookiecutter.lmnr_pipelines_dir_name}}/engine/engine.py,sha256=kCY6J7oQpm3f9YCYY2ZBzM_9bUv_XYTCRD_uFa6PLWQ,9640
 lmnr/cli/{{cookiecutter.lmnr_pipelines_dir_name}}/engine/state.py,sha256=wTx7jAv7b63-8k34cYfQp_DJxhCCOYT_qRHkmnZfnc0,1686
 lmnr/cli/{{cookiecutter.lmnr_pipelines_dir_name}}/engine/task.py,sha256=ware5VIrZvluHH3mpH6h7G0YDk5L0buSQ7s09za4Fro,1200
 lmnr/cli/{{cookiecutter.lmnr_pipelines_dir_name}}/pipelines/{{cookiecutter.pipeline_dir_name}}/__init__.py,sha256=bsfbNUBYtKv37qzc_GLhSAzBam2lnowP_dlr8pubhcg,80
-lmnr/cli/{{cookiecutter.lmnr_pipelines_dir_name}}/pipelines/{{cookiecutter.pipeline_dir_name}}/nodes/functions.py,sha256=2fLTN6PbhlSPEcrHBSSJESmg__UNMk4Ivs5mz0LqnG8,5169
+lmnr/cli/{{cookiecutter.lmnr_pipelines_dir_name}}/pipelines/{{cookiecutter.pipeline_dir_name}}/nodes/functions.py,sha256=Bwu8p7m16NAyt9wC0DTQL0MrHbM44WylLs5wTLwSxBM,7845
 lmnr/cli/{{cookiecutter.lmnr_pipelines_dir_name}}/pipelines/{{cookiecutter.pipeline_dir_name}}/{{cookiecutter.pipeline_dir_name}}.py,sha256=WG-ZMofPpGXCx5jdWVry3_XBzcKjqn8ZycFSiWEOBPg,2858
 lmnr/cli/{{cookiecutter.lmnr_pipelines_dir_name}}/types.py,sha256=iWuflMV7TiaBPs6-B-BlrovvWpZgHGGHK0v8rSqER7A,997
-lmnr/sdk/endpoint.py,sha256=tT6-w-mwbh4BAwnj5G0pCVE_Sz8EUzZmpBtacm_T2pE,6359
-lmnr/sdk/remote_debugger.py,sha256=pNSmmk-KAAOUygKQAeTB1DbB3vvlpHL8JLwPBpBhBwA,3892
-lmnr/types.py,sha256=Y41GGiCVemBUCTKedCwDhxjdsoO8uENr3LRwfdp9MI0,1771
-lmnr-0.2.5.dist-info/LICENSE,sha256=67b_wJHVV1CBaWkrKFWU1wyqTPSdzH77Ls-59631COg,10411
-lmnr-0.2.5.dist-info/METADATA,sha256=BgYsaYsxprvRZ5-RxZk8DQdSTKmF51hOvU27hg_0wn4,5492
-lmnr-0.2.5.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-lmnr-0.2.5.dist-info/entry_points.txt,sha256=Qg7ZRax4k-rcQsZ26XRYQ8YFSBiyY2PNxYfq4a6PYXI,41
-lmnr-0.2.5.dist-info/RECORD,,
+lmnr/sdk/endpoint.py,sha256=0HjcxMUcJz-klFZO2f5xtTaoLjcaEb8vrJ_YldTWUc8,7467
+lmnr/sdk/remote_debugger.py,sha256=vCpMz7y3uboOi81qEwr8z3fhQ2H1y2YtJAxXVtb6uCA,5141
+lmnr/types.py,sha256=OR9xRAQ5uTTwpJTDL_e3jZqxYJWvyX96CCoxr3oo94g,2112
+lmnr-0.2.7.dist-info/LICENSE,sha256=67b_wJHVV1CBaWkrKFWU1wyqTPSdzH77Ls-59631COg,10411
+lmnr-0.2.7.dist-info/METADATA,sha256=Ya1KVPAiyGxAZybuXSum8wmy4l-SnyYvaTTrvQ7uZRU,5427
+lmnr-0.2.7.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+lmnr-0.2.7.dist-info/entry_points.txt,sha256=Qg7ZRax4k-rcQsZ26XRYQ8YFSBiyY2PNxYfq4a6PYXI,41
+lmnr-0.2.7.dist-info/RECORD,,