lightningrod-ai 0.1.6__py3-none-any.whl

lightningrod/_generated/models/rollout.py

@@ -0,0 +1,130 @@
+from __future__ import annotations
+
+from collections.abc import Mapping
+from typing import TYPE_CHECKING, Any, 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.rollout_parsed_output_type_0 import RolloutParsedOutputType0
+
+
+T = TypeVar("T", bound="Rollout")
+
+
+@_attrs_define
+class Rollout:
+    """A single model rollout/response.
+
+    Attributes:
+        model_name (str):
+        content (str):
+        parsed_output (None | RolloutParsedOutputType0 | Unset):
+        reasoning (None | str | Unset):
+    """
+
+    model_name: str
+    content: str
+    parsed_output: None | RolloutParsedOutputType0 | Unset = UNSET
+    reasoning: None | str | Unset = UNSET
+    additional_properties: dict[str, Any] = _attrs_field(init=False, factory=dict)
+
+    def to_dict(self) -> dict[str, Any]:
+        from ..models.rollout_parsed_output_type_0 import RolloutParsedOutputType0
+
+        model_name = self.model_name
+
+        content = self.content
+
+        parsed_output: dict[str, Any] | None | Unset
+        if isinstance(self.parsed_output, Unset):
+            parsed_output = UNSET
+        elif isinstance(self.parsed_output, RolloutParsedOutputType0):
+            parsed_output = self.parsed_output.to_dict()
+        else:
+            parsed_output = self.parsed_output
+
+        reasoning: None | str | Unset
+        if isinstance(self.reasoning, Unset):
+            reasoning = UNSET
+        else:
+            reasoning = self.reasoning
+
+        field_dict: dict[str, Any] = {}
+        field_dict.update(self.additional_properties)
+        field_dict.update(
+            {
+                "model_name": model_name,
+                "content": content,
+            }
+        )
+        if parsed_output is not UNSET:
+            field_dict["parsed_output"] = parsed_output
+        if reasoning is not UNSET:
+            field_dict["reasoning"] = reasoning
+
+        return field_dict
+
+    @classmethod
+    def from_dict(cls: type[T], src_dict: Mapping[str, Any]) -> T:
+        from ..models.rollout_parsed_output_type_0 import RolloutParsedOutputType0
+
+        d = dict(src_dict)
+        model_name = d.pop("model_name")
+
+        content = d.pop("content")
+
+        def _parse_parsed_output(data: object) -> None | RolloutParsedOutputType0 | Unset:
+            if data is None:
+                return data
+            if isinstance(data, Unset):
+                return data
+            try:
+                if not isinstance(data, dict):
+                    raise TypeError()
+                parsed_output_type_0 = RolloutParsedOutputType0.from_dict(data)
+
+                return parsed_output_type_0
+            except (TypeError, ValueError, AttributeError, KeyError):
+                pass
+            return cast(None | RolloutParsedOutputType0 | Unset, data)
+
+        parsed_output = _parse_parsed_output(d.pop("parsed_output", UNSET))
+
+        def _parse_reasoning(data: object) -> None | str | Unset:
+            if data is None:
+                return data
+            if isinstance(data, Unset):
+                return data
+            return cast(None | str | Unset, data)
+
+        reasoning = _parse_reasoning(d.pop("reasoning", UNSET))
+
+        rollout = cls(
+            model_name=model_name,
+            content=content,
+            parsed_output=parsed_output,
+            reasoning=reasoning,
+        )
+
+        rollout.additional_properties = d
+        return rollout
+
+    @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/rollout_generator.py

@@ -0,0 +1,139 @@
+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.model_config import ModelConfig
+
+
+T = TypeVar("T", bound="RolloutGenerator")
+
+
+@_attrs_define
+class RolloutGenerator:
+    """
+    Attributes:
+        models (list[ModelConfig]): Model names or ModelConfig objects
+        config_type (Literal['ROLLOUT_GENERATOR'] | Unset):  Default: 'ROLLOUT_GENERATOR'.
+        prompt_template (None | str | Unset): Prompt template with {column} placeholders. If None, uses sample.prompt
+        input_columns (list[str] | Unset): Columns to substitute into template (from meta)
+        output_schema (Any | None | Unset): Pydantic model for structured output
+    """
+
+    models: list[ModelConfig]
+    config_type: Literal["ROLLOUT_GENERATOR"] | Unset = "ROLLOUT_GENERATOR"
+    prompt_template: None | str | Unset = UNSET
+    input_columns: list[str] | Unset = UNSET
+    output_schema: Any | None | Unset = UNSET
+    additional_properties: dict[str, Any] = _attrs_field(init=False, factory=dict)
+
+    def to_dict(self) -> dict[str, Any]:
+        models = []
+        for models_item_data in self.models:
+            models_item = models_item_data.to_dict()
+            models.append(models_item)
+
+        config_type = self.config_type
+
+        prompt_template: None | str | Unset
+        if isinstance(self.prompt_template, Unset):
+            prompt_template = UNSET
+        else:
+            prompt_template = self.prompt_template
+
+        input_columns: list[str] | Unset = UNSET
+        if not isinstance(self.input_columns, Unset):
+            input_columns = self.input_columns
+
+        output_schema: Any | None | Unset
+        if isinstance(self.output_schema, Unset):
+            output_schema = UNSET
+        else:
+            output_schema = self.output_schema
+
+        field_dict: dict[str, Any] = {}
+        field_dict.update(self.additional_properties)
+        field_dict.update(
+            {
+                "models": models,
+            }
+        )
+        if config_type is not UNSET:
+            field_dict["config_type"] = config_type
+        if prompt_template is not UNSET:
+            field_dict["prompt_template"] = prompt_template
+        if input_columns is not UNSET:
+            field_dict["input_columns"] = input_columns
+        if output_schema is not UNSET:
+            field_dict["output_schema"] = output_schema
+
+        return field_dict
+
+    @classmethod
+    def from_dict(cls: type[T], src_dict: Mapping[str, Any]) -> T:
+        from ..models.model_config import ModelConfig
+
+        d = dict(src_dict)
+        models = []
+        _models = d.pop("models")
+        for models_item_data in _models:
+            models_item = ModelConfig.from_dict(models_item_data)
+
+            models.append(models_item)
+
+        config_type = cast(Literal["ROLLOUT_GENERATOR"] | Unset, d.pop("config_type", UNSET))
+        if config_type != "ROLLOUT_GENERATOR" and not isinstance(config_type, Unset):
+            raise ValueError(f"config_type must match const 'ROLLOUT_GENERATOR', got '{config_type}'")
+
+        def _parse_prompt_template(data: object) -> None | str | Unset:
+            if data is None:
+                return data
+            if isinstance(data, Unset):
+                return data
+            return cast(None | str | Unset, data)
+
+        prompt_template = _parse_prompt_template(d.pop("prompt_template", UNSET))
+
+        input_columns = cast(list[str], d.pop("input_columns", UNSET))
+
+        def _parse_output_schema(data: object) -> Any | None | Unset:
+            if data is None:
+                return data
+            if isinstance(data, Unset):
+                return data
+            return cast(Any | None | Unset, data)
+
+        output_schema = _parse_output_schema(d.pop("output_schema", UNSET))
+
+        rollout_generator = cls(
+            models=models,
+            config_type=config_type,
+            prompt_template=prompt_template,
+            input_columns=input_columns,
+            output_schema=output_schema,
+        )
+
+        rollout_generator.additional_properties = d
+        return rollout_generator
+
+    @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/rollout_parsed_output_type_0.py

@@ -0,0 +1,46 @@
+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="RolloutParsedOutputType0")
+
+
+@_attrs_define
+class RolloutParsedOutputType0:
+    """ """
+
+    additional_properties: dict[str, Any] = _attrs_field(init=False, factory=dict)
+
+    def to_dict(self) -> dict[str, Any]:
+        field_dict: dict[str, Any] = {}
+        field_dict.update(self.additional_properties)
+
+        return field_dict
+
+    @classmethod
+    def from_dict(cls: type[T], src_dict: Mapping[str, Any]) -> T:
+        d = dict(src_dict)
+        rollout_parsed_output_type_0 = cls()
+
+        rollout_parsed_output_type_0.additional_properties = d
+        return rollout_parsed_output_type_0
+
+    @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/sample.py

@@ -0,0 +1,323 @@
+from __future__ import annotations
+
+from collections.abc import Mapping
+from typing import TYPE_CHECKING, Any, 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.forward_looking_question import ForwardLookingQuestion
+    from ..models.label import Label
+    from ..models.news_context import NewsContext
+    from ..models.question import Question
+    from ..models.rag_context import RAGContext
+    from ..models.rollout import Rollout
+    from ..models.sample_meta import SampleMeta
+    from ..models.seed import Seed
+
+
+T = TypeVar("T", bound="Sample")
+
+
+@_attrs_define
+class Sample:
+    """
+    Attributes:
+        seed (None | Seed | Unset):
+        question (ForwardLookingQuestion | None | Question | Unset):
+        label (Label | None | Unset):
+        prompt (None | str | Unset):
+        context (list[NewsContext | RAGContext] | None | Unset):
+        rollouts (list[Rollout] | None | Unset):
+        meta (SampleMeta | Unset):
+        is_valid (bool | Unset):  Default: True.
+    """
+
+    seed: None | Seed | Unset = UNSET
+    question: ForwardLookingQuestion | None | Question | Unset = UNSET
+    label: Label | None | Unset = UNSET
+    prompt: None | str | Unset = UNSET
+    context: list[NewsContext | RAGContext] | None | Unset = UNSET
+    rollouts: list[Rollout] | None | Unset = UNSET
+    meta: SampleMeta | Unset = UNSET
+    is_valid: bool | Unset = True
+    additional_properties: dict[str, Any] = _attrs_field(init=False, factory=dict)
+
+    def to_dict(self) -> dict[str, Any]:
+        from ..models.forward_looking_question import ForwardLookingQuestion
+        from ..models.label import Label
+        from ..models.news_context import NewsContext
+        from ..models.question import Question
+        from ..models.seed import Seed
+
+        seed: dict[str, Any] | None | Unset
+        if isinstance(self.seed, Unset):
+            seed = UNSET
+        elif isinstance(self.seed, Seed):
+            seed = self.seed.to_dict()
+        else:
+            seed = self.seed
+
+        question: dict[str, Any] | None | Unset
+        if isinstance(self.question, Unset):
+            question = UNSET
+        elif isinstance(self.question, ForwardLookingQuestion):
+            question = self.question.to_dict()
+        elif isinstance(self.question, Question):
+            question = self.question.to_dict()
+        else:
+            question = self.question
+
+        label: dict[str, Any] | None | Unset
+        if isinstance(self.label, Unset):
+            label = UNSET
+        elif isinstance(self.label, Label):
+            label = self.label.to_dict()
+        else:
+            label = self.label
+
+        prompt: None | str | Unset
+        if isinstance(self.prompt, Unset):
+            prompt = UNSET
+        else:
+            prompt = self.prompt
+
+        context: list[dict[str, Any]] | None | Unset
+        if isinstance(self.context, Unset):
+            context = UNSET
+        elif isinstance(self.context, list):
+            context = []
+            for context_type_0_item_data in self.context:
+                context_type_0_item: dict[str, Any]
+                if isinstance(context_type_0_item_data, NewsContext):
+                    context_type_0_item = context_type_0_item_data.to_dict()
+                else:
+                    context_type_0_item = context_type_0_item_data.to_dict()
+
+                context.append(context_type_0_item)
+
+        else:
+            context = self.context
+
+        rollouts: list[dict[str, Any]] | None | Unset
+        if isinstance(self.rollouts, Unset):
+            rollouts = UNSET
+        elif isinstance(self.rollouts, list):
+            rollouts = []
+            for rollouts_type_0_item_data in self.rollouts:
+                rollouts_type_0_item = rollouts_type_0_item_data.to_dict()
+                rollouts.append(rollouts_type_0_item)
+
+        else:
+            rollouts = self.rollouts
+
+        meta: dict[str, Any] | Unset = UNSET
+        if not isinstance(self.meta, Unset):
+            meta = self.meta.to_dict()
+
+        is_valid = self.is_valid
+
+        field_dict: dict[str, Any] = {}
+        field_dict.update(self.additional_properties)
+        field_dict.update({})
+        if seed is not UNSET:
+            field_dict["seed"] = seed
+        if question is not UNSET:
+            field_dict["question"] = question
+        if label is not UNSET:
+            field_dict["label"] = label
+        if prompt is not UNSET:
+            field_dict["prompt"] = prompt
+        if context is not UNSET:
+            field_dict["context"] = context
+        if rollouts is not UNSET:
+            field_dict["rollouts"] = rollouts
+        if meta is not UNSET:
+            field_dict["meta"] = meta
+        if is_valid is not UNSET:
+            field_dict["is_valid"] = is_valid
+
+        return field_dict
+
+    @classmethod
+    def from_dict(cls: type[T], src_dict: Mapping[str, Any]) -> T:
+        from ..models.forward_looking_question import ForwardLookingQuestion
+        from ..models.label import Label
+        from ..models.news_context import NewsContext
+        from ..models.question import Question
+        from ..models.rag_context import RAGContext
+        from ..models.rollout import Rollout
+        from ..models.sample_meta import SampleMeta
+        from ..models.seed import Seed
+
+        d = dict(src_dict)
+
+        def _parse_seed(data: object) -> None | Seed | Unset:
+            if data is None:
+                return data
+            if isinstance(data, Unset):
+                return data
+            try:
+                if not isinstance(data, dict):
+                    raise TypeError()
+                seed_type_0 = Seed.from_dict(data)
+
+                return seed_type_0
+            except (TypeError, ValueError, AttributeError, KeyError):
+                pass
+            return cast(None | Seed | Unset, data)
+
+        seed = _parse_seed(d.pop("seed", UNSET))
+
+        def _parse_question(data: object) -> ForwardLookingQuestion | None | Question | Unset:
+            if data is None:
+                return data
+            if isinstance(data, Unset):
+                return data
+            try:
+                if not isinstance(data, dict):
+                    raise TypeError()
+                question_type_0_type_0 = ForwardLookingQuestion.from_dict(data)
+
+                return question_type_0_type_0
+            except (TypeError, ValueError, AttributeError, KeyError):
+                pass
+            try:
+                if not isinstance(data, dict):
+                    raise TypeError()
+                question_type_0_type_1 = Question.from_dict(data)
+
+                return question_type_0_type_1
+            except (TypeError, ValueError, AttributeError, KeyError):
+                pass
+            return cast(ForwardLookingQuestion | None | Question | Unset, data)
+
+        question = _parse_question(d.pop("question", UNSET))
+
+        def _parse_label(data: object) -> Label | None | Unset:
+            if data is None:
+                return data
+            if isinstance(data, Unset):
+                return data
+            try:
+                if not isinstance(data, dict):
+                    raise TypeError()
+                label_type_0 = Label.from_dict(data)
+
+                return label_type_0
+            except (TypeError, ValueError, AttributeError, KeyError):
+                pass
+            return cast(Label | None | Unset, data)
+
+        label = _parse_label(d.pop("label", UNSET))
+
+        def _parse_prompt(data: object) -> None | str | Unset:
+            if data is None:
+                return data
+            if isinstance(data, Unset):
+                return data
+            return cast(None | str | Unset, data)
+
+        prompt = _parse_prompt(d.pop("prompt", UNSET))
+
+        def _parse_context(data: object) -> list[NewsContext | RAGContext] | None | Unset:
+            if data is None:
+                return data
+            if isinstance(data, Unset):
+                return data
+            try:
+                if not isinstance(data, list):
+                    raise TypeError()
+                context_type_0 = []
+                _context_type_0 = data
+                for context_type_0_item_data in _context_type_0:
+
+                    def _parse_context_type_0_item(data: object) -> NewsContext | RAGContext:
+                        try:
+                            if not isinstance(data, dict):
+                                raise TypeError()
+                            context_type_0_item_type_0 = NewsContext.from_dict(data)
+
+                            return context_type_0_item_type_0
+                        except (TypeError, ValueError, AttributeError, KeyError):
+                            pass
+                        if not isinstance(data, dict):
+                            raise TypeError()
+                        context_type_0_item_type_1 = RAGContext.from_dict(data)
+
+                        return context_type_0_item_type_1
+
+                    context_type_0_item = _parse_context_type_0_item(context_type_0_item_data)
+
+                    context_type_0.append(context_type_0_item)
+
+                return context_type_0
+            except (TypeError, ValueError, AttributeError, KeyError):
+                pass
+            return cast(list[NewsContext | RAGContext] | None | Unset, data)
+
+        context = _parse_context(d.pop("context", UNSET))
+
+        def _parse_rollouts(data: object) -> list[Rollout] | None | Unset:
+            if data is None:
+                return data
+            if isinstance(data, Unset):
+                return data
+            try:
+                if not isinstance(data, list):
+                    raise TypeError()
+                rollouts_type_0 = []
+                _rollouts_type_0 = data
+                for rollouts_type_0_item_data in _rollouts_type_0:
+                    rollouts_type_0_item = Rollout.from_dict(rollouts_type_0_item_data)
+
+                    rollouts_type_0.append(rollouts_type_0_item)
+
+                return rollouts_type_0
+            except (TypeError, ValueError, AttributeError, KeyError):
+                pass
+            return cast(list[Rollout] | None | Unset, data)
+
+        rollouts = _parse_rollouts(d.pop("rollouts", UNSET))
+
+        _meta = d.pop("meta", UNSET)
+        meta: SampleMeta | Unset
+        if isinstance(_meta, Unset):
+            meta = UNSET
+        else:
+            meta = SampleMeta.from_dict(_meta)
+
+        is_valid = d.pop("is_valid", UNSET)
+
+        sample = cls(
+            seed=seed,
+            question=question,
+            label=label,
+            prompt=prompt,
+            context=context,
+            rollouts=rollouts,
+            meta=meta,
+            is_valid=is_valid,
+        )
+
+        sample.additional_properties = d
+        return sample
+
+    @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/sample_meta.py

@@ -0,0 +1,46 @@
+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="SampleMeta")
+
+
+@_attrs_define
+class SampleMeta:
+    """ """
+
+    additional_properties: dict[str, Any] = _attrs_field(init=False, factory=dict)
+
+    def to_dict(self) -> dict[str, Any]:
+        field_dict: dict[str, Any] = {}
+        field_dict.update(self.additional_properties)
+
+        return field_dict
+
+    @classmethod
+    def from_dict(cls: type[T], src_dict: Mapping[str, Any]) -> T:
+        d = dict(src_dict)
+        sample_meta = cls()
+
+        sample_meta.additional_properties = d
+        return sample_meta
+
+    @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