"use strict";
Object.defineProperty(exports, "__esModule", { value: true });
exports.PredictionFunctions = void 0;
const gql_1 = require("../gql");
const BaseFunctions_1 = require("./BaseFunctions");
class PredictionFunctions extends BaseFunctions_1.BaseFunctions {
    constructor(client) {
        super(client);
    }
    // a method to create a prediction
    async create(openaiKwargs, prompt, modelBackendName, useCache = true, labelParser, metadata) {
        const query = (0, gql_1.graphql)(`
      query Prediction($openaiKwargs: OpenAIKwargs!, $prompt: String!, $modelBackendName: String!, $useCache: Boolean!, $labelParser: String, $metadata: [MetadataInput!]) {
        prediction(openaiKwargs: $openaiKwargs, prompt: $prompt, modelBackendName: $modelBackendName, useCache: $useCache, labelParser: $labelParser, metadata: $metadata) {
          completion
          id
          labels {
            id
            confidence
            name
            value
          }
          prompt
          labelParser
          metadata {
            key
            value
          }
        }
      }
    `);
        const vars = { openaiKwargs, prompt, modelBackendName, useCache, labelParser, metadata };
        return this.request(query, vars);
    }
    async countPredictions(where) {
        const query = (0, gql_1.graphql)(`
      query CountPredictions($where: PredictionWhereInput) {
        aggregatePrediction(where: $where) {
          _count {
            _all
          }
        }
      }
    `);
        let vars = {};
        if (where) {
            vars.where = where;
        }
        return this.request(query, vars);
    }
    async getLabelsForPrediction(predictionId) {
        const query = (0, gql_1.graphql)(`
      query GetLabelsForPrediction($where: LabelWhereInput) {
        labels(where: $where) {
          id
          value
          confidence
          predictionId
        }
      }
    `);
        return this.request(query, { where: { predictionId: { equals: predictionId } } });
    }
    async listPredictions(options) {
        const query = (0, gql_1.graphql)(`
    query ListPredictions($where: PredictionWhereInput, $orderBy: [PredictionOrderByWithRelationInput!], $cursor: PredictionWhereUniqueInput, $take: Int, $skip: Int, $distinct: [PredictionScalarFieldEnum!]) {
      predictions(where: $where, orderBy: $orderBy, cursor: $cursor, take: $take, skip: $skip, distinct: $distinct) {
        id
        completion
        createdAt
        _count {
          labels
          datasets
        }
        labels {
          id
          name
          value
          confidence
        }
        prompt
        requestArgs
        response
        updatedAt
        userId
        modelId
      }
    }
    `);
        return this.request(query, options);
    }
    async getPrediction(predictionId) {
        const query = (0, gql_1.graphql)(`
        query GetPrediction($id: Int!) {
          getPrediction(where: { id: $id }) {
            completion
            createdAt
            id
            _count {
              labels
              datasets
            }
            model {
              name
            }
            prompt
          }
        }
      `);
        return this.request(query, { id: predictionId });
    }
    async addPredictionToDataset(predictionId, datasetId) {
        const mutation = (0, gql_1.graphql)(`
      mutation UpdateOnePrediction($predictionId: Int!, $datasetId: Int!) {
        updateOnePrediction(where: { id: $predictionId }, data: { datasets: { connect: { id: $datasetId } } }) {
          id
          _count {
            datasets
          }
        }
      }
    `);
        return this.request(mutation, { predictionId, datasetId });
    }
    async addPredictionsToDataset(predictionIds, datasetId) {
        const mutation = (0, gql_1.graphql)(`
      mutation Mutation($datasetId: Int!, $predictionIds: [PredictionWhereUniqueInput!]!) {
        updateOneDataset(where: { id: $datasetId }, data: { predictions: { connect: $predictionIds } }) {
          id
          _count {
            predictions
          }
        }
      }
    `);
        return this.request(mutation, { predictionIds: predictionIds.map(id => ({ id })), datasetId });
    }
}
exports.PredictionFunctions = PredictionFunctions;
