lmnr 0.1.1__py3-none-any.whl → 0.1.3__py3-none-any.whl

lmnr/__init__.py

@@ -1 +1,5 @@
-from .endpoint import Laminar
+from .sdk.endpoint import Laminar
+from .types import (
+    ChatMessage, EndpointRunError, EndpointRunResponse,
+    NodeInput, SDKError,
+)

lmnr/cli/__main__.py

@@ -0,0 +1,4 @@
+from .cli import cli
+
+if __name__ == "__main__":
+    cli()

lmnr/cli/cli.py

@@ -0,0 +1,97 @@
+import requests
+from dotenv import load_dotenv
+import os
+import click
+import logging
+from cookiecutter.main import cookiecutter
+from importlib import resources as importlib_resources
+from pydantic.alias_generators import to_pascal
+
+from .parser.parser import runnable_graph_to_template_vars
+import lmnr
+
+logger = logging.getLogger(__name__)
+
+
+@click.group()
+@click.version_option()
+def cli():
+    "CLI for Laminar AI Engine"
+
+
+@cli.command(name="pull")
+@click.argument("pipeline_name")
+@click.argument("pipeline_version_name")
+@click.option(
+    "-p",
+    "--project-api-key",
+    help="Project API key",
+)
+@click.option(
+    "-l",
+    "--loglevel",
+    help="Sets logging level",
+)
+def pull(pipeline_name, pipeline_version_name, project_api_key, loglevel):
+    loglevel_str_to_val = {
+        "DEBUG": logging.DEBUG,
+        "INFO": logging.INFO,
+        "WARNING": logging.WARNING,
+        "ERROR": logging.ERROR,
+        "CRITICAL": logging.CRITICAL,
+    }
+    logging.basicConfig()
+    logging.getLogger().setLevel(loglevel_str_to_val.get(loglevel, logging.WARNING))
+
+    project_api_key = project_api_key or os.environ.get("LMNR_PROJECT_API_KEY")
+    if not project_api_key:
+        load_dotenv()
+        project_api_key = os.environ.get("LMNR_PROJECT_API_KEY")
+    if not project_api_key:
+        raise ValueError("LMNR_PROJECT_API_KEY is not set")
+
+    headers = {"Authorization": f"Bearer {project_api_key}"}
+    params = {
+        "pipelineName": pipeline_name,
+        "pipelineVersionName": pipeline_version_name,
+    }
+    res = requests.get(
+        "https://api.lmnr.ai/v1/pipeline-version-by-name",
+        headers=headers,
+        params=params,
+    )
+    if res.status_code != 200:
+        try:
+            res_json = res.json()
+        except Exception:
+            raise ValueError(
+                f"Error in fetching pipeline version: {res.status_code}\n{res.text}"
+            )
+        raise ValueError(
+            f"Error in fetching pipeline version: {res.status_code}\n{res_json}"
+        )
+
+    pipeline_version = res.json()
+
+    class_name = to_pascal(pipeline_name.replace(" ", "_"))
+
+    context = {
+        "pipeline_name": pipeline_name,
+        "pipeline_version_name": pipeline_version_name,
+        "class_name": class_name,
+        # _tasks starts from underscore because we don't want it to be templated
+        # some tasks contains LLM nodes which have prompts
+        # which we don't want to be rendered by cookiecutter
+        "_tasks": runnable_graph_to_template_vars(pipeline_version["runnableGraph"]),
+    }
+
+    logger.info(f"Context:\n{context}")
+    cookiecutter(
+        str(importlib_resources.files(lmnr)),
+        output_dir=".",
+        config_file=None,
+        extra_context=context,
+        directory="cli",
+        no_input=True,
+        overwrite_if_exists=True,
+    )

lmnr/cli/cookiecutter.json

@@ -0,0 +1,9 @@
+{
+    "lmnr_pipelines_dir_name": "lmnr_engine",
+    "pipeline_name": "Laminar Pipeline",
+    "pipeline_dir_name": "{{ cookiecutter['pipeline_name'].lower().replace('-', '_').replace(' ', '_') }}",
+    "class_name": "LaminarPipeline",
+    "pipeline_version_name": "main",
+    "_tasks": {},
+    "_jinja2_env_vars": {"lstrip_blocks": true, "trim_blocks": true}
+}

lmnr/cli/parser/nodes/__init__.py

@@ -0,0 +1,50 @@
+from abc import ABCMeta, abstractmethod
+from dataclasses import dataclass
+from typing import Optional
+import uuid
+
+
+HandleType = str  # "String" | "ChatMessageList" | "Any"
+
+
+@dataclass
+class Handle:
+    id: uuid.UUID
+    name: Optional[str]
+    type: HandleType
+
+    @classmethod
+    def from_dict(cls, dict: dict) -> "Handle":
+        return cls(
+            id=uuid.UUID(dict["id"]),
+            name=(dict["name"] if "name" in dict else None),
+            type=dict["type"],
+        )
+
+
+@abstractmethod
+class NodeFunctions(metaclass=ABCMeta):
+    @abstractmethod
+    def handles_mapping(
+        self, output_handle_id_to_node_name: dict[str, str]
+    ) -> list[tuple[str, str]]:
+        """
+        Returns a list of tuples mapping from this node's input
+        handle name to the unique name of the previous node.
+
+        Assumes previous node has only one output.
+        """
+        pass
+
+    @abstractmethod
+    def node_type(self) -> str:
+        pass
+
+    @abstractmethod
+    def config(self) -> dict:
+        """
+        Returns a dictionary of node-specific configuration.
+
+        E.g. prompt and model name for LLM node.
+        """
+        pass

lmnr/cli/parser/nodes/types.py

@@ -0,0 +1,156 @@
+from dataclasses import dataclass
+from typing import Any, Optional, Union
+import uuid
+from lmnr.cli.parser.nodes import Handle, HandleType, NodeFunctions
+from lmnr.cli.parser.utils import map_handles
+from lmnr.types import NodeInput, ChatMessage
+
+def node_input_from_json(json_val: Any) -> NodeInput:
+    if isinstance(json_val, str):
+        return json_val
+    elif isinstance(json_val, list):
+        return [ChatMessage.model_validate(msg) for msg in json_val]
+    else:
+        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]
+
+
+def node_from_dict(node_dict: dict) -> Node:
+    if node_dict["type"] == "Input":
+        return InputNode(
+            id=uuid.UUID(node_dict["id"]),
+            name=node_dict["name"],
+            outputs=[Handle.from_dict(handle) for handle in node_dict["outputs"]],
+            input=node_input_from_json(node_dict["input"]),
+            input_type=node_dict["inputType"],
+        )
+    elif node_dict["type"] == "Output":
+        return OutputNode(
+            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()
+            },
+        )
+    elif node_dict["type"] == "LLM":
+        return LLMNode(
+            id=uuid.UUID(node_dict["id"]),
+            name=node_dict["name"],
+            inputs=[Handle.from_dict(handle) for handle in node_dict["inputs"]],
+            dynamic_inputs=[
+                Handle.from_dict(handle) for handle in node_dict["dynamicInputs"]
+            ],
+            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()
+            },
+            prompt=node_dict["prompt"],
+            model=node_dict["model"],
+            model_params=(
+                node_dict["modelParams"] if "modelParams" in node_dict else None
+            ),
+            stream=False,
+            # TODO: Implement structured output
+            structured_output_enabled=False,
+            structured_output_max_retries=3,
+            structured_output_schema=None,
+            structured_output_schema_target=None,
+        )
+    else:
+        raise ValueError(f"Node type {node_dict['type']} not supported")

lmnr/cli/parser/parser.py

@@ -0,0 +1,58 @@
+from .nodes.types import node_from_dict
+
+
+def runnable_graph_to_template_vars(graph: dict) -> dict:
+    """
+    Convert a runnable graph to template vars to be rendered in a cookiecutter context.
+    """
+    node_id_to_node_name = {}
+    output_handle_id_to_node_name: dict[str, str] = {}
+    for node in graph["nodes"].values():
+        node_id_to_node_name[node["id"]] = node["name"]
+        for handle in node["outputs"]:
+            output_handle_id_to_node_name[handle["id"]] = node["name"]
+
+    tasks = []
+    for node_obj in graph["nodes"].values():
+        node = node_from_dict(node_obj)
+        handles_mapping = node.handles_mapping(output_handle_id_to_node_name)
+        node_type = node.node_type()
+        tasks.append(
+            {
+                "name": node.name,
+                "function_name": f"run_{node.name}",
+                "node_type": node_type,
+                "handles_mapping": handles_mapping,
+                # since we map from to to from, all 'to's won't repeat
+                "input_handle_names": [
+                    handle_name for (handle_name, _) in handles_mapping
+                ],
+                "handle_args": ", ".join(
+                    [
+                        f"{handle_name}: NodeInput"
+                        for (handle_name, _) in handles_mapping
+                    ]
+                ),
+                "prev": [],
+                "next": [],
+                "config": node.config(),
+            }
+        )
+
+    for to, from_ in graph["pred"].items():
+        # TODO: Make "tasks" a hashmap from node id (as str!) to task
+        to_task = [task for task in tasks if task["name"] == node_id_to_node_name[to]][
+            0
+        ]
+        from_tasks = []
+        for f in from_:
+            from_tasks.append(
+                [task for task in tasks if task["name"] == node_id_to_node_name[f]][0]
+            )
+
+        for from_task in from_tasks:
+            to_task["prev"].append(from_task["name"])
+            from_task["next"].append(node_id_to_node_name[to])
+
+    # Return as a hashmap due to cookiecutter limitations, investigate later.
+    return {task["name"]: task for task in tasks}

lmnr/cli/parser/utils.py

@@ -0,0 +1,25 @@
+# Convert a list of handles to a map of input handle names
+# to their respective values
+import uuid
+from .nodes import Handle
+
+
+def map_handles(
+    inputs: list[Handle],
+    inputs_mappings: dict[uuid.UUID, uuid.UUID],
+    output_handle_id_to_node_name: dict[str, str],
+) -> list[tuple[str, str]]:
+    mapping = []
+
+    for to, from_ in inputs_mappings.items():
+        for input in inputs:
+            if input.id == to:
+                mapping.append((input.name, from_))
+                break
+        else:
+            raise ValueError(f"Input handle {to} not found in inputs")
+
+    return [
+        (input_name, output_handle_id_to_node_name[str(output_id)])
+        for input_name, output_id in mapping
+    ]

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

@@ -0,0 +1 @@
+from .engine import Engine

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

@@ -0,0 +1,14 @@
+from dataclasses import dataclass
+from typing import Union
+
+from lmnr_engine.types import NodeInput
+
+
+@dataclass
+class RunOutput:
+    status: str  # "Success" | "Termination"  TODO: Turn into Enum
+    output: Union[NodeInput, None]
+
+
+class NodeRunError(Exception):
+    pass

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

@@ -0,0 +1,261 @@
+from concurrent.futures import ThreadPoolExecutor
+import datetime
+import logging
+from typing import Optional
+import uuid
+from dataclasses import dataclass
+import queue
+
+from .task import Task
+from .action import NodeRunError, RunOutput
+from .state import State
+from lmnr_engine.types import Message, NodeInput
+
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class ScheduledTask:
+    status: str  # "Task" | "Err"  TODO: Use an enum
+    task_name: Optional[str]
+
+
+class RunError(Exception):
+    outputs: dict[str, Message]
+
+
+@dataclass
+class Engine:
+    tasks: dict[str, Task]
+    active_tasks: set[str]
+    depths: dict[str, int]
+    outputs: dict[str, Message]
+    env: dict[str, str]
+    thread_pool_executor: ThreadPoolExecutor
+    # TODO: Store thread pool executor's Futures here to have control
+    # over them (e.g. cancel them)
+
+    @classmethod
+    def new(
+        cls, thread_pool_executor: ThreadPoolExecutor, env: dict[str, str] = {}
+    ) -> "Engine":
+        return cls(
+            tasks={},
+            active_tasks=set(),
+            depths={},
+            outputs={},
+            env=env,
+            thread_pool_executor=thread_pool_executor,
+        )
+
+    @classmethod
+    def with_tasks(
+        cls,
+        tasks: list[Task],
+        thread_pool_executor: ThreadPoolExecutor,
+        env: dict[str, str] = {},
+    ) -> "Engine":
+        dag = cls.new(thread_pool_executor, env=env)
+
+        for task in tasks:
+            dag.tasks[task.name] = task
+            dag.depths[task.name] = 0
+
+        return dag
+
+    def override_inputs(self, inputs: dict[str, NodeInput]) -> None:
+        for task in self.tasks.values():
+            # TODO: Check that it's the Input type task
+            if not task.prev:
+                task.value = inputs[task.name]
+
+    def run(self, inputs: dict[str, NodeInput]) -> dict[str, Message]:
+        self.override_inputs(inputs)
+
+        q = queue.Queue()
+
+        input_tasks = []
+        for task in self.tasks.values():
+            if len(task.prev) == 0:
+                input_tasks.append(task.name)
+
+        for task_id in input_tasks:
+            q.put(ScheduledTask(status="Task", task_name=task_id))
+
+        while True:
+            logger.info("Waiting for task from queue")
+            scheduled_task: ScheduledTask = q.get()
+            logger.info(f"Got task from queue: {scheduled_task}")
+            if scheduled_task.status == "Err":
+                # TODO: Abort all other threads
+                raise RunError(self.outputs)
+
+            task: Task = self.tasks[scheduled_task.task_name]  # type: ignore
+            logger.info(f"Task next: {task.next}")
+
+            if not task.next:
+                try:
+                    fut = self.execute_task(task, q)
+                    fut.result()
+                    if not self.active_tasks:
+                        return self.outputs
+                except Exception:
+                    raise RunError(self.outputs)
+            else:
+                self.execute_task(task, q)
+
+    def execute_task_inner(
+        self,
+        task: Task,
+        queue: queue.Queue,
+    ) -> None:
+        task_id = task.name
+        next = task.next
+        input_states = task.input_states
+        active_tasks = self.active_tasks
+        tasks = self.tasks
+        depths = self.depths
+        depth = depths[task.name]
+        outputs = self.outputs
+
+        inputs: dict[str, NodeInput] = {}
+        input_messages = []
+
+        # Wait for inputs for this task to be set
+        for handle_name, input_state in input_states.items():
+            logger.info(f"Task {task_id} waiting for semaphore for {handle_name}")
+            input_state.semaphore.acquire()
+            logger.info(f"Task {task_id} acquired semaphore for {handle_name}")
+
+            # Set the outputs of predecessors as inputs of the current
+            output = input_state.get_state()
+            # If at least one of the inputs is termination,
+            # also terminate this task early and set its state to termination
+            if output.status == "Termination":
+                return
+            message = output.get_out()
+
+            inputs[handle_name] = message.value
+            input_messages.append(message)
+
+        start_time = datetime.datetime.now()
+
+        try:
+            if callable(task.value):
+                res = task.value(**inputs, _env=self.env)
+            else:
+                res = RunOutput(status="Success", output=task.value)
+
+            if res.status == "Success":
+                id = uuid.uuid4()
+                state = State.new(
+                    Message(
+                        id=id,
+                        value=res.output,  # type: ignore
+                        start_time=start_time,
+                        end_time=datetime.datetime.now(),
+                    )
+                )
+            else:
+                assert res.status == "Termination"
+                state = State.termination()
+
+            is_termination = state.is_termination()
+            logger.info(f"Task {task_id} executed")
+
+            # remove the task from active tasks once it's done
+            if task_id in active_tasks:
+                active_tasks.remove(task_id)
+
+            if depth > 0:
+                # propagate reset once we enter the loop
+                # TODO: Implement this for cycles
+                raise NotImplementedError()
+
+            if depth == 10:
+                # TODO: Implement this for cycles
+                raise NotImplementedError()
+
+            if not next:
+                # if there are no next tasks, we can terminate the graph
+                outputs[task.name] = state.get_out()
+
+            # push next tasks to the channel only if
+            # the current task is not a termination
+            for next_task_name in next:
+                # we set the inputs of the next tasks
+                # to the outputs of the current task
+                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]
+                    next_state.set_state_and_permits(state, 1)
+
+                # push next tasks to the channel only if the task is not active
+                # and current task is not a termination
+                if not (next_task_name in active_tasks) and not is_termination:
+                    active_tasks.add(next_task_name)
+                    queue.put(
+                        ScheduledTask(
+                            status="Task",
+                            task_name=next_task_name,
+                        )
+                    )
+
+            # increment depth of the finished task
+            depths[task_id] = depth + 1
+        except NodeRunError as e:
+            logger.exception(f"Execution failed [id: {task_id}]")
+
+            error = Message(
+                id=uuid.uuid4(),
+                value=str(e),
+                start_time=start_time,
+                end_time=datetime.datetime.now(),
+            )
+
+            outputs[task.name] = error
+
+            # terminate entire graph by sending err task
+            queue.put(
+                ScheduledTask(
+                    status="Err",
+                    task_name=None,
+                )
+            )
+
+        except Exception:
+            logger.exception(f"Execution failed [id: {task_id}]")
+            error = Message(
+                id=uuid.uuid4(),
+                value="Unexpected server error",
+                start_time=start_time,
+                end_time=datetime.datetime.now(),
+            )
+            outputs[task.name] = error
+            queue.put(
+                ScheduledTask(
+                    status="Err",
+                    task_name=None,
+                )
+            )
+
+    def execute_task(
+        self,
+        task: Task,
+        queue: queue.Queue,
+    ):
+        return self.thread_pool_executor.submit(
+            self.execute_task_inner,
+            task,
+            queue,
+        )