Merge branch 'fix-litellm-logprobs-client' into 'main'

fix: update litellm logprobs after upstream bugfix

See merge request gitlab-org/modelops/applied-ml/code-suggestions/ai-assist!1958

ai_gateway/models/litellm.py

@@ -348,13 +348,8 @@ class LiteLlmTextGenModel(TextGenModelBase):
         score = 10**5  # default high value if model doesn't provide score
 
         # For fireworks/qwen, use logprob of first token as score
-        # using original_response - see https://github.com/BerriAI/litellm/issues/7974
-        if self.provider == KindModelProvider.FIREWORKS and suggestion.get(
-            "_hidden_params"
-        ):
-            score = suggestion._hidden_params["original_response"]["choices"][0][
-                "logprobs"
-            ].token_logprobs[0]
+        if self.provider == KindModelProvider.FIREWORKS:
+            score = suggestion.choices[0].logprobs.token_logprobs[0]
 
         return TextGenModelOutput(
             text=self._extract_suggestion_text(suggestion),

conftest.py

@@ -375,11 +375,9 @@ def mock_litellm_acompletion():
                 AsyncMock(
                     message=AsyncMock(content="Test response"),
                     text="Test text completion response",
+                    logprobs= AsyncMock(token_logprobs=[999]),
                 ),
             ],
-            _hidden_params={
-                "original_response": {"choices": [{"logprobs": AsyncMock(token_logprobs=[999])}]}
-                },
             usage=AsyncMock(completion_tokens=999),
         )