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

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

@@ -0,0 +1,69 @@
+import threading
+from dataclasses import dataclass
+from typing import Union
+
+from lmnr_engine.types import Message
+
+
+@dataclass
+class State:
+    status: str  # "Success", "Empty", "Termination"  # TODO: Turn into Enum
+    message: Union[Message, None]
+
+    @classmethod
+    def new(cls, val: Message) -> "State":
+        return cls(
+            status="Success",
+            message=val,
+        )
+
+    @classmethod
+    def empty(cls) -> "State":
+        return cls(
+            status="Empty",
+            message=Message.empty(),
+        )
+
+    @classmethod
+    def termination(cls) -> "State":
+        return cls(
+            status="Termination",
+            message=None,
+        )
+
+    def is_success(self) -> bool:
+        return self.status == "Success"
+
+    def is_termination(self) -> bool:
+        return self.status == "Termination"
+
+    def get_out(self) -> Message:
+        if self.message is None:
+            raise ValueError("Cannot get message from a termination state")
+
+        return self.message
+
+
+@dataclass
+class ExecState:
+    output: State
+    semaphore: threading.Semaphore
+
+    @classmethod
+    def new(cls) -> "ExecState":
+        return cls(
+            output=State.empty(),
+            semaphore=threading.Semaphore(0),
+        )
+
+    # Assume this is called by the caller who doesn't need to acquire semaphore
+    def set_state(self, output: State):
+        self.output = output
+
+    # Assume the caller is smart to call this after acquiring the semaphore
+    def get_state(self) -> State:
+        return self.output
+
+    def set_state_and_permits(self, output: State, permits: int):
+        self.output = output
+        self.semaphore.release(permits)

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

@@ -0,0 +1,38 @@
+from typing import Callable, Union
+
+from .action import RunOutput
+from .state import ExecState
+from lmnr_engine.types import NodeInput
+
+
+class Task:
+    # unique identifier
+    name: str
+    # mapping from current node's handle name to previous node's unique name
+    # assumes nodes have only one output
+    handles_mapping: list[tuple[str, str]]
+    # Value or a function that returns a value
+    # Usually a function which waits for inputs from previous nodes
+    value: Union[NodeInput, Callable[..., RunOutput]]  # TODO: Type this fully
+    # unique node names of previous nodes
+    prev: list[str]
+    # unique node names of next nodes
+    next: list[str]
+    input_states: dict[str, ExecState]
+
+    def __init__(
+        self,
+        name: str,
+        handles_mapping: list[tuple[str, str]],
+        value: Union[NodeInput, Callable[..., RunOutput]],
+        prev: list[str],
+        next: list[str],
+    ) -> None:
+        self.name = name
+        self.handles_mapping = handles_mapping
+        self.value = value
+        self.prev = prev
+        self.next = next
+        self.input_states = {
+            handle_name: ExecState.new() for (handle_name, _) in self.handles_mapping
+        }

lmnr/cli/{{cookiecutter.lmnr_pipelines_dir_name}}/pipelines/{{cookiecutter.pipeline_dir_name}}/__init__.py

@@ -0,0 +1 @@
+from .{{ cookiecutter.pipeline_dir_name }} import {{ cookiecutter.class_name }}

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

@@ -0,0 +1,149 @@
+import requests
+import json
+
+from lmnr_engine.engine.action import NodeRunError, RunOutput
+from lmnr_engine.types import ChatMessage, NodeInput
+
+
+{% for task in cookiecutter._tasks.values() %}
+{% if task.node_type == "LLM" %}
+def {{task.function_name}}({{ task.handle_args }}, _env: dict[str, str]) -> RunOutput:
+    {% set chat_messages_found = false %}
+    {% for input_handle_name in task.input_handle_names %}
+    {% if input_handle_name == 'chat_messages' %}
+    {% set chat_messages_found = true %}
+    {% endif %}
+    {% endfor %}
+
+    {% if chat_messages_found %}
+    input_chat_messages = chat_messages
+    {% else %}
+    input_chat_messages = []
+    {% endif %}
+
+    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.
+    rendered_prompt = rendered_prompt.replace("{{'{{'}}{{prompt_variable}}{{'}}'}}", {{prompt_variable}})  # type: ignore
+    {% endfor %}
+
+    {% if task.config.model_params == none %}
+    params = {}
+    {% else %}
+    params = json.loads(
+        """{{task.config.model_params}}"""
+    )
+    {% endif %}
+
+    messages = [ChatMessage(role="system", content=rendered_prompt)]
+    messages.extend(input_chat_messages)
+
+    {% if task.config.provider == "openai" %}
+    message_jsons = [
+        {"role": message.role, "content": message.content} for message in messages
+    ]
+
+    data = {
+        "model": "{{task.config.model}}",
+        "messages": message_jsons,
+    }
+    data.update(params)
+
+    headers = {
+        "Content-Type": "application/json",
+        "Authorization": f"Bearer {_env['OPENAI_API_KEY']}",
+    }
+    res = requests.post(
+        "https://api.openai.com/v1/chat/completions", json=data, headers=headers
+    )
+
+    if res.status_code != 200:
+        res_json = res.json()
+        raise NodeRunError(f'OpenAI completions request failed: {res_json["error"]["message"]}')
+
+    chat_completion = res.json()
+
+    completion_message = chat_completion["choices"][0]["message"]["content"]
+
+    meta_log = {}
+    meta_log["node_chunk_id"] = None  # TODO: Add node chunk id
+    meta_log["model"] = "{{task.config.model}}"
+    meta_log["prompt"] = rendered_prompt
+    meta_log["input_message_count"] = len(messages)
+    meta_log["input_token_count"] = chat_completion["usage"]["prompt_tokens"]
+    meta_log["output_token_count"] = chat_completion["usage"]["completion_tokens"]
+    meta_log["total_token_count"] = (
+        chat_completion["usage"]["prompt_tokens"] + chat_completion["usage"]["completion_tokens"]
+    )
+    meta_log["approximate_cost"] = None  # TODO: Add approximate cost
+    {% elif task.config.provider == "anthropic" %}
+    data = {
+        "model": "{{task.config.model}}",
+        "max_tokens": 4096,
+    }
+    data.update(params)
+
+    # 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 = [
+            {"role": message.role, "content": message.content} for message in messages
+        ]
+        data["messages"] = message_jsons
+    else:
+        data["system"] = messages[0].content
+        message_jsons = [
+            {"role": message.role, "content": message.content} for message in messages[1:]
+        ]
+        data["messages"] = message_jsons
+
+    headers = {
+        "Content-Type": "application/json",
+        "X-Api-Key": _env['ANTHROPIC_API_KEY'],
+        "Anthropic-Version": "2023-06-01",
+    }
+    res = requests.post(
+        "https://api.anthropic.com/v1/messages", json=data, headers=headers
+    )
+
+    if res.status_code != 200:
+        raise NodeRunError(f"Anthropic message request failed: {res.text}")
+
+    chat_completion = res.json()
+
+    completion_message = chat_completion["content"][0]["text"]
+
+    meta_log = {}
+    meta_log["node_chunk_id"] = None  # TODO: Add node chunk id
+    meta_log["model"] = "{{task.config.model}}"
+    meta_log["prompt"] = rendered_prompt
+    meta_log["input_message_count"] = len(messages)
+    meta_log["input_token_count"] = chat_completion["usage"]["input_tokens"]
+    meta_log["output_token_count"] = chat_completion["usage"]["output_tokens"]
+    meta_log["total_token_count"] = (
+        chat_completion["usage"]["input_tokens"] + chat_completion["usage"]["output_tokens"]
+    )
+    meta_log["approximate_cost"] = None  # TODO: Add approximate cost
+    {% else %}
+    {% endif %}
+
+    return RunOutput(status="Success", output=completion_message)
+
+
+{% elif task.node_type == "Output" %}
+def {{task.function_name}}(output: NodeInput, _env: dict[str, str]) -> RunOutput:
+    return RunOutput(status="Success", output=output)
+
+
+{% elif task.node_type == "Input" %}
+{# Do nothing for Input tasks #}
+{% else %}
+def {{task.function_name}}(output: NodeInput, _env: dict[str, str]) -> RunOutput:
+    return RunOutput(status="Success", output=output)
+
+
+{% endif %}
+{% endfor %}
+# Other functions can be added here

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

@@ -0,0 +1,87 @@
+from concurrent.futures import ThreadPoolExecutor
+from dataclasses import dataclass
+import logging
+from typing import Optional, Union
+
+from lmnr_engine.types import ChatMessageList, ChatMessage
+from lmnr_engine.engine import Engine
+{% set function_names = cookiecutter._tasks.values() | selectattr('node_type', '!=', 'Input') | map(attribute='function_name') | join(', ') %}
+from .nodes.functions import {{ function_names }}
+from lmnr_engine.engine.task import Task
+
+
+logger = logging.getLogger(__name__)
+
+
+class PipelineRunnerError(Exception):
+    pass
+
+
+@dataclass
+class PipelineRunOutput:
+    value: Union[str, ChatMessageList]
+
+
+# This class is not imported in other files and can be renamed to desired name
+class {{ cookiecutter.class_name }}:
+    thread_pool_executor: ThreadPoolExecutor
+
+    def __init__(
+        self, thread_pool_executor: Optional[ThreadPoolExecutor] = None
+    ) -> None:
+        # Set max workers to hard-coded value for now
+        self.thread_pool_executor = (
+            ThreadPoolExecutor(max_workers=10)
+            if thread_pool_executor is None
+            else thread_pool_executor
+        )
+
+    def run(
+        self,
+        inputs: dict[str, Union[str, list]],
+        env: dict[str, str] = {},
+    ) -> dict[str, PipelineRunOutput]:
+        """
+        Run the pipeline with the given graph
+
+        Raises:
+            PipelineRunnerError: if there is an error running the pipeline
+        """
+        logger.info("Running pipeline {{ cookiecutter.pipeline_name }}, pipeline_version: {{ cookiecutter.pipeline_version_name }}")
+
+        run_inputs = {}
+        for inp_name, inp in inputs.items():
+            if isinstance(inp, str):
+                run_inputs[inp_name] = inp
+            else:
+                assert isinstance(inp, list), f"Invalid input type: {type(inp)}"
+                run_inputs[inp_name] = [ChatMessage.from_dict(msg) for msg in inp]
+
+        tasks = []
+        {% for task in cookiecutter._tasks.values() %}
+        tasks.append(
+            Task(
+                name="{{ task.name }}",
+                value={{ "''" if task.node_type == "Input" else task.function_name }},
+                handles_mapping={{ task.handles_mapping }},
+                prev=[
+                    {% for prev in task.prev %}
+                    "{{ prev }}",
+                    {% endfor %}
+                ],
+                next=[
+                    {% for next in task.next %}
+                    "{{ next }}",
+                    {% endfor %}
+                ],
+            )
+        )
+        {% endfor %}
+        engine = Engine.with_tasks(tasks, self.thread_pool_executor, env=env)
+
+        # TODO: Raise PipelineRunnerError with node_errors
+        run_res = engine.run(run_inputs)
+        return {
+            name: PipelineRunOutput(value=output.value)
+            for name, output in run_res.items()
+        }

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

@@ -0,0 +1,50 @@
+from dataclasses import dataclass
+from typing import Union
+import uuid
+import datetime
+
+
+@dataclass
+class ChatMessage:
+    role: str
+    content: str
+
+    @classmethod
+    def from_dict(cls, data: dict) -> "ChatMessage":
+        return cls(
+            role=data["role"],
+            content=data["content"],
+        )
+
+
+NodeInput = Union[str, list[ChatMessage]]  # TODO: Add conditioned value
+
+
+@dataclass
+class Message:
+    id: uuid.UUID
+    # output value of producing node in form of NodeInput
+    # for the following consumer
+    value: NodeInput
+    # all input messages to this node; accumulates previous messages too
+    # input_messages: list["Message"]
+    start_time: datetime.datetime
+    end_time: datetime.datetime
+    # node_id: uuid.UUID
+    # node_name: str
+    # node_type: str
+    # all node per-run metadata that needs to be logged at the end of execution
+    # meta_log: MetaLog | None
+
+    @classmethod
+    def empty(cls) -> "Message":
+        return cls(
+            id=uuid.uuid4(),
+            value="",
+            # input_messages=[],
+            start_time=datetime.datetime.now(),
+            end_time=datetime.datetime.now(),
+            # node_id=uuid.uuid4(),
+            # node_name="",
+            # node_type="",
+        )

lmnr/sdk/endpoint.py

@@ -0,0 +1,166 @@
+import json
+from pydantic.alias_generators import to_snake
+import pydantic
+import requests
+from ..types import (
+    EndpointRunError, EndpointRunResponse, NodeInput, EndpointRunRequest,
+    ToolCall, SDKError
+)
+from typing import Callable, Optional
+from websockets.sync.client import connect
+
+class Laminar:
+    project_api_key: Optional[str] = None
+    def __init__(self, project_api_key: str):
+        """Initialize the Laminar object with your project API key
+
+        Args:
+            project_api_key (str):
+                Project api key. Generate or view your keys
+                in the project settings in the Laminar dashboard.
+        """        
+        self.project_api_key = project_api_key
+        self.url = 'https://api.lmnr.ai/v2/endpoint/run'
+        self.ws_url = 'ws://api.lmnr.ai/v2/endpoint/ws'
+    
+    def run (
+        self,
+        endpoint: str,
+        inputs: dict[str, NodeInput],
+        env: dict[str, str] = {},
+        metadata: dict[str, str] = {},
+        tools: list[Callable[..., NodeInput]] = [],
+    ) -> EndpointRunResponse:
+        """Runs the endpoint with the given inputs
+
+        Args:
+            endpoint (str): name of the Laminar endpoint
+            inputs (dict[str, NodeInput]):
+                inputs to the endpoint's target pipeline.
+                Keys in the dictionary must match input node names
+            env (dict[str, str], optional):
+                Environment variables for the pipeline execution.
+                Defaults to {}.
+            metadata (dict[str, str], optional):
+                any custom metadata to be stored
+                with execution trace. Defaults to {}.
+            tools (list[Callable[..., NodeInput]], optional):
+                List of callable functions the execution can call as tools.
+                If specified and non-empty, a bidirectional communication
+                with Laminar API through websocket will be established.
+                Defaults to [].
+
+        Returns:
+            EndpointRunResponse: response object containing the outputs
+        
+        Raises:
+            ValueError: if project API key is not set
+            EndpointRunError: if the endpoint run fails
+            SDKError: if an error occurs on client side during the execution
+        """        
+        if self.project_api_key is None:
+            raise ValueError(
+                'Please initialize the Laminar object with' 
+                ' your project API key'
+            )
+        if tools:
+            return self._run_websocket(endpoint, inputs, env, metadata, tools)
+        return self._run(endpoint, inputs, env, metadata)
+    
+    def _run(
+        self,
+        endpoint: str,
+        inputs: dict[str, NodeInput],
+        env: dict[str, str] = {},
+        metadata: dict[str, str] = {}
+    ) -> EndpointRunResponse:
+        try:
+            request = EndpointRunRequest(
+                inputs = inputs,
+                endpoint = endpoint,
+                env = env,
+                metadata = metadata
+            )
+        except Exception as e:
+            raise ValueError(f'Invalid request: {e}')
+        response = requests.post(
+            self.url, 
+            json=json.loads(request.model_dump_json()),
+            headers={'Authorization': f'Bearer {self.project_api_key}'}
+        )
+        if response.status_code != 200:
+            raise EndpointRunError(response)
+        try:
+            resp_json = response.json()
+            keys = list(resp_json.keys())
+            for key in keys:
+                value = resp_json[key]
+                del resp_json[key]
+                resp_json[to_snake(key)] = value
+            return EndpointRunResponse(**resp_json)
+        except:
+            raise EndpointRunError(response)
+
+    def _run_websocket(
+        self,
+        endpoint: str,
+        inputs: dict[str, NodeInput],
+        env: dict[str, str] = {},
+        metadata: dict[str, str] = {},
+        tools: list[Callable[..., NodeInput]] = [],
+    ) -> EndpointRunResponse:
+        try:
+            request = EndpointRunRequest(
+                inputs = inputs,
+                endpoint = endpoint,
+                env = env,
+                metadata = metadata
+            )
+        except Exception as e:
+            raise ValueError(f'Invalid request: {e}')
+
+        with connect(
+            "ws://localhost:8000/v2/endpoint/ws",
+            additional_headers={
+                'Authorization': f'Bearer {self.project_api_key}'
+            }
+        ) as websocket:
+            websocket.send(request.model_dump_json())
+
+            while True:
+                message = websocket.recv()
+                try:
+                    tool_call = ToolCall.model_validate_json(message)
+                    matching_tools = [
+                        tool for tool in tools
+                        if tool.__name__ == tool_call.function.name
+                    ]
+                    if not matching_tools:
+                        raise SDKError(
+                            f'Tool {tool_call.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
+                        arguments = {}
+                        try:
+                            arguments = json.loads(tool_call.function.arguments)
+                        except:
+                            pass
+                        response = tool(**arguments)
+                        websocket.send(json.dumps(response))
+                except pydantic.ValidationError as e:
+                    message_json = json.loads(message)
+                    keys = list(message_json.keys())
+                    for key in keys:
+                        value = message_json[key]
+                        del message_json[key]
+                        message_json[to_snake(key)] = value
+                    result = EndpointRunResponse.model_validate(message_json)
+                    websocket.close()
+                    return result
+                except Exception:
+                    websocket.close()
+                    raise SDKError('Error communicating to backend through websocket')

lmnr/types.py

@@ -0,0 +1,93 @@
+import requests
+import pydantic
+from typing import Any, Union, Optional
+
+class ChatMessage(pydantic.BaseModel):
+    """Chat message object
+
+    Attributes:
+        role (str):
+            Role of the message sender.
+            Can be 'user', 'assistant', or 'system'.
+        content (str): Message content
+    """    
+    role: str
+    content: str
+
+"""NodeInput is a common type for values shared 
+between nodes in Laminar pipelines."""
+NodeInput = Union[str, list[ChatMessage]] # TypeAlias
+
+class EndpointRunRequest(pydantic.BaseModel):
+    inputs: dict[str, NodeInput]
+    endpoint: str
+    env: dict[str, str] = pydantic.Field(default_factory=dict)
+    metadata: dict[str, str] = pydantic.Field(default_factory=dict)
+
+class EndpointRunResponse(pydantic.BaseModel):
+    """Response object from endpoint run
+
+    Attributes:
+        outputs (dict[str, dict[str, NodeInput]]):
+            Dictionary of output names and their values.
+            Each value is a dictionary with the following keys:
+            - 'value': Output value
+        run_id (str): Stringified UUID of the run. Useful to find traces.
+    """    
+    outputs: dict[str, dict[str, NodeInput]]
+    run_id: str
+
+    def get_output(self, output_name: str) -> Optional[NodeInput]:
+        """utility to extract the output value by node name
+
+        Args:
+            output_name (str): must match the output node name
+
+        Returns:
+            Optional[NodeInput]: value of the output node if found
+        """        
+        output = self.outputs.get(output_name)
+        return output.get('value') if output else None
+
+class EndpointRunError(Exception):
+    error_code: str
+    error_message: str
+
+    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']
+            super().__init__(self.error_message)
+        except:
+            super().__init__(response.text)
+    
+    def __str__(self) -> str:
+        try:
+            return str({
+                'error_code': self.error_code,
+                'error_message': self.error_message
+            })
+        except:
+            return super().__str__()
+
+class SDKError(Exception):
+    error_message: str
+
+    def __init__(self, error_message: str):
+        self.error_message = error_message
+        super().__init__(self.error_message)
+    
+    def __str__(self) -> str:
+        return super().__str__()
+
+class ToolFunctionCall(pydantic.BaseModel):
+    name: str
+    arguments: str
+
+class ToolCall(pydantic.BaseModel):
+    id: str
+    type: str
+    function: ToolFunctionCall
+
+ToolResponse = NodeInput # TypeAlias