alita-sdk 0.3.345__tar.gz → 0.3.347__tar.gz

{alita_sdk-0.3.345 → alita_sdk-0.3.347}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: alita_sdk
-Version: 0.3.345
+Version: 0.3.347
 Summary: SDK for building langchain agents using resources from Alita
 Author-email: Artem Rozumenko <artyom.rozumenko@gmail.com>, Mikalai Biazruchka <mikalai_biazruchka@epam.com>, Roman Mitusov <roman_mitusov@epam.com>, Ivan Krakhmaliuk <lifedj27@gmail.com>, Artem Dubrovskiy <ad13box@gmail.com>
 License-Expression: Apache-2.0

{alita_sdk-0.3.345 → alita_sdk-0.3.347}/alita_sdk/runtime/clients/client.py

@@ -85,7 +85,12 @@ class AlitaClient:
             # This loop iterates over each key-value pair in the arguments dictionary,
             # and if a value is a Pydantic object, it replaces it with its dictionary representation using .dict().
             for arg_name, arg_value in params.get('params', {}).get('arguments', {}).items():
-                if hasattr(arg_value, "dict") and callable(arg_value.dict):
+                if isinstance(arg_value, list):
+                    params['params']['arguments'][arg_name] = [
+                        item.dict() if hasattr(item, "dict") and callable(item.dict) else item
+                        for item in arg_value
+                    ]
+                elif hasattr(arg_value, "dict") and callable(arg_value.dict):
                     params['params']['arguments'][arg_name] = arg_value.dict()
             #
             response = requests.post(url, headers=self.headers, json=params, verify=False)

{alita_sdk-0.3.345 → alita_sdk-0.3.347}/alita_sdk/runtime/langchain/document_loaders/AlitaPDFLoader.py

@@ -1,6 +1,6 @@
 import pymupdf
 import fitz
-from langchain_community.document_loaders import PyPDFLoader
+from langchain_community.document_loaders import PyPDFium2Loader
 
 from .ImageParser import ImageParser
 from .utils import perform_llm_prediction_for_image_bytes, create_temp_file
@@ -23,6 +23,7 @@ class AlitaPDFLoader:
         self.headers = kwargs.get('headers', None)
         self.extraction_mode = kwargs.get('extraction_mode', "plain")
         self.extraction_kwargs = kwargs.get('extraction_kwargs', None)
+        self.images_parser=ImageParser(llm=self.llm, prompt=self.prompt)
 
     def get_content(self):
         if hasattr(self, 'file_path'):
@@ -119,13 +120,13 @@ class AlitaPDFLoader:
             return self._load_docs()
 
     def _load_docs(self):
-        docs = PyPDFLoader(file_path=self.file_path,
-                        password=self.password,
-                        headers=self.headers,
-                        extract_images=self.extract_images,
-                        extraction_mode=self.extraction_mode,
-                        images_parser=ImageParser(llm=self.llm, prompt=self.prompt),
-                        extraction_kwargs=self.extraction_kwargs).load()
+        docs = PyPDFium2Loader(
+                file_path = self.file_path,
+                password=self.password,
+                headers=self.headers,
+                extract_images = self.extract_images,
+                images_parser = ImageParser(llm=self.llm, prompt=self.prompt),
+            ).load()
         for doc in docs:
             doc.metadata['chunk_id'] = doc.metadata['page']
         return docs

alita_sdk-0.3.347/alita_sdk/runtime/langchain/document_loaders/ImageParser.py

@@ -0,0 +1,30 @@
+from typing import Iterator
+
+from langchain_community.document_loaders.parsers.images import BaseImageBlobParser
+from langchain_core.documents import Document
+from langchain_core.documents.base import Blob
+
+from alita_sdk.runtime.langchain.document_loaders.AlitaImageLoader import AlitaImageLoader
+
+class ImageParser(BaseImageBlobParser):
+
+    def __init__(self, **kwargs):
+        self.llm = kwargs.get('llm')
+        self.prompt = kwargs.get('prompt')
+
+    def lazy_parse(self, blob: Blob) -> Iterator[Document]:
+        try:
+            yield from super().lazy_parse(blob)
+        except Exception:
+            yield Document(page_content="[Image: Unknown]")
+
+    def _analyze_image(self, img) -> str:
+        from io import BytesIO
+
+        byte_stream = BytesIO()
+        img.save(byte_stream, format='PNG')
+        image_bytes = byte_stream.getvalue()
+        try:
+            return AlitaImageLoader(file_content=image_bytes, file_name="image.png", prompt=self.prompt, llm=self.llm).get_content()
+        except Exception:
+            return "Image: unknown"

{alita_sdk-0.3.345 → alita_sdk-0.3.347}/alita_sdk/runtime/tools/artifact.py

@@ -1,7 +1,7 @@
 import hashlib
-import json
 import logging
-from typing import Any, Optional, Generator
+import re
+from typing import Any, Optional, Generator, List
 
 from langchain_core.documents import Document
 from langchain_core.tools import ToolException
@@ -59,18 +59,54 @@ class ArtifactWrapper(NonCodeIndexerToolkit):
     def create_new_bucket(self, bucket_name: str, expiration_measure = "weeks", expiration_value = 1):
         return self.artifact.client.create_bucket(bucket_name, expiration_measure, expiration_value)
 
+    def _index_tool_params(self):
+        return {
+            'include_extensions': (Optional[List[str]], Field(
+                description="List of file extensions to include when processing: i.e. ['*.png', '*.jpg']. "
+                            "If empty, all files will be processed (except skip_extensions).",
+                default=[])),
+            'skip_extensions': (Optional[List[str]], Field(
+                description="List of file extensions to skip when processing: i.e. ['*.png', '*.jpg']",
+                default=[])),
+        }
+
     def _base_loader(self, **kwargs) -> Generator[Document, None, None]:
+        self._log_tool_event(message=f"Loading the files from artifact's bucket. {kwargs=}", tool_name="loader")
         try:
-            all_files = self.list_files(self.bucket, False)
+            all_files = self.list_files(self.bucket, False)['rows']
         except Exception as e:
             raise ToolException(f"Unable to extract files: {e}")
 
-        for file in all_files['rows']:
+        include_extensions = kwargs.get('include_extensions', [])
+        skip_extensions = kwargs.get('skip_extensions', [])
+        self._log_tool_event(message=f"Files filtering started. Include extensions: {include_extensions}. "
+                                     f"Skip extensions: {skip_extensions}", tool_name="loader")
+        # show the progress of filtering
+        total_files = len(all_files) if isinstance(all_files, list) else 0
+        filtered_files_count = 0
+        for file in all_files:
+            filtered_files_count += 1
+            if filtered_files_count % 10 == 0 or filtered_files_count == total_files:
+                self._log_tool_event(message=f"Files filtering progress: {filtered_files_count}/{total_files}",
+                                     tool_name="loader")
+            file_name = file['name']
+
+            # Check if file should be skipped based on skip_extensions
+            if any(re.match(pattern.replace('*', '.*') + '$', file_name, re.IGNORECASE)
+                   for pattern in skip_extensions):
+                continue
+
+            # Check if file should be included based on include_extensions
+            # If include_extensions is empty, process all files (that weren't skipped)
+            if include_extensions and not (any(re.match(pattern.replace('*', '.*') + '$', file_name, re.IGNORECASE)
+                                               for pattern in include_extensions)):
+                continue
+
             metadata = {
                 ("updated_on" if k == "modified" else k): str(v)
                 for k, v in file.items()
             }
-            metadata['id'] = self.get_hash_from_bucket_and_file_name(self.bucket, file['name'])
+            metadata['id'] = self.get_hash_from_bucket_and_file_name(self.bucket, file_name)
             yield Document(page_content="", metadata=metadata)
 
     def get_hash_from_bucket_and_file_name(self, bucket, file_name):

{alita_sdk-0.3.345 → alita_sdk-0.3.347}/alita_sdk/tools/figma/api_wrapper.py

@@ -22,7 +22,7 @@ GLOBAL_REMOVE = []
 GLOBAL_DEPTH_START = 4
 GLOBAL_DEPTH_END = 6
 EXTRA_PARAMS = (
-    Optional[Dict[str, Union[str, int, None]]],
+    Optional[Dict[str, Union[str, int, List, None]]],
     Field(
         description=(
             "Additional parameters for customizing response processing:\n"

{alita_sdk-0.3.345 → alita_sdk-0.3.347}/alita_sdk/tools/github/github_client.py

@@ -105,7 +105,7 @@ class GitHubClient(BaseModel):
                     self._github_repo_instance = None
             except Exception as e:
                 # Only raise when accessed, not during initialization
-                return ToolException(e)
+                raise ToolException(e)
         return self._github_repo_instance
 
     @model_validator(mode='before')

{alita_sdk-0.3.345 → alita_sdk-0.3.347}/alita_sdk/tools/testrail/api_wrapper.py

@@ -298,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",
@@ -330,29 +342,74 @@ class TestrailAPIWrapper(NonCodeIndexerToolkit):
         cls._client = TestRailAPI(url, email, password)
         return super().validate_toolkit(values)
 
-    def _validate_suite_mode_requirements(self, project_id: str, suite_id: Optional[str] = None) -> None:
+    def _is_suite_id_required(self, project_id: str) -> bool:
         """
-        Validate if project requires suite_id when in multiple suite mode.
-        
+        Returns True if project requires suite_id (multiple suite or baselines mode), otherwise False.
         Args:
             project_id: The TestRail project ID to check
-            suite_id: The suite ID if provided (optional)
-            custom_error_msg: Custom error message to use (optional)
-        
-        Raises:
-            ToolException: If project is in multiple suite mode and no suite_id is provided
         """
-        if suite_id:
-            return  # No validation needed if suite_id is already provided
-            
         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)
-            if suite_mode == 3:
-                raise ToolException("Project is in multiple suite mode, please provide suite_id to extract test cases.")
-        except StatusCodeError as e:
-            logger.warning(f"Unable to check project suite mode: {e}")
+            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.
@@ -442,17 +499,10 @@ class TestrailAPIWrapper(NonCodeIndexerToolkit):
         invalid_keys = [key for key in keys if key not in SUPPORTED_KEYS]
 
         try:
-            # Check if project requires suite_id for multiple suite mode
-            self._validate_suite_mode_requirements(
-                project_id=project_id,
-                suite_id=suite_id
-            )
-            
-            extracted_cases = self._client.cases.get_cases(project_id=project_id, suite_id=suite_id)
-            # support old versions of testrail_api
-            cases = extracted_cases.get("cases") if isinstance(extracted_cases, dict) else extracted_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 = [
@@ -506,31 +556,35 @@ class TestrailAPIWrapper(NonCodeIndexerToolkit):
                 )
             self._log_tool_event(message=f"Extract test cases per filter {params}", tool_name='get_cases_by_filter')
             
-            # Check if project requires suite_id when not provided in params
-            suite_id_in_params = params.get('suite_id', None)
-            self._validate_suite_mode_requirements(
-                project_id=project_id,
-                suite_id=str(suite_id_in_params) if suite_id_in_params else None
-            )
+            # 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
             
-            extracted_cases = self._client.cases.get_cases(
-                project_id=project_id, **params
+            # 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="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:
+            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)
@@ -578,6 +632,40 @@ class TestrailAPIWrapper(NonCodeIndexerToolkit):
         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,
@@ -589,24 +677,13 @@ class TestrailAPIWrapper(NonCodeIndexerToolkit):
         self._include_attachments = kwargs.get('include_attachments', False)
         self._skip_attachment_extensions = kwargs.get('skip_attachment_extensions', [])
 
-        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
-
         try:
-            # Check if project requires suite_id when not provided
-            self._validate_suite_mode_requirements(project_id=project_id, suite_id=suite_id)
-
-            if suite_id:
-                resp = self._client.cases.get_cases(project_id=project_id, suite_id=int(suite_id))
-            else:
-                resp = self._client.cases.get_cases(project_id=project_id)
-            
-            cases = _extract_cases_from_response(resp)
-            
+            # 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:
@@ -810,6 +887,12 @@ class TestrailAPIWrapper(NonCodeIndexerToolkit):
                 "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-0.3.345 → alita_sdk-0.3.347}/alita_sdk.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: alita_sdk
-Version: 0.3.345
+Version: 0.3.347
 Summary: SDK for building langchain agents using resources from Alita
 Author-email: Artem Rozumenko <artyom.rozumenko@gmail.com>, Mikalai Biazruchka <mikalai_biazruchka@epam.com>, Roman Mitusov <roman_mitusov@epam.com>, Ivan Krakhmaliuk <lifedj27@gmail.com>, Artem Dubrovskiy <ad13box@gmail.com>
 License-Expression: Apache-2.0

{alita_sdk-0.3.345 → alita_sdk-0.3.347}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "alita_sdk"
-version = "0.3.345"
+version = "0.3.347"
 description = "SDK for building langchain agents using resources from Alita"
 readme = "README.md"
 requires-python = ">=3.10"

alita_sdk-0.3.345/alita_sdk/runtime/langchain/document_loaders/ImageParser.py

@@ -1,17 +0,0 @@
-from langchain_community.document_loaders.parsers.images import BaseImageBlobParser
-
-from alita_sdk.runtime.langchain.document_loaders.AlitaImageLoader import AlitaImageLoader
-
-class ImageParser(BaseImageBlobParser):
-
-    def __init__(self, **kwargs):
-        self.llm = kwargs.get('llm')
-        self.prompt = kwargs.get('prompt')
-
-    def _analyze_image(self, img) -> str:
-        from io import BytesIO
-
-        byte_stream = BytesIO()
-        img.save(byte_stream, format='PNG')
-        image_bytes = byte_stream.getvalue()
-        return AlitaImageLoader(file_content=image_bytes, file_name="image.png", prompt=self.prompt, llm=self.llm).get_content()