alita-sdk 0.3.274__tar.gz → 0.3.276__tar.gz

{alita_sdk-0.3.274 → alita_sdk-0.3.276}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: alita_sdk
-Version: 0.3.274
+Version: 0.3.276
 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 <lifedjik@gmail.com>, Artem Dubrovskiy <ad13box@gmail.com>
 License-Expression: Apache-2.0

{alita_sdk-0.3.274 → alita_sdk-0.3.276}/alita_sdk/tools/ado/test_plan/test_plan_wrapper.py

@@ -1,7 +1,7 @@
 import json
 import logging
 import xml.etree.ElementTree as ET
-from typing import Generator, Literal, Optional
+from typing import Generator, Literal, Optional, List
 
 from azure.devops.connection import Connection
 from azure.devops.v7_0.test_plan.models import TestPlanCreateParams, TestSuiteCreateParams, \
@@ -359,6 +359,7 @@ class TestPlanApiWrapper(NonCodeIndexerToolkit):
             test_cases = self._client.get_test_case_list(self.project, plan_id, suite_id)
             return [test_case.as_dict() for test_case in test_cases]
         except Exception as e:
+            self._log_tool_event(f"Error getting test cases: {e}", 'get_test_cases')
             logger.error(f"Error getting test cases: {e}")
             return ToolException(f"Error getting test cases: {e}")
 
@@ -398,7 +399,7 @@ class TestPlanApiWrapper(NonCodeIndexerToolkit):
         """Return the parameters for indexing data."""
         return {
             "plan_id": (str, Field(description="ID of the test plan for which test cases are requested")),
-            "suite_ids": (str, Field(description="List of test suite IDs for which test cases are requested (can be empty)")),
+            "suite_ids": (Optional[List[str]], Field(description='List of test suite IDs for which test cases are requested (can be empty). Example: ["2", "23"]', default=[])),
             'chunking_tool':(Literal['html'], Field(description="Name of chunking tool", default='html'))
         }
 

{alita_sdk-0.3.274 → alita_sdk-0.3.276}/alita_sdk/tools/base_indexer_toolkit.py

@@ -126,6 +126,11 @@ class BaseIndexerToolkit(VectorStoreWrapperBase):
         """
         return {}
 
+    def _remove_metadata_keys(self) -> List[str]:
+        """ Returns a list of metadata keys to be removed from documents before indexing.
+        Override this method in subclasses to provide specific keys to remove."""
+        return [IndexerKeywords.CONTENT_IN_BYTES.value, IndexerKeywords.CONTENT_FILE_NAME.value]
+
     def _base_loader(self, **kwargs) -> Generator[Document, None, None]:
         """ Loads documents from a source, processes them,
         and returns a list of Document objects with base metadata: id and created_on."""
@@ -158,11 +163,14 @@ class BaseIndexerToolkit(VectorStoreWrapperBase):
         documents = self._extend_data(documents) # update content of not-reduced base document if needed (for sharepoint and similar)
         documents = self._collect_dependencies(documents) # collect dependencies for base documents
         documents = self._apply_loaders_chunkers(documents, chunking_tool, chunking_config)
+        list_documents = list(documents)
+        self._clean_metadata(list_documents)
+
         #
-        return self._save_index(list(documents), collection_suffix=collection_suffix, progress_step=progress_step)
+        return self._save_index(list_documents, collection_suffix=collection_suffix, progress_step=progress_step)
     
     def _apply_loaders_chunkers(self, documents: Generator[Document, None, None], chunking_tool: str=None, chunking_config=None) -> Generator[Document, None, None]:
-        from alita_sdk.tools.chunkers import __all__ as chunkers
+        from ..tools.chunkers import __all__ as chunkers
 
         if chunking_config is None:
             chunking_config = {}
@@ -203,6 +211,12 @@ class BaseIndexerToolkit(VectorStoreWrapperBase):
                 dep.metadata[IndexerKeywords.PARENT.value] = document.metadata.get('id', None)
                 yield dep
 
+    def _clean_metadata(self, documents: list[Document]):
+        for document in documents:
+            remove_keys = self._remove_metadata_keys()
+            for key in remove_keys:
+                document.metadata.pop(key, None)
+
     def _reduce_duplicates(
             self,
             documents: Generator[Any, None, None],

{alita_sdk-0.3.274 → alita_sdk-0.3.276}/alita_sdk/tools/bitbucket/api_wrapper.py

@@ -11,7 +11,7 @@ from .bitbucket_constants import create_pr_data
 from .cloud_api_wrapper import BitbucketCloudApi, BitbucketServerApi
 from pydantic.fields import PrivateAttr
 
-from ..elitea_base import BaseCodeToolApiWrapper, extend_with_vector_tools
+from ..elitea_base import BaseCodeToolApiWrapper
 
 logger = logging.getLogger(__name__)
 
@@ -399,7 +399,6 @@ class BitbucketAPIWrapper(BaseCodeToolApiWrapper):
         except Exception as e:
             return f"Failed to read file {file_path}: {str(e)}"
 
-    @extend_with_vector_tools
     def get_available_tools(self):
         return [
             {
@@ -473,5 +472,5 @@ class BitbucketAPIWrapper(BaseCodeToolApiWrapper):
                 "ref": self.add_pull_request_comment,
                 "description": self.add_pull_request_comment.__doc__ or "Add a comment to a pull request in the repository.",
                 "args_schema": AddPullRequestCommentModel,
-            },
-        ]
+            }
+        ] + self._get_vector_search_tools()

{alita_sdk-0.3.274 → alita_sdk-0.3.276}/alita_sdk/tools/figma/api_wrapper.py

@@ -1,4 +1,3 @@
-import base64
 import functools
 import json
 import logging
@@ -12,7 +11,8 @@ from langchain_core.documents import Document
 from langchain_core.tools import ToolException
 from pydantic import Field, PrivateAttr, create_model, model_validator, SecretStr
 
-from ..elitea_base import BaseVectorStoreToolApiWrapper, extend_with_vector_tools
+from ..non_code_indexer_toolkit import NonCodeIndexerToolkit
+from ..utils.available_tools_decorator import extend_with_parent_available_tools
 from ..utils.content_parser import load_content_from_bytes
 
 GLOBAL_LIMIT = 10000
@@ -229,46 +229,87 @@ class ArgsSchema(Enum):
     )
 
 
-class FigmaApiWrapper(BaseVectorStoreToolApiWrapper):
+class FigmaApiWrapper(NonCodeIndexerToolkit):
     token: Optional[SecretStr] = Field(default=None)
     oauth2: Optional[SecretStr] = Field(default=None)
     global_limit: Optional[int] = Field(default=GLOBAL_LIMIT)
     global_regexp: Optional[str] = Field(default=None)
     _client: Optional[FigmaPy] = PrivateAttr()
 
-    def _base_loader(self, project_id: Optional[str] = None, file_keys: Optional[List[str]] = None, **kwargs) -> Generator[Document, None, None]:
-        files = []
-        if project_id:
-            files = json.loads(self.get_project_files(project_id)).get('files', [])
-            for file in files:
-                yield Document(page_content=json.dumps(file), metadata={
-                    'id': file.get('key', ''),
-                    'file_key': file.get('key', ''),
-                    'name': file.get('name', ''),
-                    'updated_on': file.get('last_modified', '')
-                })
-        elif file_keys:
-            for file_key in file_keys:
+    def _base_loader(
+            self,
+            project_id: Optional[str] = None,
+            file_keys_include: Optional[List[str]] = None,
+            file_keys_exclude: Optional[List[str]] = None,
+            node_ids_include: Optional[List[str]] = None,
+            node_ids_exclude: Optional[List[str]] = None,
+            node_types_include: Optional[List[str]] = None,
+            node_types_exclude: Optional[List[str]] = None,
+            **kwargs
+    ) -> Generator[Document, None, None]:
+        # If both include and exclude are provided, use only include
+        if file_keys_include:
+            for file_key in file_keys_include:
                 file = self._client.get_file(file_key)
+                if not file:
+                    raise ToolException(f"Unexpected error while retrieving file {file_key}. Probably file is under editing. Try again later.")
                 metadata = {
                     'id': file_key,
                     'file_key': file_key,
                     'name': file.name,
-                    'updated_on': file.last_modified
+                    'updated_on': file.last_modified,
+                    'figma_pages_include': node_ids_include or [],
+                    'figma_pages_exclude': node_ids_exclude or [],
+                    'figma_nodes_include': node_types_include or [],
+                    'figma_nodes_exclude': node_types_exclude or []
                 }
                 yield Document(page_content=json.dumps(metadata), metadata=metadata)
+        elif project_id:
+            files = json.loads(self.get_project_files(project_id)).get('files', [])
+            for file in files:
+                if file_keys_exclude and file.get('key', '') in file_keys_exclude:
+                    continue
+                yield Document(page_content=json.dumps(file), metadata={
+                    'id': file.get('key', ''),
+                    'file_key': file.get('key', ''),
+                    'name': file.get('name', ''),
+                    'updated_on': file.get('last_modified', ''),
+                    'figma_pages_include': node_ids_include or [],
+                    'figma_pages_exclude': node_ids_exclude or [],
+                    'figma_nodes_include': node_types_include or [],
+                    'figma_nodes_exclude': node_types_exclude or []
+                })
+        elif file_keys_exclude or node_ids_exclude:
+            raise ValueError("Excludes without parent (project_id or file_keys_include) do not make sense.")
+        else:
+            raise ValueError("You must provide at least project_id or file_keys_include.")
 
     def _process_document(self, document: Document) -> Generator[Document, None, None]:
         file_key = document.metadata.get('id', '')
         #
-        node_ids = []
-        children = self._client.get_file(file_key).document.get('children', [])
-        if children:
-            nodes = children[0].get('children', [])
-            node_ids = [node['id'] for node in nodes if 'id' in node]
-        images = self._client.get_file_images(file_key, node_ids).images or {}
+        figma_pages = self._client.get_file(file_key).document.get('children', [])
+        node_ids_include = document.metadata.pop('figma_pages_include', [])
+        node_ids_exclude = document.metadata.pop('figma_pages_exclude', [])
+        node_types_include = [t.lower() for t in document.metadata.pop('figma_nodes_include', [])]
+        node_types_exclude = [t.lower() for t in document.metadata.pop('figma_nodes_exclude', [])]
+        if node_ids_include:
+            figma_pages = [node for node in figma_pages if ('id' in node and node['id'].replace(':', '-') in node_ids_include)]
+        elif node_ids_exclude:
+            figma_pages = [node for node in figma_pages if ('id' in node and node['id'].replace(':', '-') not in node_ids_exclude)]
+        node_ids = [
+            child['id']
+            for page in figma_pages
+            if 'children' in page
+            for child in page['children']
+            if 'id' in child
+            and (
+                (node_types_include and child.get('type').lower() in node_types_include)
+                or (node_types_exclude and child.get('type').lower() not in node_types_exclude)
+                or (not node_types_include and not node_types_exclude)
+            )
+        ]
 
-        # iterate over images values
+        images = self._client.get_file_images(file_key, node_ids).images or {}
         for node_id, image_url in images.items():
             if not image_url:
                 logging.warning(f"Image URL not found for node_id {node_id} in file {file_key}. Skipping.")
@@ -284,20 +325,40 @@ class FigmaApiWrapper(BaseVectorStoreToolApiWrapper):
                     yield Document(
                                     page_content=page_content,
                                     metadata={
+                                        'id': file_key,
+                                        'updated_on': document.metadata.get('updated_on', ''),
                                         'file_key': file_key,
                                         'node_id': node_id,
                                         'image_url': image_url
                                     }
                                 )
 
+    def _remove_metadata_keys(self):
+        return super()._remove_metadata_keys() + ['figma_pages_include', 'figma_pages_exclude', 'figma_nodes_include', 'figma_nodes_exclude']
+
     def _index_tool_params(self):
         """Return the parameters for indexing data."""
         return {
             "project_id": (Optional[str], Field(
-                description="ID of the project to list files from: i.e. '55391681'",
+                description="ID of the project to list files from: i.e. 55391681'",
+                default=None)),
+            'file_keys_include': (Optional[List[str]], Field(
+                description="List of file keys to include in index if project_id is not provided: i.e. ['Fp24FuzPwH0L74ODSrCnQo', 'jmhAr6q78dJoMRqt48zisY']",
                 default=None)),
-            'file_keys': (Optional[List[str]], Field(
-                description='List of file keys to index: i.e. ["Fp24FuzPwH0L74ODSrCnQo", "jmhAr6q78dJoMRqt48zisY"]',
+            'file_keys_exclude': (Optional[List[str]], Field(
+                description="List of file keys to exclude from index. It is applied only if project_id is provided and file_keys_include is not provided: i.e. ['Fp24FuzPwH0L74ODSrCnQo', 'jmhAr6q78dJoMRqt48zisY']",
+                default=None)),
+            'node_ids_include': (Optional[List[str]], Field(
+                description="List of top-level nodes (pages) in file to include in index. It is node-id from figma url: i.e. ['123-56', '7651-9230'].",
+                default=None)),
+            'node_ids_exclude': (Optional[List[str]], Field(
+                description="List of top-level nodes (pages) in file to exclude from index. It is applied only if node_ids_include is not provided. It is node-id from figma url: i.e. ['Fp24FuzPwH0L74ODSrCnQo', 'jmhAr6q78dJoMRqt48zisY']",
+                default=None)),
+            'node_types_include': (Optional[List[str]], Field(
+                description="List type of nodes to include in index: i.e. ['FRAME', 'COMPONENT', 'RECTANGLE', 'COMPONENT_SET', 'INSTANCE', 'VECTOR', ...].",
+                default=None)),
+            'node_types_exclude': (Optional[List[str]], Field(
+                description="List type of nodes to exclude from index. It is applied only if node_types_include is not provided: i.e. ['FRAME', 'COMPONENT', 'RECTANGLE', 'COMPONENT_SET', 'INSTANCE', 'VECTOR', ...]",
                 default=None))
         }
 
@@ -328,6 +389,11 @@ class FigmaApiWrapper(BaseVectorStoreToolApiWrapper):
             logging.error(msg)
             raise ToolException(msg)
 
+    @model_validator(mode='before')
+    @classmethod
+    def check_before(cls, values):
+        return super().validate_toolkit(values)
+
     @model_validator(mode="after")
     @classmethod
     def validate_toolkit(cls, values):
@@ -510,7 +576,7 @@ class FigmaApiWrapper(BaseVectorStoreToolApiWrapper):
         """Retrieves all files for a specified project ID from Figma."""
         return self._client.get_project_files(project_id)
 
-    @extend_with_vector_tools
+    @extend_with_parent_available_tools
     def get_available_tools(self):
         return [
             {

{alita_sdk-0.3.274 → alita_sdk-0.3.276}/alita_sdk.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: alita_sdk
-Version: 0.3.274
+Version: 0.3.276
 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 <lifedjik@gmail.com>, Artem Dubrovskiy <ad13box@gmail.com>
 License-Expression: Apache-2.0

{alita_sdk-0.3.274 → alita_sdk-0.3.276}/pyproject.toml

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