alita-sdk 0.3.209__py3-none-any.whl → 0.3.211__py3-none-any.whl

alita_sdk/tools/testrail/api_wrapper.py

@@ -4,10 +4,11 @@ from typing import Dict, List, Optional, Union, Any, Generator
 
 import pandas as pd
 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 ..elitea_base import BaseVectorStoreToolApiWrapper, BaseIndexParams
+from ..elitea_base import BaseVectorStoreToolApiWrapper, extend_with_vector_tools
 from langchain_core.documents import Document
 
 from ...runtime.utils.utils import IndexerKeywords
@@ -289,20 +290,6 @@ updateCase = create_model(
     ),
 )
 
-# Schema for indexing TestRail data into vector store
-indexData = create_model(
-    "indexData",
-    __base__=BaseIndexParams,
-    project_id=(str, Field(description="TestRail project ID to index data from")),
-    suite_id=(Optional[str], Field(default=None, description="Optional TestRail suite ID to filter test cases")),
-    section_id=(Optional[int], Field(default=None, description="Optional section ID to filter test cases")),
-    title_keyword=(Optional[str], Field(default=None, description="Optional keyword to filter test cases by title")),
-    progress_step=(Optional[int],
-                   Field(default=None, ge=0, le=100, description="Optional step size for progress reporting during indexing")),
-    clean_index=(Optional[bool],
-                       Field(default=False, description="Optional flag to enforce clean existing index before indexing new data")),
-)
-
 SUPPORTED_KEYS = {
     "id", "title", "section_id", "template_id", "type_id", "priority_id", "milestone_id",
     "refs", "created_by", "created_on", "updated_by", "updated_on", "estimate",
@@ -317,14 +304,6 @@ class TestrailAPIWrapper(BaseVectorStoreToolApiWrapper):
     password: Optional[SecretStr] = None,
     email: Optional[str] = None,
     _client: Optional[TestRailAPI] = PrivateAttr() # Private attribute for the TestRail client
-    llm: Any = None
-
-    connection_string: Optional[SecretStr] = None
-    collection_name: Optional[str] = None
-    embedding_model: Optional[str] = "HuggingFaceEmbeddings"
-    embedding_model_params: Optional[Dict[str, Any]] = {"model_name": "sentence-transformers/all-MiniLM-L6-v2"}
-    vectorstore_type: Optional[str] = "PGVector"
-
 
     @model_validator(mode="before")
     @classmethod
@@ -490,7 +469,8 @@ class TestrailAPIWrapper(BaseVectorStoreToolApiWrapper):
                 project_id=project_id, **params
             )
 
-            cases = extracted_cases.get("cases")
+            # support old versions of testrail_api
+            cases = extracted_cases.get("cases") if isinstance(extracted_cases, dict) else extracted_cases
 
             if cases is None:
                 return ToolException("No test cases found in the extracted data.")
@@ -554,7 +534,8 @@ class TestrailAPIWrapper(BaseVectorStoreToolApiWrapper):
     def _base_loader(self, project_id: str,
                      suite_id: Optional[str] = None,
                      section_id: Optional[int] = None,
-                     title_keyword: Optional[str] = None
+                     title_keyword: Optional[str] = None,
+                     **kwargs: Any
                      ) -> Generator[Document, None, None]:
         try:
             if suite_id:
@@ -577,7 +558,7 @@ class TestrailAPIWrapper(BaseVectorStoreToolApiWrapper):
                 'title': case.get('title', ''),
                 'suite_id': suite_id or case.get('suite_id', ''),
                 'id': str(case.get('id', '')),
-                'updated_on': case.get('updated_on') or -1,
+                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,
                 'priority': case.get('priority_id') or -1,
@@ -588,22 +569,6 @@ class TestrailAPIWrapper(BaseVectorStoreToolApiWrapper):
                 'entity_type': 'test_case',
             })
 
-    def index_data(
-            self,
-            project_id: str,
-            suite_id: Optional[str] = None,
-            collection_suffix: str = "",
-            section_id: Optional[int] = None,
-            title_keyword: Optional[str] = None,
-            progress_step: Optional[int] = None,
-            clean_index: Optional[bool] = False
-    ):
-        """Load TestRail test cases into the vector store."""
-        docs = self._base_loader(project_id, suite_id, section_id, title_keyword)
-        embedding = get_embeddings(self.embedding_model, self.embedding_model_params)
-        vs = self._init_vector_store(collection_suffix, embeddings=embedding)
-        return vs.index_documents(docs, progress_step=progress_step, clean_index=clean_index)
-
     def _process_document(self, document: Document) -> Generator[Document, None, None]:
         """
         Process an existing base document to extract relevant metadata for full document preparation.
@@ -626,16 +591,15 @@ class TestrailAPIWrapper(BaseVectorStoreToolApiWrapper):
 
             # process each attachment to extract its content
             for attachment in attachments:
-                attachment_id = attachment['id']
+                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),
                                                      metadata={
                                                          'project_id': base_data.get('project_id', ''),
-                                                         IndexerKeywords.PARENT.value: case_id,
-                                                         'id': attachment_id,
+                                                         'id': str(attachment_id),
+                                                         IndexerKeywords.PARENT.value: str(case_id),
                                                          'filename': attachment['filename'],
                                                          'filetype': attachment['filetype'],
                                                          'created_on': attachment['created_on'],
@@ -663,10 +627,20 @@ class TestrailAPIWrapper(BaseVectorStoreToolApiWrapper):
             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)
+            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
 
+    def _index_tool_params(self):
+        return {
+            'project_id': (str, Field(description="TestRail project ID to index data from")),
+            'suite_id': (Optional[str],
+                         Field(default=None, description="Optional TestRail suite ID to filter test cases")),
+            'section_id': (Optional[int], Field(default=None, description="Optional section ID to filter test cases")),
+        }
+
     def _to_markup(self, data: List[Dict], output_format: str) -> str:
         """
         Converts the given data into the specified format: 'json', 'csv', or 'markdown'.
@@ -694,6 +668,7 @@ class TestrailAPIWrapper(BaseVectorStoreToolApiWrapper):
         if output_format == "markdown":
             return df.to_markdown(index=False)
 
+    @extend_with_vector_tools
     def get_available_tools(self):
         tools = [
             {
@@ -731,14 +706,6 @@ class TestrailAPIWrapper(BaseVectorStoreToolApiWrapper):
                 "ref": self.update_case,
                 "description": self.update_case.__doc__,
                 "args_schema": updateCase,
-            },
-            {
-                "name": "index_data",
-                "ref": self.index_data,
-                "description": self.index_data.__doc__,
-                "args_schema": indexData,
             }
         ]
-        # Add vector search from base
-        tools.extend(self._get_vector_search_tools())
         return tools

alita_sdk/tools/utils/content_parser.py

@@ -61,7 +61,7 @@ IMAGE_EXTENSIONS = ['jpg', 'jpeg', 'png', 'gif', 'bmp', 'tiff', 'webp', 'svg']
 
 
 def parse_file_content(file_name=None, file_content=None, is_capture_image: bool = False, page_number: int = None,
-                       sheet_name: str = None, llm=None, file_path: str = None):
+                       sheet_name: str = None, llm=None, file_path: str = None, excel_by_sheets: bool = False):
     """Parse the content of a file based on its type and return the parsed content.
 
     Args:
@@ -91,7 +91,7 @@ def parse_file_content(file_name=None, file_content=None, is_capture_image: bool
     elif file_name.endswith('.docx'):
         return read_docx_from_bytes(file_content)
     elif file_name.endswith('.xlsx') or file_name.endswith('.xls'):
-        return parse_excel(file_content, sheet_name)
+        return parse_excel(file_content, sheet_name, excel_by_sheets)
     elif file_name.endswith('.pdf'):
         return parse_pdf(file_content, page_number, is_capture_image, llm)
     elif file_name.endswith('.pptx'):
@@ -109,17 +109,26 @@ def parse_txt(file_content):
     except Exception as e:
         return ToolException(f"Error decoding file content: {e}")
 
-def parse_excel(file_content, sheet_name = None):
+def parse_excel(file_content, sheet_name = None, return_by_sheets: bool = False):
     try:
         excel_file = io.BytesIO(file_content)
         if sheet_name:
             return parse_sheet(excel_file, sheet_name)
         dfs = pd.read_excel(excel_file, sheet_name=sheet_name)
-        result = []
-        for sheet_name, df in dfs.items():
-            df.fillna('', inplace=True)
-            result.append(f"=== Sheet: {sheet_name} ===\n{df.to_string(index=False)}")
-        return "\n\n".join(result)
+
+        if return_by_sheets:
+            result = {}
+            for sheet_name, df in dfs.items():
+                df.fillna('', inplace=True)
+                result[sheet_name] = df.to_dict(orient='records')
+            return result
+        else:
+            result = []
+            for sheet_name, df in dfs.items():
+                df.fillna('', inplace=True)
+                string_content = df.to_string(index=False)
+                result.append(f"====== Sheet name: {sheet_name} ======\n{string_content}")
+            return "\n\n".join(result)
     except Exception as e:
         return ToolException(f"Error reading Excel file: {e}")
 
@@ -194,6 +203,8 @@ def describe_image(image):
     return "\n[Picture: " + processor.decode(out[0], skip_special_tokens=True) + "]\n"
 
 def __perform_llm_prediction_for_image(llm, image: bytes, image_format='png', prompt=image_processing_prompt) -> str:
+    if not llm:
+        raise ToolException("LLM is not provided for image processing.")
     base64_string = bytes_to_base64(image)
     result = llm.invoke([
         HumanMessage(
@@ -207,6 +218,59 @@ def __perform_llm_prediction_for_image(llm, image: bytes, image_format='png', pr
     ])
     return f"\n[Image description: {result.content}]\n"
 
+# TODO: review usage of this function alongside with functions above
+def load_content(file_path: str, extension: str = None, loader_extra_config: dict = None, llm = None) -> str:
+    """
+    Loads the content of a file based on its extension using a configured loader.
+    """
+    try:
+        from ...runtime.langchain.document_loaders.constants import loaders_map
+
+        if not extension:
+            extension = file_path.split('.')[-1].lower()
+
+        loader_config = loaders_map.get(extension)
+        if not loader_config:
+            logger.warning(f"No loader found for file extension: {extension}. File: {file_path}")
+            return ""
+
+        loader_cls = loader_config['class']
+        loader_kwargs = loader_config['kwargs']
+
+        if loader_extra_config:
+            loader_kwargs.update(loader_extra_config)
+        if loader_config['is_multimodal_processing'] and llm:
+            loader_kwargs.update({'llm': llm})
+
+        loader = loader_cls(file_path, **loader_kwargs)
+        documents = loader.load()
+
+        page_contents = [doc.page_content for doc in documents]
+        return "\n".join(page_contents)
+    except Exception as e:
+        error_message = f"Error loading attachment: {str(e)}"
+        logger.warning(f"{error_message} for file {file_path}")
+        return ""
+
+def load_content_from_bytes(file_content: bytes, extension: str = None, loader_extra_config: dict = None, llm = None) -> str:
+    """Loads the content of a file from bytes based on its extension using a configured loader."""
+
+    import tempfile
+
+    # Automatic cleanup with context manager
+    with tempfile.NamedTemporaryFile(mode='w+b', delete=True) as temp_file:
+        # Write data to temp file
+        temp_file.write(file_content)
+        temp_file.flush()  # Ensure data is written
+
+        # Get the file path for operations
+        temp_path = temp_file.name
+
+        # Perform your operations
+        return load_content(temp_path, extension, loader_extra_config, llm)
+
+
+
 def file_to_bytes(filepath):
     """
     Reads a file and returns its content as a bytes object.

alita_sdk/tools/xray/__init__.py

@@ -20,7 +20,14 @@ def get_tools(tool):
         client_secret=tool['settings'].get('client_secret', None),
         limit=tool['settings'].get('limit', 20),
         verify_ssl=tool['settings'].get('verify_ssl', True),
-        toolkit_name=tool.get('toolkit_name')
+        toolkit_name=tool.get('toolkit_name'),
+
+        # indexer settings
+        connection_string=tool['settings'].get('connection_string', None),
+        collection_name=f"{tool.get('toolkit_name')}_{str(tool['id'])}",
+        embedding_model="HuggingFaceEmbeddings",
+        embedding_model_params={"model_name": "sentence-transformers/all-MiniLM-L6-v2"},
+        vectorstore_type="PGVector"
     ).get_tools()