From eb069c7e4612108923697c9a8f20dd0b2e5489f0 Mon Sep 17 00:00:00 2001
From: geraud <gtwin@bourdin.name>
Date: Thu, 24 Oct 2024 12:47:30 +0200
Subject: [PATCH] add timeout param.

---
 examples/socket_timeout.php | 46 +++++++++++++++++++++++++++++++++++++
 src/MistralClient.php       | 40 +++++++++++++++++++++++++++++---
 2 files changed, 83 insertions(+), 3 deletions(-)
 create mode 100755 examples/socket_timeout.php

diff --git a/examples/socket_timeout.php b/examples/socket_timeout.php
new file mode 100755
index 0000000..8710a06
--- /dev/null
+++ b/examples/socket_timeout.php
@@ -0,0 +1,46 @@
+<?php
+require_once __DIR__ . '/../vendor/autoload.php';
+
+use Partitech\PhpMistral\MistralClient;
+use Partitech\PhpMistral\MistralClientException;
+use Partitech\PhpMistral\Messages;
+use Symfony\Component\HttpClient\Exception\TransportException;
+
+// export MISTRAL_API_KEY=
+$apiKey = getenv('MISTRAL_API_KEY');
+$client = new MistralClient(apiKey: $apiKey, timeout:0.1);
+$messages = new Messages();
+$messages->addUserMessage('What is the best French cheese?');
+
+$params = [
+    'model' => 'mistral-large-latest',
+    'temperature' => 0.7,
+    'top_p' => 1,
+    'max_tokens' => null,
+    'safe_prompt' => false,
+    'random_seed' => null
+];
+
+try {
+    foreach ($client->chatStream($messages, $params) as $chunk) {
+        echo $chunk->getChunk();
+    }
+} catch (MistralClientException $e) {
+    echo $e->getMessage();
+    exit(1);
+} catch (TransportException $e) {
+    echo 'Idle timeout reached' . PHP_EOL;
+
+}
+
+$client->setTimeout(null);
+try {
+    foreach ($client->chatStream($messages, $params) as $chunk) {
+        echo $chunk->getChunk();
+    }
+} catch (MistralClientException $e) {
+    echo $e->getMessage();
+    exit(1);
+} catch (TransportException $e) {
+    echo 'Idle timeout reached' . PHP_EOL;;
+}
\ No newline at end of file
diff --git a/src/MistralClient.php b/src/MistralClient.php
index 140ab54..e54506c 100644
--- a/src/MistralClient.php
+++ b/src/MistralClient.php
@@ -6,9 +6,11 @@
 
 use Generator;
 use KnpLabs\JsonSchema\ObjectSchema;
+use Symfony\Component\HttpClient\Exception\TransportException;
 use Symfony\Component\HttpClient\HttpClient;
 use Symfony\Component\HttpClient\Retry\GenericRetryStrategy;
 use Symfony\Component\HttpClient\RetryableHttpClient;
+use Symfony\Contracts\HttpClient\Exception\TransportExceptionInterface;
 use Symfony\Contracts\HttpClient\HttpClientInterface;
 use Symfony\Contracts\HttpClient\ResponseInterface;
 use Throwable;
@@ -43,9 +45,11 @@ class MistralClient
     protected string $url;
     private HttpClientInterface $httpClient;
     private string $mode;
+    private null|int|float $timeout = null; // null = default_socket_timeout
 
-    public function __construct(string $apiKey, string $url = self::ENDPOINT)
+    public function __construct(string $apiKey, string $url = self::ENDPOINT, int|float $timeout = null)
     {
+        $this->setTimeout($timeout);
         $this->httpClient = new RetryableHttpClient(
             HttpClient::create(),
             new GenericRetryStrategy(self::RETRY_STATUS_CODES, 500, 2)
@@ -83,11 +87,21 @@ protected function request(
         bool   $stream = false
     ): array|ResponseInterface
     {
+        $params = [
+            'json' => $request,
+            'headers' => ['Authorization' => 'Bearer ' . $this->apiKey,],
+            'buffer' => $stream,
+        ];
+
+        if(!is_null($this->getTimeout())){
+            $params['timeout'] = $this->getTimeout();
+        }
+
         try {
             $response = $this->httpClient->request(
                 $method,
                 $this->url . '/' . $path,
-                ['json' => $request, 'headers' => ['Authorization' => 'Bearer ' . $this->apiKey,], 'buffer' => $stream,]
+                $params
             );
         } catch (Throwable $e) {
             throw new MistralClientException($e->getMessage(), $e->getCode(), $e);
@@ -127,7 +141,7 @@ public function fim(string $prompt, ?string $suffix, array $params = []): Respon
     }
 
     /**
-     * @throws MistralClientException
+     * @throws MistralClientException|TransportExceptionInterface
      */
     public function fimStream(string $prompt, ?string $suffix, array $params = []): Generator
     {
@@ -327,6 +341,7 @@ protected function makeChatCompletionRequest(Messages $messages, array $params,
 
     /**
      * @throws MistralClientException
+     * @throws TransportExceptionInterface
      */
     public function chatStream(Messages $messages, array $params = []): Generator
     {
@@ -346,10 +361,17 @@ public function embeddings(array $datas): array
     }
 
 
+    /**
+     * @throws MistralClientException
+     * @throws TransportExceptionInterface
+     */
     public function getStream($stream): Generator
     {
         $response = null;
         foreach ($this->httpClient->stream($stream) as $chunk) {
+            if ($chunk->isTimeout()) {
+                throw new TransportException('Stream is closed');
+            }
             try {
                 $chunk = trim($chunk->getContent());
             } catch (Throwable $e) {
@@ -378,4 +400,16 @@ public function getStream($stream): Generator
             }
         }
     }
+
+    public function setTimeout(null|int|float $timeout): self
+    {
+        $this->timeout = $timeout;
+
+        return $this;
+    }
+
+    public function getTimeout(): null|int|float
+    {
+        return $this->timeout;
+    }
 }