alita-sdk 0.3.162__py3-none-any.whl → 0.3.164__py3-none-any.whl

alita_sdk/runtime/langchain/assistant.py

@@ -61,8 +61,8 @@ class Assistant:
                                 "Review toolkits configuration or use pipeline as master agent.")
 
         # configure memory store if memory tool is defined
-        # memory_tool = next((tool for tool in data['tools'] if tool['type'] == 'memory'), None)
-        # self._configure_store(memory_tool)
+        memory_tool = next((tool for tool in data['tools'] if tool['type'] == 'memory'), None)
+        self._configure_store(memory_tool)
         
         # Lazy import to avoid circular dependency
         from ..toolkits.tools import get_tools

alita_sdk/runtime/langchain/store_manager.py

@@ -1,6 +1,8 @@
 import threading
 import atexit
 import logging
+from urllib.parse import urlparse, unquote
+
 from psycopg import Connection
 from langgraph.store.postgres import PostgresStore
 
@@ -18,11 +20,30 @@ class StoreManager:
                     cls._instance._stores = {}
         return cls._instance
 
+    def _parse_connection_string(self, conn_str: str) -> dict:
+        """
+        Parse the connection string from SQLAlchemy style to args dict.
+        """
+        if conn_str.startswith("postgresql+psycopg://"):
+            url = conn_str[len("postgresql+psycopg://"):]
+
+        parsed = urlparse(f"//{url}")
+
+        return {
+            "user": unquote(parsed.username) if parsed.username else None,
+            "password": unquote(parsed.password) if parsed.password else None,
+            "host": parsed.hostname,
+            "port": parsed.port,
+            "dbname": parsed.path.lstrip("/") if parsed.path else None
+        }
+
     def get_store(self, conn_str: str) -> PostgresStore:
         store = self._stores.get(conn_str)
         if store is None:
             logger.info(f"Creating new PostgresStore for connection: {conn_str}")
-            conn = Connection.connect(conn_str, autocommit=True, prepare_threshold=0)
+            conn_params = self._parse_connection_string(conn_str)
+            conn_params.update({'autocommit': True, 'prepare_threshold': 0})
+            conn = Connection.connect(**conn_params)
             store = PostgresStore(conn)
             store.setup()
             self._stores[conn_str] = store

alita_sdk/runtime/toolkits/tools.py

@@ -95,7 +95,7 @@ def get_tools(tools_list: list, alita_client, llm, memory_store: BaseStore = Non
     # Add community tools
     tools += community_tools(tools_list, alita_client, llm)
     # Add alita tools
-    tools += alita_tools(tools_list, alita_client, llm)
+    tools += alita_tools(tools_list, alita_client, llm, memory_store)
     # Add MCP tools
     tools += _mcp_tools(tools_list, alita_client)
     

alita_sdk/tools/__init__.py

@@ -1,5 +1,8 @@
 import logging
 from importlib import import_module
+from typing import Optional
+
+from langgraph.store.base import BaseStore
 
 logger = logging.getLogger(__name__)
 
@@ -74,13 +77,15 @@ _safe_import_tool('carrier', 'carrier', 'get_tools', 'AlitaCarrierToolkit')
 _safe_import_tool('ocr', 'ocr', 'get_tools', 'OCRToolkit')
 _safe_import_tool('pptx', 'pptx', 'get_tools', 'PPTXToolkit')
 _safe_import_tool('postman', 'postman', 'get_tools', 'PostmanToolkit')
+_safe_import_tool('memory', 'memory', 'get_tools', 'MemoryToolkit')
+_safe_import_tool('zephyr_squad', 'zephyr_squad', 'get_tools', 'ZephyrSquadToolkit')
 
 # Log import summary
 available_count = len(AVAILABLE_TOOLS)
 total_attempted = len(AVAILABLE_TOOLS) + len(FAILED_IMPORTS)
 logger.info(f"Tool imports completed: {available_count}/{total_attempted} successful")
 
-def get_tools(tools_list, alita, llm, *args, **kwargs):
+def get_tools(tools_list, alita, llm, store: Optional[BaseStore] = None, *args, **kwargs):
     tools = []
     for tool in tools_list:
         # validate tool name syntax - it cannot be started with _
@@ -90,6 +95,7 @@ def get_tools(tools_list, alita, llm, *args, **kwargs):
         
         tool['settings']['alita'] = alita
         tool['settings']['llm'] = llm
+        tool['settings']['store'] = store
         tool_type = tool['type']
         
         # Check if tool is available and has get_tools function

alita_sdk/tools/carrier/api_wrapper.py

@@ -67,6 +67,10 @@ class CarrierAPIWrapper(BaseModel):
     def run_test(self, test_id: str, json_body):
         return self._client.run_test(test_id, json_body)
 
+    def run_ui_test(self, test_id: str, json_body):
+        """Run a UI test with the given test ID and JSON body."""
+        return self._client.run_ui_test(test_id, json_body)
+
     def get_engagements_list(self) -> List[Dict[str, Any]]:
         return self._client.get_engagements_list()
 
@@ -76,17 +80,21 @@ class CarrierAPIWrapper(BaseModel):
     def get_report_file_log(self, bucket: str, file_name: str):
         return self._client.get_report_file_log(bucket, file_name)
 
-    def upload_excel_report(self, bucket_name: str, excel_report_name: str):
-        return self._client.upload_excel_report(bucket_name, excel_report_name)
+    def upload_file(self, bucket_name: str, file_name: str):
+        return self._client.upload_file(bucket_name, file_name)
 
     def get_ui_reports_list(self) -> List[Dict[str, Any]]:
         """Get list of UI test reports from the Carrier platform."""
-        return self._client.get_ui_tests_list()
+        return self._client.get_ui_reports_list()
 
     def get_ui_tests_list(self) -> List[Dict[str, Any]]:
         """Get list of UI tests from the Carrier platform."""
         return self._client.get_ui_tests_list()
 
+    def get_locations(self) -> Dict[str, Any]:
+        """Get list of available locations/cloud settings from the Carrier platform."""
+        return self._client.get_locations()
+
     def get_ui_report_links(self, uid: str) -> list:
         """Get all unique file_names for a given UI report UID, ending with .html, without #index=*, and only unique values."""
         endpoint = f"api/v1/ui_performance/results/{self.project_id}/{uid}?sort=loop&order=asc"
@@ -120,4 +128,68 @@ class CarrierAPIWrapper(BaseModel):
             return [prefix + name for name in sorted_names]
         except Exception as e:
             logger.error(f"Failed to fetch UI report links: {e}")
-            return []
+            return []
+
+    def update_ui_test(self, test_id: str, json_body) -> Dict[str, Any]:
+        """Update UI test configuration and schedule."""
+        return self._client.update_ui_test(test_id, json_body)
+
+    def get_ui_test_details(self, test_id: str) -> Dict[str, Any]:
+        """Get detailed UI test configuration by test ID."""
+        return self._client.get_ui_test_details(test_id)
+
+    def create_ui_test(self, json_body: Dict[str, Any]) -> Dict[str, Any]:
+        """Create a new UI test."""
+        return self._client.create_ui_test(json_body)
+
+    def cancel_ui_test(self, test_id: str) -> Dict[str, Any]:
+        """Cancel a UI test by setting its status to Canceled."""
+        return self._client.cancel_ui_test(test_id)
+
+    def get_ui_report_json_files(self, uid: str) -> list:
+        """Get all JSON file names for a given UI report UID for Excel processing."""
+        endpoint = f"api/v1/ui_performance/results/{self.project_id}/{uid}?sort=loop&order=asc"
+        try:
+            response = self._client.request('get', endpoint)
+            file_names = set()
+            
+            def clean_file_name(file_name):
+                # Extract JSON files only and clean the names
+                if file_name.endswith('.json'):
+                    return file_name
+                return None
+            
+            # If the response is a dict with lists as values, flatten all file_names from all values
+            if isinstance(response, dict):
+                for value in response.values():
+                    if isinstance(value, list):
+                        for item in value:
+                            file_name = item.get("file_name")
+                            if file_name:
+                                clean_name = clean_file_name(file_name)
+                                if clean_name:
+                                    file_names.add(clean_name)
+            elif isinstance(response, list):
+                for item in response:
+                    file_name = item.get("file_name")
+                    if file_name:
+                        clean_name = clean_file_name(file_name)
+                        if clean_name:
+                            file_names.add(clean_name)
+            
+            sorted_names = sorted(file_names)
+            prefix = f"https://platform.getcarrier.io/api/v1/artifacts/artifact/default/{self.project_id}/reports/"
+            return [prefix + name for name in sorted_names]
+        except Exception as e:
+            logger.error(f"Failed to fetch UI report JSON files: {e}")
+            return []
+
+    def download_ui_report_json(self, bucket: str, file_name: str) -> str:
+        """Download UI report JSON file content."""
+        endpoint = f"api/v1/artifacts/artifact/default/{self.project_id}/{bucket}/{file_name}"
+        try:
+            response = self._client.request('get', endpoint)
+            return response
+        except Exception as e:
+            logger.error(f"Failed to download UI report JSON: {e}")
+            return None

alita_sdk/tools/carrier/backend_reports_tool.py

@@ -1,6 +1,8 @@
 import logging
 from datetime import datetime
 import json
+import zipfile
+from itertools import islice
 import traceback
 from typing import Type
 from langchain_core.tools import BaseTool, ToolException
@@ -78,14 +80,31 @@ class GetReportByIDTool(BaseTool):
 
     def _run(self, report_id: str):
         try:
-            reports = self.api_wrapper.get_reports_list()
-            report_data = {}
-            for report in reports:
-                if report_id == str(report["id"]):
-                    report_data = report
-                    break
-
-            return json.dumps(report_data)
+            report, test_log_file_path, errors_log_file_path = self.api_wrapper.get_report_file_name(report_id)
+            try:
+                with open(errors_log_file_path, mode='r') as f:
+                    # Use islice to read up to 100 lines
+                    errors = list(islice(f, 100))
+                report["errors_log"] = errors
+                # Archive with errors log file path
+                zip_file_path = f'/tmp/{report["build_id"]}_error_log_archive.zip'
+
+                # Create zip archive
+                with zipfile.ZipFile(zip_file_path, 'w', zipfile.ZIP_DEFLATED) as zipf:
+                    arcname = os.path.basename(errors_log_file_path)
+                    zipf.write(errors_log_file_path, arcname)
+
+                bucket_name = report["name"].replace("_", "").replace(" ", "").lower()
+                self.api_wrapper.upload_file(bucket_name, zip_file_path)
+                report["link_to_errors_file"] = f"{self.api_wrapper.url.rstrip('/')}/api/v1/artifacts/artifact/default/" \
+                                                f"{self.api_wrapper.project_id}/{bucket_name}/" \
+                                                f"{zip_file_path.replace('/tmp/', '')}"
+            except Exception as e:
+                logger.error(e)
+                report["errors_log"] = []
+                report["link_to_errors_file"] = "link is not available"
+
+            return json.dumps(report)
         except Exception:
             stacktrace = traceback.format_exc()
             logger.error(f"Error downloading reports: {stacktrace}")
@@ -164,17 +183,17 @@ class CreateExcelReportTool(BaseTool):
 
     def _process_report_by_id(self, report_id, parameters):
         """Process report using report ID."""
-        report, file_path = self.api_wrapper.get_report_file_name(report_id)
+        report, test_log_file_path, errors_log_file_path = self.api_wrapper.get_report_file_name(report_id)
         carrier_report = f"{self.api_wrapper.url.rstrip('/')}/-/performance/backend/results?result_id={report_id}"
         lg_type = report.get("lg_type")
         excel_report_file_name = f'/tmp/reports_test_results_{report["build_id"]}_excel_report.xlsx'
         bucket_name = report["name"].replace("_", "").replace(" ", "").lower()
 
-        result_stats_j = self._parse_report(file_path, lg_type, parameters["think_time"], is_absolute_file_path=True)
+        result_stats_j = self._parse_report(test_log_file_path, lg_type, parameters["think_time"], is_absolute_file_path=True)
         calc_thr_j = self._calculate_thresholds(result_stats_j, parameters)
 
         return self._generate_and_upload_report(
-            result_stats_j, carrier_report, calc_thr_j, parameters, excel_report_file_name, bucket_name, file_path
+            result_stats_j, carrier_report, calc_thr_j, parameters, excel_report_file_name, bucket_name, test_log_file_path
         )
 
     def _process_report_by_file(self, bucket, file_name, parameters):
@@ -233,7 +252,7 @@ class CreateExcelReportTool(BaseTool):
         excel_reporter_object.prepare_headers_and_titles()
         excel_reporter_object.write_to_excel(result_stats_j, carrier_report, calc_thr_j, parameters["pct"])
 
-        self.api_wrapper.upload_excel_report(bucket_name, excel_report_file_name)
+        self.api_wrapper.upload_file(bucket_name, excel_report_file_name)
 
         # Clean up
         self._cleanup(file_path, excel_report_file_name)

alita_sdk/tools/carrier/backend_tests_tool.py

@@ -80,22 +80,28 @@ class RunTestByIDTool(BaseTool):
     description: str = "Execute test plan from the Carrier platform."
     args_schema: Type[BaseModel] = create_model(
         "RunTestByIdInput",
-        test_id=(str, Field(default="", description="Test id to execute")),
+        test_id=(str, Field(default=None, description="Test id to execute")),
+        name=(str, Field(default=None, description="Test name to execute")),
         test_parameters=(dict, Field(default=None, description="Test parameters to override")),
     )
 
-    def _run(self, test_id: str, test_parameters=None):
+    def _run(self, test_id=None, name=None, test_parameters=None):
         try:
+            if not test_id and not name:
+                return {"message": "Please provide test id or test name to start"}
+
             # Fetch test data
             tests = self.api_wrapper.get_tests_list()
-            test_data = {}
-            for test in tests:
-                if test_id == str(test["id"]):
-                    test_data = test
-                    break
+
+            # Find the test data based on test_id or name
+            test_data = next(
+                (test for test in tests if
+                 (test_id and str(test["id"]) == test_id) or (name and str(test["name"]) == name)),
+                None
+            )
 
             if not test_data:
-                raise ValueError(f"Test with id {test_id} not found.")
+                raise ValueError(f"Test with id {test_id} or name {name} not found.")
 
             # Default test parameters
             default_test_parameters = test_data.get("test_parameters", [])

alita_sdk/tools/carrier/cancel_ui_test_tool.py

@@ -0,0 +1,178 @@
+import logging
+import re
+import traceback
+from typing import Type
+from langchain_core.tools import BaseTool, ToolException
+from pydantic.fields import Field
+from pydantic import create_model, BaseModel
+from .api_wrapper import CarrierAPIWrapper
+
+logger = logging.getLogger(__name__)
+
+
+class CancelUITestTool(BaseTool):
+    api_wrapper: CarrierAPIWrapper = Field(..., description="Carrier API Wrapper instance")
+    name: str = "cancel_ui_test"
+    description: str = "Cancel a UI test or show available tests to cancel in the Carrier platform."
+    args_schema: Type[BaseModel] = create_model(
+        "CancelUITestInput",
+        **{
+            "message": (str, Field(description="User input message (e.g., 'Cancel UI test' or 'Cancel UI test 12345')")),
+        }
+    )
+
+    def _run(self, message: str):
+        try:
+            # Parse the message to extract test ID if provided
+            test_id = self._extract_test_id(message)
+            
+            if test_id:
+                # User provided a specific test ID to cancel
+                return self._cancel_specific_test(test_id)
+            else:
+                # User didn't provide ID, show available tests to cancel
+                return self._show_cancelable_tests()
+                
+        except Exception as e:
+            logger.error(f"Error in cancel UI test: {e}")
+            raise ToolException(f"Failed to process cancel UI test request: {str(e)}")
+
+    def _extract_test_id(self, message: str) -> str:
+        """Extract test ID from user message if present."""
+        # Look for patterns like "Cancel UI test 12345" or "cancel ui test 12345"
+        match = re.search(r'cancel\s+ui\s+test\s+(\d+)', message.lower())
+        if match:
+            return match.group(1)
+        return ""
+
+    def _show_cancelable_tests(self) -> str:
+        """Show list of tests that can be canceled."""
+        try:
+            # Get all UI reports/tests
+            reports = self.api_wrapper.get_ui_reports_list()
+            
+            if not reports:
+                return "❌ **No UI tests found.**"
+            
+            # Filter tests that can be canceled (not in final states)
+            final_states = {"Canceled", "Finished", "Failed"}
+            cancelable_tests = []
+            
+            for report in reports:
+                test_status = report.get("test_status", {})
+                status = test_status.get("status", "Unknown")
+                
+                # Check if status is not in final states
+                if status not in final_states:
+                    cancelable_tests.append({
+                        "id": report.get("id"),
+                        "name": report.get("name", "Unknown"),
+                        "status": status,
+                        "percentage": test_status.get("percentage", 0),
+                        "description": test_status.get("description", "")
+                    })
+            
+            if not cancelable_tests:
+                return """# ℹ️ No Tests Available for Cancellation
+
+All UI tests are already in final states (Canceled, Finished, or Failed).
+
+## 🔍 To cancel a specific test:
+Use the command: `Cancel UI test <test_id>`
+
+Example: `Cancel UI test 12345`"""
+            
+            # Build the response message
+            response = """# 🚫 UI Tests Available for Cancellation
+
+The following tests are currently running and can be canceled:
+
+## 📋 Active Tests:
+"""
+            
+            for test in cancelable_tests:
+                response += f"""
+### 🔸 Test ID: `{test['id']}`
+- **Name:** `{test['name']}`
+- **Status:** `{test['status']}`
+- **Progress:** {test['percentage']}%
+- **Description:** {test['description']}
+"""
+            
+            response += """
+## 🚫 To cancel a specific test:
+Use the command: `Cancel UI test <test_id>`
+
+Example: `Cancel UI test 12345`"""
+            
+            return response
+            
+        except Exception as e:
+            logger.error(f"Error fetching cancelable tests: {e}")
+            return f"❌ **Error fetching tests:** {str(e)}"
+
+    def _cancel_specific_test(self, test_id: str) -> str:
+        """Cancel a specific UI test by ID."""
+        try:
+            # First, get the current status of the test
+            reports = self.api_wrapper.get_ui_reports_list()
+            target_test = None
+            
+            for report in reports:
+                if str(report.get("id")) == test_id:
+                    target_test = report
+                    break
+            
+            if not target_test:
+                return f"❌ **Test with ID `{test_id}` not found.**"
+            
+            # Check if test can be canceled
+            test_status = target_test.get("test_status", {})
+            current_status = test_status.get("status", "Unknown")
+            final_states = {"Canceled", "Finished", "Failed"}
+            
+            if current_status in final_states:
+                return f"""# ❌ Cannot Cancel Test
+
+## Test Information:
+- **Test ID:** `{test_id}`
+- **Name:** `{target_test.get('name', 'Unknown')}`
+- **Current Status:** `{current_status}`
+
+## 🚫 Reason:
+This test cannot be canceled because it is already in a final state (`{current_status}`).
+
+Only tests with status **not** in `Canceled`, `Finished`, or `Failed` can be canceled."""
+            
+            # Attempt to cancel the test
+            try:
+                cancel_response = self.api_wrapper.cancel_ui_test(test_id)
+                
+                return f"""# ✅ UI Test Canceled Successfully!
+
+## Test Information:
+- **Test ID:** `{test_id}`
+- **Name:** `{target_test.get('name', 'Unknown')}`
+- **Previous Status:** `{current_status}`
+- **New Status:** `Canceled`
+
+## 🎯 Result:
+The test has been successfully canceled and will stop executing."""
+                
+            except Exception as cancel_error:
+                logger.error(f"Error canceling test {test_id}: {cancel_error}")
+                return f"""# ❌ Failed to Cancel Test
+
+## Test Information:
+- **Test ID:** `{test_id}`
+- **Name:** `{target_test.get('name', 'Unknown')}`
+- **Current Status:** `{current_status}`
+
+## 🚫 Error:
+{str(cancel_error)}
+
+Please check the test ID and try again."""
+                
+        except Exception as e:
+            logger.error(f"Error canceling specific test {test_id}: {e}")
+            return f"❌ **Error processing cancellation for test `{test_id}`: {str(e)}**"

alita_sdk/tools/carrier/carrier_sdk.py

@@ -85,6 +85,11 @@ class CarrierClient(BaseModel):
         endpoint = f"api/v1/backend_performance/test/{self.credentials.project_id}/{test_id}"
         return self.request('post', endpoint, json=json_body).get("result_id", "")
 
+    def run_ui_test(self, test_id: str, json_body):
+        """Run a UI test with the given test ID and JSON body."""
+        endpoint = f"api/v1/ui_performance/test/{self.credentials.project_id}/{test_id}"
+        return self.request('post', endpoint, json=json_body).get("result_id", "")
+
     def get_engagements_list(self) -> List[Dict[str, Any]]:
         endpoint = f"api/v1/engagements/engagements/{self.credentials.project_id}"
         return self.request('get', endpoint).get("items", [])
@@ -123,15 +128,17 @@ class CarrierClient(BaseModel):
         for file_name in file_list:
             if file_name.startswith(report_archive_prefix) and "excel_report" not in file_name:
                 report_files_list.append(file_name)
-        file_path = self.download_and_merge_reports(report_files_list, lg_type, bucket_name, extract_to)
+        test_log_file_path, errors_log_file_path = self.download_and_merge_reports(report_files_list, lg_type, bucket_name, extract_to)
 
-        return report_info, file_path
+        return report_info, test_log_file_path, errors_log_file_path
 
-    def download_and_merge_reports(self, report_files_list: list, lg_type: str, bucket: str, extract_to: str = "/tmp") -> str:
+    def download_and_merge_reports(self, report_files_list: list, lg_type: str, bucket: str, extract_to: str = "/tmp"):
         if lg_type == "jmeter":
             summary_log_file_path = f"summary_{bucket}_jmeter.jtl"
+            error_log_file_path = f"error_{bucket}_jmeter.log"
         else:
             summary_log_file_path = f"summary_{bucket}_simulation.log"
+            error_log_file_path = f"error_{bucket}_simulation.log"
         extracted_reports = []
         for each in report_files_list:
             endpoint = f"api/v1/artifacts/artifact/{self.credentials.project_id}/{bucket}/{each}"
@@ -153,10 +160,21 @@ class CarrierClient(BaseModel):
                 os.remove(local_file_path)
             extracted_reports.append(extract_dir)
 
-        # get files from extract_dirs and merge to summary_log_file_path
+        # get files from extract_dirs and merge to summary_log_file_path and error_log_file_path
         self.merge_log_files(summary_log_file_path, extracted_reports, lg_type)
+        try:
+            self.merge_error_files(error_log_file_path, extracted_reports)
+        except Exception as e:
+            logger.error(f"Failed to merge errors log: {e}")
+
+        # Clean up
+        for each in extracted_reports:
+            try:
+                shutil.rmtree(each)
+            except Exception as e:
+                logger.error(e)
 
-        return summary_log_file_path
+        return summary_log_file_path, error_log_file_path
 
     def merge_log_files(self, summary_file, extracted_reports, lg_type):
         with open(summary_file, mode='w') as summary:
@@ -173,11 +191,14 @@ class CarrierClient(BaseModel):
                     else:
                         # Skip the first line (header) for subsequent files
                         summary.writelines(lines[1:])
-        for each in extracted_reports:
-            try:
-                shutil.rmtree(each)
-            except Exception as e:
-                logger.error(e)
+
+    def merge_error_files(self, error_file, extracted_reports):
+        with open(error_file, mode='w') as summary_errors:
+            for i, log_file in enumerate(extracted_reports):
+                report_file = f"{log_file}/simulation-errors.log"
+                with open(report_file, mode='r') as f:
+                    lines = f.readlines()
+                    summary_errors.writelines(lines)
 
     def get_report_file_log(self, bucket: str, file_name: str):
         bucket_endpoint = f"api/v1/artifacts/artifact/default/{self.credentials.project_id}/{bucket}/{file_name}"
@@ -190,10 +211,10 @@ class CarrierClient(BaseModel):
             f.write(response.content)
         return file_path
 
-    def upload_excel_report(self, bucket_name: str, excel_report_name: str):
+    def upload_file(self, bucket_name: str, file_name: str):
         upload_url = f'api/v1/artifacts/artifacts/{self.credentials.project_id}/{bucket_name}'
         full_url = f"{self.credentials.url.rstrip('/')}/{upload_url.lstrip('/')}"
-        files = {'file': open(excel_report_name, 'rb')}
+        files = {'file': open(file_name, 'rb')}
         headers = {'Authorization': f'bearer {self.credentials.token}'}
         s3_config = {'integration_id': 1, 'is_local': False}
         requests.post(full_url, params=s3_config, allow_redirects=True, files=files, headers=headers)
@@ -204,6 +225,69 @@ class CarrierClient(BaseModel):
         return self.request('get', endpoint).get("rows", [])
 
     def get_ui_reports_list(self) -> List[Dict[str, Any]]:
-            """Get list of UI test reports from the Carrier platform."""
-            endpoint = f"api/v1/ui_performance/reports/{self.credentials.project_id}"
-            return self.request('get', endpoint).get("rows", [])
+        """Get list of UI test reports from the Carrier platform."""
+        endpoint = f"api/v1/ui_performance/reports/{self.credentials.project_id}"
+        return self.request('get', endpoint).get("rows", [])
+
+    def get_locations(self) -> Dict[str, Any]:
+        """Get list of available locations/cloud settings from the Carrier platform."""
+        endpoint = f"api/v1/shared/locations/{self.credentials.project_id}"
+        return self.request('get', endpoint)
+
+    def update_ui_test(self, test_id: str, json_body) -> Dict[str, Any]:
+        """Update UI test configuration and schedule."""
+        endpoint = f"api/v1/ui_performance/test/{self.credentials.project_id}/{test_id}"
+        return self.request('put', endpoint, json=json_body)
+
+    def get_ui_test_details(self, test_id: str) -> Dict[str, Any]:
+        """Get detailed UI test configuration by test ID."""
+        endpoint = f"api/v1/ui_performance/test/{self.credentials.project_id}/{test_id}"
+        return self.request('get', endpoint)
+
+    def create_ui_test(self, json_body: Dict[str, Any]) -> Dict[str, Any]:
+        """Create a new UI test."""
+        endpoint = f"api/v1/ui_performance/tests/{self.credentials.project_id}"
+
+        # Print full JSON POST body for debugging
+        print("=" * 60)
+        print("DEBUG: Full JSON POST body for create_ui_test:")
+        print("=" * 60)
+        print(json.dumps(json_body, indent=2))
+        print("=" * 60)
+
+        # Use multipart/form-data with data field containing the JSON body
+        form_data = {'data': json.dumps(json_body)}
+
+        # Temporarily remove Content-Type header to let requests set it for multipart
+        original_headers = self.session.headers.copy()
+        if 'Content-Type' in self.session.headers:
+            del self.session.headers['Content-Type']
+
+        try:
+            full_url = f"{self.credentials.url.rstrip('/')}/{endpoint.lstrip('/')}"
+            response = self.session.post(full_url, data=form_data)
+            response.raise_for_status()
+            return response.json()
+        except requests.HTTPError as http_err:
+            logger.error(f"HTTP {response.status_code} error: {response.text[:500]}")
+            raise CarrierAPIError(f"Request to {full_url} failed with status {response.status_code}")
+        except json.JSONDecodeError:
+            logger.error(f"Response was not valid JSON. Body:\n{response.text[:500]}")
+            raise CarrierAPIError("Server returned non-JSON response")
+        finally:
+            # Restore original headers
+            self.session.headers.update(original_headers)
+
+    def cancel_ui_test(self, test_id: str) -> Dict[str, Any]:
+        """Cancel a UI test by setting its status to Canceled."""
+        endpoint = f"api/v1/ui_performance/report_status/{self.credentials.project_id}/{test_id}"
+
+        cancel_body = {
+            "test_status": {
+                "status": "Canceled",
+                "percentage": 100,
+                "description": "Test was canceled"
+            }
+        }
+
+        return self.request('put', endpoint, json=cancel_body)