lightningrod-ai 0.1.6__py3-none-any.whl

lightningrod/_generated/models/web_search_labeler.py

@@ -0,0 +1,120 @@
+from __future__ import annotations
+
+from collections.abc import Mapping
+from typing import TYPE_CHECKING, Any, Literal, TypeVar, cast
+
+from attrs import define as _attrs_define
+from attrs import field as _attrs_field
+
+from ..types import UNSET, Unset
+
+if TYPE_CHECKING:
+    from ..models.answer_type import AnswerType
+
+
+T = TypeVar("T", bound="WebSearchLabeler")
+
+
+@_attrs_define
+class WebSearchLabeler:
+    """
+    Attributes:
+        config_type (Literal['WEB_SEARCH_LABELER'] | Unset): Type of transform configuration Default:
+            'WEB_SEARCH_LABELER'.
+        confidence_threshold (float | Unset): Minimum confidence threshold for including questions Default: 0.9.
+        answer_type (AnswerType | None | Unset): The type of answer expected, used to guide the labeler
+        resolve_redirects (bool | Unset): Resolve redirect URLs to actual destinations Default: False.
+    """
+
+    config_type: Literal["WEB_SEARCH_LABELER"] | Unset = "WEB_SEARCH_LABELER"
+    confidence_threshold: float | Unset = 0.9
+    answer_type: AnswerType | None | Unset = UNSET
+    resolve_redirects: bool | Unset = False
+    additional_properties: dict[str, Any] = _attrs_field(init=False, factory=dict)
+
+    def to_dict(self) -> dict[str, Any]:
+        from ..models.answer_type import AnswerType
+
+        config_type = self.config_type
+
+        confidence_threshold = self.confidence_threshold
+
+        answer_type: dict[str, Any] | None | Unset
+        if isinstance(self.answer_type, Unset):
+            answer_type = UNSET
+        elif isinstance(self.answer_type, AnswerType):
+            answer_type = self.answer_type.to_dict()
+        else:
+            answer_type = self.answer_type
+
+        resolve_redirects = self.resolve_redirects
+
+        field_dict: dict[str, Any] = {}
+        field_dict.update(self.additional_properties)
+        field_dict.update({})
+        if config_type is not UNSET:
+            field_dict["config_type"] = config_type
+        if confidence_threshold is not UNSET:
+            field_dict["confidence_threshold"] = confidence_threshold
+        if answer_type is not UNSET:
+            field_dict["answer_type"] = answer_type
+        if resolve_redirects is not UNSET:
+            field_dict["resolve_redirects"] = resolve_redirects
+
+        return field_dict
+
+    @classmethod
+    def from_dict(cls: type[T], src_dict: Mapping[str, Any]) -> T:
+        from ..models.answer_type import AnswerType
+
+        d = dict(src_dict)
+        config_type = cast(Literal["WEB_SEARCH_LABELER"] | Unset, d.pop("config_type", UNSET))
+        if config_type != "WEB_SEARCH_LABELER" and not isinstance(config_type, Unset):
+            raise ValueError(f"config_type must match const 'WEB_SEARCH_LABELER', got '{config_type}'")
+
+        confidence_threshold = d.pop("confidence_threshold", UNSET)
+
+        def _parse_answer_type(data: object) -> AnswerType | None | Unset:
+            if data is None:
+                return data
+            if isinstance(data, Unset):
+                return data
+            try:
+                if not isinstance(data, dict):
+                    raise TypeError()
+                answer_type_type_0 = AnswerType.from_dict(data)
+
+                return answer_type_type_0
+            except (TypeError, ValueError, AttributeError, KeyError):
+                pass
+            return cast(AnswerType | None | Unset, data)
+
+        answer_type = _parse_answer_type(d.pop("answer_type", UNSET))
+
+        resolve_redirects = d.pop("resolve_redirects", UNSET)
+
+        web_search_labeler = cls(
+            config_type=config_type,
+            confidence_threshold=confidence_threshold,
+            answer_type=answer_type,
+            resolve_redirects=resolve_redirects,
+        )
+
+        web_search_labeler.additional_properties = d
+        return web_search_labeler
+
+    @property
+    def additional_keys(self) -> list[str]:
+        return list(self.additional_properties.keys())
+
+    def __getitem__(self, key: str) -> Any:
+        return self.additional_properties[key]
+
+    def __setitem__(self, key: str, value: Any) -> None:
+        self.additional_properties[key] = value
+
+    def __delitem__(self, key: str) -> None:
+        del self.additional_properties[key]
+
+    def __contains__(self, key: str) -> bool:
+        return key in self.additional_properties

lightningrod/_generated/py.typed

@@ -0,0 +1 @@
+# Marker file for PEP 561

lightningrod/_generated/types.py

@@ -0,0 +1,54 @@
+"""Contains some shared types for properties"""
+
+from collections.abc import Mapping, MutableMapping
+from http import HTTPStatus
+from typing import IO, BinaryIO, Generic, Literal, TypeVar
+
+from attrs import define
+
+
+class Unset:
+    def __bool__(self) -> Literal[False]:
+        return False
+
+
+UNSET: Unset = Unset()
+
+# The types that `httpx.Client(files=)` can accept, copied from that library.
+FileContent = IO[bytes] | bytes | str
+FileTypes = (
+    # (filename, file (or bytes), content_type)
+    tuple[str | None, FileContent, str | None]
+    # (filename, file (or bytes), content_type, headers)
+    | tuple[str | None, FileContent, str | None, Mapping[str, str]]
+)
+RequestFiles = list[tuple[str, FileTypes]]
+
+
+@define
+class File:
+    """Contains information for file uploads"""
+
+    payload: BinaryIO
+    file_name: str | None = None
+    mime_type: str | None = None
+
+    def to_tuple(self) -> FileTypes:
+        """Return a tuple representation that httpx will accept for multipart/form-data"""
+        return self.file_name, self.payload, self.mime_type
+
+
+T = TypeVar("T")
+
+
+@define
+class Response(Generic[T]):
+    """A response from an endpoint"""
+
+    status_code: HTTPStatus
+    content: bytes
+    headers: MutableMapping[str, str]
+    parsed: T | None
+
+
+__all__ = ["UNSET", "File", "FileTypes", "RequestFiles", "Response", "Unset"]

lightningrod/client.py

@@ -0,0 +1,48 @@
+from typing import List
+
+from lightningrod._generated.client import AuthenticatedClient
+from lightningrod._generated.models.sample import Sample
+from lightningrod.datasets.client import DatasetSamplesClient, DatasetsClient
+from lightningrod.datasets.dataset import Dataset
+from lightningrod.files.client import FilesClient
+from lightningrod.filesets.client import FileSetsClient
+from lightningrod.organization.client import OrganizationsClient
+from lightningrod.transforms.client import TransformsClient
+
+
+class LightningRod:
+    """
+    Python SDK for the Lightning Rod API.
+    
+    Args:
+        api_key: Your Lightning Rod API key
+        base_url: Base URL for the API (defaults to production)
+    
+    Example:
+        >>> lr = LightningRod(api_key="your-api-key")
+        >>> config = QuestionPipeline(...)
+        >>> dataset = lr.transforms.run(config)
+        >>> samples = dataset.to_samples()
+    """
+    
+    def __init__(
+        self,
+        api_key: str,
+        base_url: str = "https://api.lightningrod.ai/api/public/v1"
+    ):
+        self.api_key: str = api_key
+        self.base_url: str = base_url.rstrip("/")
+        self._generated_client: AuthenticatedClient = AuthenticatedClient(
+            base_url=self.base_url,
+            token=api_key,
+            prefix="Bearer",
+            auth_header_name="Authorization",
+        )
+        
+        self._dataset_samples: DatasetSamplesClient = DatasetSamplesClient(self._generated_client)
+        self.transforms: TransformsClient = TransformsClient(self._generated_client, self._dataset_samples)
+        self.datasets: DatasetsClient = DatasetsClient(self._generated_client, self._dataset_samples)
+        self.organization: OrganizationsClient = OrganizationsClient(self._generated_client)
+         # TODO(filesets): Enable when filesets are publicly supported
+        # self.files: FilesClient = FilesClient(self._generated_client)
+        # self.filesets: FileSetsClient = FileSetsClient(self._generated_client, self.files)

lightningrod/datasets/__init__.py

@@ -0,0 +1,4 @@
+from lightningrod.datasets.client import DatasetsClient, DatasetSamplesClient
+from lightningrod.datasets.dataset import Dataset, AsyncDataset
+
+__all__ = ["DatasetsClient", "DatasetSamplesClient", "Dataset", "AsyncDataset"]

lightningrod/datasets/client.py

@@ -0,0 +1,174 @@
+from typing import List, Optional
+
+from lightningrod._generated.models import (
+    HTTPValidationError,
+    UploadSamplesRequest,
+)
+from lightningrod._generated.models.sample import Sample
+from lightningrod._generated.api.datasets import (
+    create_dataset_datasets_post,
+    get_dataset_datasets_dataset_id_get,
+    get_dataset_samples_datasets_dataset_id_samples_get,
+    upload_samples_datasets_dataset_id_samples_post,
+)
+from lightningrod._generated.types import Unset
+from lightningrod._generated.client import AuthenticatedClient
+from lightningrod.datasets.dataset import Dataset
+from lightningrod._errors import handle_response_error
+
+
+class DatasetSamplesClient:
+    def __init__(self, client: AuthenticatedClient):
+        self._client: AuthenticatedClient = client
+    
+    def list(self, dataset_id: str) -> List[Sample]:
+        samples: List[Sample] = []
+        cursor: Optional[str] = None
+        
+        while True:
+            response = get_dataset_samples_datasets_dataset_id_samples_get.sync_detailed(
+                dataset_id=dataset_id,
+                client=self._client,
+                limit=100,
+                cursor=cursor,
+            )
+            
+            parsed = handle_response_error(response, "fetch samples")
+            
+            samples.extend(parsed.samples)
+            
+            if not parsed.has_more:
+                break
+            if isinstance(parsed.next_cursor, Unset) or parsed.next_cursor is None:
+                break
+            cursor = str(parsed.next_cursor)
+        
+        return samples
+    
+    def upload(
+        self,
+        dataset_id: str,
+        samples: List[Sample],
+    ) -> None:
+        """
+        Upload samples to an existing dataset.
+        
+        Args:
+            dataset_id: ID of the dataset to upload samples to
+            samples: List of Sample objects to upload
+            
+        Example:
+            >>> lr = LightningRod(api_key="your-api-key")
+            >>> samples = [Sample(seed=Seed(...), ...), ...]
+            >>> lr.datasets.upload(samples)
+        """
+        request = UploadSamplesRequest(samples=samples)
+        
+        response = upload_samples_datasets_dataset_id_samples_post.sync_detailed(
+            dataset_id=dataset_id,
+            client=self._client,
+            body=request,
+        )
+        
+        handle_response_error(response, "upload samples")
+
+
+class DatasetsClient:
+    def __init__(self, client: AuthenticatedClient, dataset_samples_client: DatasetSamplesClient):
+        self._client: AuthenticatedClient = client
+        self._dataset_samples_client: DatasetSamplesClient = dataset_samples_client
+    
+    def create(self) -> Dataset:
+        """
+        Create a new empty dataset.
+        
+        Returns:
+            Dataset object representing the newly created dataset
+            
+        Example:
+            >>> lr = LightningRod(api_key="your-api-key")
+            >>> dataset = lr.datasets.create()
+            >>> print(f"Created dataset: {dataset.id}")
+        """
+        response = create_dataset_datasets_post.sync_detailed(
+            client=self._client,
+        )
+        
+        create_result = handle_response_error(response, "create dataset")
+        
+        dataset_response = get_dataset_datasets_dataset_id_get.sync_detailed(
+            dataset_id=create_result.id,
+            client=self._client,
+        )
+        dataset_result = handle_response_error(dataset_response, "get dataset")
+        
+        return Dataset(
+            id=dataset_result.id,
+            num_rows=dataset_result.num_rows,
+            datasets_client=self._dataset_samples_client
+        )
+    
+    def create_from_samples(
+        self,
+        samples: List[Sample],
+        batch_size: int = 1000,
+    ) -> Dataset:
+        """
+        Create a new dataset and upload samples to it.
+        
+        This is a convenience method that creates a dataset and uploads all samples
+        in batches. Useful for creating input datasets from a collection of seeds.
+        
+        Args:
+            samples: List of Sample objects to upload
+            batch_size: Number of samples to upload per batch (default: 1000)
+            
+        Returns:
+            Dataset object with all samples uploaded
+            
+        Example:
+            >>> lr = LightningRod(api_key="your-api-key")
+            >>> samples = [Sample(seed=Seed(...), ...), ...]
+            >>> dataset = lr.datasets.create_from_samples(samples, batch_size=1000)
+            >>> print(f"Created dataset with {dataset.num_rows} samples")
+        """
+        dataset = self.create()
+        
+        for i in range(0, len(samples), batch_size):
+            batch = samples[i:i + batch_size]
+            self._dataset_samples_client.upload(dataset.id, batch)
+        
+        dataset_response = get_dataset_datasets_dataset_id_get.sync_detailed(
+            dataset_id=dataset.id,
+            client=self._client,
+        )
+        dataset_result = handle_response_error(dataset_response, "refresh dataset")
+        
+        dataset.num_rows = dataset_result.num_rows
+        return dataset
+    
+    def get(self, dataset_id: str) -> Dataset:
+        """
+        Get a dataset by ID.
+        
+        Args:
+            dataset_id: ID of the dataset to retrieve
+            
+        Returns:
+            Dataset object
+            
+        Example:
+            >>> lr = LightningRod(api_key="your-api-key")
+            >>> dataset = lr.datasets.get("dataset-id-here")
+        """
+        dataset_response = get_dataset_datasets_dataset_id_get.sync_detailed(
+            dataset_id=dataset_id,
+            client=self._client,
+        )
+        dataset_result = handle_response_error(dataset_response, "get dataset")
+        
+        return Dataset(
+            id=dataset_result.id,
+            num_rows=dataset_result.num_rows,
+            datasets_client=self._dataset_samples_client
+        )

lightningrod/datasets/dataset.py

@@ -0,0 +1,255 @@
+from typing import List, Optional, Dict, Any, TYPE_CHECKING
+import asyncio
+
+from lightningrod._generated.models.sample import Sample
+from lightningrod._generated.models.forward_looking_question import ForwardLookingQuestion
+from lightningrod._generated.models.question import Question
+from lightningrod._generated.models.news_context import NewsContext
+from lightningrod._generated.models.rag_context import RAGContext
+from lightningrod._generated.models.sample_meta import SampleMeta
+from lightningrod._generated.types import UNSET, Unset
+
+# avoid circular import
+if TYPE_CHECKING:
+    from lightningrod.datasets.client import DatasetSamplesClient
+
+class Dataset:
+    """
+    Represents a dataset in Lightning Rod.
+    
+    A dataset contains rows of sample data. Use this class to access 
+    dataset metadata and download the actual samples.
+    
+    Note: Datasets should only be created through LightningRod methods,
+    not instantiated directly.
+    
+    Attributes:
+        id: Unique identifier for the dataset
+        num_rows: Number of rows in the dataset
+    
+    Example:
+        >>> lr = LightningRod(api_key="your-api-key")
+        >>> config = QuestionPipeline(...)
+        >>> dataset = lr.transforms.run(config)
+        >>> samples = dataset.to_samples()
+        >>> print(f"Dataset has {len(samples)} samples")
+    """
+    
+    def __init__(
+        self,
+        id: str,
+        num_rows: int,
+        datasets_client: "DatasetSamplesClient"
+    ):
+        self.id: str = id
+        self.num_rows: int = num_rows
+        self._datasets_client: "DatasetSamplesClient" = datasets_client
+        self._samples: Optional[List[Sample]] = None
+    
+    def download(self) -> List[Sample]:
+        """
+        Download all samples from the dataset via the paginated API.
+        
+        Returns:
+            List of Sample objects
+        
+        Example:
+            >>> lr = LightningRod(api_key="your-api-key")
+            >>> dataset = lr.transforms.run(config)
+            >>> samples = dataset.download()
+            >>> for sample in samples:
+            ...     print(sample.seed.seed_text)
+        """
+        self._samples = self._datasets_client.list(self.id)
+        return self._samples
+
+    def samples(self) -> List[Sample]:
+        """
+        Get all samples from the dataset. 
+        Automatically downloads the samples if they haven't been downloaded yet.
+        
+        Returns:
+            List of Sample objects
+        """
+        if not self._samples:
+            self.download()
+        return self._samples
+
+    def to_samples(self) -> List[Sample]:
+        """
+        Download all samples from the dataset via the paginated API.
+        
+        Returns:
+            List of Sample objects
+        
+        Example:
+            >>> lr = LightningRod(api_key="your-api-key")
+            >>> config = QuestionPipeline(...)
+            >>> dataset = lr.transforms.run(config)
+            >>> samples = dataset.to_samples()
+            >>> for sample in samples:
+            ...     print(sample.seed.seed_text)
+        """
+        return self.samples()
+
+    def flattened(self) -> List[Dict[str, Any]]:
+        """
+        Convert all samples to a list of dictionaries.
+        Automatically downloads the samples if they haven't been downloaded yet.
+        
+        Handles different question types (Question, ForwardLookingQuestion) and
+        extracts relevant fields from labels, seeds, and prompts.
+        
+        Returns:
+            List of dictionaries, each representing a sample row
+        
+        Example:
+            >>> lr = LightningRod(api_key="your-api-key")
+            >>> config = QuestionPipeline(...)
+            >>> dataset = lr.transforms.run(config)
+            >>> rows = dataset.flattened()
+            >>> import pandas as pd
+            >>> df = pd.DataFrame(rows)
+        """
+        samples = self.samples()
+        return [self._sample_to_dict(sample) for sample in samples]
+
+    def _sample_to_dict(self, sample: Sample) -> Dict[str, Any]:
+        row: Dict[str, Any] = {}
+        
+        if sample.question and not isinstance(sample.question, Unset):
+            if isinstance(sample.question, ForwardLookingQuestion):
+                row['question.question_text'] = sample.question.question_text
+                row['question.date_close'] = sample.question.date_close.isoformat()
+                row['question.event_date'] = sample.question.event_date.isoformat()
+                row['question.resolution_criteria'] = sample.question.resolution_criteria
+                if sample.question.prediction_date is not None and not isinstance(sample.question.prediction_date, Unset):
+                    row['question.prediction_date'] = sample.question.prediction_date.isoformat()
+            elif isinstance(sample.question, Question):
+                row['question.question_text'] = sample.question.question_text
+            else:
+                question_text = getattr(sample.question, 'question_text', None)
+                if question_text is not None:
+                    row['question.question_text'] = question_text
+        
+        if sample.label and not isinstance(sample.label, Unset):
+            row['label.label'] = sample.label.label
+            row['label.label_confidence'] = sample.label.label_confidence
+            if sample.label.resolution_date is not None and not isinstance(sample.label.resolution_date, Unset):
+                row['label.resolution_date'] = sample.label.resolution_date.isoformat()
+            if sample.label.reasoning is not None and not isinstance(sample.label.reasoning, Unset):
+                row['label.reasoning'] = sample.label.reasoning
+            if sample.label.answer_sources is not None and not isinstance(sample.label.answer_sources, Unset):
+                row['label.answer_sources'] = sample.label.answer_sources
+        
+        if sample.prompt and not isinstance(sample.prompt, Unset):
+            row['prompt'] = sample.prompt
+        
+        if sample.seed and not isinstance(sample.seed, Unset):
+            row['seed.seed_text'] = sample.seed.seed_text
+            if sample.seed.url is not None and not isinstance(sample.seed.url, Unset):
+                row['seed.url'] = sample.seed.url
+            if sample.seed.seed_creation_date is not None and not isinstance(sample.seed.seed_creation_date, Unset):
+                row['seed.seed_creation_date'] = sample.seed.seed_creation_date.isoformat()
+            if sample.seed.search_query is not None and not isinstance(sample.seed.search_query, Unset):
+                row['seed.search_query'] = sample.seed.search_query
+        
+        if sample.is_valid is not None and not isinstance(sample.is_valid, Unset):
+            row['is_valid'] = sample.is_valid
+        
+        if sample.context is not None and not isinstance(sample.context, Unset):
+            for idx, ctx in enumerate(sample.context):
+                if isinstance(ctx, NewsContext):
+                    row[f'context.{idx}.rendered_context'] = ctx.rendered_context
+                    row[f'context.{idx}.search_query'] = ctx.search_query
+                    row[f'context.{idx}.context_type'] = ctx.context_type
+                elif isinstance(ctx, RAGContext):
+                    row[f'context.{idx}.rendered_context'] = ctx.rendered_context
+                    row[f'context.{idx}.document_id'] = ctx.document_id
+                    row[f'context.{idx}.context_type'] = ctx.context_type
+        
+        if sample.meta is not None and not isinstance(sample.meta, Unset):
+            if isinstance(sample.meta, SampleMeta):
+                for key, value in sample.meta.additional_properties.items():
+                    row[f'meta.{key}'] = value
+        
+        if sample.additional_properties:
+            for key, value in sample.additional_properties.items():
+                row[f'additional_properties.{key}'] = value
+        
+        return row
+
+class AsyncDataset:
+    """
+    Async wrapper for Dataset.
+    
+    This class provides an async interface to Dataset operations by running
+    the synchronous operations in a thread pool using asyncio.to_thread.
+    
+    Note: AsyncDatasets should only be created through AsyncLightningRod methods,
+    not instantiated directly.
+    
+    Attributes:
+        id: Unique identifier for the dataset
+        num_rows: Number of rows in the dataset
+    
+    Example:
+        >>> lr = AsyncLightningRod(api_key="your-api-key")
+        >>> config = QuestionPipeline(...)
+        >>> dataset = await lr.transforms.run(config)
+        >>> samples = await dataset.to_samples()
+        >>> print(f"Dataset has {len(samples)} samples")
+    """
+    
+    def __init__(self, sync_dataset: Dataset):
+        self._sync_dataset: Dataset = sync_dataset
+    
+    @property
+    def id(self) -> str:
+        return self._sync_dataset.id
+    
+    @property
+    def num_rows(self) -> int:
+        return self._sync_dataset.num_rows
+    
+    async def to_samples(self) -> List[Sample]:
+        """
+        Download all samples from the dataset via the paginated API.
+        
+        All operations are run in a thread pool to avoid blocking the event loop.
+        
+        Returns:
+            List of Sample objects
+        
+        Example:
+            >>> lr = AsyncLightningRod(api_key="your-api-key")
+            >>> config = QuestionPipeline(...)
+            >>> dataset = await lr.transforms.run(config)
+            >>> samples = await dataset.to_samples()
+            >>> for sample in samples:
+            ...     print(sample.seed.seed_text)
+        """
+        return await asyncio.to_thread(self._sync_dataset.to_samples)
+
+    async def flattened(self) -> List[Dict[str, Any]]:
+        """
+        Convert all samples to a list of dictionaries.
+        Automatically downloads the samples if they haven't been downloaded yet.
+        
+        All operations are run in a thread pool to avoid blocking the event loop.
+        
+        Handles different question types (Question, ForwardLookingQuestion) and
+        extracts relevant fields from labels, seeds, and prompts.
+        
+        Returns:
+            List of dictionaries, each representing a sample row
+        
+        Example:
+            >>> lr = AsyncLightningRod(api_key="your-api-key")
+            >>> config = QuestionPipeline(...)
+            >>> dataset = await lr.transforms.run(config)
+            >>> rows = await dataset.flattened()
+            >>> import pandas as pd
+            >>> df = pd.DataFrame(rows)
+        """
+        return await asyncio.to_thread(self._sync_dataset.flattened)