class Vellum::Client

def execute_prompt(inputs:, prompt_deployment_id: nil, prompt_deployment_name: nil, release_tag: nil,
                   external_id: nil, expand_meta: nil, raw_overrides: nil, expand_raw: nil, metadata: nil, request_options: nil)
  response = @request_client.conn.post do |req|
    req.options.timeout = request_options.timeout_in_seconds unless request_options&.timeout_in_seconds.nil?
    req.headers["X_API_KEY"] = request_options.api_key unless request_options&.api_key.nil?
    req.headers = { **req.headers, **(request_options&.additional_headers || {}) }.compact
    req.body = {
      **(request_options&.additional_body_parameters || {}),
      inputs: inputs,
      prompt_deployment_id: prompt_deployment_id,
      prompt_deployment_name: prompt_deployment_name,
      release_tag: release_tag,
      external_id: external_id,
      expand_meta: expand_meta,
      raw_overrides: raw_overrides,
      expand_raw: expand_raw,
      metadata: metadata
    }.compact
    req.url "#{@request_client.default_environment[:Predict]}/v1/execute-prompt"
  end
  ExecutePromptResponse.from_json(json_object: response.body)
end

def execute_workflow(inputs:, workflow_deployment_id: nil, workflow_deployment_name: nil, release_tag: nil,
                     external_id: nil, request_options: nil)
  response = @request_client.conn.post do |req|
    req.options.timeout = request_options.timeout_in_seconds unless request_options&.timeout_in_seconds.nil?
    req.headers["X_API_KEY"] = request_options.api_key unless request_options&.api_key.nil?
    req.headers = { **req.headers, **(request_options&.additional_headers || {}) }.compact
    req.body = {
      **(request_options&.additional_body_parameters || {}),
      inputs: inputs,
      workflow_deployment_id: workflow_deployment_id,
      workflow_deployment_name: workflow_deployment_name,
      release_tag: release_tag,
      external_id: external_id
    }.compact
    req.url "#{@request_client.default_environment[:Predict]}/v1/execute-workflow"
  end
  ExecuteWorkflowResponse.from_json(json_object: response.body)
end

def generate(requests:, deployment_id: nil, deployment_name: nil, options: nil, request_options: nil)
  response = @request_client.conn.post do |req|
    req.options.timeout = request_options.timeout_in_seconds unless request_options&.timeout_in_seconds.nil?
    req.headers["X_API_KEY"] = request_options.api_key unless request_options&.api_key.nil?
    req.headers = { **req.headers, **(request_options&.additional_headers || {}) }.compact
    req.body = {
      **(request_options&.additional_body_parameters || {}),
      deployment_id: deployment_id,
      deployment_name: deployment_name,
      requests: requests,
      options: options
    }.compact
    req.url "#{@request_client.default_environment[:Predict]}/v1/generate"
  end
  GenerateResponse.from_json(json_object: response.body)
end

def initialize(api_key:, environment: Environment::PRODUCTION, max_retries: nil, timeout_in_seconds: nil)
  @request_client = RequestClient.new(environment: environment, max_retries: max_retries,
                                      timeout_in_seconds: timeout_in_seconds, api_key: api_key)
  @deployments = DeploymentsClient.new(request_client: @request_client)
  @document_indexes = DocumentIndexesClient.new(request_client: @request_client)
  @documents = DocumentsClient.new(request_client: @request_client)
  @folder_entities = FolderEntitiesClient.new(request_client: @request_client)
  @model_versions = ModelVersionsClient.new(request_client: @request_client)
  @registered_prompts = RegisteredPromptsClient.new(request_client: @request_client)
  @sandboxes = SandboxesClient.new(request_client: @request_client)
  @test_suite_runs = TestSuiteRunsClient.new(request_client: @request_client)
  @test_suites = TestSuitesClient.new(request_client: @request_client)
  @workflow_deployments = WorkflowDeploymentsClient.new(request_client: @request_client)
end

def search(query:, index_id: nil, index_name: nil, options: nil, request_options: nil)
  response = @request_client.conn.post do |req|
    req.options.timeout = request_options.timeout_in_seconds unless request_options&.timeout_in_seconds.nil?
    req.headers["X_API_KEY"] = request_options.api_key unless request_options&.api_key.nil?
    req.headers = { **req.headers, **(request_options&.additional_headers || {}) }.compact
    req.body = {
      **(request_options&.additional_body_parameters || {}),
      index_id: index_id,
      index_name: index_name,
      query: query,
      options: options
    }.compact
    req.url "#{@request_client.default_environment[:Predict]}/v1/search"
  end
  SearchResponse.from_json(json_object: response.body)
end

def submit_completion_actuals(actuals:, deployment_id: nil, deployment_name: nil, request_options: nil)
  @request_client.conn.post do |req|
    req.options.timeout = request_options.timeout_in_seconds unless request_options&.timeout_in_seconds.nil?
    req.headers["X_API_KEY"] = request_options.api_key unless request_options&.api_key.nil?
    req.headers = { **req.headers, **(request_options&.additional_headers || {}) }.compact
    req.body = {
      **(request_options&.additional_body_parameters || {}),
      deployment_id: deployment_id,
      deployment_name: deployment_name,
      actuals: actuals
    }.compact
    req.url "#{@request_client.default_environment[:Predict]}/v1/submit-completion-actuals"
  end
end

def submit_workflow_execution_actuals(actuals:, execution_id: nil, external_id: nil, request_options: nil)
  @request_client.conn.post do |req|
    req.options.timeout = request_options.timeout_in_seconds unless request_options&.timeout_in_seconds.nil?
    req.headers["X_API_KEY"] = request_options.api_key unless request_options&.api_key.nil?
    req.headers = { **req.headers, **(request_options&.additional_headers || {}) }.compact
    req.body = {
      **(request_options&.additional_body_parameters || {}),
      actuals: actuals,
      execution_id: execution_id,
      external_id: external_id
    }.compact
    req.url "#{@request_client.default_environment[:Predict]}/v1/submit-workflow-execution-actuals"
  end
end