Elep

<?php

namespace OtomaticAi\Api\Flux;

use Exception;
use OtomaticAi\Api\Ideogram\Exceptions\ApiException;
use OtomaticAi\Models\Usages\Usage;
use OtomaticAi\Utils\Settings;
use OtomaticAi\Vendors\GuzzleHttp\Client as HttpClient;
use OtomaticAi\Vendors\Illuminate\Support\Arr;
use OtomaticAi\Vendors\GuzzleHttp\Pool;
use OtomaticAi\Vendors\GuzzleHttp\Psr7\Request;
use OtomaticAi\Vendors\GuzzleHttp\Psr7\Response;

class PoolClient
{
    /**
     * The Flux api endpoint
     *
     * @var string
     */
    private string $endpoint = "https://api.bfl.ai/v1";

/**
     * The GuzzleHttp client
     *
     * @var HttpClient
     */
    private HttpClient $client;

/**
     * The auth key
     *
     * @var HttpClient
     */
    private string $key;

// cache for store tasks data
    public array $tasks = [];

/**
     * Create a new Flux Api client
     *
     * @param string|null $key
     * @throws Exception
     */
    function __construct(string $key = null)
    {
        // get the key
        if (empty($key)) {
            $key = Settings::get("api.flux.api_key");
        }
        if (empty($key)) {
            throw new Exception("No Flux Api Key provided.");
        }

$this->key = $key;

// create the http client
        $this->client = new HttpClient(
            [
                "base_uri" => rtrim($this->endpoint, "/") . '/',
                "headers" => [
                    "X-Key" => $this->key,
                    "Accept" => "application/json",
                    "Content-Type" => "application/json",
                ],
                "timeout" => 600
            ]
        );
    }

/**
     * Call the create task api endpoint
     *
     * @param array $settings
     * @return array
     * @throws Exception
     */
    public function createTasks(array $settings = []): array
    {
        $requests = [];

foreach ($settings as $payload) {
            // make the payload
            if ($payload["model"] === "flux-kontext-pro" || $payload["model"] === "flux-kontext-max") {
                $payload = array_merge([
                    "aspect_ratio" => "21:9",
                    "prompt_upsampling" => true,
                ], $payload);
            } else {
                $payload = array_merge([
                    "width" => 1344,
                    "height" => 768,
                    "prompt_upsampling" => true,
                ], $payload);
            }

// make the request
            $requests[] = [
                "request" => new Request(
                    'POST',
                    $payload["model"],
                    [],
                    json_encode($payload)
                ),
                "after" => function ($result) use ($payload) {

$id = Arr::get($result, "id");

$this->tasks[$id] = [
                        "id" => $id,
                        "payload" => $payload,
                        "result" => null,
                    ];

// store the usage
                    Usage::create([
                        "provider" => "flux",
                        "payload" => [
                            "model" => $payload["model"],
                            "artifacts" => 1,
                        ]
                    ]);
                }
            ];
        }

// call the api endpoint
        $responses = $this->run($requests);

return $responses;
    }

/**
     * Call the get result api endpoint
     *
     * @param array $pollingUrls
     * @return array
     * @throws Exception
     */
    public function getResults(array $pollingUrls): array
    {
        $requests = [];
        foreach ($pollingUrls as $pollingUrl) {
            $id = $this->getIdFromUrl($pollingUrl);
            $status = Arr::get($this->tasks, $id . ".result.status", null);
            if (empty($status) || $status === "Pending") {
                $requests[] = [
                    "request" => new Request(
                        'GET',
                        $pollingUrl
                    )
                ];
            }
        }

// call the api endpoint
        $responses = $this->run($requests);

foreach ($responses as $response) {
            if (isset($this->tasks[$response["id"]])) {
                $this->tasks[$response["id"]]["result"] = $response;
            }
        }

$output = [];
        foreach ($pollingUrls as $pollingUrl) {
            $id = $this->getIdFromUrl($pollingUrl);
            $output[] = Arr::get($this->tasks, $id . ".result");
        }

return $output;
    }

/**
     * Perform an api request
     *
     * @param array $requests
     * @param stirng $uri
     * @param array $options
     * @return array
     * @throws Exception
     */
    private function run(array $requests)
    {
        $output = [];

// get the http requests
        $httpRequests = Arr::pluck($requests, "request");

// create the pool
        $pool = new Pool($this->client, $httpRequests, [
            'concurrency' => 10,
            'fulfilled' => function (Response $response, $index) use (&$output, $requests) {
                $response = $response->getBody()->getContents();
                $response = json_decode($response, true);
                $output[$index] = $response;

// call after callback
                if (!empty($requests[$index]["after"]) && is_callable($requests[$index]["after"])) {
                    $requests[$index]["after"]($response);
                }
            },
            'rejected' => function (Exception $reason, $index) {
                $e = ApiException::make($reason);
                throw $e;
            },
        ]);

// call the pool
        $promise = $pool->promise();
        $promise->wait();

ksort($output);
        $output = array_values($output);

return $output;
    }

private function getIdFromUrl(string $url): ?string
    {
        $query = parse_url($url, PHP_URL_QUERY);

if (!$query) {
            return null;
        }

parse_str($query, $params);

return Arr::get($params, "id");
    }
}