trismik 0.9.12__py3-none-any.whl → 1.0.0__py3-none-any.whl

trismik/adaptive_test.py

@@ -1,669 +0,0 @@
-"""
-Trismik adaptive test runner.
-
-This module provides both synchronous and asynchronous interfaces for running
-Trismik tests. The async implementation is the core, with sync methods wrapping
-the async ones.
-"""
-
-import asyncio
-from typing import Any, Callable, Dict, List, Literal, Optional, Union, overload
-
-import nest_asyncio
-from tqdm.auto import tqdm
-
-from trismik.client_async import TrismikAsyncClient
-from trismik.settings import evaluation_settings
-from trismik.types import (
-    AdaptiveTestScore,
-    TrismikAdaptiveTestState,
-    TrismikClassicEvalRequest,
-    TrismikClassicEvalResponse,
-    TrismikDataset,
-    TrismikItem,
-    TrismikMeResponse,
-    TrismikProject,
-    TrismikReplayRequest,
-    TrismikReplayRequestItem,
-    TrismikRunMetadata,
-    TrismikRunResults,
-)
-
-
-class AdaptiveTest:
-    """
-    Trismik test runner with both sync and async interfaces.
-
-    This class provides both synchronous and asynchronous interfaces for
-    running Trismik tests. The async implementation is the core, with sync
-    methods wrapping the async ones.
-    """
-
-    def __init__(
-        self,
-        item_processor: Callable[[TrismikItem], Any],
-        client: Optional[TrismikAsyncClient] = None,
-        api_key: Optional[str] = None,
-        max_items: int = evaluation_settings["max_iterations"],
-    ) -> None:
-        """
-        Initialize a new Trismik runner.
-
-        Args:
-            item_processor (Callable[[TrismikItem], Any]): Function to process
-              test items. For async usage, this should be an async function.
-            client (Optional[TrismikAsyncClient]): Trismik async client to use
-              for requests. If not provided, a new one will be created.
-            api_key (Optional[str]): API key to use if a new client is created.
-            max_items (int): Maximum number of items to process. Default is 60.
-
-        Raises:
-            ValueError: If both client and api_key are provided.
-            TrismikApiError: If API request fails.
-        """
-        if client and api_key:
-            raise ValueError(
-                "Either 'client' or 'api_key' should be provided, not both."
-            )
-        self._item_processor = item_processor
-        if client:
-            self._client = client
-        else:
-            self._client = TrismikAsyncClient(api_key=api_key)
-        self._max_items = max_items
-        self._loop = None
-
-    def _get_loop(self) -> asyncio.AbstractEventLoop:
-        """
-        Get or create an event loop, handling nested loops if needed.
-
-        Returns:
-            asyncio.AbstractEventLoop: The event loop to use.
-        """
-        try:
-            loop = asyncio.get_event_loop()
-        except RuntimeError:
-            # No event loop in this thread, create one
-            loop = asyncio.new_event_loop()
-            asyncio.set_event_loop(loop)
-
-        # Allow nested event loops (needed for Jupyter, etc)
-        nest_asyncio.apply(loop)
-        return loop
-
-    def list_datasets(self) -> List[TrismikDataset]:
-        """
-        Get a list of available datasets synchronously.
-
-        Returns:
-            List[TrismikDataset]: List of available datasets.
-
-        Raises:
-            TrismikApiError: If API request fails.
-        """
-        loop = self._get_loop()
-        return loop.run_until_complete(self.list_datasets_async())
-
-    async def list_datasets_async(self) -> List[TrismikDataset]:
-        """
-        Get a list of available datasets asynchronously.
-
-        Returns:
-            List[TrismikDataset]: List of available datasets.
-
-        Raises:
-            TrismikApiError: If API request fails.
-        """
-        return await self._client.list_datasets()
-
-    def me(self) -> TrismikMeResponse:
-        """
-        Get current user information synchronously.
-
-        Returns:
-            TrismikMeResponse: User information including validity and payload.
-
-        Raises:
-            TrismikApiError: If API request fails.
-        """
-        loop = self._get_loop()
-        return loop.run_until_complete(self.me_async())
-
-    async def me_async(self) -> TrismikMeResponse:
-        """
-        Get current user information asynchronously.
-
-        Returns:
-            TrismikMeResponse: User information including validity and payload.
-
-        Raises:
-            TrismikApiError: If API request fails.
-        """
-        return await self._client.me()
-
-    def create_project(
-        self,
-        name: str,
-        team_id: Optional[str] = None,
-        description: Optional[str] = None,
-    ) -> TrismikProject:
-        """
-        Create a new project synchronously.
-
-        Args:
-            name (str): Name of the project.
-            team_id (str): ID of the team to create the
-                project in.
-            description (Optional[str]): Optional description of the project.
-
-        Returns:
-            TrismikProject: Created project information.
-
-        Raises:
-            TrismikValidationError: If the request fails validation.
-            TrismikApiError: If API request fails.
-        """
-        loop = self._get_loop()
-        return loop.run_until_complete(
-            self.create_project_async(name, team_id, description)
-        )
-
-    async def create_project_async(
-        self,
-        name: str,
-        team_id: Optional[str] = None,
-        description: Optional[str] = None,
-    ) -> TrismikProject:
-        """
-        Create a new project asynchronously.
-
-        Args:
-            name (str): Name of the project.
-            team_id (str): ID of the team to create the
-                project in.
-            description (Optional[str]): Optional description of the project.
-
-        Returns:
-            TrismikProject: Created project information.
-
-        Raises:
-            TrismikValidationError: If the request fails validation.
-            TrismikApiError: If API request fails.
-        """
-        return await self._client.create_project(name, team_id, description)
-
-    @overload
-    def run(  # noqa: E704
-        self,
-        test_id: str,
-        project_id: str,
-        experiment: str,
-        run_metadata: TrismikRunMetadata,
-        return_dict: Literal[True],
-        with_responses: bool = False,
-    ) -> Dict[str, Any]: ...
-
-    @overload
-    def run(  # noqa: E704
-        self,
-        test_id: str,
-        project_id: str,
-        experiment: str,
-        run_metadata: TrismikRunMetadata,
-        return_dict: Literal[False],
-        with_responses: bool = False,
-    ) -> TrismikRunResults: ...
-
-    def run(
-        self,
-        test_id: str,
-        project_id: str,
-        experiment: str,
-        run_metadata: TrismikRunMetadata,
-        return_dict: bool = True,
-        with_responses: bool = False,
-    ) -> Union[TrismikRunResults, Dict[str, Any]]:
-        """
-        Run a test synchronously.
-
-        Args:
-            test_id (str): ID of the test to run.
-            project_id (str): ID of the project.
-            experiment (str): Name of the experiment.
-            run_metadata (TrismikRunMetadata): Metadata for the
-              run.
-            return_dict (bool): If True, return results as a dictionary instead
-              of TrismikRunResults object. Defaults to True.
-            with_responses (bool): If True, responses will be included with
-              the results.
-
-        Returns:
-            Union[TrismikRunResults, Dict[str, Any]]: Either TrismikRunResults
-              object or dictionary representation based on return_dict
-              parameter.
-
-        Raises:
-            TrismikApiError: If API request fails.
-            NotImplementedError: If with_responses = True (not yet implemented).
-        """
-        loop = self._get_loop()
-        if return_dict:
-            return loop.run_until_complete(
-                self.run_async(
-                    test_id,
-                    project_id,
-                    experiment,
-                    run_metadata,
-                    True,
-                    with_responses,
-                )
-            )
-        else:
-            return loop.run_until_complete(
-                self.run_async(
-                    test_id,
-                    project_id,
-                    experiment,
-                    run_metadata,
-                    False,
-                    with_responses,
-                )
-            )
-
-    @overload
-    async def run_async(  # noqa: E704
-        self,
-        test_id: str,
-        project_id: str,
-        experiment: str,
-        run_metadata: TrismikRunMetadata,
-        return_dict: Literal[True],
-        with_responses: bool = False,
-    ) -> Dict[str, Any]: ...
-
-    @overload
-    async def run_async(  # noqa: E704
-        self,
-        test_id: str,
-        project_id: str,
-        experiment: str,
-        run_metadata: TrismikRunMetadata,
-        return_dict: Literal[False],
-        with_responses: bool = False,
-    ) -> TrismikRunResults: ...
-
-    async def run_async(
-        self,
-        test_id: str,
-        project_id: str,
-        experiment: str,
-        run_metadata: TrismikRunMetadata,
-        return_dict: bool = True,
-        with_responses: bool = False,
-    ) -> Union[TrismikRunResults, Dict[str, Any]]:
-        """
-        Run a test asynchronously.
-
-        Args:
-            test_id: ID of the test to run.
-            project_id: ID of the project.
-            experiment: Name of the experiment.
-            run_metadata: Metadata for the run.
-            return_dict: If True, return results as a dictionary instead
-              of TrismikRunResults object. Defaults to True.
-            with_responses: If True, responses will be included with
-              the results.
-
-        Returns:
-            Union[TrismikRunResults, Dict[str, Any]]: Either TrismikRunResults
-              object or dictionary representation based on return_dict
-              parameter.
-
-        Raises:
-            TrismikApiError: If API request fails.
-            NotImplementedError: If with_responses = True (not yet implemented).
-        """
-        if with_responses:
-            raise NotImplementedError(
-                "with_responses is not yet implemented for the new API flow"
-            )
-
-        # Start run and get first item
-        start_response = await self._client.start_run(
-            test_id, project_id, experiment, run_metadata
-        )
-
-        # Initialize state tracking
-        states: List[TrismikAdaptiveTestState] = []
-        run_id = start_response.run_info.id
-
-        # Add initial state
-        states.append(
-            TrismikAdaptiveTestState(
-                run_id=run_id,
-                state=start_response.state,
-                completed=start_response.completed,
-            )
-        )
-
-        # Run the test and get last state
-        last_state = await self._run_async(
-            run_id, start_response.next_item, states
-        )
-
-        if not last_state:
-            raise RuntimeError(
-                "Test run completed but no final state was captured"
-            )
-
-        score = AdaptiveTestScore(
-            theta=last_state.state.thetas[-1],
-            std_error=last_state.state.std_error_history[-1],
-        )
-
-        results = TrismikRunResults(run_id, score=score)
-
-        if return_dict:
-            return {
-                "run_id": results.run_id,
-                "score": (
-                    {
-                        "theta": results.score.theta,
-                        "std_error": results.score.std_error,
-                    }
-                    if results.score
-                    else None
-                ),
-                "responses": results.responses,
-            }
-        else:
-            return results
-
-    @overload
-    def run_replay(  # noqa: E704
-        self,
-        previous_run_id: str,
-        run_metadata: TrismikRunMetadata,
-        return_dict: Literal[True],
-        with_responses: bool = False,
-    ) -> Dict[str, Any]: ...
-
-    @overload
-    def run_replay(  # noqa: E704
-        self,
-        previous_run_id: str,
-        run_metadata: TrismikRunMetadata,
-        return_dict: Literal[False],
-        with_responses: bool = False,
-    ) -> TrismikRunResults: ...
-
-    def run_replay(
-        self,
-        previous_run_id: str,
-        run_metadata: TrismikRunMetadata,
-        return_dict: bool = True,
-        with_responses: bool = False,
-    ) -> Union[TrismikRunResults, Dict[str, Any]]:
-        """
-        Replay the exact sequence of questions from a previous run.
-
-        Wraps the run_replay_async method.
-
-        Args:
-            previous_run_id: ID of a previous run to replay.
-            run_metadata: Metadata for the replay run.
-            return_dict: If True, return results as a dictionary instead
-              of TrismikRunResults object. Defaults to True.
-            with_responses: If True, responses will be included
-             with the results.
-
-        Returns:
-            Union[TrismikRunResults, Dict[str, Any]]: Either TrismikRunResults
-              object or dictionary representation based on return_dict
-              parameter.
-
-        Raises:
-            TrismikApiError: If API request fails.
-        """
-        loop = self._get_loop()
-        if return_dict:
-            return loop.run_until_complete(
-                self.run_replay_async(
-                    previous_run_id,
-                    run_metadata,
-                    True,
-                    with_responses,
-                )
-            )
-        else:
-            return loop.run_until_complete(
-                self.run_replay_async(
-                    previous_run_id,
-                    run_metadata,
-                    False,
-                    with_responses,
-                )
-            )
-
-    @overload
-    async def run_replay_async(  # noqa: E704
-        self,
-        previous_run_id: str,
-        run_metadata: TrismikRunMetadata,
-        return_dict: Literal[True],
-        with_responses: bool = False,
-    ) -> Dict[str, Any]: ...
-
-    @overload
-    async def run_replay_async(  # noqa: E704
-        self,
-        previous_run_id: str,
-        run_metadata: TrismikRunMetadata,
-        return_dict: Literal[False],
-        with_responses: bool = False,
-    ) -> TrismikRunResults: ...
-
-    async def run_replay_async(
-        self,
-        previous_run_id: str,
-        run_metadata: TrismikRunMetadata,
-        return_dict: bool = True,
-        with_responses: bool = False,
-    ) -> Union[TrismikRunResults, Dict[str, Any]]:
-        """
-        Replay the exact sequence of questions from a previous run.
-
-        Args:
-            previous_run_id: ID of a previous run to replay.
-            run_metadata: Metadata for the run.
-            return_dict: If True, return results as a dictionary instead
-              of TrismikRunResults object. Defaults to True.
-            with_responses: If True, responses will be included
-              with the results.
-
-        Returns:
-            Union[TrismikRunResults, Dict[str, Any]]: Either TrismikRunResults
-              object or dictionary representation based on return_dict
-              parameter.
-
-        Raises:
-            TrismikApiError: If API request fails.
-        """
-        # Get the original run summary to access dataset and responses
-        original_summary = await self._client.run_summary(previous_run_id)
-
-        # Build replay request by processing each item in the original order
-        replay_items = []
-        with tqdm(
-            total=len(original_summary.dataset), desc="Running replay..."
-        ) as pbar:
-            for item in original_summary.dataset:
-                # Handle both sync and async item processors
-                if asyncio.iscoroutinefunction(self._item_processor):
-                    response = await self._item_processor(item)
-                else:
-                    response = self._item_processor(item)
-
-                # Create replay request item
-                replay_item = TrismikReplayRequestItem(
-                    itemId=item.id, itemChoiceId=response
-                )
-                replay_items.append(replay_item)
-                pbar.update(1)
-
-        # Create replay request
-        replay_request = TrismikReplayRequest(responses=replay_items)
-
-        # Submit replay with metadata
-        replay_response = await self._client.submit_replay(
-            previous_run_id, replay_request, run_metadata
-        )
-
-        # Create score from replay response
-        score = AdaptiveTestScore(
-            theta=replay_response.state.thetas[-1],
-            std_error=replay_response.state.std_error_history[-1],
-        )
-
-        # Return results with optional responses
-        if with_responses:
-            results = TrismikRunResults(
-                run_id=replay_response.id,
-                score=score,
-                responses=replay_response.responses,
-            )
-        else:
-            results = TrismikRunResults(run_id=replay_response.id, score=score)
-
-        if return_dict:
-            return {
-                "run_id": results.run_id,
-                "score": (
-                    {
-                        "theta": results.score.theta,
-                        "std_error": results.score.std_error,
-                    }
-                    if results.score
-                    else None
-                ),
-                "responses": (
-                    [
-                        {
-                            "dataset_item_id": resp.dataset_item_id,
-                            "value": resp.value,
-                            "correct": resp.correct,
-                        }
-                        for resp in results.responses
-                    ]
-                    if results.responses
-                    else None
-                ),
-            }
-        else:
-            return results
-
-    async def _run_async(
-        self,
-        run_id: str,
-        first_item: Optional[TrismikItem],
-        states: List[TrismikAdaptiveTestState],
-    ) -> Optional[TrismikAdaptiveTestState]:
-        """
-        Run a test asynchronously.
-
-        Args:
-            run_id (str): ID of the run to execute.
-            first_item (Optional[TrismikItem]): First item from run start.
-            states (List[TrismikAdaptiveTestState]): List to accumulate states.
-
-        Returns:
-            Optional[TrismikAdaptiveTestState]: Last state of the run.
-
-        Raises:
-            TrismikApiError: If API request fails.
-        """
-        item = first_item
-        with tqdm(total=self._max_items, desc="Evaluating") as pbar:
-            while item is not None:
-                # Handle both sync and async item processors
-                if asyncio.iscoroutinefunction(self._item_processor):
-                    response = await self._item_processor(item)
-                else:
-                    response = self._item_processor(item)
-
-                # Continue run with response
-                continue_response = await self._client.continue_run(
-                    run_id, response
-                )
-
-                # Update state tracking
-                states.append(
-                    TrismikAdaptiveTestState(
-                        run_id=run_id,
-                        state=continue_response.state,
-                        completed=continue_response.completed,
-                    )
-                )
-
-                pbar.update(1)
-
-                if continue_response.completed:
-                    pbar.total = pbar.n  # Update total to current position
-                    pbar.refresh()
-                    break
-
-                item = continue_response.next_item
-
-        last_state = states[-1] if states else None
-
-        return last_state
-
-    def submit_classic_eval(
-        self, classic_eval_request: TrismikClassicEvalRequest
-    ) -> TrismikClassicEvalResponse:
-        """
-        Submit a classic evaluation run with pre-computed results synchronously.
-
-        Args:
-            classic_eval_request (TrismikClassicEvalRequest): Request containing
-                project info, dataset, model outputs, and metrics.
-
-        Returns:
-            TrismikClassicEvalResponse: Response from the classic evaluation
-                endpoint.
-
-        Raises:
-            TrismikPayloadTooLargeError: If the request payload exceeds the
-                server's size limit.
-            TrismikValidationError: If the request fails validation.
-            TrismikApiError: If API request fails.
-        """
-        loop = self._get_loop()
-        return loop.run_until_complete(
-            self.submit_classic_eval_async(classic_eval_request)
-        )
-
-    async def submit_classic_eval_async(
-        self, classic_eval_request: TrismikClassicEvalRequest
-    ) -> TrismikClassicEvalResponse:
-        """
-        Submit a classic evaluation run with pre-computed results async.
-
-        This method allows you to submit pre-computed model outputs and metrics
-        for evaluation without running an interactive test.
-
-        Args:
-            classic_eval_request (TrismikClassicEvalRequest): Request containing
-                project info, dataset, model outputs, and metrics.
-
-        Returns:
-            TrismikClassicEvalResponse: Response from the classic evaluation
-                endpoint.
-
-        Raises:
-            TrismikPayloadTooLargeError: If the request payload exceeds the
-                server's size limit.
-            TrismikValidationError: If the request fails validation.
-            TrismikApiError: If API request fails.
-        """
-        return await self._client.submit_classic_eval(classic_eval_request)