refactor(litellm): Remove prefix from model names

tymondesigns · tymondesigns · commit 064cb71206a7 · 2025-06-04T07:49:21.000+01:00
diff --git a/phpstan.dist.neon b/phpstan.dist.neon
@@ -5,8 +5,5 @@ parameters:
   level: 10
   paths:
     - src
-  excludePaths:
-    analyse:
-      - src/Providers/LiteLLMModelInfoProvider.php
   tmpDir: .phpstan-cache
   checkBenevolentUnionTypes: true
diff --git a/src/Providers/LiteLLMModelInfoProvider.php b/src/Providers/LiteLLMModelInfoProvider.php
@@ -4,7 +4,6 @@
 
 namespace Cortex\ModelInfo\Providers;
 
-use JsonException;
 use SensitiveParameter;
 use Psr\SimpleCache\CacheInterface;
 use Cortex\ModelInfo\Data\ModelInfo;
@@ -17,6 +16,10 @@
 use Cortex\ModelInfo\Providers\Concerns\ChecksSupport;
 use Cortex\ModelInfo\Providers\Concerns\MakesRequests;
 
+/**
+ * @phpstan-type LiteLLMModelInfoResponse array{litellm_provider: string, mode: string, max_input_tokens: ?int, max_output_tokens: ?int, input_cost_per_token: ?float, output_cost_per_token: ?float, supports_response_schema: ?bool, supports_function_calling: ?bool, supports_vision: ?bool, supports_tool_choice: ?bool, supports_reasoning: ?bool, supports_web_search: ?bool, supports_prompt_caching: ?bool, supports_audio_input: ?bool, supports_audio_output: ?bool, deprecation_date?: ?string}
+ * @phpstan-type LiteLLMModelInfoResponseWithName array{name: string, litellm_provider: string, mode: string, max_input_tokens: ?int, max_output_tokens: ?int, input_cost_per_token: ?float, output_cost_per_token: ?float, supports_response_schema: ?bool, supports_function_calling: ?bool, supports_vision: ?bool, supports_tool_choice: ?bool, supports_reasoning: ?bool, supports_web_search: ?bool, supports_prompt_caching: ?bool, supports_audio_input: ?bool, supports_audio_output: ?bool, deprecation_date?: ?string}
+ */
 class LiteLLMModelInfoProvider implements ModelInfoProvider
 {
     use ChecksSupport;
@@ -52,18 +55,17 @@ public function getModels(ModelProvider $modelProvider): array
     {
         $this->checkSupportOrFail($modelProvider);
 
-        $body = $this->getStaticResponse();
+        $body = $this->getResponse();
 
-        $models = array_filter(
+        $models = self::removePrefixFromModelName(array_filter(
             $body,
             fn(array $model): bool => $model['litellm_provider'] === $modelProvider->value,
-        );
+        ));
 
-        return array_values(array_map(
-            fn(array $modelInfo, string $model): ModelInfo => self::mapModelInfo($modelProvider, $model, $modelInfo),
+        return array_map(
+            fn(array $modelInfo): ModelInfo => self::mapModelInfo($modelProvider, $modelInfo),
             $models,
-            array_keys($models),
-        ));
+        );
     }
 
     /**
@@ -73,50 +75,59 @@ public function getModelInfo(ModelProvider $modelProvider, string $model): Model
     {
         $this->checkSupportOrFail($modelProvider);
 
-        $body = $this->getStaticResponse();
+        $body = $this->getResponse();
 
-        $models = array_filter(
+        $models = self::removePrefixFromModelName(array_filter(
             $body,
             fn(array $model): bool => $model['litellm_provider'] === $modelProvider->value,
-        );
+        ));
 
         $modelInfo = array_values(
-            array_filter(
-                $models,
-                fn(string $key): bool => $key === $model,
-                ARRAY_FILTER_USE_KEY,
-            ),
+            array_filter($models, fn(array $modelInfo): bool => $modelInfo['name'] === $model),
         )[0] ?? null;
 
         if ($modelInfo === null) {
             throw new ModelInfoException('Model not found');
         }
 
-        return self::mapModelInfo($modelProvider, $model, $modelInfo);
+        return self::mapModelInfo($modelProvider, $modelInfo);
     }
 
-    protected static function mapModelInfo(
-        ModelProvider $modelProvider,
-        string $model,
-        array $modelInfo,
-    ): ModelInfo {
+    /**
+     * @param array<array-key, LiteLLMModelInfoResponse> $models
+     *
+     * @return array<array-key, LiteLLMModelInfoResponseWithName>
+     */
+    protected static function removePrefixFromModelName(array $models): array
+    {
+        return array_map(fn(array $modelInfo, string $modelName): array => [
+            'name' => str_replace($modelInfo['litellm_provider'] . '/', '', $modelName),
+            ...$modelInfo,
+        ], $models, array_keys($models));
+    }
+
+    /**
+     * @param LiteLLMModelInfoResponseWithName $modelInfo
+     */
+    protected static function mapModelInfo(ModelProvider $modelProvider, array $modelInfo): ModelInfo
+    {
         return new ModelInfo(
-            name: $model,
+            name: $modelInfo['name'],
             provider: $modelProvider,
-            type: self::mapModelType($modelInfo['mode'] ?? ''),
+            type: self::mapModelType($modelInfo['mode']),
             maxInputTokens: $modelInfo['max_input_tokens'] ?? null,
             maxOutputTokens: $modelInfo['max_output_tokens'] ?? null,
             inputCostPerToken: $modelInfo['input_cost_per_token'] ?? 0.0,
             outputCostPerToken: $modelInfo['output_cost_per_token'] ?? 0.0,
             features: self::getFeatures($modelInfo),
-            isDeprecated: isset($modelInfo['deprecation_date']) && $modelInfo['deprecation_date'] !== null,
+            isDeprecated: isset($modelInfo['deprecation_date']),
         );
     }
 
     /**
-     * @param array<string, mixed> $info
+     * @param LiteLLMModelInfoResponse $info
      *
-     * @return array<int, \Cortex\Enums\ModelFeature>
+     * @return array<array-key, \Cortex\ModelInfo\Enums\ModelFeature>
      */
     protected static function getFeatures(array $info): array
     {
@@ -177,39 +188,48 @@ protected static function mapModelType(string $type): ModelType
     }
 
     /**
-     * @throws \Cortex\Exceptions\ModelInfoException
+     * @throws \Cortex\ModelInfo\Exceptions\ModelInfoException
      *
-     * @return array<string, mixed>
+     * @return array<array-key, LiteLLMModelInfoResponse>
      */
     protected function getStaticResponse(): array
     {
-        $request = self::discoverHttpRequestFactory()->createRequest('GET', self::LITELLM_STATIC_URL);
-        $response = $this->httpClient->sendRequest($request);
-
-        if ($response->getStatusCode() !== 200) {
-            throw new ModelInfoException('Failed to get model info');
-        }
+        $request = self::discoverHttpRequestFactoryOrFail()
+            ->createRequest('GET', self::LITELLM_STATIC_URL);
 
-        try {
-            return json_decode($response->getBody()->getContents(), true, flags: JSON_THROW_ON_ERROR);
-        } catch (JsonException $jsonException) {
-            throw new ModelInfoException('Failed to decode model info', previous: $jsonException);
-        }
+        // @phpstan-ignore return.type
+        return $this->getJsonResponse($request);
     }
 
     /**
      * @throws \Cortex\ModelInfo\Exceptions\ModelInfoException
      *
-     * @return array<string, mixed>
+     * @return array<array-key, LiteLLMModelInfoResponse>
      */
     protected function getApiModelsResponse(): array
     {
-        $request = self::discoverHttpRequestFactory()
+        $request = self::discoverHttpRequestFactoryOrFail()
             ->createRequest('GET', $this->host . '/v1/models');
 
+        // @phpstan-ignore return.type
         return $this->getJsonResponse($request);
     }
 
+    /**
+     * @throws \Cortex\ModelInfo\Exceptions\ModelInfoException
+     *
+     * @return array<array-key, LiteLLMModelInfoResponse>
+     */
+    protected function getResponse(): array
+    {
+        return $this->shouldUseStaticResponse()
+            ? $this->getStaticResponse()
+            : $this->getApiModelsResponse();
+    }
+
+    /**
+     * Determine if the static response should be used.
+     */
     protected function shouldUseStaticResponse(): bool
     {
         return $this->host === null || $this->apiKey === null;
diff --git a/src/Providers/OllamaModelInfoProvider.php b/src/Providers/OllamaModelInfoProvider.php
@@ -62,19 +62,20 @@ public function getModelInfo(ModelProvider $modelProvider, string $model): Model
 
         return self::mapModelInfo(
             $this->getModelInfoResponse($model),
+            $model,
         );
     }
 
     /**
      * @param OllamaModelInfoResponse|OllamaModelsResponse $modelResponseBody
      */
-    protected static function mapModelInfo(array $modelResponseBody): ModelInfo
+    protected static function mapModelInfo(array $modelResponseBody, ?string $modelName = null): ModelInfo
     {
         $modelInfo = $modelResponseBody['model_info'] ?? [];
         $capabilities = $modelResponseBody['capabilities'] ?? [];
 
         return new ModelInfo(
-            name: $modelResponseBody['name'],
+            name: $modelName ?? $modelResponseBody['name'],
             provider: ModelProvider::Ollama,
             type: self::getModelType($capabilities),
             maxInputTokens: self::getMaxInputTokens($modelInfo),
diff --git a/tests/Unit/Providers/LiteLLMModelInfoProviderTest.php b/tests/Unit/Providers/LiteLLMModelInfoProviderTest.php
@@ -26,15 +26,34 @@
                 'mode' => 'chat',
             ],
         ])),
+        new Response(body: json_encode([
+            'xai/grok-2-latest' => [
+                'litellm_provider' => 'xai',
+                'mode' => 'chat',
+            ],
+        ])),
     );
 
     $provider = new LiteLLMModelInfoProvider(
         httpClient: $client,
     );
 
-    $models = $provider->getModels(ModelProvider::OpenAI);
+    $openAIModels = $provider->getModels(ModelProvider::OpenAI);
+    $xAIModels = $provider->getModels(ModelProvider::XAI);
+
+    expect($openAIModels)->toBeArray()
+        ->toHaveCount(2)
+        ->toContainOnlyInstancesOf(ModelInfo::class);
+
+    expect($openAIModels[0]->name)->toBe('gpt-4o');
+    expect($openAIModels[1]->name)->toBe('gpt-3.5-turbo');
+
+    expect($xAIModels)->toBeArray()
+        ->toHaveCount(1)
+        ->toContainOnlyInstancesOf(ModelInfo::class);
 
-    expect($models)->toBeArray()->toContainOnlyInstancesOf(ModelInfo::class);
+    // prefix is removed from model name
+    expect($xAIModels[0]->name)->toBe('grok-2-latest');
 });
 
 test('it can get the model info', function (): void {
diff --git a/tests/Unit/Providers/OllamaModelInfoProviderTest.php b/tests/Unit/Providers/OllamaModelInfoProviderTest.php
@@ -57,7 +57,7 @@
         new Response(body: json_encode([
             'name' => 'mistral-small3.1',
             'model_info' => [
-                'mock.context_length' => 1024 * 8,
+                'mock.context_length' => 1024 * 128,
             ],
             'capabilities' => [
                 'completion',
@@ -76,7 +76,7 @@
     expect($modelInfo->name)->toBe('mistral-small3.1');
     expect($modelInfo->provider)->toBe(ModelProvider::Ollama);
     expect($modelInfo->type)->toBe(ModelType::Chat);
-    expect($modelInfo->maxInputTokens)->toBe(8000);
+    expect($modelInfo->maxInputTokens)->toBe(128000);
     expect($modelInfo->maxOutputTokens)->toBeNull();
     expect($modelInfo->inputCostPerToken)->toBe(0.0);
     expect($modelInfo->outputCostPerToken)->toBe(0.0);
