pygeai 0.1.6__py3-none-any.whl → 0.6.0b15__py3-none-any.whl

pygeai/tests/cli/commands/test_chat.py

@@ -0,0 +1,130 @@
+import unittest
+from unittest.mock import patch, MagicMock
+from pygeai.cli.commands.chat import show_help, get_chat_completion, chat_with_iris, chat_with_agent
+from pygeai.core.common.exceptions import MissingRequirementException, WrongArgumentError
+from pygeai.cli.commands import Option
+
+
+class TestChatCommands(unittest.TestCase):
+    """
+    python -m unittest pygeai.tests.cli.commands.test_chat.TestChatCommands
+    """
+
+    @patch('pygeai.cli.commands.chat.ChatClient')
+    @patch('pygeai.cli.commands.common.get_messages')
+    def test_get_chat_completion_success(self, mock_get_messages, mock_chat_client):
+        mock_instance = mock_chat_client.return_value
+        mock_instance.chat_completion.return_value = "Chat response"
+        mock_get_messages.return_value = [{"role": "user", "content": "Hello"}]
+
+        option_list = [
+            (Option("model", ["--model"], "", True), "test-model"),
+            (Option("messages", ["--messages"], "", True), '{"role": "user", "content": "Hello"}')
+        ]
+
+        get_chat_completion(option_list)
+        mock_chat_client.assert_called_once()
+        mock_instance.chat_completion.assert_called_once()
+
+    def test_get_chat_completion_missing_requirements(self):
+        option_list = [
+            (Option("model", ["--model"], "", True), "test-model")
+        ]
+        with self.assertRaises(MissingRequirementException):
+            get_chat_completion(option_list)
+
+    def test_get_chat_completion_invalid_messages_format(self):
+        option_list = [
+            (Option("model", ["--model"], "", True), "test-model"),
+            (Option("messages", ["--messages"], "", True), "invalid-json")
+        ]
+        with self.assertRaises(WrongArgumentError):
+            get_chat_completion(option_list)
+
+    @patch('pygeai.cli.commands.chat.Iris')
+    @patch('prompt_toolkit.PromptSession.prompt')
+    def test_chat_with_iris_success(self, mock_prompt, mock_iris):
+        mock_prompt.side_effect = ["Hello", "exit"]
+        mock_iris_instance = mock_iris.return_value
+        mock_iris_instance.stream_answer.return_value = iter(["Hi there!"])
+
+        chat_with_iris()
+        mock_iris.assert_called_once()
+        mock_iris_instance.stream_answer.assert_called()
+
+    @patch('pygeai.cli.commands.chat.get_project_id')
+    @patch('pygeai.cli.commands.chat.AgentClient')
+    @patch('pygeai.cli.commands.chat.AgentChatSession')
+    @patch('prompt_toolkit.PromptSession.prompt')
+    def test_chat_with_agent_success(self, mock_prompt, mock_agent_session, mock_agent_client, mock_get_project_id):
+        mock_prompt.side_effect = ["Hello", "exit"]
+        mock_session_instance = mock_agent_session.return_value
+        mock_session_instance.get_answer.return_value = "Hello, I'm your agent!"
+        mock_session_instance.stream_answer.return_value = iter(["Hi there!"])
+        
+        mock_get_project_id.return_value = "test-project-id"
+        mock_client_instance = mock_agent_client.return_value
+        mock_client_instance.get_agent.return_value = {"agentId": "123", "agentName": "test-agent"}
+
+        option_list = [
+            (Option("agent_name", ["--agent-name"], "", True), "test-agent")
+        ]
+
+        chat_with_agent(option_list)
+        mock_agent_session.assert_called_once_with("test-agent")
+        mock_session_instance.stream_answer.assert_called()
+
+    def test_chat_with_agent_missing_agent_name(self):
+        option_list = []
+        with self.assertRaises(MissingRequirementException):
+            chat_with_agent(option_list)
+
+    @patch('pygeai.cli.commands.chat.get_project_id')
+    @patch('pygeai.cli.commands.chat.AgentClient')
+    @patch('pygeai.cli.commands.chat.AgentChatSession')
+    @patch('prompt_toolkit.PromptSession.prompt')
+    @patch('builtins.open', new_callable=MagicMock)
+    @patch('json.load')
+    def test_chat_with_agent_restore_session_success(self, mock_json_load, mock_open, mock_prompt, mock_agent_session, mock_agent_client, mock_get_project_id):
+        mock_prompt.side_effect = ["Hello", "exit"]
+        mock_session_instance = mock_agent_session.return_value
+        mock_session_instance.get_answer.return_value = "Hello, I'm your agent!"
+        mock_session_instance.stream_answer.return_value = iter(["Hi there!"])
+        mock_json_load.return_value = [{"role": "user", "content": "Hi"}, {"role": "assistant", "content": "Hello!"}]
+        
+        mock_get_project_id.return_value = "test-project-id"
+        mock_client_instance = mock_agent_client.return_value
+        mock_client_instance.get_agent.return_value = {"agentId": "123", "agentName": "test-agent"}
+
+        option_list = [
+            (Option("agent_name", ["--agent-name"], "", True), "test-agent"),
+            (Option("restore_session", ["--restore-session"], "", True), "session.json")
+        ]
+
+        chat_with_agent(option_list)
+        mock_json_load.assert_called_once()
+
+    @patch('pygeai.cli.commands.chat.get_project_id')
+    @patch('pygeai.cli.commands.chat.AgentClient')
+    @patch('pygeai.cli.commands.chat.AgentChatSession')
+    @patch('prompt_toolkit.PromptSession.prompt')
+    @patch('builtins.open', new_callable=MagicMock)
+    @patch('json.dump')
+    def test_chat_with_agent_save_session_success(self, mock_json_dump, mock_open, mock_prompt, mock_agent_session, mock_agent_client, mock_get_project_id):
+        mock_prompt.side_effect = ["Hello", "exit"]
+        mock_session_instance = mock_agent_session.return_value
+        mock_session_instance.get_answer.return_value = "Hello, I'm your agent!"
+        mock_session_instance.stream_answer.return_value = iter(["Hi there!"])
+        
+        mock_get_project_id.return_value = "test-project-id"
+        mock_client_instance = mock_agent_client.return_value
+        mock_client_instance.get_agent.return_value = {"agentId": "123", "agentName": "test-agent"}
+
+        option_list = [
+            (Option("agent_name", ["--agent-name"], "", True), "test-agent"),
+            (Option("save_session", ["--save-session"], "", True), "session.json")
+        ]
+
+        chat_with_agent(option_list)
+        mock_json_dump.assert_called()
+

pygeai/tests/cli/commands/test_common.py

@@ -0,0 +1,350 @@
+import unittest
+from unittest.mock import patch
+
+from pygeai.cli.commands.common import (
+    get_llm_settings,
+    get_welcome_data,
+    get_messages,
+    get_boolean_value,
+    get_penalty_float_value,
+    get_search_options,
+    get_index_options,
+    get_welcome_data_feature_list,
+    get_welcome_data_example_prompt
+)
+from pygeai.core.common.exceptions import WrongArgumentError
+
+
+class TestCommon(unittest.TestCase):
+    """
+    python -m unittest pygeai.tests.cli.commands.test_common.TestCommon
+    """
+    def test_get_llm_settings_all_values(self):
+        result = get_llm_settings(
+            provider_name="OpenAI",
+            model_name="gpt-3.5-turbo",
+            temperature=0.7,
+            max_tokens=1000
+        )
+        expected = {
+            "providerName": "OpenAI",
+            "modelName": "gpt-3.5-turbo",
+            "temperature": 0.7,
+            "maxTokens": 1000
+        }
+        self.assertEqual(result, expected)
+
+    def test_get_llm_settings_partial_values(self):
+        result = get_llm_settings(
+            provider_name="OpenAI",
+            model_name="",
+            temperature=0.0,
+            max_tokens=0
+        )
+        expected = {"providerName": "OpenAI"}
+        self.assertEqual(result, expected)
+
+    def test_get_welcome_data_full(self):
+        feature_list = [{"title": "Feature1", "description": "Desc1"}]
+        examples_list = [{"title": "Example1", "description": "Desc1", "prompt_text": "Prompt1"}]
+        result = get_welcome_data(
+            welcome_data_title="Welcome",
+            welcome_data_description="Description",
+            feature_list=feature_list,
+            examples_prompt_list=examples_list
+        )
+        expected = {
+            "title": "Welcome",
+            "description": "Description",
+            "features": [{"title": "Feature1", "description": "Desc1"}],
+            "examplesPrompt": [{"title": "Example1", "description": "Desc1", "promptText": "Prompt1"}]
+        }
+        self.assertEqual(result, expected)
+
+    def test_get_welcome_data_empty_lists(self):
+        result = get_welcome_data(
+            welcome_data_title="Welcome",
+            welcome_data_description="Description",
+            feature_list=[],
+            examples_prompt_list=[]
+        )
+        expected = {
+            "title": "Welcome",
+            "description": "Description",
+            "features": [],
+            "examplesPrompt": []
+        }
+        self.assertEqual(result, expected)
+
+    def test_get_welcome_data_invalid_feature(self):
+        feature_list = [{"title": "Feature1"}]
+        with self.assertRaises(KeyError):
+            get_welcome_data(
+                welcome_data_title="Welcome",
+                welcome_data_description="Description",
+                feature_list=feature_list,
+                examples_prompt_list=[]
+            )
+
+    def test_get_welcome_data_invalid_example(self):
+        examples_list = [{"title": "Example1", "description": "Desc1"}]
+        with self.assertRaises(KeyError):
+            get_welcome_data(
+                welcome_data_title="Welcome",
+                welcome_data_description="Description",
+                feature_list=[],
+                examples_prompt_list=examples_list
+            )
+
+    def test_get_messages_valid(self):
+        message_list = [{"role": "user", "content": "Hello"}]
+        result = get_messages(message_list)
+        expected = [{"role": "user", "content": "Hello"}]
+        self.assertEqual(result, expected)
+
+    def test_get_messages_empty(self):
+        result = get_messages([])
+        self.assertEqual(result, [])
+
+    def test_get_messages_invalid(self):
+        message_list = [{"role": "user"}]
+        with self.assertRaises(KeyError):
+            get_messages(message_list)
+
+    def test_get_boolean_value_valid(self):
+        self.assertTrue(get_boolean_value("1"))
+        self.assertFalse(get_boolean_value("0"))
+
+    def test_get_boolean_value_invalid(self):
+        with self.assertRaises(WrongArgumentError) as context:
+            get_boolean_value("x")
+        self.assertEqual(str(context.exception), "Possible values are 0 or 1, for off and on, respectively.")
+
+    def test_get_penalty_float_value_valid(self):
+        penalty_value = get_penalty_float_value("0.0")
+        self.assertEqual(penalty_value, 0.0)
+
+    def test_get_penalty_float_value_out_of_range(self):
+        with self.assertRaises(WrongArgumentError) as context:
+            get_penalty_float_value("2.5")
+        self.assertEqual(str(context.exception), "If defined, penalty must be a number between -2.0 and 2.0")
+
+    def test_get_penalty_float_value_invalid(self):
+        with self.assertRaises(WrongArgumentError) as context:
+            get_penalty_float_value("invalid")
+        self.assertEqual(str(context.exception), "If defined, penalty must be a number between -2.0 and 2.0")
+
+    def test_get_search_options_full(self):
+        result = get_search_options(
+            history_count=5,
+            llm_cache=True,
+            llm_frequency_penalty=0.5,
+            llm_max_tokens=1000,
+            llm_model_name="gpt-3.5-turbo",
+            llm_n=1,
+            llm_presence_penalty=0.2,
+            llm_provider="OpenAI",
+            llm_stream=False,
+            llm_temperature=0.7,
+            llm_top_p=0.9,
+            llm_type={"key": "value"},
+            llm_verbose=True,
+            search_k=10,
+            search_type="similarity",
+            search_fetch_k=20,
+            search_lambda=0.5,
+            search_prompt="test prompt",
+            search_return_source_documents=True,
+            search_score_threshold=0.8,
+            search_template="test template",
+            retriever_type="vectorStore",
+            retriever_search_type="hybrid",
+            retriever_step="all",
+            retriever_prompt="custom prompt"
+        )
+        self.assertEqual(result["history_count"], 5)
+        self.assertEqual(result["llm"]["cache"], True)
+        self.assertEqual(result["llm"]["frequencyPenalty"], 0.5)
+        self.assertEqual(result["llm"]["maxTokens"], 1000)
+        self.assertEqual(result["llm"]["modelName"], "gpt-3.5-turbo")
+        self.assertEqual(result["search"]["k"], 10)
+        self.assertEqual(result["search"]["type"], "similarity")
+        self.assertEqual(result["retriever"]["type"], "vectorStore")
+        self.assertEqual(result["retriever"]["prompt"], "custom prompt")
+
+    def test_get_search_options_partial(self):
+        result = get_search_options(
+            history_count=0,
+            llm_cache=False,
+            llm_frequency_penalty=0.0,
+            llm_max_tokens=0,
+            llm_model_name="",
+            llm_n=0,
+            llm_presence_penalty=0.0,
+            llm_provider="",
+            llm_stream=False,
+            llm_temperature=0.0,
+            llm_top_p=0.0,
+            llm_type={},
+            llm_verbose=False,
+            search_k=0,
+            search_type="",
+            search_fetch_k=0,
+            search_lambda=0.0,
+            search_prompt="",
+            search_return_source_documents=False,
+            search_score_threshold=0.0,
+            search_template="",
+            retriever_type="",
+            retriever_search_type="",
+            retriever_step="",
+            retriever_prompt=""
+        )
+        self.assertIsNone(result["history_count"])
+        self.assertEqual(result["llm"], {"cache": False, "stream": False, "verbose": False})
+        self.assertEqual(result["search"], {"returnSourceDocuments": False})
+        self.assertEqual(result["retriever"], {})
+
+    def test_get_index_options_full(self):
+        result = get_index_options(
+            chunk_overlap=50,
+            chunk_size=500,
+            use_parent_document=True,
+            child_k=5.0,
+            child_chunk_size=100.0,
+            child_chunk_overlap=20.0
+        )
+        expected = {
+            "chunks": {"chunkOverlap": 50, "chunkSize": 500},
+            "useParentDocument": True,
+            "childDocument": {
+                "childK": 5.0,
+                "child": {"chunkSize": 100.0, "chunkOverlap": 20.0}
+            }
+        }
+        self.assertEqual(result, expected)
+
+    def test_get_index_options_partial(self):
+        result = get_index_options(
+            chunk_overlap=0,
+            chunk_size=0,
+            use_parent_document=False,
+            child_k=0.0,
+            child_chunk_size=0.0,
+            child_chunk_overlap=0.0
+        )
+        expected = {
+            "chunks": {"chunkOverlap": 0, "chunkSize": 0},
+            "useParentDocument": False,
+            "childDocument": {}
+        }
+        self.assertEqual(result, expected)
+
+    def test_get_welcome_data_feature_list_valid_dict(self):
+        feature_list = []
+        option_arg = '{"title": "Feature1", "description": "Desc1"}'
+        result = get_welcome_data_feature_list(feature_list, option_arg)
+        expected = [{"title": "Feature1", "description": "Desc1"}]
+        self.assertEqual(result, expected)
+
+    def test_get_welcome_data_feature_list_valid_list(self):
+        feature_list = []
+        option_arg = '[{"title": "Feature1", "description": "Desc1"}]'
+        result = get_welcome_data_feature_list(feature_list, option_arg)
+        expected = [{"title": "Feature1", "description": "Desc1"}]
+        self.assertEqual(result, expected)
+
+    def test_get_welcome_data_feature_list_invalid_json(self):
+        feature_list = []
+        option_arg = "invalid_json"
+        with self.assertRaises(WrongArgumentError) as context:
+            get_welcome_data_feature_list(feature_list, option_arg)
+        self.assertIn("Features must be a JSON string", str(context.exception))
+
+    def test_get_welcome_data_example_prompt_valid_dict(self):
+        examples_list = []
+        option_arg = '{"title": "Example1", "description": "Desc1", "prompt_text": "Prompt1"}'
+        result = get_welcome_data_example_prompt(examples_list, option_arg)
+        expected = [{"title": "Example1", "description": "Desc1", "prompt_text": "Prompt1"}]
+        self.assertEqual(result, expected)
+
+    def test_get_welcome_data_example_prompt_valid_list(self):
+        examples_list = []
+        option_arg = '[{"title": "Example1", "description": "Desc1", "prompt_text": "Prompt1"}]'
+        result = get_welcome_data_example_prompt(examples_list, option_arg)
+        expected = [{"title": "Example1", "description": "Desc1", "prompt_text": "Prompt1"}]
+        self.assertEqual(result, expected)
+
+    def test_get_welcome_data_example_prompt_invalid_json(self):
+        examples_list = []
+        option_arg = "invalid_json"
+        with self.assertRaises(WrongArgumentError) as context:
+            get_welcome_data_example_prompt(examples_list, option_arg)
+        self.assertIn("Example prompt text must be a JSON string", str(context.exception))
+
+    def test_get_llm_settings_empty_values(self):
+        result = get_llm_settings(
+            provider_name="",
+            model_name="",
+            temperature=0.0,
+            max_tokens=0
+        )
+        expected = {}
+        self.assertEqual(result, expected)
+
+    def test_get_penalty_float_value_boundary_values(self):
+        self.assertEqual(get_penalty_float_value("2.0"), 2.0)
+        self.assertEqual(get_penalty_float_value("-2.0"), -2.0)
+
+    def test_get_search_options_only_history_count(self):
+        result = get_search_options(
+            history_count=3,
+            llm_cache=False,
+            llm_frequency_penalty=0.0,
+            llm_max_tokens=0,
+            llm_model_name="",
+            llm_n=0,
+            llm_presence_penalty=0.0,
+            llm_provider="",
+            llm_stream=False,
+            llm_temperature=0.0,
+            llm_top_p=0.0,
+            llm_type={},
+            llm_verbose=False,
+            search_k=0,
+            search_type="",
+            search_fetch_k=0,
+            search_lambda=0.0,
+            search_prompt="",
+            search_return_source_documents=False,
+            search_score_threshold=0.0,
+            search_template="",
+            retriever_type="",
+            retriever_search_type="",
+            retriever_step="",
+            retriever_prompt=""
+        )
+        self.assertEqual(result["history_count"], 3)
+        self.assertEqual(result["llm"], {"cache": False, "stream": False, "verbose": False})
+        self.assertEqual(result["search"], {"returnSourceDocuments": False})
+        self.assertEqual(result["retriever"], {})
+
+    def test_get_welcome_data_feature_list_append_to_existing(self):
+        feature_list = [{"title": "Feature1", "description": "Desc1"}]
+        option_arg = '{"title": "Feature2", "description": "Desc2"}'
+        result = get_welcome_data_feature_list(feature_list, option_arg)
+        expected = [
+            {"title": "Feature1", "description": "Desc1"},
+            {"title": "Feature2", "description": "Desc2"}
+        ]
+        self.assertEqual(result, expected)
+
+    def test_get_welcome_data_example_prompt_append_to_existing(self):
+        examples_list = [{"title": "Example1", "description": "Desc1", "prompt_text": "Prompt1"}]
+        option_arg = '{"title": "Example2", "description": "Desc2", "prompt_text": "Prompt2"}'
+        result = get_welcome_data_example_prompt(examples_list, option_arg)
+        expected = [
+            {"title": "Example1", "description": "Desc1", "prompt_text": "Prompt1"},
+            {"title": "Example2", "description": "Desc2", "prompt_text": "Prompt2"}
+        ]
+        self.assertEqual(result, expected)

pygeai/tests/cli/commands/test_embeddings.py

@@ -0,0 +1,132 @@
+import unittest
+from unittest.mock import patch, MagicMock
+
+from pygeai.cli.commands.embeddings import show_help, generate_embeddings, embeddings_commands, generate_embeddings_options
+from pygeai.cli.commands import Option, Command
+from pygeai.core.common.exceptions import MissingRequirementException, WrongArgumentError
+
+
+class TestEmbeddingsCommands(unittest.TestCase):
+    """
+    python -m unittest pygeai.tests.cli.commands.test_embeddings.TestEmbeddingsCommands
+    """
+    def setUp(self):
+        self.mock_console = MagicMock()
+        self.mock_client = MagicMock()
+        self.mock_client.generate_embeddings.return_value = {
+            "model": "test-model",
+            "object": "list",
+            "data": [
+                {
+                    "object": "embedding",
+                    "index": 0,
+                    "embedding": [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7]
+                }
+            ],
+            "usage": {
+                "prompt_tokens": 10,
+                "total_tokens": 10
+            }
+        }
+        self.option_list_with_all_params = [
+            (Option("model", ["--model"], "", True), "test-model"),
+            (Option("input", ["--input"], "", True), "input1"),
+            (Option("input", ["--input"], "", True), "input2"),
+            (Option("encoding_format", ["--encoding-format"], "", True), "float"),
+            (Option("dimensions", ["--dimensions"], "", True), "512"),
+            (Option("user", ["--user"], "", True), "user123"),
+            (Option("input_type", ["--input-type"], "", True), "text"),
+            (Option("timeout", ["--timeout"], "", True), "30"),
+            (Option("cache", ["--cache"], "", True), "1")
+        ]
+        self.option_list_minimal = [
+            (Option("model", ["--model"], "", True), "test-model"),
+            (Option("input", ["--input"], "", True), "input1")
+        ]
+
+    @patch('pygeai.cli.commands.embeddings.Console.write_stdout')
+    @patch('pygeai.cli.commands.embeddings.EmbeddingsClient')
+    def test_generate_embeddings_success_all_params(self, mock_client_class, mock_write_stdout):
+        mock_client_class.return_value = self.mock_client
+
+        generate_embeddings(self.option_list_with_all_params)
+
+        mock_client_class.assert_called_once()
+        self.mock_client.generate_embeddings.assert_called_once_with(
+            input_list=["input1", "input2"],
+            model="test-model",
+            encoding_format="float",
+            dimensions=512,
+            user="user123",
+            input_type="text",
+            timeout=30,
+            cache=True
+        )
+        mock_write_stdout.assert_called_once()
+
+    @patch('pygeai.cli.commands.embeddings.Console.write_stdout')
+    @patch('pygeai.cli.commands.embeddings.EmbeddingsClient')
+    def test_generate_embeddings_success_minimal_params(self, mock_client_class, mock_write_stdout):
+        mock_client_class.return_value = self.mock_client
+
+        generate_embeddings(self.option_list_minimal)
+
+        mock_client_class.assert_called_once()
+        self.mock_client.generate_embeddings.assert_called_once_with(
+            input_list=["input1"],
+            model="test-model",
+            encoding_format=None,
+            dimensions=None,
+            user=None,
+            input_type=None,
+            timeout=None,
+            cache=None
+        )
+        mock_write_stdout.assert_called_once()
+
+    def test_generate_embeddings_missing_model(self):
+        option_list_no_model = [
+            (Option("input", ["--input"], "", True), "input1")
+        ]
+
+        with self.assertRaises(MissingRequirementException):
+            generate_embeddings(option_list_no_model)
+
+    def test_generate_embeddings_missing_input(self):
+        option_list_no_input = [
+            (Option("model", ["--model"], "", True), "test-model")
+        ]
+
+        with self.assertRaises(MissingRequirementException):
+            generate_embeddings(option_list_no_input)
+
+    def test_generate_embeddings_invalid_cache_value(self):
+        option_list_invalid_cache = [
+            (Option("model", ["--model"], "", True), "test-model"),
+            (Option("input", ["--input"], "", True), "input1"),
+            (Option("cache", ["--cache"], "", True), "invalid")
+        ]
+
+        with self.assertRaises(WrongArgumentError):
+            generate_embeddings(option_list_invalid_cache)
+
+    def test_generate_embeddings_invalid_dimensions(self):
+        option_list_invalid_dimensions = [
+            (Option("model", ["--model"], "", True), "test-model"),
+            (Option("input", ["--input"], "", True), "input1"),
+            (Option("dimensions", ["--dimensions"], "", True), "invalid")
+        ]
+
+        with self.assertRaises(WrongArgumentError):
+            generate_embeddings(option_list_invalid_dimensions)
+
+    def test_generate_embeddings_invalid_timeout(self):
+        option_list_invalid_timeout = [
+            (Option("model", ["--model"], "", True), "test-model"),
+            (Option("input", ["--input"], "", True), "input1"),
+            (Option("timeout", ["--timeout"], "", True), "invalid")
+        ]
+
+        with self.assertRaises(WrongArgumentError):
+            generate_embeddings(option_list_invalid_timeout)
+