Skip to content

Commit

Permalink
Update Mistral and TogetherAI tests [release]
Browse files Browse the repository at this point in the history
  • Loading branch information
AjayP13 committed Aug 2, 2024
1 parent 8e546d2 commit 6994b30
Showing 1 changed file with 3 additions and 1 deletion.
4 changes: 3 additions & 1 deletion src/tests/llms/test_llms.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@
import multiprocessing
import os
import sys
import typing
import uuid
from collections import Counter, defaultdict
from concurrent.futures import ThreadPoolExecutor
Expand Down Expand Up @@ -2988,6 +2989,7 @@ def test_get_max_context_length(self, create_datadreamer):
)

@pytest.mark.order("last")
@typing.no_type_check
def test_run(self, create_datadreamer, mocker):
def chat_mocked(**kwargs):
from mistralai.models.chat_completion import ChatCompletionResponse
Expand All @@ -3014,7 +3016,7 @@ def chat_mocked(**kwargs):
"total_tokens": 0,
},
}
return ChatCompletionResponse(**response) # type: ignore[arg-type]
return ChatCompletionResponse(**response)

with create_datadreamer():
llm = MistralAI("mistral-tiny", api_key="fakeapikey")
Expand Down

0 comments on commit 6994b30

Please sign in to comment.