lightningrod-ai 0.1.6__py3-none-any.whl

lightningrod/_generated/client.py

@@ -0,0 +1,268 @@
+import ssl
+from typing import Any
+
+import httpx
+from attrs import define, evolve, field
+
+
+@define
+class Client:
+    """A class for keeping track of data related to the API
+
+    The following are accepted as keyword arguments and will be used to construct httpx Clients internally:
+
+        ``base_url``: The base URL for the API, all requests are made to a relative path to this URL
+
+        ``cookies``: A dictionary of cookies to be sent with every request
+
+        ``headers``: A dictionary of headers to be sent with every request
+
+        ``timeout``: The maximum amount of a time a request can take. API functions will raise
+        httpx.TimeoutException if this is exceeded.
+
+        ``verify_ssl``: Whether or not to verify the SSL certificate of the API server. This should be True in production,
+        but can be set to False for testing purposes.
+
+        ``follow_redirects``: Whether or not to follow redirects. Default value is False.
+
+        ``httpx_args``: A dictionary of additional arguments to be passed to the ``httpx.Client`` and ``httpx.AsyncClient`` constructor.
+
+
+    Attributes:
+        raise_on_unexpected_status: Whether or not to raise an errors.UnexpectedStatus if the API returns a
+            status code that was not documented in the source OpenAPI document. Can also be provided as a keyword
+            argument to the constructor.
+    """
+
+    raise_on_unexpected_status: bool = field(default=False, kw_only=True)
+    _base_url: str = field(alias="base_url")
+    _cookies: dict[str, str] = field(factory=dict, kw_only=True, alias="cookies")
+    _headers: dict[str, str] = field(factory=dict, kw_only=True, alias="headers")
+    _timeout: httpx.Timeout | None = field(default=None, kw_only=True, alias="timeout")
+    _verify_ssl: str | bool | ssl.SSLContext = field(default=True, kw_only=True, alias="verify_ssl")
+    _follow_redirects: bool = field(default=False, kw_only=True, alias="follow_redirects")
+    _httpx_args: dict[str, Any] = field(factory=dict, kw_only=True, alias="httpx_args")
+    _client: httpx.Client | None = field(default=None, init=False)
+    _async_client: httpx.AsyncClient | None = field(default=None, init=False)
+
+    def with_headers(self, headers: dict[str, str]) -> "Client":
+        """Get a new client matching this one with additional headers"""
+        if self._client is not None:
+            self._client.headers.update(headers)
+        if self._async_client is not None:
+            self._async_client.headers.update(headers)
+        return evolve(self, headers={**self._headers, **headers})
+
+    def with_cookies(self, cookies: dict[str, str]) -> "Client":
+        """Get a new client matching this one with additional cookies"""
+        if self._client is not None:
+            self._client.cookies.update(cookies)
+        if self._async_client is not None:
+            self._async_client.cookies.update(cookies)
+        return evolve(self, cookies={**self._cookies, **cookies})
+
+    def with_timeout(self, timeout: httpx.Timeout) -> "Client":
+        """Get a new client matching this one with a new timeout configuration"""
+        if self._client is not None:
+            self._client.timeout = timeout
+        if self._async_client is not None:
+            self._async_client.timeout = timeout
+        return evolve(self, timeout=timeout)
+
+    def set_httpx_client(self, client: httpx.Client) -> "Client":
+        """Manually set the underlying httpx.Client
+
+        **NOTE**: This will override any other settings on the client, including cookies, headers, and timeout.
+        """
+        self._client = client
+        return self
+
+    def get_httpx_client(self) -> httpx.Client:
+        """Get the underlying httpx.Client, constructing a new one if not previously set"""
+        if self._client is None:
+            self._client = httpx.Client(
+                base_url=self._base_url,
+                cookies=self._cookies,
+                headers=self._headers,
+                timeout=self._timeout,
+                verify=self._verify_ssl,
+                follow_redirects=self._follow_redirects,
+                **self._httpx_args,
+            )
+        return self._client
+
+    def __enter__(self) -> "Client":
+        """Enter a context manager for self.client—you cannot enter twice (see httpx docs)"""
+        self.get_httpx_client().__enter__()
+        return self
+
+    def __exit__(self, *args: Any, **kwargs: Any) -> None:
+        """Exit a context manager for internal httpx.Client (see httpx docs)"""
+        self.get_httpx_client().__exit__(*args, **kwargs)
+
+    def set_async_httpx_client(self, async_client: httpx.AsyncClient) -> "Client":
+        """Manually set the underlying httpx.AsyncClient
+
+        **NOTE**: This will override any other settings on the client, including cookies, headers, and timeout.
+        """
+        self._async_client = async_client
+        return self
+
+    def get_async_httpx_client(self) -> httpx.AsyncClient:
+        """Get the underlying httpx.AsyncClient, constructing a new one if not previously set"""
+        if self._async_client is None:
+            self._async_client = httpx.AsyncClient(
+                base_url=self._base_url,
+                cookies=self._cookies,
+                headers=self._headers,
+                timeout=self._timeout,
+                verify=self._verify_ssl,
+                follow_redirects=self._follow_redirects,
+                **self._httpx_args,
+            )
+        return self._async_client
+
+    async def __aenter__(self) -> "Client":
+        """Enter a context manager for underlying httpx.AsyncClient—you cannot enter twice (see httpx docs)"""
+        await self.get_async_httpx_client().__aenter__()
+        return self
+
+    async def __aexit__(self, *args: Any, **kwargs: Any) -> None:
+        """Exit a context manager for underlying httpx.AsyncClient (see httpx docs)"""
+        await self.get_async_httpx_client().__aexit__(*args, **kwargs)
+
+
+@define
+class AuthenticatedClient:
+    """A Client which has been authenticated for use on secured endpoints
+
+    The following are accepted as keyword arguments and will be used to construct httpx Clients internally:
+
+        ``base_url``: The base URL for the API, all requests are made to a relative path to this URL
+
+        ``cookies``: A dictionary of cookies to be sent with every request
+
+        ``headers``: A dictionary of headers to be sent with every request
+
+        ``timeout``: The maximum amount of a time a request can take. API functions will raise
+        httpx.TimeoutException if this is exceeded.
+
+        ``verify_ssl``: Whether or not to verify the SSL certificate of the API server. This should be True in production,
+        but can be set to False for testing purposes.
+
+        ``follow_redirects``: Whether or not to follow redirects. Default value is False.
+
+        ``httpx_args``: A dictionary of additional arguments to be passed to the ``httpx.Client`` and ``httpx.AsyncClient`` constructor.
+
+
+    Attributes:
+        raise_on_unexpected_status: Whether or not to raise an errors.UnexpectedStatus if the API returns a
+            status code that was not documented in the source OpenAPI document. Can also be provided as a keyword
+            argument to the constructor.
+        token: The token to use for authentication
+        prefix: The prefix to use for the Authorization header
+        auth_header_name: The name of the Authorization header
+    """
+
+    raise_on_unexpected_status: bool = field(default=False, kw_only=True)
+    _base_url: str = field(alias="base_url")
+    _cookies: dict[str, str] = field(factory=dict, kw_only=True, alias="cookies")
+    _headers: dict[str, str] = field(factory=dict, kw_only=True, alias="headers")
+    _timeout: httpx.Timeout | None = field(default=None, kw_only=True, alias="timeout")
+    _verify_ssl: str | bool | ssl.SSLContext = field(default=True, kw_only=True, alias="verify_ssl")
+    _follow_redirects: bool = field(default=False, kw_only=True, alias="follow_redirects")
+    _httpx_args: dict[str, Any] = field(factory=dict, kw_only=True, alias="httpx_args")
+    _client: httpx.Client | None = field(default=None, init=False)
+    _async_client: httpx.AsyncClient | None = field(default=None, init=False)
+
+    token: str
+    prefix: str = "Bearer"
+    auth_header_name: str = "Authorization"
+
+    def with_headers(self, headers: dict[str, str]) -> "AuthenticatedClient":
+        """Get a new client matching this one with additional headers"""
+        if self._client is not None:
+            self._client.headers.update(headers)
+        if self._async_client is not None:
+            self._async_client.headers.update(headers)
+        return evolve(self, headers={**self._headers, **headers})
+
+    def with_cookies(self, cookies: dict[str, str]) -> "AuthenticatedClient":
+        """Get a new client matching this one with additional cookies"""
+        if self._client is not None:
+            self._client.cookies.update(cookies)
+        if self._async_client is not None:
+            self._async_client.cookies.update(cookies)
+        return evolve(self, cookies={**self._cookies, **cookies})
+
+    def with_timeout(self, timeout: httpx.Timeout) -> "AuthenticatedClient":
+        """Get a new client matching this one with a new timeout configuration"""
+        if self._client is not None:
+            self._client.timeout = timeout
+        if self._async_client is not None:
+            self._async_client.timeout = timeout
+        return evolve(self, timeout=timeout)
+
+    def set_httpx_client(self, client: httpx.Client) -> "AuthenticatedClient":
+        """Manually set the underlying httpx.Client
+
+        **NOTE**: This will override any other settings on the client, including cookies, headers, and timeout.
+        """
+        self._client = client
+        return self
+
+    def get_httpx_client(self) -> httpx.Client:
+        """Get the underlying httpx.Client, constructing a new one if not previously set"""
+        if self._client is None:
+            self._headers[self.auth_header_name] = f"{self.prefix} {self.token}" if self.prefix else self.token
+            self._client = httpx.Client(
+                base_url=self._base_url,
+                cookies=self._cookies,
+                headers=self._headers,
+                timeout=self._timeout,
+                verify=self._verify_ssl,
+                follow_redirects=self._follow_redirects,
+                **self._httpx_args,
+            )
+        return self._client
+
+    def __enter__(self) -> "AuthenticatedClient":
+        """Enter a context manager for self.client—you cannot enter twice (see httpx docs)"""
+        self.get_httpx_client().__enter__()
+        return self
+
+    def __exit__(self, *args: Any, **kwargs: Any) -> None:
+        """Exit a context manager for internal httpx.Client (see httpx docs)"""
+        self.get_httpx_client().__exit__(*args, **kwargs)
+
+    def set_async_httpx_client(self, async_client: httpx.AsyncClient) -> "AuthenticatedClient":
+        """Manually set the underlying httpx.AsyncClient
+
+        **NOTE**: This will override any other settings on the client, including cookies, headers, and timeout.
+        """
+        self._async_client = async_client
+        return self
+
+    def get_async_httpx_client(self) -> httpx.AsyncClient:
+        """Get the underlying httpx.AsyncClient, constructing a new one if not previously set"""
+        if self._async_client is None:
+            self._headers[self.auth_header_name] = f"{self.prefix} {self.token}" if self.prefix else self.token
+            self._async_client = httpx.AsyncClient(
+                base_url=self._base_url,
+                cookies=self._cookies,
+                headers=self._headers,
+                timeout=self._timeout,
+                verify=self._verify_ssl,
+                follow_redirects=self._follow_redirects,
+                **self._httpx_args,
+            )
+        return self._async_client
+
+    async def __aenter__(self) -> "AuthenticatedClient":
+        """Enter a context manager for underlying httpx.AsyncClient—you cannot enter twice (see httpx docs)"""
+        await self.get_async_httpx_client().__aenter__()
+        return self
+
+    async def __aexit__(self, *args: Any, **kwargs: Any) -> None:
+        """Exit a context manager for underlying httpx.AsyncClient (see httpx docs)"""
+        await self.get_async_httpx_client().__aexit__(*args, **kwargs)

lightningrod/_generated/errors.py

@@ -0,0 +1,16 @@
+"""Contains shared errors types that can be raised from API functions"""
+
+
+class UnexpectedStatus(Exception):
+    """Raised by api functions when the response status an undocumented status and Client.raise_on_unexpected_status is True"""
+
+    def __init__(self, status_code: int, content: bytes):
+        self.status_code = status_code
+        self.content = content
+
+        super().__init__(
+            f"Unexpected status code: {status_code}\n\nResponse content:\n{content.decode(errors='ignore')}"
+        )
+
+
+__all__ = ["UnexpectedStatus"]

lightningrod/_generated/models/__init__.py

@@ -0,0 +1,147 @@
+"""Contains all the data models used in inputs/outputs"""
+
+from .answer_type import AnswerType
+from .answer_type_enum import AnswerTypeEnum
+from .balance_response import BalanceResponse
+from .chat_completion_request import ChatCompletionRequest
+from .chat_completion_response import ChatCompletionResponse
+from .chat_message import ChatMessage
+from .choice import Choice
+from .create_dataset_response import CreateDatasetResponse
+from .create_file_set_file_request import CreateFileSetFileRequest
+from .create_file_set_file_request_metadata_type_0 import CreateFileSetFileRequestMetadataType0
+from .create_file_set_request import CreateFileSetRequest
+from .create_file_upload_request import CreateFileUploadRequest
+from .create_file_upload_response import CreateFileUploadResponse
+from .create_file_upload_response_metadata_type_0 import CreateFileUploadResponseMetadataType0
+from .create_transform_job_request import CreateTransformJobRequest
+from .dataset_metadata import DatasetMetadata
+from .estimate_cost_request import EstimateCostRequest
+from .estimate_cost_response import EstimateCostResponse
+from .event_usage_summary import EventUsageSummary
+from .file_set import FileSet
+from .file_set_file import FileSetFile
+from .file_set_file_metadata_type_0 import FileSetFileMetadataType0
+from .file_set_query_seed_generator import FileSetQuerySeedGenerator
+from .file_set_seed_generator import FileSetSeedGenerator
+from .filter_criteria import FilterCriteria
+from .forward_looking_question import ForwardLookingQuestion
+from .forward_looking_question_generator import ForwardLookingQuestionGenerator
+from .gdelt_seed_generator import GdeltSeedGenerator
+from .http_validation_error import HTTPValidationError
+from .job_usage import JobUsage
+from .job_usage_by_step_type_0 import JobUsageByStepType0
+from .label import Label
+from .list_file_set_files_response import ListFileSetFilesResponse
+from .list_file_sets_response import ListFileSetsResponse
+from .llm_model_usage_summary import LLMModelUsageSummary
+from .mock_transform_config import MockTransformConfig
+from .mock_transform_config_metadata_additions import MockTransformConfigMetadataAdditions
+from .model_config import ModelConfig
+from .model_source_type import ModelSourceType
+from .news_context import NewsContext
+from .news_context_generator import NewsContextGenerator
+from .news_seed_generator import NewsSeedGenerator
+from .paginated_samples_response import PaginatedSamplesResponse
+from .pipeline_metrics_response import PipelineMetricsResponse
+from .question import Question
+from .question_and_label_generator import QuestionAndLabelGenerator
+from .question_generator import QuestionGenerator
+from .question_pipeline import QuestionPipeline
+from .question_renderer import QuestionRenderer
+from .rag_context import RAGContext
+from .response_message import ResponseMessage
+from .rollout import Rollout
+from .rollout_generator import RolloutGenerator
+from .rollout_parsed_output_type_0 import RolloutParsedOutputType0
+from .sample import Sample
+from .sample_meta import SampleMeta
+from .seed import Seed
+from .step_cost_breakdown import StepCostBreakdown
+from .transform_job import TransformJob
+from .transform_job_status import TransformJobStatus
+from .transform_step_metrics_response import TransformStepMetricsResponse
+from .transform_type import TransformType
+from .upload_samples_request import UploadSamplesRequest
+from .upload_samples_response import UploadSamplesResponse
+from .usage import Usage
+from .usage_summary import UsageSummary
+from .usage_summary_events import UsageSummaryEvents
+from .usage_summary_llm_by_model import UsageSummaryLlmByModel
+from .validate_sample_response import ValidateSampleResponse
+from .validation_error import ValidationError
+from .web_search_labeler import WebSearchLabeler
+
+__all__ = (
+    "AnswerType",
+    "AnswerTypeEnum",
+    "BalanceResponse",
+    "ChatCompletionRequest",
+    "ChatCompletionResponse",
+    "ChatMessage",
+    "Choice",
+    "CreateDatasetResponse",
+    "CreateFileSetFileRequest",
+    "CreateFileSetFileRequestMetadataType0",
+    "CreateFileSetRequest",
+    "CreateFileUploadRequest",
+    "CreateFileUploadResponse",
+    "CreateFileUploadResponseMetadataType0",
+    "CreateTransformJobRequest",
+    "DatasetMetadata",
+    "EstimateCostRequest",
+    "EstimateCostResponse",
+    "EventUsageSummary",
+    "FileSet",
+    "FileSetFile",
+    "FileSetFileMetadataType0",
+    "FileSetQuerySeedGenerator",
+    "FileSetSeedGenerator",
+    "FilterCriteria",
+    "ForwardLookingQuestion",
+    "ForwardLookingQuestionGenerator",
+    "GdeltSeedGenerator",
+    "HTTPValidationError",
+    "JobUsage",
+    "JobUsageByStepType0",
+    "Label",
+    "ListFileSetFilesResponse",
+    "ListFileSetsResponse",
+    "LLMModelUsageSummary",
+    "MockTransformConfig",
+    "MockTransformConfigMetadataAdditions",
+    "ModelConfig",
+    "ModelSourceType",
+    "NewsContext",
+    "NewsContextGenerator",
+    "NewsSeedGenerator",
+    "PaginatedSamplesResponse",
+    "PipelineMetricsResponse",
+    "Question",
+    "QuestionAndLabelGenerator",
+    "QuestionGenerator",
+    "QuestionPipeline",
+    "QuestionRenderer",
+    "RAGContext",
+    "ResponseMessage",
+    "Rollout",
+    "RolloutGenerator",
+    "RolloutParsedOutputType0",
+    "Sample",
+    "SampleMeta",
+    "Seed",
+    "StepCostBreakdown",
+    "TransformJob",
+    "TransformJobStatus",
+    "TransformStepMetricsResponse",
+    "TransformType",
+    "UploadSamplesRequest",
+    "UploadSamplesResponse",
+    "Usage",
+    "UsageSummary",
+    "UsageSummaryEvents",
+    "UsageSummaryLlmByModel",
+    "ValidateSampleResponse",
+    "ValidationError",
+    "WebSearchLabeler",
+)

lightningrod/_generated/models/answer_type.py

@@ -0,0 +1,129 @@
+from __future__ import annotations
+
+from collections.abc import Mapping
+from typing import Any, TypeVar, cast
+
+from attrs import define as _attrs_define
+from attrs import field as _attrs_field
+
+from ..models.answer_type_enum import AnswerTypeEnum
+from ..types import UNSET, Unset
+
+T = TypeVar("T", bound="AnswerType")
+
+
+@_attrs_define
+class AnswerType:
+    """
+    Attributes:
+        answer_type (AnswerTypeEnum):
+        answer_format_instruction (None | str | Unset): Instructions describing how the answer should be formatted and
+            given. If not set, uses default based on answer_type.
+        labeler_instruction (None | str | Unset): Custom instructions for the labeler. If not set, uses default based on
+            answer_type.
+        question_generation_instruction (None | str | Unset): Custom instructions for generating questions of this type.
+            If not set, uses default based on answer_type.
+    """
+
+    answer_type: AnswerTypeEnum
+    answer_format_instruction: None | str | Unset = UNSET
+    labeler_instruction: None | str | Unset = UNSET
+    question_generation_instruction: None | str | Unset = UNSET
+    additional_properties: dict[str, Any] = _attrs_field(init=False, factory=dict)
+
+    def to_dict(self) -> dict[str, Any]:
+        answer_type = self.answer_type.value
+
+        answer_format_instruction: None | str | Unset
+        if isinstance(self.answer_format_instruction, Unset):
+            answer_format_instruction = UNSET
+        else:
+            answer_format_instruction = self.answer_format_instruction
+
+        labeler_instruction: None | str | Unset
+        if isinstance(self.labeler_instruction, Unset):
+            labeler_instruction = UNSET
+        else:
+            labeler_instruction = self.labeler_instruction
+
+        question_generation_instruction: None | str | Unset
+        if isinstance(self.question_generation_instruction, Unset):
+            question_generation_instruction = UNSET
+        else:
+            question_generation_instruction = self.question_generation_instruction
+
+        field_dict: dict[str, Any] = {}
+        field_dict.update(self.additional_properties)
+        field_dict.update(
+            {
+                "answer_type": answer_type,
+            }
+        )
+        if answer_format_instruction is not UNSET:
+            field_dict["answer_format_instruction"] = answer_format_instruction
+        if labeler_instruction is not UNSET:
+            field_dict["labeler_instruction"] = labeler_instruction
+        if question_generation_instruction is not UNSET:
+            field_dict["question_generation_instruction"] = question_generation_instruction
+
+        return field_dict
+
+    @classmethod
+    def from_dict(cls: type[T], src_dict: Mapping[str, Any]) -> T:
+        d = dict(src_dict)
+        answer_type = AnswerTypeEnum(d.pop("answer_type"))
+
+        def _parse_answer_format_instruction(data: object) -> None | str | Unset:
+            if data is None:
+                return data
+            if isinstance(data, Unset):
+                return data
+            return cast(None | str | Unset, data)
+
+        answer_format_instruction = _parse_answer_format_instruction(d.pop("answer_format_instruction", UNSET))
+
+        def _parse_labeler_instruction(data: object) -> None | str | Unset:
+            if data is None:
+                return data
+            if isinstance(data, Unset):
+                return data
+            return cast(None | str | Unset, data)
+
+        labeler_instruction = _parse_labeler_instruction(d.pop("labeler_instruction", UNSET))
+
+        def _parse_question_generation_instruction(data: object) -> None | str | Unset:
+            if data is None:
+                return data
+            if isinstance(data, Unset):
+                return data
+            return cast(None | str | Unset, data)
+
+        question_generation_instruction = _parse_question_generation_instruction(
+            d.pop("question_generation_instruction", UNSET)
+        )
+
+        answer_type = cls(
+            answer_type=answer_type,
+            answer_format_instruction=answer_format_instruction,
+            labeler_instruction=labeler_instruction,
+            question_generation_instruction=question_generation_instruction,
+        )
+
+        answer_type.additional_properties = d
+        return answer_type
+
+    @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/models/answer_type_enum.py

@@ -0,0 +1,11 @@
+from enum import Enum
+
+
+class AnswerTypeEnum(str, Enum):
+    BINARY = "BINARY"
+    CONTINUOUS = "CONTINUOUS"
+    FREE_RESPONSE = "FREE_RESPONSE"
+    MULTIPLE_CHOICE = "MULTIPLE_CHOICE"
+
+    def __str__(self) -> str:
+        return str(self.value)

lightningrod/_generated/models/balance_response.py

@@ -0,0 +1,61 @@
+from __future__ import annotations
+
+from collections.abc import Mapping
+from typing import Any, TypeVar
+
+from attrs import define as _attrs_define
+from attrs import field as _attrs_field
+
+T = TypeVar("T", bound="BalanceResponse")
+
+
+@_attrs_define
+class BalanceResponse:
+    """
+    Attributes:
+        balance_dollars (float):
+    """
+
+    balance_dollars: float
+    additional_properties: dict[str, Any] = _attrs_field(init=False, factory=dict)
+
+    def to_dict(self) -> dict[str, Any]:
+        balance_dollars = self.balance_dollars
+
+        field_dict: dict[str, Any] = {}
+        field_dict.update(self.additional_properties)
+        field_dict.update(
+            {
+                "balance_dollars": balance_dollars,
+            }
+        )
+
+        return field_dict
+
+    @classmethod
+    def from_dict(cls: type[T], src_dict: Mapping[str, Any]) -> T:
+        d = dict(src_dict)
+        balance_dollars = d.pop("balance_dollars")
+
+        balance_response = cls(
+            balance_dollars=balance_dollars,
+        )
+
+        balance_response.additional_properties = d
+        return balance_response
+
+    @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