alita-sdk 0.3.257__py3-none-any.whl → 0.3.562__py3-none-any.whl

alita_sdk/tools/testrail/api_wrapper.py

@@ -1,20 +1,19 @@
 import json
 import logging
-from typing import Dict, List, Optional, Union, Any, Generator
+from typing import Dict, List, Literal, Optional, Union, Any, Generator
 
 import pandas as pd
+from langchain_core.documents import Document
 from langchain_core.tools import ToolException
 from openai import BadRequestError
 from pydantic import SecretStr, create_model, model_validator
 from pydantic.fields import Field, PrivateAttr
 from testrail_api import StatusCodeError, TestRailAPI
 
-from ..chunkers.code.constants import get_file_extension
-from ..elitea_base import BaseVectorStoreToolApiWrapper, extend_with_vector_tools
-from langchain_core.documents import Document
-
+from ..chunkers.code.constants import get_file_extension, image_extensions
+from ..non_code_indexer_toolkit import NonCodeIndexerToolkit
+from ..utils.available_tools_decorator import extend_with_parent_available_tools
 from ...runtime.utils.utils import IndexerKeywords
-from ..utils.content_parser import parse_file_content
 
 try:
     from alita_sdk.runtime.langchain.interfaces.llm_processor import get_embeddings
@@ -118,6 +117,13 @@ getCases = create_model(
             description="A list of case field keys to include in the data output. If None, defaults to ['title', 'id'].",
         ),
     ),
+    suite_id=(Optional[str],
+              Field(
+                  default=None,
+                  description="[Optional] Suite id for test cases extraction in case "
+                              "project is in multiple suite mode (setting 3)",
+              ),
+              ),
 )
 
 getCasesByFilter = create_model(
@@ -292,6 +298,18 @@ updateCase = create_model(
     ),
 )
 
+getSuites = create_model(
+    "getSuites",
+    project_id=(str, Field(description="Project id")),
+    output_format=(
+        str,
+        Field(
+            default="json",
+            description="Desired output format. Supported values: 'json', 'csv', 'markdown'. Defaults to 'json'.",
+        ),
+    ),
+)
+
 SUPPORTED_KEYS = {
     "id", "title", "section_id", "template_id", "type_id", "priority_id", "milestone_id",
     "refs", "created_by", "created_on", "updated_by", "updated_on", "estimate",
@@ -301,7 +319,7 @@ SUPPORTED_KEYS = {
 }
 
 
-class TestrailAPIWrapper(BaseVectorStoreToolApiWrapper):
+class TestrailAPIWrapper(NonCodeIndexerToolkit):
     url: str
     password: Optional[SecretStr] = None,
     email: Optional[str] = None,
@@ -322,7 +340,76 @@ class TestrailAPIWrapper(BaseVectorStoreToolApiWrapper):
         password = values.get("password")
         email = values.get("email")
         cls._client = TestRailAPI(url, email, password)
-        return values
+        return super().validate_toolkit(values)
+
+    def _is_suite_id_required(self, project_id: str) -> bool:
+        """
+        Returns True if project requires suite_id (multiple suite or baselines mode), otherwise False.
+        Args:
+            project_id: The TestRail project ID to check
+        """
+        try:
+            project = self._client.projects.get_project(project_id=project_id)
+            # 1 for single suite mode, 2 for single suite + baselines, 3 for multiple suites
+            suite_mode = project.get('suite_mode', 1)
+            return suite_mode == 2 or suite_mode == 3
+        except StatusCodeError:
+            return False
+
+    def _fetch_cases_with_suite_handling(
+        self, 
+        project_id: str, 
+        suite_id: Optional[str] = None, 
+        **api_params
+    ) -> List[Dict]:
+        """
+        Unified method to fetch test cases with proper TestRail suite mode handling.
+        
+        Args:
+            project_id: The TestRail project ID
+            suite_id: Optional suite ID to filter by
+            **api_params: Additional parameters to pass to the get_cases API call
+            
+        Returns:
+            List of test case dictionaries
+        """
+        def _extract_cases_from_response(response):
+            """Extract cases from API response, supporting both old and new testrail_api versions."""
+            return response.get('cases', []) if isinstance(response, dict) else response
+
+        suite_required = self._is_suite_id_required(project_id=project_id)
+        all_cases = []
+
+        if suite_required:
+            # Suite modes 2 & 3: Require suite_id parameter
+            if suite_id:
+                response = self._client.cases.get_cases(
+                    project_id=project_id, suite_id=int(suite_id), **api_params
+                )
+                cases_from_suite = _extract_cases_from_response(response)
+                all_cases.extend(cases_from_suite)
+            else:
+                suites = self._get_raw_suites(project_id)
+                suite_ids = [suite['id'] for suite in suites if 'id' in suite]
+                for current_suite_id in suite_ids:
+                    try:
+                        response = self._client.cases.get_cases(
+                            project_id=project_id, suite_id=int(current_suite_id), **api_params
+                        )
+                        cases_from_suite = _extract_cases_from_response(response)
+                        all_cases.extend(cases_from_suite)
+                    except StatusCodeError:
+                        continue
+        else:
+            # Suite mode 1: Can fetch all cases directly without suite_id
+            try:
+                response = self._client.cases.get_cases(project_id=project_id, **api_params)
+                cases_from_project = _extract_cases_from_response(response)
+                all_cases.extend(cases_from_project)
+            except StatusCodeError as e:
+                logger.warning(f"Unable to fetch cases at project level: {e}")
+
+        return all_cases
 
     def add_cases(self, add_test_cases_data: str):
         """Adds new test cases into Testrail per defined parameters.
@@ -390,7 +477,8 @@ class TestrailAPIWrapper(BaseVectorStoreToolApiWrapper):
         return f"Extracted test case:\n{str(extracted_case)}"
 
     def get_cases(
-        self, project_id: str, output_format: str = "json", keys: Optional[List[str]] = None
+        self, project_id: str, output_format: str = "json", keys: Optional[List[str]] = None,
+            suite_id: Optional[str] = None
     ) -> Union[str, ToolException]:
         """
         Extracts a list of test cases in the specified format: `json`, `csv`, or `markdown`.
@@ -411,10 +499,10 @@ class TestrailAPIWrapper(BaseVectorStoreToolApiWrapper):
         invalid_keys = [key for key in keys if key not in SUPPORTED_KEYS]
 
         try:
-            extracted_cases = self._client.cases.get_cases(project_id=project_id)
-            cases = extracted_cases.get("cases")
+            # Use unified suite handling method
+            cases = self._fetch_cases_with_suite_handling(project_id=project_id, suite_id=suite_id)
 
-            if cases is None:
+            if not cases:
                 return ToolException("No test cases found in the extracted data.")
 
             extracted_cases_data = [
@@ -467,24 +555,36 @@ class TestrailAPIWrapper(BaseVectorStoreToolApiWrapper):
                     "json_case_arguments must be a JSON string or dictionary."
                 )
             self._log_tool_event(message=f"Extract test cases per filter {params}", tool_name='get_cases_by_filter')
-            extracted_cases = self._client.cases.get_cases(
-                project_id=project_id, **params
+            
+            # Extract suite_id from params for unified handling
+            suite_id_in_params = params.pop('suite_id', None)
+            suite_id = str(suite_id_in_params) if suite_id_in_params else None
+            
+            # Use unified suite handling method with remaining filter parameters
+            cases = self._fetch_cases_with_suite_handling(
+                project_id=project_id, 
+                suite_id=suite_id, 
+                **params
             )
-            self._log_tool_event(message=f"Test cases were extracted", tool_name='get_cases_by_filter')
-            # support old versions of testrail_api
-            cases = extracted_cases.get("cases") if isinstance(extracted_cases, dict) else extracted_cases
 
-            if cases is None:
+            self._log_tool_event(message="Test cases were extracted", tool_name='get_cases_by_filter')
+
+            if not cases:
                 return ToolException("No test cases found in the extracted data.")
 
             if keys is None:
                 return self._to_markup(cases, output_format)
 
-            extracted_cases_data = [
-                {key: case.get(key, "N/A") for key in keys} for case in cases
-            ]
+            extracted_cases_data = []
+            for case in cases:
+                case_dict = {}
+                for key in keys:
+                    if key in case:
+                        case_dict[key] = case[key]
+                if case_dict:
+                    extracted_cases_data.append(case_dict)
 
-            if extracted_cases_data is None:
+            if not extracted_cases_data:
                 return ToolException("No valid test case data found to format.")
 
             result = self._to_markup(extracted_cases_data, output_format)
@@ -532,47 +632,86 @@ class TestrailAPIWrapper(BaseVectorStoreToolApiWrapper):
         return (
             f"Test case #{case_id} has been updated at '{updated_case['updated_on']}')"
         )
+    
+    def _get_raw_suites(self, project_id: str) -> List[Dict]:
+        """
+        Internal helper to get raw suite data from TestRail API.
+        Handles both old and new testrail_api response formats.
+        """
+        suites_response = self._client.suites.get_suites(project_id=project_id)
+        if isinstance(suites_response, dict) and 'suites' in suites_response:
+            return suites_response['suites']
+        else:
+            return suites_response if isinstance(suites_response, list) else []
+    
+    def get_suites(self, project_id: str, output_format: str = "json",) -> Union[str, ToolException]:
+        """Extracts a list of test suites for a given project from Testrail"""
+        try:
+            suites = self._get_raw_suites(project_id)
+
+            if not suites:
+                return ToolException("No test suites found for the specified project.")
+
+            suite_dicts = []
+            for suite in suites:
+                if isinstance(suite, dict):
+                    suite_dict = {}
+                    for field in ["id", "name", "description", "project_id", "is_baseline", 
+                                 "completed_on", "url", "is_master", "is_completed"]:
+                        if field in suite:
+                            suite_dict[field] = suite[field]
+                    suite_dicts.append(suite_dict)
+                else:
+                    suite_dicts.append({"suite": str(suite)})
+            return self._to_markup(suite_dicts, output_format)
+        except StatusCodeError as e:
+            return ToolException(f"Unable to extract test suites: {e}")
 
     def _base_loader(self, project_id: str,
                      suite_id: Optional[str] = None,
                      section_id: Optional[int] = None,
                      title_keyword: Optional[str] = None,
+                     chunking_tool: str = None,
                      **kwargs: Any
                      ) -> Generator[Document, None, None]:
         self._include_attachments = kwargs.get('include_attachments', False)
         self._skip_attachment_extensions = kwargs.get('skip_attachment_extensions', [])
 
         try:
-            if suite_id:
-                resp = self._client.cases.get_cases(project_id=project_id, suite_id=int(suite_id))
-                cases = resp.get('cases', [])
-            else:
-                resp = self._client.cases.get_cases(project_id=project_id)
-                cases = resp.get('cases', [])
+            # Use unified suite handling method
+            cases = self._fetch_cases_with_suite_handling(project_id=project_id, suite_id=suite_id)
         except StatusCodeError as e:
             raise ToolException(f"Unable to extract test cases: {e}")
-            # Apply filters
+
+        # Apply filters
         if section_id is not None:
             cases = [case for case in cases if case.get('section_id') == section_id]
         if title_keyword is not None:
             cases = [case for case in cases if title_keyword.lower() in case.get('title', '').lower()]
 
         for case in cases:
-            yield Document(page_content=json.dumps(case), metadata={
+            metadata = {
                 'project_id': project_id,
                 'title': case.get('title', ''),
                 'suite_id': suite_id or case.get('suite_id', ''),
                 'id': str(case.get('id', '')),
                 IndexerKeywords.UPDATED_ON.value: case.get('updated_on') or -1,
                 'labels': [lbl['title'] for lbl in case.get('labels', [])],
-                'type': case.get('type_id') or -1,
+                'type': "testrail_test_case",
                 'priority': case.get('priority_id') or -1,
                 'milestone': case.get('milestone_id') or -1,
                 'estimate': case.get('estimate') or '',
                 'automation_type': case.get('custom_automation_type') or -1,
                 'section_id': case.get('section_id') or -1,
                 'entity_type': 'test_case',
-            })
+            }
+            if chunking_tool:
+                # content is in metadata for chunking tool post-processing
+                metadata[IndexerKeywords.CONTENT_IN_BYTES.value] = json.dumps(case).encode("utf-8")
+                page_content = ""
+            else:
+                page_content = json.dumps(case)
+            yield Document(page_content=page_content, metadata=metadata)
 
     def _process_document(self, document: Document) -> Generator[Document, None, None]:
         """
@@ -592,37 +731,56 @@ class TestrailAPIWrapper(BaseVectorStoreToolApiWrapper):
                 return
 
             # get base data from the document required to extract attachments and other metadata
-            base_data = json.loads(document.page_content)
+            base_data = document.metadata
             case_id = base_data.get("id")
 
             # get a list of attachments for the case
-            attachments = self._client.attachments.get_attachments_for_case_bulk(case_id=case_id)
+            attachments_response = self._client.attachments.get_attachments_for_case(case_id=case_id)
+
+            # Extract attachments from response - handle both old and new API response formats
+            if isinstance(attachments_response, dict) and 'attachments' in attachments_response:
+                attachments = attachments_response['attachments']
+            else:
+                attachments = attachments_response if isinstance(attachments_response, list) else []
 
             # process each attachment to extract its content
             for attachment in attachments:
-                if get_file_extension(attachment['filename']) in self._skip_attachment_extensions:
+                attachment_name = attachment.get('filename') or attachment.get('name')
+                attachment['filename'] = attachment_name
+                
+                # Handle filetype: use existing field if present, otherwise extract from filename
+                if 'filetype' not in attachment or not attachment['filetype']:
+                    file_extension = get_file_extension(attachment_name)
+                    attachment['filetype'] = file_extension.lstrip('.')
+                
+                # Handle is_image: use existing field if present, otherwise check file extension
+                if 'is_image' not in attachment:
+                    file_extension = get_file_extension(attachment_name)
+                    attachment['is_image'] = file_extension in image_extensions
+                
+                if get_file_extension(attachment_name) in self._skip_attachment_extensions:
                     logger.info(f"Skipping attachment {attachment['filename']} with unsupported extension.")
                     continue
 
                 attachment_id = f"attach_{attachment['id']}"
-                # add attachment id to metadata of parent
-                document.metadata.setdefault(IndexerKeywords.DEPENDENT_DOCS.value, []).append(attachment_id)
-                # TODO: pass it to chunkers
-                yield Document(page_content=self._process_attachment(attachment),
+                file_name, content_bytes = self._process_attachment(attachment)
+                yield Document(page_content="",
                                                      metadata={
                                                          'project_id': base_data.get('project_id', ''),
+                                                         'updated_on': base_data.get('updated_on', ''),
                                                          'id': str(attachment_id),
-                                                         IndexerKeywords.PARENT.value: str(case_id),
                                                          'filename': attachment['filename'],
                                                          'filetype': attachment['filetype'],
                                                          'created_on': attachment['created_on'],
                                                          'entity_type': 'test_case_attachment',
                                                          'is_image': attachment['is_image'],
+                                                         IndexerKeywords.CONTENT_FILE_NAME.value: file_name,
+                                                         IndexerKeywords.CONTENT_IN_BYTES.value: content_bytes
                                                      })
         except json.JSONDecodeError as e:
             raise ToolException(f"Failed to decode JSON from document: {e}")
 
-    def _process_attachment(self, attachment: Dict[str, Any]) -> str:
+    def _process_attachment(self, attachment: Dict[str, Any]) -> tuple[Any, bytes]:
         """
         Processes an attachment to extract its content.
 
@@ -633,18 +791,21 @@ class TestrailAPIWrapper(BaseVectorStoreToolApiWrapper):
             str: string description of the attachment.
         """
 
-        page_content = "This filetype is not supported."
         if attachment['filetype'] == 'txt' :
-            page_content =  self._client.get(endpoint=f"get_attachment/{attachment['id']}")
+            response = self._client.get(endpoint=f"get_attachment/{attachment['id']}")
+            if hasattr(response, "content"):
+                return attachment['filename'], response.content
+            elif isinstance(response, str):
+                return attachment['filename'], response.encode("utf-8")
         else:
             try:
                 attachment_path = self._client.attachments.get_attachment(attachment_id=attachment['id'], path=f"./{attachment['filename']}")
-                page_content = parse_file_content(file_name=attachment['filename'], file_content=attachment_path.read_bytes(), llm=self.llm, is_capture_image=True)
+                return attachment['filename'], attachment_path.read_bytes()
             except BadRequestError as ai_e:
                 logger.error(f"Unable to parse page's content with type: {attachment['filetype']} due to AI service issues: {ai_e}")
             except Exception as e:
                 logger.error(f"Unable to parse page's content with type: {attachment['filetype']}: {e}")
-        return page_content
+        return attachment['filename'], b"This filetype is not supported."
 
     def _index_tool_params(self):
         return {
@@ -658,6 +819,7 @@ class TestrailAPIWrapper(BaseVectorStoreToolApiWrapper):
             'skip_attachment_extensions': (Optional[List[str]], Field(
                 description="List of file extensions to skip when processing attachments: i.e. ['.png', '.jpg']",
                 default=[])),
+            'chunking_tool':(Literal['json', ''], Field(description="Name of chunking tool", default='json'))
         }
 
     def _to_markup(self, data: List[Dict], output_format: str) -> str:
@@ -687,7 +849,7 @@ class TestrailAPIWrapper(BaseVectorStoreToolApiWrapper):
         if output_format == "markdown":
             return df.to_markdown(index=False)
 
-    @extend_with_vector_tools
+    @extend_with_parent_available_tools
     def get_available_tools(self):
         tools = [
             {
@@ -725,6 +887,12 @@ class TestrailAPIWrapper(BaseVectorStoreToolApiWrapper):
                 "ref": self.update_case,
                 "description": self.update_case.__doc__,
                 "args_schema": updateCase,
+            },
+            {
+                "name": "get_suites",
+                "ref": self.get_suites,
+                "description": self.get_suites.__doc__,
+                "args_schema": getSuites,
             }
         ]
         return tools

alita_sdk/tools/utils/__init__.py

@@ -7,6 +7,8 @@ import requests
 from pydantic import create_model, Field
 
 
+# DEPRECATED: Tool names no longer use prefixes
+# Kept for backward compatibility only
 TOOLKIT_SPLITTER = "___"
 TOOL_NAME_LIMIT = 64
 
@@ -22,10 +24,13 @@ def clean_string(s: str, max_length: int = 0):
 
 
 def get_max_toolkit_length(selected_tools: Any):
-    """Calculates the maximum length of the toolkit name based on the selected tools per toolkit."""
-
-    longest_tool_name_length = max(len(tool_name) for tool_name in selected_tools.keys())
-    return TOOL_NAME_LIMIT - longest_tool_name_length - len(TOOLKIT_SPLITTER)
+    """DEPRECATED: Calculates the maximum length of the toolkit name.
+    
+    This function is deprecated as tool names no longer use prefixes.
+    Returns a fixed value for backward compatibility.
+    """
+    # Return a reasonable default since we no longer use prefixes
+    return 50
 
 
 def parse_list(list_str: str = None) -> List[str]:
@@ -55,6 +60,8 @@ def parse_type(type_str):
     """Parse a type string into an actual Python type."""
     try:
         # Evaluate the type string using builtins and imported modules
+        if type_str == 'number':
+            type_str = 'int'
         return eval(type_str, {**vars(builtins), **globals()})
     except Exception as e:
         print(f"Error parsing type: {e}")
@@ -95,3 +102,20 @@ def check_connection_response(check_fun):
         else:
             return f"Service Unreachable: return code {response.status_code}"
     return _wrapper
+
+
+def make_json_serializable(obj):
+    if isinstance(obj, BaseModel):
+        return obj.model_dump()
+    if isinstance(obj, dict):
+        return {k: make_json_serializable(v) for k, v in obj.items()}
+    if isinstance(obj, list):
+        return [make_json_serializable(i) for i in obj]
+    if isinstance(obj, bool):
+        return bool(obj)
+    if isinstance(obj, (str, int, float)) or obj is None:
+        return obj
+    # Fallback: handle objects that look like booleans but were not caught above
+    if str(obj) in ("True", "False"):
+        return str(obj) == "True"
+    return str(obj)