codemie-test-harness 0.1.184__py3-none-any.whl → 0.1.197__py3-none-any.whl

codemie_test_harness/tests/llm/assistants/test_llm.py

@@ -1,8 +1,13 @@
 import pytest
+from codemie_sdk.models.assistant import ToolKitDetails, ToolDetails
 from hamcrest import assert_that, has_item
 from codemie_test_harness.tests.enums.model_types import ModelTypes
+from codemie_test_harness.tests.enums.tools import Toolkit, FileManagementTool
+from codemie_test_harness.tests.test_data.file_test_data import file_test_data
 from codemie_test_harness.tests.test_data.llm_test_data import MODEL_RESPONSES
+from codemie_test_harness.tests.utils.base_utils import get_random_name
 from codemie_test_harness.tests.utils.client_factory import get_client
+from codemie_test_harness.tests.utils.constants import FILES_PATH
 from codemie_test_harness.tests.utils.env_resolver import get_environment
 from codemie_test_harness.tests.utils.pytest_utils import check_mark
 
@@ -36,14 +41,25 @@ def pytest_generate_tests(metafunc):
 @pytest.mark.api
 @pytest.mark.smoke
 def test_assistant_with_different_models(
-    llm_utils, assistant_utils, model_type, similarity_check
+    llm_utils, assistant_utils, model_type, similarity_check, filesystem_integration
 ):
     assert_that(
         [row.base_name for row in llm_utils.list_llm_models()],
         has_item(model_type),
         f"{model_type} is missing in backend response",
     )
-    assistant = assistant_utils.create_assistant(model_type)
+
+    tool = ToolKitDetails(
+        toolkit=Toolkit.FILE_MANAGEMENT,
+        tools=[
+            ToolDetails(
+                name=FileManagementTool.GENERATE_IMAGE, settings=filesystem_integration
+            )
+        ],
+        settings=filesystem_integration,
+    )
+
+    assistant = assistant_utils.create_assistant(model_type, toolkits=[tool])
     response = assistant_utils.ask_assistant(assistant, "Just say one word: 'Hello'")
 
     if model_type in [ModelTypes.DEEPSEEK_R1, ModelTypes.RLAB_QWQ_32B]:
@@ -91,3 +107,41 @@ def test_assistant_with_different_models_with_temperature_parameter(
     if model_type in [ModelTypes.DEEPSEEK_R1, ModelTypes.RLAB_QWQ_32B]:
         response = "\n".join(response.split("\n")[-3:])
     similarity_check.check_similarity(response, "Hello")
+
+
+@pytest.mark.assistant
+@pytest.mark.llm
+@pytest.mark.api
+@pytest.mark.smoke
+def test_assistant_with_different_models_with_datasource_attached(
+    llm_utils,
+    assistant_utils,
+    model_type,
+    similarity_check,
+    datasource_utils,
+    default_embedding_llm,
+    kb_context,
+):
+    assert_that(
+        [row.base_name for row in llm_utils.list_llm_models()],
+        has_item(model_type),
+        f"{model_type} is missing in backend response",
+    )
+
+    file_name = file_test_data[2][0]
+
+    datasource = datasource_utils.create_file_datasource(
+        name=get_random_name(),
+        description=f"[Autotest] {file_name} with {default_embedding_llm.base_name} embedding model",
+        files=[str(FILES_PATH / file_name)],
+        embeddings_model=default_embedding_llm.base_name,
+    )
+
+    assistant = assistant_utils.create_assistant(
+        model_type, context=[kb_context(datasource)]
+    )
+    response = assistant_utils.ask_assistant(assistant, "Just say one word: 'Hello'")
+
+    if model_type in [ModelTypes.DEEPSEEK_R1, ModelTypes.RLAB_QWQ_32B]:
+        response = "\n".join(response.split("\n")[-3:])
+    similarity_check.check_similarity(response, "Hello")