# Generated by ariadne-codegen
# Source: vals/graphql/

from typing import Any, Dict, List, Optional, Union

from .add_all_tests_to_queue_single import AddAllTestsToQueueSingle
from .add_batch_tests import AddBatchTests
from .add_or_remove_users_to_run_review import AddOrRemoveUsersToRunReview
from .async_base_client import AsyncBaseClient
from .base_model import UNSET, UnsetType
from .batch_add_question_answer_pairs import BatchAddQuestionAnswerPairs
from .create_or_update_test_suite import CreateOrUpdateTestSuite
from .create_question_answer_set import CreateQuestionAnswerSet
from .delete_test_suite import DeleteTestSuite
from .enums import RunStatus
from .get_custom_metric import GetCustomMetric
from .get_default_parameters import GetDefaultParameters
from .get_run_dataframe import GetRunDataframe
from .get_run_status import GetRunStatus
from .get_single_run_review import GetSingleRunReview
from .get_test_data import GetTestData
from .get_test_suite_data import GetTestSuiteData
from .get_test_suites_with_count import GetTestSuitesWithCount
from .get_user_options import GetUserOptions
from .input_types import (
    CheckInputType,
    LocalEvalUploadInputType,
    ParameterInputType,
    QuestionAnswerPairInputType,
    TestMutationInfo,
    TestResultReviewOptionsInput,
    TestReviewFilterOptionsInput,
)
from .list_projects import ListProjects
from .list_question_answer_pairs import ListQuestionAnswerPairs
from .list_runs import ListRuns
from .mark_question_answer_set_as_complete import MarkQuestionAnswerSetAsComplete
from .poll_custom_metric_task import PollCustomMetricTask
from .pull_run import PullRun
from .pull_test_results_with_count import PullTestResultsWithCount
from .remove_old_tests import RemoveOldTests
from .rerun_tests import RerunTests
from .set_archived_status_custom_metrics import SetArchivedStatusCustomMetrics
from .set_custom_metrics import SetCustomMetrics
from .single_test_result_reviews_with_count import SingleTestResultReviewsWithCount
from .start_custom_metric_task import StartCustomMetricTask
from .start_run import StartRun
from .stop_run import StopRun
from .unfinished_tests import UnfinishedTests
from .update_custom_metric_result import UpdateCustomMetricResult
from .update_global_checks import UpdateGlobalChecks
from .update_run_status import UpdateRunStatus
from .upload_local_evaluation import UploadLocalEvaluation
from .upsert_custom_metric import UpsertCustomMetric


def gql(q: str) -> str:
    return q


class Client(AsyncBaseClient):
    async def upsert_custom_metric(
        self,
        project_id: str,
        name: str,
        file_id: str,
        update_past: bool,
        id: Union[Optional[str], UnsetType] = UNSET,
        description: Union[Optional[str], UnsetType] = UNSET,
        **kwargs: Any
    ) -> UpsertCustomMetric:
        query = gql(
            """
            mutation upsertCustomMetric($id: String, $project_id: String!, $name: String!, $description: String, $file_id: String!, $update_past: Boolean!) {
              upsertCustomMetric(
                id: $id
                projectId: $project_id
                name: $name
                description: $description
                fileId: $file_id
                updatePast: $update_past
              ) {
                metric {
                  id
                }
              }
            }
            """
        )
        variables: Dict[str, object] = {
            "id": id,
            "project_id": project_id,
            "name": name,
            "description": description,
            "file_id": file_id,
            "update_past": update_past,
        }
        response = await self.execute(
            query=query,
            operation_name="upsertCustomMetric",
            variables=variables,
            **kwargs
        )
        data = self.get_data(response)
        return UpsertCustomMetric.model_validate(data)

    async def set_archived_status_custom_metrics(
        self, ids: List[str], archived: bool, **kwargs: Any
    ) -> SetArchivedStatusCustomMetrics:
        query = gql(
            """
            mutation setArchivedStatusCustomMetrics($ids: [String!]!, $archived: Boolean!) {
              setArchivedStatusCustomMetrics(ids: $ids, archived: $archived) {
                affectedIds
              }
            }
            """
        )
        variables: Dict[str, object] = {"ids": ids, "archived": archived}
        response = await self.execute(
            query=query,
            operation_name="setArchivedStatusCustomMetrics",
            variables=variables,
            **kwargs
        )
        data = self.get_data(response)
        return SetArchivedStatusCustomMetrics.model_validate(data)

    async def start_custom_metric_task(
        self, run_id: str, file_id: str, **kwargs: Any
    ) -> StartCustomMetricTask:
        query = gql(
            """
            mutation startCustomMetricTask($run_id: String!, $file_id: String!) {
              startCustomMetricTask(runId: $run_id, fileId: $file_id, temporary: false) {
                messageId
              }
            }
            """
        )
        variables: Dict[str, object] = {"run_id": run_id, "file_id": file_id}
        response = await self.execute(
            query=query,
            operation_name="startCustomMetricTask",
            variables=variables,
            **kwargs
        )
        data = self.get_data(response)
        return StartCustomMetricTask.model_validate(data)

    async def get_custom_metric(self, id: str, **kwargs: Any) -> GetCustomMetric:
        query = gql(
            """
            query getCustomMetric($id: String!) {
              customMetric(id: $id) {
                id
                project {
                  slug
                }
                archived
                name
                description
                fileId
              }
            }
            """
        )
        variables: Dict[str, object] = {"id": id}
        response = await self.execute(
            query=query, operation_name="getCustomMetric", variables=variables, **kwargs
        )
        data = self.get_data(response)
        return GetCustomMetric.model_validate(data)

    async def poll_custom_metric_task(
        self, message_id: str, **kwargs: Any
    ) -> PollCustomMetricTask:
        query = gql(
            """
            query pollCustomMetricTask($messageId: String!) {
              pollCustomMetricTask(messageId: $messageId) {
                status
                result
                error
                traceback
              }
            }
            """
        )
        variables: Dict[str, object] = {"messageId": message_id}
        response = await self.execute(
            query=query,
            operation_name="pollCustomMetricTask",
            variables=variables,
            **kwargs
        )
        data = self.get_data(response)
        return PollCustomMetricTask.model_validate(data)

    async def get_run_dataframe(self, run_id: str, **kwargs: Any) -> GetRunDataframe:
        query = gql(
            """
            query GetRunDataframe($runId: String!) {
              getRunDataframe(runId: $runId) {
                id
                tags
                input
                inputContext
                output
                outputContext
                rightAnswer
                refusedToAnswer
                isRephrasal
                beenRephrased
                fileIds
                operator
                criteria
                eval
                cont
                feedback
                isGlobal
                modifiers
              }
            }
            """
        )
        variables: Dict[str, object] = {"runId": run_id}
        response = await self.execute(
            query=query, operation_name="GetRunDataframe", variables=variables, **kwargs
        )
        data = self.get_data(response)
        return GetRunDataframe.model_validate(data)

    async def add_all_tests_to_queue_single(
        self,
        run_id: str,
        test_result_review_options: List[TestResultReviewOptionsInput],
        number_of_reviews: int,
        rereview_auto_eval: bool,
        project_id: str,
        template_ids: Union[Optional[List[str]], UnsetType] = UNSET,
        **kwargs: Any
    ) -> AddAllTestsToQueueSingle:
        query = gql(
            """
            mutation AddAllTestsToQueueSingle($runId: String!, $templateIds: [String!], $testResultReviewOptions: [TestResultReviewOptionsInput]!, $numberOfReviews: Int!, $rereviewAutoEval: Boolean!, $projectId: String!) {
              addAllSingleTestReviewToQueue(
                runId: $runId
                templateIds: $templateIds
                testResultReviewOptions: $testResultReviewOptions
                numberOfReviews: $numberOfReviews
                rereviewAutoEval: $rereviewAutoEval
                projectId: $projectId
              ) {
                singleRunReview {
                  id
                }
              }
            }
            """
        )
        variables: Dict[str, object] = {
            "runId": run_id,
            "templateIds": template_ids,
            "testResultReviewOptions": test_result_review_options,
            "numberOfReviews": number_of_reviews,
            "rereviewAutoEval": rereview_auto_eval,
            "projectId": project_id,
        }
        response = await self.execute(
            query=query,
            operation_name="AddAllTestsToQueueSingle",
            variables=variables,
            **kwargs
        )
        data = self.get_data(response)
        return AddAllTestsToQueueSingle.model_validate(data)

    async def get_single_run_review(
        self, run_review_id: str, project_id: str, **kwargs: Any
    ) -> GetSingleRunReview:
        query = gql(
            """
            query GetSingleRunReview($runReviewId: String!, $projectId: String!) {
              singleRunReview(runReviewId: $runReviewId, projectId: $projectId) {
                id
                createdBy
                createdAt
                status
                passRateHumanEval
                agreementRateAutoEval
                agreementRateHumanEval
                flaggedRate
                completedTime
                numberOfReviews
                assignedReviewers
                rereviewAutoEval
                project {
                  id
                }
                run {
                  id
                }
              }
            }
            """
        )
        variables: Dict[str, object] = {
            "runReviewId": run_review_id,
            "projectId": project_id,
        }
        response = await self.execute(
            query=query,
            operation_name="GetSingleRunReview",
            variables=variables,
            **kwargs
        )
        data = self.get_data(response)
        return GetSingleRunReview.model_validate(data)

    async def single_test_result_reviews_with_count(
        self,
        run_id: str,
        filter_options: Union[
            Optional[TestReviewFilterOptionsInput], UnsetType
        ] = UNSET,
        **kwargs: Any
    ) -> SingleTestResultReviewsWithCount:
        query = gql(
            """
            query SingleTestResultReviewsWithCount($runId: String!, $filterOptions: TestReviewFilterOptionsInput) {
              testResultReviewsWithCount(runId: $runId, filterOptions: $filterOptions) {
                count
                singleTestResults {
                  reviewedBy
                  ...TestResultFields
                  singleTestReviews {
                    id
                    completedBy
                    feedback
                    completedAt
                    startedAt
                    createdBy
                    status
                    perCheckTestReview {
                      binaryHumanEval
                      isFlagged
                    }
                    testResult {
                      id
                      resultJson {
                        ...ResultJsonElementTypeFields
                      }
                    }
                    customReviewValues {
                      template {
                        id
                        name
                        instructions
                        optional
                        categories
                        type
                        minValue
                        maxValue
                      }
                      value
                    }
                  }
                }
              }
            }

            fragment CheckModifierOutputTypeFields on CheckModifierOutputType {
              optional
              severity
              displayMetrics
              examples {
                ...ExampleOutputTypeFields
              }
              extractor
              conditional {
                ...ConditionalCheckOutputTypeFields
              }
              category
            }

            fragment ConditionalCheckOutputTypeFields on ConditionalCheckOutputType {
              operator
              criteria
            }

            fragment ExampleOutputTypeFields on ExampleOutputType {
              type
              text
            }

            fragment MetadataOutputTypeFields on MetadataOutputType {
              inTokens
              outTokens
              durationSeconds
            }

            fragment ResultJsonElementTypeFields on ResultJsonElementType {
              operator
              criteria
              modifiers {
                ...CheckModifierOutputTypeFields
              }
              autoEval
              feedback
              evalCont
              isGlobal
              severity
              customOperator {
                rubric {
                  label
                  score
                }
              }
            }

            fragment TestFragment on TestType {
              id
              inputUnderTest
              context
              tags
              crossVersionId
              goldenOutput
              fileIds
              checks {
                operator
                criteria
                modifiers {
                  optional
                  severity
                  examples {
                    type
                    text
                  }
                  extractor
                  conditional {
                    operator
                    criteria
                  }
                  category
                  displayMetrics
                }
              }
              testSuite {
                id
              }
            }

            fragment TestResultFields on TestResultType {
              id
              llmOutput
              passPercentage
              passPercentageWithWeight
              resultJson {
                ...ResultJsonElementTypeFields
              }
              qaPair {
                context
                outputContext
                errorMessage
              }
              test {
                ...TestFragment
              }
              metadata {
                ...MetadataOutputTypeFields
              }
            }
            """
        )
        variables: Dict[str, object] = {
            "runId": run_id,
            "filterOptions": filter_options,
        }
        response = await self.execute(
            query=query,
            operation_name="SingleTestResultReviewsWithCount",
            variables=variables,
            **kwargs
        )
        data = self.get_data(response)
        return SingleTestResultReviewsWithCount.model_validate(data)

    async def get_user_options(self, project_id: str, **kwargs: Any) -> GetUserOptions:
        query = gql(
            """
            query GetUserOptions($projectId: String!) {
              userEmails(projectId: $projectId)
            }
            """
        )
        variables: Dict[str, object] = {"projectId": project_id}
        response = await self.execute(
            query=query, operation_name="GetUserOptions", variables=variables, **kwargs
        )
        data = self.get_data(response)
        return GetUserOptions.model_validate(data)

    async def add_or_remove_users_to_run_review(
        self,
        single_run_review_id: str,
        project_id: str,
        test_result_review_options: List[TestResultReviewOptionsInput],
        **kwargs: Any
    ) -> AddOrRemoveUsersToRunReview:
        query = gql(
            """
            mutation AddOrRemoveUsersToRunReview($singleRunReviewId: String!, $projectId: String!, $testResultReviewOptions: [TestResultReviewOptionsInput]!) {
              updateAssignedReviewers(
                singleRunReviewId: $singleRunReviewId
                projectId: $projectId
                testResultReviewOptions: $testResultReviewOptions
              ) {
                newUsersAddedToQueue
              }
            }
            """
        )
        variables: Dict[str, object] = {
            "singleRunReviewId": single_run_review_id,
            "projectId": project_id,
            "testResultReviewOptions": test_result_review_options,
        }
        response = await self.execute(
            query=query,
            operation_name="AddOrRemoveUsersToRunReview",
            variables=variables,
            **kwargs
        )
        data = self.get_data(response)
        return AddOrRemoveUsersToRunReview.model_validate(data)

    async def list_projects(
        self, offset: int, limit: int, search: str, **kwargs: Any
    ) -> ListProjects:
        query = gql(
            """
            query listProjects($offset: Int!, $limit: Int!, $search: String!) {
              projectsWithCount(
                filterOptions: {offset: $offset, limit: $limit, archived: false, search: $search}
              ) {
                projects {
                  id
                  name
                  slug
                  created
                  isDefault
                }
                count
              }
            }
            """
        )
        variables: Dict[str, object] = {
            "offset": offset,
            "limit": limit,
            "search": search,
        }
        response = await self.execute(
            query=query, operation_name="listProjects", variables=variables, **kwargs
        )
        data = self.get_data(response)
        return ListProjects.model_validate(data)

    async def create_question_answer_set(
        self,
        test_suite_id: str,
        question_answer_pairs: List[QuestionAnswerPairInputType],
        parameters: Any,
        project_id: str,
        model_id: str,
        run_name: Union[Optional[str], UnsetType] = UNSET,
        **kwargs: Any
    ) -> CreateQuestionAnswerSet:
        query = gql(
            """
            mutation CreateQuestionAnswerSet($testSuiteId: String!, $questionAnswerPairs: [QuestionAnswerPairInputType!]!, $parameters: GenericScalar!, $projectId: String!, $modelId: String!, $runName: String) {
              createQuestionAnswerSet(
                testSuiteId: $testSuiteId
                questionAnswerPairs: $questionAnswerPairs
                parameters: $parameters
                projectId: $projectId
                modelId: $modelId
                runName: $runName
                createRun: true
              ) {
                questionAnswerSet {
                  id
                }
                runId
              }
            }
            """
        )
        variables: Dict[str, object] = {
            "testSuiteId": test_suite_id,
            "questionAnswerPairs": question_answer_pairs,
            "parameters": parameters,
            "projectId": project_id,
            "modelId": model_id,
            "runName": run_name,
        }
        response = await self.execute(
            query=query,
            operation_name="CreateQuestionAnswerSet",
            variables=variables,
            **kwargs
        )
        data = self.get_data(response)
        return CreateQuestionAnswerSet.model_validate(data)

    async def batch_add_question_answer_pairs(
        self,
        question_answer_set_id: str,
        question_answer_pairs: List[QuestionAnswerPairInputType],
        **kwargs: Any
    ) -> BatchAddQuestionAnswerPairs:
        query = gql(
            """
            mutation BatchAddQuestionAnswerPairs($questionAnswerSetId: String!, $questionAnswerPairs: [QuestionAnswerPairInputType!]!) {
              batchAddQuestionAnswerPairs(
                questionAnswerSetId: $questionAnswerSetId
                questionAnswerPairs: $questionAnswerPairs
              ) {
                questionAnswerPairs {
                  id
                  inputUnderTest
                  llmOutput
                  fileIds
                  context
                  outputContext
                  test {
                    id
                  }
                }
              }
            }
            """
        )
        variables: Dict[str, object] = {
            "questionAnswerSetId": question_answer_set_id,
            "questionAnswerPairs": question_answer_pairs,
        }
        response = await self.execute(
            query=query,
            operation_name="BatchAddQuestionAnswerPairs",
            variables=variables,
            **kwargs
        )
        data = self.get_data(response)
        return BatchAddQuestionAnswerPairs.model_validate(data)

    async def mark_question_answer_set_as_complete(
        self, question_answer_set_id: str, **kwargs: Any
    ) -> MarkQuestionAnswerSetAsComplete:
        query = gql(
            """
            mutation MarkQuestionAnswerSetAsComplete($questionAnswerSetId: String!) {
              markQuestionAnswerSetAsComplete(questionAnswerSetId: $questionAnswerSetId) {
                questionAnswerSet {
                  id
                }
              }
            }
            """
        )
        variables: Dict[str, object] = {"questionAnswerSetId": question_answer_set_id}
        response = await self.execute(
            query=query,
            operation_name="MarkQuestionAnswerSetAsComplete",
            variables=variables,
            **kwargs
        )
        data = self.get_data(response)
        return MarkQuestionAnswerSetAsComplete.model_validate(data)

    async def list_question_answer_pairs(
        self,
        qa_set_id: str,
        offset: Union[Optional[int], UnsetType] = UNSET,
        limit: Union[Optional[int], UnsetType] = UNSET,
        **kwargs: Any
    ) -> ListQuestionAnswerPairs:
        query = gql(
            """
            query ListQuestionAnswerPairs($qaSetId: String!, $offset: Int, $limit: Int) {
              questionAnswerPairsWithCount(
                qaSetId: $qaSetId
                filterOptions: {offset: $offset, limit: $limit}
              ) {
                questionAnswerPairs {
                  id
                  inputUnderTest
                  llmOutput
                  context
                  outputContext
                  metadata {
                    inTokens
                    outTokens
                    durationSeconds
                  }
                  fileIds
                  localEvals {
                    id
                    score
                    feedback
                    createdAt
                  }
                  test {
                    id
                  }
                }
                count
              }
            }
            """
        )
        variables: Dict[str, object] = {
            "qaSetId": qa_set_id,
            "offset": offset,
            "limit": limit,
        }
        response = await self.execute(
            query=query,
            operation_name="ListQuestionAnswerPairs",
            variables=variables,
            **kwargs
        )
        data = self.get_data(response)
        return ListQuestionAnswerPairs.model_validate(data)

    async def start_run(
        self,
        test_suite_id: str,
        typed_parameters: ParameterInputType,
        qa_set_id: Union[Optional[str], UnsetType] = UNSET,
        run_name: Union[Optional[str], UnsetType] = UNSET,
        run_id: Union[Optional[str], UnsetType] = UNSET,
        **kwargs: Any
    ) -> StartRun:
        query = gql(
            """
            mutation startRun($test_suite_id: String!, $typed_parameters: ParameterInputType!, $qa_set_id: String = null, $run_name: String = null, $run_id: String = null) {
              startRun(
                testSuiteId: $test_suite_id
                typedParameters: $typed_parameters
                qaSetId: $qa_set_id
                runName: $run_name
                runId: $run_id
              ) {
                runId
              }
            }
            """
        )
        variables: Dict[str, object] = {
            "test_suite_id": test_suite_id,
            "typed_parameters": typed_parameters,
            "qa_set_id": qa_set_id,
            "run_name": run_name,
            "run_id": run_id,
        }
        response = await self.execute(
            query=query, operation_name="startRun", variables=variables, **kwargs
        )
        data = self.get_data(response)
        return StartRun.model_validate(data)

    async def update_run_status(
        self, run_id: str, status: RunStatus, **kwargs: Any
    ) -> UpdateRunStatus:
        query = gql(
            """
            mutation updateRunStatus($run_id: String!, $status: RunStatus!) {
              updateRunStatus(runId: $run_id, status: $status) {
                run {
                  runId
                  status
                }
              }
            }
            """
        )
        variables: Dict[str, object] = {"run_id": run_id, "status": status}
        response = await self.execute(
            query=query, operation_name="updateRunStatus", variables=variables, **kwargs
        )
        data = self.get_data(response)
        return UpdateRunStatus.model_validate(data)

    async def stop_run(self, run_id: str, status: RunStatus, **kwargs: Any) -> StopRun:
        query = gql(
            """
            mutation stopRun($run_id: String!, $status: RunStatus!) {
              stopRun(runId: $run_id, status: $status) {
                success
              }
            }
            """
        )
        variables: Dict[str, object] = {"run_id": run_id, "status": status}
        response = await self.execute(
            query=query, operation_name="stopRun", variables=variables, **kwargs
        )
        data = self.get_data(response)
        return StopRun.model_validate(data)

    async def update_custom_metric_result(
        self,
        run_id: str,
        local: bool,
        custom_metric_result: float,
        custom_metric_id: Union[Optional[str], UnsetType] = UNSET,
        custom_metric_name: Union[Optional[str], UnsetType] = UNSET,
        **kwargs: Any
    ) -> UpdateCustomMetricResult:
        query = gql(
            """
            mutation updateCustomMetricResult($runId: String!, $local: Boolean!, $custom_metric_result: Float!, $custom_metric_id: String, $custom_metric_name: String) {
              updateCustomMetricResult(
                runId: $runId
                local: $local
                customMetricResult: $custom_metric_result
                customMetricId: $custom_metric_id
                customMetricName: $custom_metric_name
              ) {
                id
              }
            }
            """
        )
        variables: Dict[str, object] = {
            "runId": run_id,
            "local": local,
            "custom_metric_result": custom_metric_result,
            "custom_metric_id": custom_metric_id,
            "custom_metric_name": custom_metric_name,
        }
        response = await self.execute(
            query=query,
            operation_name="updateCustomMetricResult",
            variables=variables,
            **kwargs
        )
        data = self.get_data(response)
        return UpdateCustomMetricResult.model_validate(data)

    async def get_run_status(self, run_id: str, **kwargs: Any) -> GetRunStatus:
        query = gql(
            """
            query GetRunStatus($runId: String!) {
              run(runId: $runId) {
                status
              }
            }
            """
        )
        variables: Dict[str, object] = {"runId": run_id}
        response = await self.execute(
            query=query, operation_name="GetRunStatus", variables=variables, **kwargs
        )
        data = self.get_data(response)
        return GetRunStatus.model_validate(data)

    async def pull_run(self, run_id: str, **kwargs: Any) -> PullRun:
        query = gql(
            """
            query PullRun($runId: String!) {
              run(runId: $runId) {
                qaSet {
                  id
                }
                runId
                singlePendingOrCompletedRunReview {
                  id
                }
                passPercentage
                status
                textSummary
                timestamp
                completedAt
                archived
                name
                parameters {
                  evalModel
                  maximumThreads
                  runConfidenceEvaluation
                  heavyweightFactor
                  createTextSummary
                  modelUnderTest
                  temperature
                  maxOutputTokens
                  systemPrompt
                  customParameters
                }
                passRate {
                  value
                  error
                }
                successRate {
                  value
                  error
                }
                customMetrics {
                  metricId
                  name
                  passRate
                  error
                }
                testSuite {
                  id
                  title
                }
                project {
                  id
                  slug
                }
              }
            }
            """
        )
        variables: Dict[str, object] = {"runId": run_id}
        response = await self.execute(
            query=query, operation_name="PullRun", variables=variables, **kwargs
        )
        data = self.get_data(response)
        return PullRun.model_validate(data)

    async def pull_test_results_with_count(
        self,
        run_id: str,
        offset: Union[Optional[int], UnsetType] = UNSET,
        limit: Union[Optional[int], UnsetType] = UNSET,
        **kwargs: Any
    ) -> PullTestResultsWithCount:
        query = gql(
            """
            query PullTestResultsWithCount($runId: String!, $offset: Int, $limit: Int) {
              testResultsWithCount(
                runId: $runId
                filterOptions: {offset: $offset, limit: $limit}
              ) {
                testResults {
                  id
                  llmOutput
                  passPercentage
                  passPercentageWithWeight
                  resultJson {
                    ...ResultJsonElementTypeFields
                  }
                  qaPair {
                    context
                    outputContext
                    errorMessage
                  }
                  test {
                    ...TestFragment
                  }
                  metadata {
                    ...MetadataOutputTypeFields
                  }
                }
                count
              }
            }

            fragment CheckModifierOutputTypeFields on CheckModifierOutputType {
              optional
              severity
              displayMetrics
              examples {
                ...ExampleOutputTypeFields
              }
              extractor
              conditional {
                ...ConditionalCheckOutputTypeFields
              }
              category
            }

            fragment ConditionalCheckOutputTypeFields on ConditionalCheckOutputType {
              operator
              criteria
            }

            fragment ExampleOutputTypeFields on ExampleOutputType {
              type
              text
            }

            fragment MetadataOutputTypeFields on MetadataOutputType {
              inTokens
              outTokens
              durationSeconds
            }

            fragment ResultJsonElementTypeFields on ResultJsonElementType {
              operator
              criteria
              modifiers {
                ...CheckModifierOutputTypeFields
              }
              autoEval
              feedback
              evalCont
              isGlobal
              severity
              customOperator {
                rubric {
                  label
                  score
                }
              }
            }

            fragment TestFragment on TestType {
              id
              inputUnderTest
              context
              tags
              crossVersionId
              goldenOutput
              fileIds
              checks {
                operator
                criteria
                modifiers {
                  optional
                  severity
                  examples {
                    type
                    text
                  }
                  extractor
                  conditional {
                    operator
                    criteria
                  }
                  category
                  displayMetrics
                }
              }
              testSuite {
                id
              }
            }
            """
        )
        variables: Dict[str, object] = {
            "runId": run_id,
            "offset": offset,
            "limit": limit,
        }
        response = await self.execute(
            query=query,
            operation_name="PullTestResultsWithCount",
            variables=variables,
            **kwargs
        )
        data = self.get_data(response)
        return PullTestResultsWithCount.model_validate(data)

    async def list_runs(
        self,
        project_id: str,
        archived: Union[Optional[bool], UnsetType] = UNSET,
        suite_id: Union[Optional[str], UnsetType] = UNSET,
        limit: Union[Optional[int], UnsetType] = UNSET,
        offset: Union[Optional[int], UnsetType] = UNSET,
        search: Union[Optional[str], UnsetType] = UNSET,
        **kwargs: Any
    ) -> ListRuns:
        query = gql(
            """
            query ListRuns($archived: Boolean, $suiteId: String, $projectId: String!, $limit: Int, $offset: Int, $search: String) {
              runsWithCount(
                filterOptions: {archived: $archived, suiteId: $suiteId, projectId: $projectId, limit: $limit, offset: $offset, sortBy: STARTED_AT, search: $search}
              ) {
                runResults {
                  runId
                  passPercentage
                  passRate {
                    value
                    error
                  }
                  successRate {
                    value
                    error
                  }
                  name
                  status
                  textSummary
                  timestamp
                  completedAt
                  archived
                  parameters {
                    evalModel
                    maximumThreads
                    runConfidenceEvaluation
                    heavyweightFactor
                    createTextSummary
                    modelUnderTest
                    temperature
                    maxOutputTokens
                    systemPrompt
                  }
                  testSuite {
                    title
                  }
                }
              }
            }
            """
        )
        variables: Dict[str, object] = {
            "archived": archived,
            "suiteId": suite_id,
            "projectId": project_id,
            "limit": limit,
            "offset": offset,
            "search": search,
        }
        response = await self.execute(
            query=query, operation_name="ListRuns", variables=variables, **kwargs
        )
        data = self.get_data(response)
        return ListRuns.model_validate(data)

    async def unfinished_tests(self, run_id: str, **kwargs: Any) -> UnfinishedTests:
        query = gql(
            """
            query UnfinishedTests($runId: String!) {
              unfinishedTests(runId: $runId) {
                unfinishedTests {
                  ...TestFragment
                }
                testSuiteSize
              }
            }

            fragment TestFragment on TestType {
              id
              inputUnderTest
              context
              tags
              crossVersionId
              goldenOutput
              fileIds
              checks {
                operator
                criteria
                modifiers {
                  optional
                  severity
                  examples {
                    type
                    text
                  }
                  extractor
                  conditional {
                    operator
                    criteria
                  }
                  category
                  displayMetrics
                }
              }
              testSuite {
                id
              }
            }
            """
        )
        variables: Dict[str, object] = {"runId": run_id}
        response = await self.execute(
            query=query, operation_name="UnfinishedTests", variables=variables, **kwargs
        )
        data = self.get_data(response)
        return UnfinishedTests.model_validate(data)

    async def create_or_update_test_suite(
        self,
        test_suite_id: str,
        title: str,
        description: str,
        project_id: str,
        **kwargs: Any
    ) -> CreateOrUpdateTestSuite:
        query = gql(
            """
            mutation createOrUpdateTestSuite($testSuiteId: String!, $title: String!, $description: String!, $projectId: String!) {
              updateTestSuite(
                testSuiteId: $testSuiteId
                title: $title
                description: $description
                projectId: $projectId
              ) {
                testSuite {
                  description
                  id
                  org
                  title
                  project {
                    slug
                  }
                }
              }
            }
            """
        )
        variables: Dict[str, object] = {
            "testSuiteId": test_suite_id,
            "title": title,
            "description": description,
            "projectId": project_id,
        }
        response = await self.execute(
            query=query,
            operation_name="createOrUpdateTestSuite",
            variables=variables,
            **kwargs
        )
        data = self.get_data(response)
        return CreateOrUpdateTestSuite.model_validate(data)

    async def add_batch_tests(
        self, test_info: List[TestMutationInfo], create_only: bool, **kwargs: Any
    ) -> AddBatchTests:
        query = gql(
            """
            mutation addBatchTests($test_info: [TestMutationInfo!]!, $createOnly: Boolean!) {
              batchUpdateTest(tests: $test_info, createOnly: $createOnly) {
                tests {
                  ...TestFragment
                }
              }
            }

            fragment TestFragment on TestType {
              id
              inputUnderTest
              context
              tags
              crossVersionId
              goldenOutput
              fileIds
              checks {
                operator
                criteria
                modifiers {
                  optional
                  severity
                  examples {
                    type
                    text
                  }
                  extractor
                  conditional {
                    operator
                    criteria
                  }
                  category
                  displayMetrics
                }
              }
              testSuite {
                id
              }
            }
            """
        )
        variables: Dict[str, object] = {
            "test_info": test_info,
            "createOnly": create_only,
        }
        response = await self.execute(
            query=query, operation_name="addBatchTests", variables=variables, **kwargs
        )
        data = self.get_data(response)
        return AddBatchTests.model_validate(data)

    async def delete_test_suite(
        self, test_suite_id: str, **kwargs: Any
    ) -> DeleteTestSuite:
        query = gql(
            """
            mutation deleteTestSuite($testSuiteId: String!) {
              deleteSuite(suiteId: $testSuiteId) {
                success
              }
            }
            """
        )
        variables: Dict[str, object] = {"testSuiteId": test_suite_id}
        response = await self.execute(
            query=query, operation_name="deleteTestSuite", variables=variables, **kwargs
        )
        data = self.get_data(response)
        return DeleteTestSuite.model_validate(data)

    async def update_global_checks(
        self, test_suite_id: str, checks: List[CheckInputType], **kwargs: Any
    ) -> UpdateGlobalChecks:
        query = gql(
            """
            mutation updateGlobalChecks($testSuiteId: String!, $checks: [CheckInputType!]!) {
              updateGlobalChecks(testSuiteId: $testSuiteId, checks: $checks) {
                success
              }
            }
            """
        )
        variables: Dict[str, object] = {"testSuiteId": test_suite_id, "checks": checks}
        response = await self.execute(
            query=query,
            operation_name="updateGlobalChecks",
            variables=variables,
            **kwargs
        )
        data = self.get_data(response)
        return UpdateGlobalChecks.model_validate(data)

    async def remove_old_tests(
        self, test_suite_id: str, in_use_tests: List[str], **kwargs: Any
    ) -> RemoveOldTests:
        query = gql(
            """
            mutation removeOldTests($testSuiteId: String!, $inUseTests: [String!]!) {
              removeUnusedTests(testSuiteId: $testSuiteId, inUseTests: $inUseTests) {
                success
              }
            }
            """
        )
        variables: Dict[str, object] = {
            "testSuiteId": test_suite_id,
            "inUseTests": in_use_tests,
        }
        response = await self.execute(
            query=query, operation_name="removeOldTests", variables=variables, **kwargs
        )
        data = self.get_data(response)
        return RemoveOldTests.model_validate(data)

    async def rerun_tests(self, run_id: str, **kwargs: Any) -> RerunTests:
        query = gql(
            """
            mutation RerunTests($runId: String!) {
              rerunFailingTests(runId: $runId) {
                success
              }
            }
            """
        )
        variables: Dict[str, object] = {"runId": run_id}
        response = await self.execute(
            query=query, operation_name="RerunTests", variables=variables, **kwargs
        )
        data = self.get_data(response)
        return RerunTests.model_validate(data)

    async def upload_local_evaluation(
        self,
        question_answer_set_id: str,
        local_evals: List[LocalEvalUploadInputType],
        **kwargs: Any
    ) -> UploadLocalEvaluation:
        query = gql(
            """
            mutation uploadLocalEvaluation($questionAnswerSetId: String!, $localEvals: [LocalEvalUploadInputType!]!) {
              uploadLocalEvaluation(
                questionAnswerSetId: $questionAnswerSetId
                localEvals: $localEvals
              ) {
                uploads {
                  id
                  score
                  feedback
                  createdAt
                }
              }
            }
            """
        )
        variables: Dict[str, object] = {
            "questionAnswerSetId": question_answer_set_id,
            "localEvals": local_evals,
        }
        response = await self.execute(
            query=query,
            operation_name="uploadLocalEvaluation",
            variables=variables,
            **kwargs
        )
        data = self.get_data(response)
        return UploadLocalEvaluation.model_validate(data)

    async def set_custom_metrics(
        self, test_suite_id: str, ids: List[str], **kwargs: Any
    ) -> SetCustomMetrics:
        query = gql(
            """
            mutation setCustomMetrics($testSuiteId: String!, $ids: [String!]!) {
              setCustomMetrics(testSuiteId: $testSuiteId, ids: $ids) {
                success
              }
            }
            """
        )
        variables: Dict[str, object] = {"testSuiteId": test_suite_id, "ids": ids}
        response = await self.execute(
            query=query,
            operation_name="setCustomMetrics",
            variables=variables,
            **kwargs
        )
        data = self.get_data(response)
        return SetCustomMetrics.model_validate(data)

    async def get_test_suite_data(
        self, suite_id: str, **kwargs: Any
    ) -> GetTestSuiteData:
        query = gql(
            """
            query getTestSuiteData($suiteId: String!) {
              testSuite(testSuiteId: $suiteId) {
                description
                id
                project {
                  slug
                }
                org
                title
                created
                globalChecks {
                  ...CheckOutputTypeFields
                }
              }
            }

            fragment CheckModifierOutputTypeFields on CheckModifierOutputType {
              optional
              severity
              displayMetrics
              examples {
                ...ExampleOutputTypeFields
              }
              extractor
              conditional {
                ...ConditionalCheckOutputTypeFields
              }
              category
            }

            fragment CheckOutputTypeFields on CheckOutputType {
              operator
              criteria
              modifiers {
                ...CheckModifierOutputTypeFields
              }
            }

            fragment ConditionalCheckOutputTypeFields on ConditionalCheckOutputType {
              operator
              criteria
            }

            fragment ExampleOutputTypeFields on ExampleOutputType {
              type
              text
            }
            """
        )
        variables: Dict[str, object] = {"suiteId": suite_id}
        response = await self.execute(
            query=query,
            operation_name="getTestSuiteData",
            variables=variables,
            **kwargs
        )
        data = self.get_data(response)
        return GetTestSuiteData.model_validate(data)

    async def get_test_data(
        self,
        suite_id: str,
        offset: Union[Optional[int], UnsetType] = UNSET,
        limit: Union[Optional[int], UnsetType] = UNSET,
        **kwargs: Any
    ) -> GetTestData:
        query = gql(
            """
            query getTestData($suiteId: String!, $offset: Int, $limit: Int) {
              testsWithCount(
                testSuiteId: $suiteId
                filterOptions: {offset: $offset, limit: $limit}
              ) {
                tests {
                  ...TestFragment
                }
                count
              }
            }

            fragment TestFragment on TestType {
              id
              inputUnderTest
              context
              tags
              crossVersionId
              goldenOutput
              fileIds
              checks {
                operator
                criteria
                modifiers {
                  optional
                  severity
                  examples {
                    type
                    text
                  }
                  extractor
                  conditional {
                    operator
                    criteria
                  }
                  category
                  displayMetrics
                }
              }
              testSuite {
                id
              }
            }
            """
        )
        variables: Dict[str, object] = {
            "suiteId": suite_id,
            "offset": offset,
            "limit": limit,
        }
        response = await self.execute(
            query=query, operation_name="getTestData", variables=variables, **kwargs
        )
        data = self.get_data(response)
        return GetTestData.model_validate(data)

    async def get_test_suites_with_count(
        self,
        project_id: str,
        offset: Union[Optional[int], UnsetType] = UNSET,
        limit: Union[Optional[int], UnsetType] = UNSET,
        search: Union[Optional[str], UnsetType] = UNSET,
        **kwargs: Any
    ) -> GetTestSuitesWithCount:
        query = gql(
            """
            query getTestSuitesWithCount($offset: Int, $limit: Int, $search: String, $projectId: String!) {
              testSuitesWithCount(
                filterOptions: {offset: $offset, limit: $limit, search: $search, projectId: $projectId}
              ) {
                testSuites {
                  id
                  title
                  description
                  created
                  creator
                  lastModifiedBy
                  lastModifiedAt
                  folder {
                    id
                    name
                  }
                }
              }
            }
            """
        )
        variables: Dict[str, object] = {
            "offset": offset,
            "limit": limit,
            "search": search,
            "projectId": project_id,
        }
        response = await self.execute(
            query=query,
            operation_name="getTestSuitesWithCount",
            variables=variables,
            **kwargs
        )
        data = self.get_data(response)
        return GetTestSuitesWithCount.model_validate(data)

    async def get_default_parameters(self, **kwargs: Any) -> GetDefaultParameters:
        query = gql(
            """
            query GetDefaultParameters {
              defaultParameters {
                evalModel
                maximumThreads
                runConfidenceEvaluation
                createTextSummary
                detectRefusals
                modelUnderTest
                temperature
                maxOutputTokens
                heavyweightFactor
                systemPrompt
              }
            }
            """
        )
        variables: Dict[str, object] = {}
        response = await self.execute(
            query=query,
            operation_name="GetDefaultParameters",
            variables=variables,
            **kwargs
        )
        data = self.get_data(response)
        return GetDefaultParameters.model_validate(data)
