lmnr 0.6.16__py3-none-any.whl → 0.7.26__py3-none-any.whl

lmnr/__init__.py

@@ -3,29 +3,21 @@ from .sdk.client.asynchronous.async_client import AsyncLaminarClient
 from .sdk.datasets import EvaluationDataset, LaminarDataset
 from .sdk.evaluations import evaluate
 from .sdk.laminar import Laminar
-from .sdk.types import (
-    AgentOutput,
-    FinalOutputChunkContent,
-    HumanEvaluator,
-    RunAgentResponseChunk,
-    StepChunkContent,
-    TracingLevel,
-)
+from .sdk.types import SessionRecordingOptions, MaskInputOptions
+from .sdk.types import HumanEvaluator
 from .sdk.decorators import observe
 from .sdk.types import LaminarSpanContext
 from .opentelemetry_lib.litellm import LaminarLiteLLMCallback
 from .opentelemetry_lib.tracing.attributes import Attributes
 from .opentelemetry_lib.tracing.instruments import Instruments
 from .opentelemetry_lib.tracing.processor import LaminarSpanProcessor
+from .opentelemetry_lib.tracing.span import LaminarSpan
 from .opentelemetry_lib.tracing.tracer import get_laminar_tracer_provider, get_tracer
-from opentelemetry.trace import use_span
 
 __all__ = [
-    "AgentOutput",
     "AsyncLaminarClient",
     "Attributes",
     "EvaluationDataset",
-    "FinalOutputChunkContent",
     "HumanEvaluator",
     "Instruments",
     "Laminar",
@@ -34,12 +26,11 @@ __all__ = [
     "LaminarLiteLLMCallback",
     "LaminarSpanContext",
     "LaminarSpanProcessor",
-    "RunAgentResponseChunk",
-    "StepChunkContent",
-    "TracingLevel",
+    "LaminarSpan",
     "get_laminar_tracer_provider",
     "get_tracer",
     "evaluate",
     "observe",
-    "use_span",
+    "SessionRecordingOptions",
+    "MaskInputOptions",
 ]

lmnr/cli/__init__.py

@@ -0,0 +1,270 @@
+from argparse import ArgumentParser, _SubParsersAction
+import asyncio
+
+from lmnr.cli.datasets import handle_datasets_command
+from lmnr.cli.evals import run_evaluation
+from lmnr.cli.rules import add_cursor_rules
+from lmnr.sdk.log import get_default_logger
+from lmnr.sdk.utils import from_env
+
+LOG = get_default_logger(__name__)
+EVAL_DIR = "evals"
+DEFAULT_DATASET_PULL_BATCH_SIZE = 100
+DEFAULT_DATASET_PUSH_BATCH_SIZE = 100
+
+
+def setup_eval_parser(subparsers: _SubParsersAction) -> None:
+    """Setup the eval subcommand parser."""
+    parser_eval = subparsers.add_parser(
+        "eval",
+        description="Run an evaluation",
+        help="Run an evaluation",
+    )
+    parser_eval.add_argument(
+        "file",
+        nargs="*",
+        help="Files or a file containing the evaluation to run. "
+        + "If no file name is provided, all evaluation files in the `evals` directory are run as long "
+        + "as they match *_eval.py or eval_*.py",
+        default=[],
+    )
+    parser_eval.add_argument(
+        "--continue-on-error",
+        action="store_true",
+        default=False,
+        help="Continue execution upon errors",
+    )
+    parser_eval.add_argument(
+        "--output-file",
+        help="Output file to write the results to. Outputs are written in JSON format.",
+        nargs="?",
+    )
+
+
+def setup_add_cursor_rules_parser(subparsers: _SubParsersAction) -> None:
+    """Setup the add-cursor-rules subcommand parser."""
+    subparsers.add_parser(
+        "add-cursor-rules",
+        description="Download laminar.mdc file and add it to .cursor/rules",
+        help="Download laminar.mdc file and add it to .cursor/rules",
+    )
+
+
+def setup_laminar_args(parser: ArgumentParser) -> None:
+    """Setup the laminar arguments parser."""
+    parser.add_argument(
+        "--project-api-key",
+        help="[Optional] Project API key to use for the command. "
+        + "If no project API key is provided, the project API key will be read "
+        + "from the environment variable LMNR_PROJECT_API_KEY.",
+        default=from_env("LMNR_PROJECT_API_KEY"),
+    )
+    parser.add_argument(
+        "--base-url",
+        help="[Optional] Base URL to use for the command. "
+        + "If no base URL is provided, the base URL will be read from the "
+        + "'LMNR_BASE_URL' environment variable or we default to 'https://api.lmnr.ai'.",
+        default=from_env("LMNR_BASE_URL") or "https://api.lmnr.ai",
+    )
+    parser.add_argument(
+        "--port",
+        help="[Optional] Port to use for the command. "
+        + "If no port is provided, the port defaults to '443'.",
+        type=int,
+    )
+
+
+def setup_datasets_list_parser(subparsers: _SubParsersAction) -> None:
+    """Setup the datasets list subcommand parser."""
+    subparsers.add_parser(
+        "list",
+        description="List datasets",
+        help="List datasets",
+    )
+
+
+def setup_datasets_push_parser(subparsers: _SubParsersAction) -> None:
+    """Setup the datasets push subcommand parser."""
+    parser_datasets_push: ArgumentParser = subparsers.add_parser(
+        "push",
+        description="Push datapoints to an existing dataset",
+        help="Push datapoints to an existing dataset",
+    )
+    parser_datasets_push.add_argument(
+        "--name",
+        "-n",
+        help="Name of the dataset to push data to. Exactly one of name or id must be provided.",
+        default=None,
+    )
+    parser_datasets_push.add_argument(
+        "--id",
+        help="ID of the dataset to push data to. Exactly one of name or id must be provided.",
+        default=None,
+    )
+    parser_datasets_push.add_argument(
+        "paths",
+        nargs="*",
+        help="Paths to the files or directories containing the data to push to the dataset. "
+        + "Supported formats: JSON, CSV, JSONL",
+    )
+    parser_datasets_push.add_argument(
+        "-r",
+        "--recursive",
+        action="store_true",
+        default=False,
+        help="Recursively read all files in the directories and their subdirectories.",
+    )
+    parser_datasets_push.add_argument(
+        "--batch-size",
+        type=int,
+        help="Batch size to push data in. If no batch size is provided, "
+        + f"data is pushed in batches of '{DEFAULT_DATASET_PUSH_BATCH_SIZE}'.",
+        default=DEFAULT_DATASET_PUSH_BATCH_SIZE,
+    )
+
+
+def setup_datasets_pull_parser(subparsers: _SubParsersAction) -> None:
+    """Setup the datasets pull subcommand parser."""
+    parser_datasets_pull: ArgumentParser = subparsers.add_parser(
+        "pull",
+        description="Pull data from a dataset",
+        help="Pull data from a dataset",
+    )
+    parser_datasets_pull.add_argument(
+        "--name",
+        "-n",
+        help="Name of the dataset to pull data from",
+        default=None,
+    )
+    parser_datasets_pull.add_argument(
+        "--id",
+        help="ID of the dataset to pull data from",
+        default=None,
+    )
+    parser_datasets_pull.add_argument(
+        "output_path",
+        help="Path to the file to save the data to. "
+        + "If no path is provided, data is printed to the console in the format "
+        + "specified by '--output-format'.",
+        nargs="?",
+    )
+    parser_datasets_pull.add_argument(
+        "--output-format",
+        choices=["json", "csv", "jsonl"],
+        help="Output format to save the data to. "
+        + "If no format is provided, it is inferred from the file extension.",
+    )
+    parser_datasets_pull.add_argument(
+        "--batch-size",
+        type=int,
+        help="Batch size to pull data in. If no batch size is provided, "
+        + f"data is pulled in batches of '{DEFAULT_DATASET_PULL_BATCH_SIZE}'.",
+        default=DEFAULT_DATASET_PULL_BATCH_SIZE,
+    )
+    parser_datasets_pull.add_argument(
+        "--limit",
+        type=int,
+        help="Limit the number of data points to pull. "
+        + "If no limit is provided, all data points are pulled.",
+    )
+    parser_datasets_pull.add_argument(
+        "--offset",
+        type=int,
+        help="Offset the number of data points to pull. "
+        + "If no offset is provided, data is pulled from the beginning.",
+    )
+
+
+def setup_datasets_create_parser(subparsers: _SubParsersAction) -> None:
+    """Setup the datasets create subcommand parser."""
+    parser_datasets_create: ArgumentParser = subparsers.add_parser(
+        "create",
+        description="Create a dataset from input files and download it in Laminar format",
+        help="Create a dataset from input files and download it in Laminar format",
+    )
+    parser_datasets_create.add_argument(
+        "name",
+        help="Name of the dataset to create",
+    )
+    parser_datasets_create.add_argument(
+        "paths",
+        nargs="+",
+        help="Paths to the files or directories containing the data to push to the dataset. "
+        + "Supported formats: JSON, CSV, JSONL",
+    )
+    parser_datasets_create.add_argument(
+        "-o",
+        "--output-file",
+        required=True,
+        help="Path to the file to save the pulled data to",
+    )
+    parser_datasets_create.add_argument(
+        "--output-format",
+        choices=["json", "csv", "jsonl"],
+        help="Output format to save the data to. "
+        + "If no format is provided, it is inferred from the output file extension.",
+    )
+    parser_datasets_create.add_argument(
+        "-r",
+        "--recursive",
+        action="store_true",
+        default=False,
+        help="Recursively read all files in the directories and their subdirectories.",
+    )
+    parser_datasets_create.add_argument(
+        "--batch-size",
+        type=int,
+        help="Batch size to push/pull data in. If no batch size is provided, "
+        + f"data is processed in batches of '{DEFAULT_DATASET_PUSH_BATCH_SIZE}'.",
+        default=DEFAULT_DATASET_PUSH_BATCH_SIZE,
+    )
+
+
+def setup_datasets_parser(subparsers: _SubParsersAction) -> None:
+    """Setup the datasets subcommand parser and its subcommands."""
+    parser_datasets: ArgumentParser = subparsers.add_parser(
+        "datasets",
+        description="Manage datasets",
+        help="Manage datasets",
+    )
+
+    setup_laminar_args(parser_datasets)
+
+    parser_datasets_subparsers = parser_datasets.add_subparsers(
+        title="command",
+        dest="command",
+    )
+
+    # Setup all dataset subcommands
+    setup_datasets_list_parser(parser_datasets_subparsers)
+    setup_datasets_push_parser(parser_datasets_subparsers)
+    setup_datasets_pull_parser(parser_datasets_subparsers)
+    setup_datasets_create_parser(parser_datasets_subparsers)
+
+
+def cli() -> None:
+    """Main CLI entry point."""
+    parser = ArgumentParser(
+        prog="lmnr",
+        description="CLI for Laminar. "
+        + "Call `lmnr [subcommand] --help` for more information on each subcommand.",
+    )
+
+    subparsers = parser.add_subparsers(title="subcommands", dest="subcommand")
+
+    # Setup all subcommand parsers
+    setup_eval_parser(subparsers)
+    setup_add_cursor_rules_parser(subparsers)
+    setup_datasets_parser(subparsers)
+
+    # Parse arguments and dispatch to appropriate handler
+    parsed = parser.parse_args()
+
+    if parsed.subcommand == "eval":
+        asyncio.run(run_evaluation(parsed))
+    elif parsed.subcommand == "add-cursor-rules":
+        add_cursor_rules()
+    elif parsed.subcommand == "datasets":
+        asyncio.run(handle_datasets_command(parsed))
+    else:
+        parser.print_help()

lmnr/cli/datasets.py

@@ -0,0 +1,371 @@
+from argparse import Namespace
+from pathlib import Path
+from typing import Any
+
+import csv
+import sys
+
+import orjson
+
+from lmnr.sdk.client.asynchronous.async_client import AsyncLaminarClient
+from lmnr.sdk.datasets.file_utils import load_from_paths, parse_paths
+from lmnr.sdk.log import get_default_logger
+from lmnr.sdk.types import Datapoint
+
+LOG = get_default_logger(__name__, verbose=False)
+DEFAULT_DATASET_PULL_BATCH_SIZE = 100
+DEFAULT_DATASET_PUSH_BATCH_SIZE = 100
+
+
+def _dump_json(data: Any, do_indent: bool = True) -> str:
+    return orjson.dumps(
+        data,
+        option=(orjson.OPT_INDENT_2 if do_indent else 0)
+        | orjson.OPT_SERIALIZE_DATACLASS
+        | orjson.OPT_SERIALIZE_UUID
+        | orjson.OPT_UTC_Z
+        | orjson.OPT_NON_STR_KEYS,
+    ).decode()
+
+
+async def _pull_all_data(
+    client: AsyncLaminarClient,
+    identifier: dict,
+    batch_size: int,
+    offset: int = 0,
+    limit: int | None = None,
+) -> list[Datapoint]:
+    """
+    Pull all data from a dataset. This function does not close the client.
+
+    Args:
+        client: The AsyncLaminarClient instance
+        identifier: Dict with either 'name' or 'id' key
+        batch_size: Size of batches to pull
+        offset: Starting offset for pulling data
+        limit: Maximum number of items to pull (None for all)
+
+    Returns:
+        List of all pulled datapoints
+    """
+    has_more = True
+    current_offset = offset
+    stop_at = offset + limit if limit else None
+
+    result: list[Datapoint] = []
+    while has_more and (stop_at is None or current_offset < stop_at):
+        data = await client.datasets.pull(
+            **identifier,
+            offset=current_offset,
+            limit=batch_size,
+        )
+        result.extend(data.items)
+        if stop_at is not None and current_offset + batch_size >= stop_at:
+            has_more = False
+        elif (
+            data.total_count is not None
+            and current_offset + batch_size >= data.total_count
+        ):
+            has_more = False
+        current_offset += batch_size
+
+    if limit is not None:
+        return result[:limit]
+    return result
+
+
+def _write_data_to_file(
+    data: list[Datapoint],
+    output_path: Path,
+    output_format: str | None = None,
+) -> bool:
+    """
+    Write datapoints to a file.
+
+    Args:
+        data: List of datapoints to write
+        output_path: Path to write the file to
+        output_format: Optional format override (json, csv, jsonl)
+
+    Returns:
+        True if successful, False otherwise
+    """
+    if output_path.is_dir():
+        LOG.error(f"Output path is a directory: {output_path}")
+        return False
+
+    # Create parent directories if they don't exist
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+
+    # Determine output format
+    format = output_format or output_path.suffix[1:]
+    if output_format and output_format != output_path.suffix[1:]:
+        LOG.warning(
+            f"Output format {output_format} does not match file extension {output_path.suffix[1:]}"
+        )
+
+    if format not in ["json", "csv", "jsonl"]:
+        LOG.error(f"Unsupported output format: {format}")
+        return False
+
+    # Write output file
+    if format == "json":
+        output_path.write_text(_dump_json([item.model_dump() for item in data]))
+    elif format == "csv":
+        if not data:
+            LOG.error("No data to write to CSV")
+            return False
+        with output_path.open("w") as f:
+            writer = csv.writer(f)
+            keys = list(data[0].model_dump().keys())
+            writer.writerow(keys)
+            for item in data:
+                writer.writerow([item.model_dump()[key] for key in keys])
+    elif format == "jsonl":
+        with output_path.open("w") as f:
+            for item in data:
+                f.write(_dump_json(item.model_dump(), do_indent=False) + "\n")
+
+    return True
+
+
+def _print_data_to_console(data: list[Datapoint], output_format: str = "json") -> bool:
+    """
+    Print datapoints to console.
+
+    Args:
+        data: List of datapoints to print
+        output_format: Format to use (json, csv, jsonl)
+
+    Returns:
+        True if successful, False otherwise
+    """
+    if output_format not in ["json", "csv", "jsonl"]:
+        LOG.error(f"Unsupported output format: {output_format}")
+        return False
+
+    if output_format == "json":
+        print(_dump_json([item.model_dump() for item in data]))
+    elif output_format == "csv":
+        if not data:
+            LOG.error("No data to print")
+            return False
+        writer = csv.writer(sys.stdout)
+        keys = list(data[0].model_dump().keys())
+        writer.writerow(keys)
+        for item in data:
+            writer.writerow([item.model_dump()[key] for key in keys])
+    elif output_format == "jsonl":
+        for item in data:
+            print(_dump_json(item.model_dump(), do_indent=False))
+    print()
+
+    return True
+
+
+async def handle_datasets_list(args: Namespace) -> None:
+    """
+    Handle datasets list command.
+
+    Lists all datasets in a formatted table.
+    """
+    client = AsyncLaminarClient(
+        project_api_key=args.project_api_key,
+        base_url=args.base_url,
+        port=args.port,
+    )
+
+    try:
+        datasets = await client.datasets.list_datasets()
+    except Exception as e:
+        LOG.error(f"Failed to list datasets: {e}")
+        return
+    finally:
+        await client.close()
+
+    if not datasets:
+        print("No datasets found.")
+        return
+
+    # Print table header
+    id_width = 36  # UUID length
+    created_at_width = 19  # YYYY-MM-DD HH:MM:SS format
+
+    print(f"\n{'ID':<{id_width}}  {'Created At':<{created_at_width}}  Name")
+    print(f"{'-' * id_width}  {'-' * created_at_width}  {'-' * 20}")
+
+    # Print each dataset row
+    for dataset in datasets:
+        created_at_str = dataset.created_at.strftime("%Y-%m-%d %H:%M:%S")
+        print(
+            f"{str(dataset.id):<{id_width}}  {created_at_str:<{created_at_width}}  {dataset.name}"
+        )
+
+    print(f"\nTotal: {len(datasets)} dataset(s)\n")
+
+
+async def handle_datasets_push(args: Namespace) -> None:
+    """
+    Handle datasets push command.
+
+    Pushes data from file(s) to an existing dataset.
+    """
+    if not args.name and not args.id:
+        LOG.error("Either name or id must be provided")
+        return
+    if args.name and args.id:
+        LOG.error("Only one of name or id must be provided")
+        return
+    identifier = {"name": args.name} if args.name else {"id": args.id}
+    client = AsyncLaminarClient(
+        project_api_key=args.project_api_key,
+        base_url=args.base_url,
+        port=args.port,
+    )
+    data = load_from_paths(parse_paths(args.paths), recursive=args.recursive)
+    if len(data) == 0:
+        LOG.warning("No data to push. Skipping")
+        return
+    try:
+        await client.datasets.push(
+            data,
+            **identifier,
+            batch_size=args.batch_size or DEFAULT_DATASET_PUSH_BATCH_SIZE,
+        )
+        LOG.info(f"Pushed {len(data)} data points to dataset {args.name or args.id}")
+    except Exception as e:
+        LOG.error(f"Failed to push dataset: {e}")
+    finally:
+        await client.close()
+
+
+async def handle_datasets_pull(args: Namespace) -> None:
+    """
+    Handle datasets pull command.
+
+    Pulls data from a dataset and saves it to a file.
+    """
+    if not args.name and not args.id:
+        LOG.error("Either name or id must be provided")
+        return
+    if args.name and args.id:
+        LOG.error("Only one of name or id must be provided")
+        return
+    identifier = {"name": args.name} if args.name else {"id": args.id}
+    client = AsyncLaminarClient(
+        project_api_key=args.project_api_key,
+        base_url=args.base_url,
+        port=args.port,
+    )
+
+    # Pull data from dataset
+    try:
+        result = await _pull_all_data(
+            client=client,
+            identifier=identifier,
+            batch_size=args.batch_size or DEFAULT_DATASET_PULL_BATCH_SIZE,
+            offset=args.offset or 0,
+            limit=args.limit,
+        )
+    except Exception as e:
+        LOG.error(f"Failed to pull dataset: {e}")
+        return
+    finally:
+        await client.close()
+
+    # Write to file or print to console
+    if args.output_path:
+        if not _write_data_to_file(
+            data=result,
+            output_path=Path(args.output_path),
+            output_format=args.output_format,
+        ):
+            return
+    else:
+        if not _print_data_to_console(
+            data=result,
+            output_format=args.output_format or "json",
+        ):
+            return
+
+
+async def handle_datasets_create(args: Namespace) -> None:
+    """
+    Handle datasets create command.
+
+    Creates a dataset from input files, pushes the data to it, and then pulls it back
+    in Laminar format to save to the output file.
+    """
+    client = AsyncLaminarClient(
+        project_api_key=args.project_api_key,
+        base_url=args.base_url,
+        port=args.port,
+    )
+
+    # Load data from input files
+    data = load_from_paths(parse_paths(args.paths), recursive=args.recursive)
+    if len(data) == 0:
+        LOG.warning("No data to push. Skipping")
+        return
+
+    # Push data to create/populate the dataset
+    LOG.info(f"Pushing {len(data)} data points to dataset '{args.name}'...")
+    try:
+        await client.datasets.push(
+            data,
+            name=args.name,
+            batch_size=args.batch_size or DEFAULT_DATASET_PUSH_BATCH_SIZE,
+            create_dataset=True,
+        )
+        LOG.info(
+            f"Successfully pushed {len(data)} data points to dataset '{args.name}'"
+        )
+    except Exception as e:
+        LOG.error(f"Failed to create dataset: {e}")
+        return
+
+    # Pull data back from the dataset
+    LOG.info(f"Pulling data from dataset '{args.name}'...")
+    try:
+        result = await _pull_all_data(
+            client=client,
+            identifier={"name": args.name},
+            batch_size=args.batch_size or DEFAULT_DATASET_PULL_BATCH_SIZE,
+            offset=0,
+            limit=None,
+        )
+    except Exception as e:
+        LOG.error(f"Failed to pull dataset after creation: {e}")
+        return
+    finally:
+        await client.close()
+
+    # Save to output file
+    if not _write_data_to_file(
+        data=result,
+        output_path=Path(args.output_file),
+        output_format=args.output_format,
+    ):
+        return
+
+    LOG.info(
+        f"Successfully created dataset '{args.name}' and saved {len(result)} data points to {args.output_file}"
+    )
+
+
+async def handle_datasets_command(args: Namespace) -> None:
+    """
+    Handle datasets subcommand dispatching.
+
+    Dispatches to the appropriate handler based on the command.
+    """
+    if args.command == "list":
+        await handle_datasets_list(args)
+    elif args.command == "push":
+        await handle_datasets_push(args)
+    elif args.command == "pull":
+        await handle_datasets_pull(args)
+    elif args.command == "create":
+        await handle_datasets_create(args)
+    else:
+        LOG.error(f"Unknown datasets command: {args.command}")