lightningrod-ai 0.1.6__py3-none-any.whl

lightningrod/__init__.py

@@ -0,0 +1,66 @@
+"""
+Lightning Rod Python SDK
+
+AI-powered forecasting dataset generation platform.
+"""
+
+from lightningrod.client import LightningRod
+from lightningrod.datasets.dataset import Dataset
+from lightningrod._generated.models import (
+    AnswerType,
+    AnswerTypeEnum,
+    TransformJob,
+    TransformJobStatus,
+    NewsSeedGenerator,
+    GdeltSeedGenerator,
+    NewsContextGenerator,
+    QuestionGenerator,
+    QuestionAndLabelGenerator,
+    ForwardLookingQuestionGenerator,
+    QuestionPipeline,
+    QuestionRenderer,
+    WebSearchLabeler,
+    FilterCriteria,
+    Sample,
+    SampleMeta,
+    Seed,
+    # TODO(filesets): Enable when filesets are publicly supported
+    # FileSetSeedGenerator,
+    # FileSetQuerySeedGenerator,
+    # CreateFileSetRequest,
+    # CreateFileSetFileRequest,
+    # CreateFileUploadResponse,
+    # FileSetFile,
+)
+
+__version__ = "0.1.6"
+__all__ = [
+    "AnswerType",
+    "AnswerTypeEnum",
+    "AnswerTypes",
+    "AsyncDataset",
+    "Dataset",
+    # TODO(filesets): Enable when filesets are publicly supported
+    # "FileSetSeedGenerator",
+    # "FileSetQuerySeedGenerator",
+    # "CreateFileSetRequest",
+    # "CreateFileSetFileRequest",
+    # "CreateFileUploadResponse",
+    # "FileSetFile",
+    "FilterCriteria",
+    "ForwardLookingQuestionGenerator",
+    "GdeltSeedGenerator",
+    "NewsContextGenerator",
+    "NewsSeedGenerator",
+    "QuestionAndLabelGenerator",
+    "QuestionGenerator",
+    "QuestionPipeline",
+    "QuestionRenderer",
+    "Sample",
+    "SampleMeta",
+    "Seed",
+    "TransformJob",
+    "TransformJobStatus",
+    "WebSearchLabeler",
+    "LightningRod",
+]

lightningrod/_display.py

@@ -0,0 +1,204 @@
+from __future__ import annotations
+
+from typing import Any, Optional
+
+from rich.console import Console, Group, RenderableType
+from rich.panel import Panel
+from rich.table import Table
+from rich.text import Text
+
+from lightningrod._generated.types import Unset
+
+
+def _is_set(value: Any) -> bool:
+    return not isinstance(value, Unset) and value is not None
+
+
+def _safe_markup(text: Optional[str]) -> Text:
+    """Parse text as rich markup, falling back to plain text if parsing fails."""
+    if text is None:
+        return Text("")
+    try:
+        return Text.from_markup(text)
+    except Exception:
+        return Text(text)
+
+
+def _format_duration(seconds: float) -> str:
+    if seconds < 60:
+        return f"{seconds:.0f}s"
+    minutes = int(seconds // 60)
+    secs = int(seconds % 60)
+    return f"{minutes}m {secs}s"
+
+
+def _build_cost_lines(job: Any) -> list[RenderableType]:
+    """Build cost info lines from job.usage. Returns empty list if no data."""
+    if not _is_set(job.usage):
+        return []
+    usage = job.usage
+    lines: list[RenderableType] = []
+
+    # Total cost
+    if _is_set(usage.current_cost_dollars):
+        lines.append(_safe_markup(f"  [bold]Total cost:[/bold] [bright_green]${usage.current_cost_dollars:.2f}[/bright_green]"))
+    if _is_set(usage.max_cost_dollars):
+        lines.append(_safe_markup(f"  [bold]Budget:[/bold]     ${usage.max_cost_dollars:.2f}"))
+    if _is_set(usage.estimated_cost_dollars):
+        lines.append(_safe_markup(f"  [bold]Estimated:[/bold]  ${usage.estimated_cost_dollars:.2f}"))
+
+    return lines
+
+
+def build_live_display(
+    metrics: Any = None,
+    job: Any = None,
+) -> RenderableType:
+    """Build the live display renderable for the polling loop."""
+    renderables: list[RenderableType] = []
+
+    renderables.append(_safe_markup("[bold bright_blue]>> Pipeline Running[/bold bright_blue]"))
+    renderables.append(Text(""))
+
+    # Cost summary from job.usage
+    if job is not None:
+        cost_lines = _build_cost_lines(job)
+        if cost_lines:
+            renderables.extend(cost_lines)
+            renderables.append(Text(""))
+
+    if metrics is None:
+        renderables.append(Text("Waiting for metrics...", style="dim italic"))
+        return Panel(
+            Group(*renderables),
+            border_style="bright_blue",
+            padding=(1, 2),
+        )
+
+    # Per-step table
+    table = Table(show_header=True, header_style="bold cyan", expand=True)
+    table.add_column("Step", style="bold", no_wrap=True)
+    table.add_column("Progress", width=20)
+    table.add_column("In", justify="right")
+    table.add_column("Out", justify="right")
+    table.add_column("Rejected", justify="right")
+    table.add_column("Errors", justify="right")
+    table.add_column("Duration", justify="right")
+
+    for step in sorted(metrics.steps, key=lambda s: s.step_index):
+        if step.progress >= 1.0:
+            status = Text("Complete", style="bold bright_green")
+        elif step.progress > 0:
+            status = Text("In progress", style="bold bright_yellow")
+        else:
+            status = Text("Pending", style="dim")
+
+        rejected_style = "bright_red" if step.rejected_count > 0 else "dim"
+        error_style = "bold bright_red" if step.error_count > 0 else "dim"
+
+        table.add_row(
+            step.transform_name,
+            status,
+            str(step.input_rows),
+            str(step.output_rows),
+            Text(str(step.rejected_count), style=rejected_style),
+            Text(str(step.error_count), style=error_style),
+            _format_duration(step.duration_seconds),
+        )
+
+    renderables.append(table)
+
+    return Panel(
+        Group(*renderables),
+        border_style="bright_blue",
+        padding=(1, 2),
+    )
+
+
+def _is_notebook() -> bool:
+    """Check if we're running inside a Jupyter notebook."""
+    try:
+        from IPython import get_ipython
+        shell = get_ipython()
+        return shell is not None and shell.__class__.__name__ == "ZMQInteractiveShell"
+    except ImportError:
+        return False
+
+
+def run_live_display(
+    poll_callback: Any,
+    poll_interval: float = 15,
+    warning_message: Optional[str] = None,
+) -> None:
+    """Run a live-updating display that polls for metrics.
+
+    Args:
+        poll_callback: A callable that returns (metrics, job, is_running) each cycle.
+            - metrics: PipelineMetricsResponse or None
+            - job: TransformJob with current status/usage
+            - is_running: bool, False to stop the loop
+        poll_interval: Seconds between polls.
+        warning_message: Optional warning to persist above the live display.
+    """
+    import time
+    console = Console()
+
+    if _is_notebook():
+        from IPython.display import clear_output
+        metrics, job, is_running = poll_callback()
+        while is_running:
+            clear_output(wait=True)
+            if warning_message:
+                display_warning(warning_message)
+            console.print(build_live_display(metrics=metrics, job=job))
+            time.sleep(poll_interval)
+            metrics, job, is_running = poll_callback()
+    else:
+        from rich.live import Live
+        with Live(
+            build_live_display(metrics=None, job=None),
+            console=console,
+            refresh_per_second=1,
+            transient=True,
+        ) as live:
+            metrics, job, is_running = poll_callback()
+            while is_running:
+                live.update(build_live_display(metrics=metrics, job=job))
+                time.sleep(poll_interval)
+                metrics, job, is_running = poll_callback()
+            # Final update
+            live.update(build_live_display(metrics=metrics, job=job))
+
+
+def display_error(message: str, title: str = "Error", job: Any = None) -> None:
+    console = Console()
+    renderables: list[RenderableType] = []
+
+    renderables.append(_safe_markup(f"[bold bright_red]>> {title}[/bold bright_red]"))
+    renderables.append(Text(""))
+    renderables.append(_safe_markup(f"[bold]{message}[/bold]"))
+
+    if job is not None:
+        cost_lines = _build_cost_lines(job)
+        if cost_lines:
+            renderables.append(Text(""))
+            renderables.extend(cost_lines)
+
+    console.print(Panel(Group(*renderables), border_style="bright_red", padding=(1, 2)))
+
+
+def display_warning(message: str, title: str = "Warning", job: Any = None) -> None:
+    console = Console()
+    renderables: list[RenderableType] = []
+
+    renderables.append(_safe_markup(f"[bold yellow]>> {title}[/bold yellow]"))
+    renderables.append(Text(""))
+    renderables.append(_safe_markup(message))
+
+    if job is not None:
+        cost_lines = _build_cost_lines(job)
+        if cost_lines:
+            renderables.append(Text(""))
+            renderables.extend(cost_lines)
+
+    console.print(Panel(Group(*renderables), border_style="yellow", padding=(1, 2)))

lightningrod/_errors.py

@@ -0,0 +1,67 @@
+import json
+from http import HTTPStatus
+from typing import Any, TypeVar
+
+from lightningrod._display import display_error
+from lightningrod._generated.models import HTTPValidationError
+from lightningrod._generated.types import Response
+
+T = TypeVar("T")
+
+
+def extract_error_message(response: Response[Any], operation: str) -> str:
+    """
+    Extract a detailed error message from a Response object.
+    
+    Args:
+        response: The Response object from sync_detailed
+        operation: Description of the operation that failed (e.g., "create dataset")
+        
+    Returns:
+        A detailed error message string
+    """
+    if isinstance(response.parsed, HTTPValidationError):
+        return f"Failed to {operation}: {response.parsed.detail}"
+    
+    if response.parsed is None:
+        status_code = response.status_code.value if isinstance(response.status_code, HTTPStatus) else response.status_code
+        
+        try:
+            error_data = json.loads(response.content.decode('utf-8'))
+            if isinstance(error_data, dict) and 'detail' in error_data:
+                detail = error_data['detail']
+                return f"Failed to {operation}: {detail} (HTTP {status_code})"
+            elif isinstance(error_data, dict):
+                return f"Failed to {operation}: {json.dumps(error_data)} (HTTP {status_code})"
+        except (json.JSONDecodeError, UnicodeDecodeError):
+            pass
+        
+        content_preview = response.content.decode('utf-8', errors='ignore')[:500]
+        if content_preview:
+            return f"Failed to {operation}: HTTP {status_code} - {content_preview}"
+        else:
+            return f"Failed to {operation}: HTTP {status_code} (no response body)"
+    
+    return f"Failed to {operation}: unexpected response format"
+
+
+def handle_response_error(response: Response[T], operation: str) -> T:
+    """
+    Validate a Response object and return the parsed response, raising an exception if there's an error.
+    
+    Args:
+        response: The Response object from sync_detailed
+        operation: Description of the operation that failed (e.g., "create dataset")
+        
+    Returns:
+        The parsed response object
+        
+    Raises:
+        Exception: If the response indicates an error (parsed is None or HTTPValidationError)
+    """
+    if response.parsed is None or isinstance(response.parsed, HTTPValidationError):
+        error_msg = extract_error_message(response, operation)
+        display_error(error_msg, title=f"API Error: {operation}")
+        raise Exception(error_msg)
+
+    return response.parsed

lightningrod/_generated/__init__.py

@@ -0,0 +1,8 @@
+"""A client library for accessing LightningRod API"""
+
+from .client import AuthenticatedClient, Client
+
+__all__ = (
+    "AuthenticatedClient",
+    "Client",
+)

lightningrod/_generated/api/__init__.py

@@ -0,0 +1 @@
+"""Contains methods for accessing the API"""

lightningrod/_generated/api/datasets/__init__.py

@@ -0,0 +1 @@
+"""Contains endpoint functions for accessing the API"""

lightningrod/_generated/api/datasets/create_dataset_datasets_post.py

@@ -0,0 +1,133 @@
+from http import HTTPStatus
+from typing import Any
+
+import httpx
+
+from ... import errors
+from ...client import AuthenticatedClient, Client
+from ...models.create_dataset_response import CreateDatasetResponse
+from ...types import Response
+
+
+def _get_kwargs() -> dict[str, Any]:
+    _kwargs: dict[str, Any] = {
+        "method": "post",
+        "url": "/datasets",
+    }
+
+    return _kwargs
+
+
+def _parse_response(*, client: AuthenticatedClient | Client, response: httpx.Response) -> CreateDatasetResponse | None:
+    if response.status_code == 201:
+        response_201 = CreateDatasetResponse.from_dict(response.json())
+
+        return response_201
+
+    if client.raise_on_unexpected_status:
+        raise errors.UnexpectedStatus(response.status_code, response.content)
+    else:
+        return None
+
+
+def _build_response(
+    *, client: AuthenticatedClient | Client, response: httpx.Response
+) -> Response[CreateDatasetResponse]:
+    return Response(
+        status_code=HTTPStatus(response.status_code),
+        content=response.content,
+        headers=response.headers,
+        parsed=_parse_response(client=client, response=response),
+    )
+
+
+def sync_detailed(
+    *,
+    client: AuthenticatedClient,
+) -> Response[CreateDatasetResponse]:
+    """Create Dataset
+
+     Create a new empty dataset
+
+    Raises:
+        errors.UnexpectedStatus: If the server returns an undocumented status code and Client.raise_on_unexpected_status is True.
+        httpx.TimeoutException: If the request takes longer than Client.timeout.
+
+    Returns:
+        Response[CreateDatasetResponse]
+    """
+
+    kwargs = _get_kwargs()
+
+    response = client.get_httpx_client().request(
+        **kwargs,
+    )
+
+    return _build_response(client=client, response=response)
+
+
+def sync(
+    *,
+    client: AuthenticatedClient,
+) -> CreateDatasetResponse | None:
+    """Create Dataset
+
+     Create a new empty dataset
+
+    Raises:
+        errors.UnexpectedStatus: If the server returns an undocumented status code and Client.raise_on_unexpected_status is True.
+        httpx.TimeoutException: If the request takes longer than Client.timeout.
+
+    Returns:
+        CreateDatasetResponse
+    """
+
+    return sync_detailed(
+        client=client,
+    ).parsed
+
+
+async def asyncio_detailed(
+    *,
+    client: AuthenticatedClient,
+) -> Response[CreateDatasetResponse]:
+    """Create Dataset
+
+     Create a new empty dataset
+
+    Raises:
+        errors.UnexpectedStatus: If the server returns an undocumented status code and Client.raise_on_unexpected_status is True.
+        httpx.TimeoutException: If the request takes longer than Client.timeout.
+
+    Returns:
+        Response[CreateDatasetResponse]
+    """
+
+    kwargs = _get_kwargs()
+
+    response = await client.get_async_httpx_client().request(**kwargs)
+
+    return _build_response(client=client, response=response)
+
+
+async def asyncio(
+    *,
+    client: AuthenticatedClient,
+) -> CreateDatasetResponse | None:
+    """Create Dataset
+
+     Create a new empty dataset
+
+    Raises:
+        errors.UnexpectedStatus: If the server returns an undocumented status code and Client.raise_on_unexpected_status is True.
+        httpx.TimeoutException: If the request takes longer than Client.timeout.
+
+    Returns:
+        CreateDatasetResponse
+    """
+
+    return (
+        await asyncio_detailed(
+            client=client,
+        )
+    ).parsed

lightningrod/_generated/api/datasets/get_dataset_datasets_dataset_id_get.py

@@ -0,0 +1,168 @@
+from http import HTTPStatus
+from typing import Any
+from urllib.parse import quote
+
+import httpx
+
+from ... import errors
+from ...client import AuthenticatedClient, Client
+from ...models.dataset_metadata import DatasetMetadata
+from ...models.http_validation_error import HTTPValidationError
+from ...types import Response
+
+
+def _get_kwargs(
+    dataset_id: str,
+) -> dict[str, Any]:
+    _kwargs: dict[str, Any] = {
+        "method": "get",
+        "url": "/datasets/{dataset_id}".format(
+            dataset_id=quote(str(dataset_id), safe=""),
+        ),
+    }
+
+    return _kwargs
+
+
+def _parse_response(
+    *, client: AuthenticatedClient | Client, response: httpx.Response
+) -> DatasetMetadata | HTTPValidationError | None:
+    if response.status_code == 200:
+        response_200 = DatasetMetadata.from_dict(response.json())
+
+        return response_200
+
+    if response.status_code == 422:
+        response_422 = HTTPValidationError.from_dict(response.json())
+
+        return response_422
+
+    if client.raise_on_unexpected_status:
+        raise errors.UnexpectedStatus(response.status_code, response.content)
+    else:
+        return None
+
+
+def _build_response(
+    *, client: AuthenticatedClient | Client, response: httpx.Response
+) -> Response[DatasetMetadata | HTTPValidationError]:
+    return Response(
+        status_code=HTTPStatus(response.status_code),
+        content=response.content,
+        headers=response.headers,
+        parsed=_parse_response(client=client, response=response),
+    )
+
+
+def sync_detailed(
+    dataset_id: str,
+    *,
+    client: AuthenticatedClient,
+) -> Response[DatasetMetadata | HTTPValidationError]:
+    """Get Dataset
+
+     Get dataset metadata including ID and row count
+
+    Args:
+        dataset_id (str):
+
+    Raises:
+        errors.UnexpectedStatus: If the server returns an undocumented status code and Client.raise_on_unexpected_status is True.
+        httpx.TimeoutException: If the request takes longer than Client.timeout.
+
+    Returns:
+        Response[DatasetMetadata | HTTPValidationError]
+    """
+
+    kwargs = _get_kwargs(
+        dataset_id=dataset_id,
+    )
+
+    response = client.get_httpx_client().request(
+        **kwargs,
+    )
+
+    return _build_response(client=client, response=response)
+
+
+def sync(
+    dataset_id: str,
+    *,
+    client: AuthenticatedClient,
+) -> DatasetMetadata | HTTPValidationError | None:
+    """Get Dataset
+
+     Get dataset metadata including ID and row count
+
+    Args:
+        dataset_id (str):
+
+    Raises:
+        errors.UnexpectedStatus: If the server returns an undocumented status code and Client.raise_on_unexpected_status is True.
+        httpx.TimeoutException: If the request takes longer than Client.timeout.
+
+    Returns:
+        DatasetMetadata | HTTPValidationError
+    """
+
+    return sync_detailed(
+        dataset_id=dataset_id,
+        client=client,
+    ).parsed
+
+
+async def asyncio_detailed(
+    dataset_id: str,
+    *,
+    client: AuthenticatedClient,
+) -> Response[DatasetMetadata | HTTPValidationError]:
+    """Get Dataset
+
+     Get dataset metadata including ID and row count
+
+    Args:
+        dataset_id (str):
+
+    Raises:
+        errors.UnexpectedStatus: If the server returns an undocumented status code and Client.raise_on_unexpected_status is True.
+        httpx.TimeoutException: If the request takes longer than Client.timeout.
+
+    Returns:
+        Response[DatasetMetadata | HTTPValidationError]
+    """
+
+    kwargs = _get_kwargs(
+        dataset_id=dataset_id,
+    )
+
+    response = await client.get_async_httpx_client().request(**kwargs)
+
+    return _build_response(client=client, response=response)
+
+
+async def asyncio(
+    dataset_id: str,
+    *,
+    client: AuthenticatedClient,
+) -> DatasetMetadata | HTTPValidationError | None:
+    """Get Dataset
+
+     Get dataset metadata including ID and row count
+
+    Args:
+        dataset_id (str):
+
+    Raises:
+        errors.UnexpectedStatus: If the server returns an undocumented status code and Client.raise_on_unexpected_status is True.
+        httpx.TimeoutException: If the request takes longer than Client.timeout.
+
+    Returns:
+        DatasetMetadata | HTTPValidationError
+    """
+
+    return (
+        await asyncio_detailed(
+            dataset_id=dataset_id,
+            client=client,
+        )
+    ).parsed