alita-sdk 0.3.204__py3-none-any.whl → 0.3.206__py3-none-any.whl

alita_sdk/tools/google/bigquery/schemas.py

@@ -0,0 +1,102 @@
+from enum import Enum
+from typing import Any, Dict, List, Optional, Union
+
+from pydantic import Field, create_model
+
+
+class ArgsSchema(Enum):
+    NoInput = create_model("NoInput")
+    GetDocuments = create_model(
+        "GetDocuments",
+        ids=(
+            Optional[List[str]],
+            Field(default=None, description="List of document IDs to retrieve."),
+        ),
+        filter=(
+            Optional[Union[Dict[str, Any], str]],
+            Field(default=None, description="Filter as dict or SQL WHERE clause."),
+        ),
+    )
+    SimilaritySearch = create_model(
+        "SimilaritySearch",
+        query=(str, Field(description="Text query to search for similar documents.")),
+        k=(int, Field(default=5, description="Number of top results to return.")),
+        filter=(
+            Optional[Union[Dict[str, Any], str]],
+            Field(default=None, description="Filter as dict or SQL WHERE clause."),
+        ),
+    )
+    BatchSearch = create_model(
+        "BatchSearch",
+        queries=(
+            Optional[List[str]],
+            Field(default=None, description="List of text queries."),
+        ),
+        embeddings=(
+            Optional[List[List[float]]],
+            Field(default=None, description="List of embedding vectors."),
+        ),
+        k=(int, Field(default=5, description="Number of top results to return.")),
+        filter=(
+            Optional[Union[Dict[str, Any], str]],
+            Field(default=None, description="Filter as dict or SQL WHERE clause."),
+        ),
+    )
+    JobStatsArgs = create_model(
+        "JobStatsArgs", job_id=(str, Field(description="BigQuery job ID."))
+    )
+    SimilaritySearchByVectorArgs = create_model(
+        "SimilaritySearchByVectorArgs",
+        embedding=(List[float], Field(description="Embedding vector.")),
+        k=(int, Field(default=5, description="Number of top results to return.")),
+    )
+    SimilaritySearchByVectorWithScoreArgs = create_model(
+        "SimilaritySearchByVectorWithScoreArgs",
+        embedding=(List[float], Field(description="Embedding vector.")),
+        filter=(
+            Optional[Union[Dict[str, Any], str]],
+            Field(default=None, description="Filter as dict or SQL WHERE clause."),
+        ),
+        k=(int, Field(default=5, description="Number of top results to return.")),
+    )
+    SimilaritySearchWithScoreArgs = create_model(
+        "SimilaritySearchWithScoreArgs",
+        query=(str, Field(description="Text query.")),
+        filter=(
+            Optional[Union[Dict[str, Any], str]],
+            Field(default=None, description="Filter as dict or SQL WHERE clause."),
+        ),
+        k=(int, Field(default=5, description="Number of top results to return.")),
+    )
+    SimilaritySearchByVectorsArgs = create_model(
+        "SimilaritySearchByVectorsArgs",
+        embeddings=(List[List[float]], Field(description="List of embedding vectors.")),
+        filter=(
+            Optional[Union[Dict[str, Any], str]],
+            Field(default=None, description="Filter as dict or SQL WHERE clause."),
+        ),
+        k=(int, Field(default=5, description="Number of top results to return.")),
+        with_scores=(bool, Field(default=False)),
+        with_embeddings=(bool, Field(default=False)),
+    )
+    ExecuteArgs = create_model(
+        "ExecuteArgs",
+        method=(str, Field(description="Name of the BigQuery client method to call.")),
+        args=(
+            Optional[List[Any]],
+            Field(default=None, description="Positional arguments for the method."),
+        ),
+        kwargs=(
+            Optional[Dict[str, Any]],
+            Field(default=None, description="Keyword arguments for the method."),
+        ),
+    )
+    CreateDeltaLakeTable = create_model(
+        "CreateDeltaLakeTable",
+        table_name=(str, Field(description="Name of the Delta Lake table to create in BigQuery.")),
+        dataset=(Optional[str], Field(default=None, description="BigQuery dataset to contain the table (defaults to self.dataset).")),
+        connection_id=(str, Field(description="Fully qualified connection ID (project.region.connection_id).")),
+        source_uris=(list, Field(description="List of GCS URIs (prefixes) for the Delta Lake table.")),
+        autodetect=(bool, Field(default=True, description="Whether to autodetect schema (default: True).")),
+        project=(Optional[str], Field(default=None, description="GCP project ID (defaults to self.project).")),
+    )

alita_sdk/tools/google/bigquery/tool.py

@@ -0,0 +1,34 @@
+from typing import Optional, Type
+
+from langchain_core.callbacks import CallbackManagerForToolRun
+from pydantic import BaseModel, field_validator, Field
+from langchain_core.tools import BaseTool
+from traceback import format_exc
+from .api_wrapper import BigQueryApiWrapper
+
+
+class BigQueryAction(BaseTool):
+    """Tool for interacting with the BigQuery API."""
+
+    api_wrapper: BigQueryApiWrapper = Field(default_factory=BigQueryApiWrapper)
+    name: str
+    mode: str = ""
+    description: str = ""
+    args_schema: Optional[Type[BaseModel]] = None
+
+    @field_validator('name', mode='before')
+    @classmethod
+    def remove_spaces(cls, v):
+        return v.replace(' ', '')
+
+    def _run(
+        self,
+        *args,
+        run_manager: Optional[CallbackManagerForToolRun] = None,
+        **kwargs,
+    ) -> str:
+        """Use the GitHub API to run an operation."""
+        try:
+            return self.api_wrapper.run(self.mode, *args, **kwargs)
+        except Exception as e:
+            return f"Error: {format_exc()}"

alita_sdk/tools/llm/llm_utils.py

@@ -10,12 +10,6 @@ def get_model(model_type: str, model_params: dict):
         return None
     if model_type in llms:
         return get_llm(model_type)(**model_params)
-    elif model_type == "Alita":
-        try:
-            from alita_sdk.llms.alita import AlitaChatModel
-        except ImportError:
-            raise RuntimeError("Alita model not found")
-        return AlitaChatModel(**model_params)
     elif model_type in chat_models:
         model = getattr(__import__("langchain_community.chat_models", fromlist=[model_type]), model_type)
         return model(**model_params)

alita_sdk/tools/openapi/__init__.py

@@ -1,12 +1,15 @@
 import json
 import re
+import logging
 from typing import List, Any, Optional, Dict
-from langchain_core.tools import BaseTool, BaseToolkit
+from langchain_core.tools import BaseTool, BaseToolkit, ToolException
 from requests_openapi import Operation, Client, Server
 
 from pydantic import create_model, Field
 from functools import partial
 
+logger = logging.getLogger(__name__)
+
 name = "openapi"
 
 def get_tools(tool):
@@ -105,11 +108,19 @@ class AlitaOpenAPIToolkit(BaseToolkit):
             c.requestor.headers.update(headers)
         tools = []
         for i in tools_set:
+
             try:
+                if not i:
+                    raise ToolException("Operation id is missing for some of declared operations.")
                 tool = c.operations[i]
+                if not isinstance(tool, Operation):
+                    raise ToolException(f"Operation {i} is not an instance of Operation class.")
                 tools.append(create_api_tool(i, tool))
-            except KeyError:
-                ...
+            except ToolException:
+                raise
+            except Exception as e:
+                logger.warning(f"Tool {i} not found in OpenAPI spec.")
+                raise ToolException(f"Cannot create API tool ({i}): \n{e}.")
         return cls(request_session=c, tools=tools)
 
     def get_tools(self):

alita_sdk/tools/sharepoint/__init__.py

@@ -14,7 +14,8 @@ def get_tools(tool):
         site_url=tool['settings'].get('site_url', None),
         client_id=tool['settings'].get('client_id', None),
         client_secret=tool['settings'].get('client_secret', None),
-        toolkit_name=tool.get('toolkit_name'))
+        toolkit_name=tool.get('toolkit_name'),
+        llm=tool['settings'].get('llm'))
             .get_tools())
 
 

alita_sdk/tools/sharepoint/api_wrapper.py

@@ -1,5 +1,6 @@
+import json
 import logging
-from typing import Optional
+from typing import Optional, List, Dict, Any
 
 from ..utils.content_parser import parse_file_content
 from langchain_core.tools import ToolException
@@ -7,7 +8,9 @@ from office365.runtime.auth.client_credential import ClientCredential
 from office365.sharepoint.client_context import ClientContext
 from pydantic import Field, PrivateAttr, create_model, model_validator, SecretStr
 
-from ..elitea_base import BaseToolApiWrapper
+from ..elitea_base import BaseToolApiWrapper, BaseIndexParams, BaseVectorStoreToolApiWrapper
+from ...runtime.langchain.interfaces.llm_processor import get_embeddings
+from langchain_core.documents import Document
 
 NoInput = create_model(
     "NoInput"
@@ -29,17 +32,36 @@ ReadDocument = create_model(
     "ReadDocument",
     path=(str, Field(description="Contains the server-relative path of a document for reading.")),
     is_capture_image=(Optional[bool], Field(description="Determines is pictures in the document should be recognized.", default=False)),
-    page_number=(Optional[int], Field(description="Specifies which page to read. If it is None, then full document will be read.", default=None))
+    page_number=(Optional[int], Field(description="Specifies which page to read. If it is None, then full document will be read.", default=None)),
+    sheet_name=(Optional[str], Field(
+                        description="Specifies which sheet to read. If it is None, then full document will be read.",
+                        default=None))
+)
+
+indexData = create_model(
+    "indexData",
+    __base__=BaseIndexParams,
+    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")),
 )
 
 
-class SharepointApiWrapper(BaseToolApiWrapper):
+class SharepointApiWrapper(BaseVectorStoreToolApiWrapper):
     site_url: str
     client_id: str = None
     client_secret: SecretStr = None
     token: SecretStr = None
     _client: Optional[ClientContext] = PrivateAttr()  # Private attribute for the office365 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
     def validate_toolkit(cls, values):
@@ -111,7 +133,8 @@ class SharepointApiWrapper(BaseToolApiWrapper):
                     'Path': file.properties['ServerRelativeUrl'],
                     'Created': file.properties['TimeCreated'],
                     'Modified': file.properties['TimeLastModified'],
-                    'Link': file.properties['LinkingUrl']
+                    'Link': file.properties['LinkingUrl'],
+                    'id': file.properties['UniqueId']
                 }
                 result.append(temp_props)
             return result if result else ToolException("Can not get files or folder is empty. Please, double check folder name and read permissions.")
@@ -119,7 +142,7 @@ class SharepointApiWrapper(BaseToolApiWrapper):
             logging.error(f"Failed to load files from sharepoint: {e}")
             return ToolException("Can not get files. Please, double check folder name and read permissions.")
 
-    def read_file(self, path, is_capture_image: bool = False, page_number: int = None):
+    def read_file(self, path, is_capture_image: bool = False, page_number: int = None, sheet_name: str=None):
         """ Reads file located at the specified server-relative path. """
         try:
             file = self._client.web.get_file_by_server_relative_path(path)
@@ -130,7 +153,42 @@ class SharepointApiWrapper(BaseToolApiWrapper):
         except Exception as e:
             logging.error(f"Failed to load file from SharePoint: {e}. Path: {path}. Please, double check file name and path.")
             return ToolException("File not found. Please, check file name and path.")
-        return parse_file_content(file.name, file_content, is_capture_image, page_number)
+        return parse_file_content(file_name=file.name,
+                                  file_content=file_content,
+                                  is_capture_image=is_capture_image,
+                                  page_number=page_number,
+                                  sheet_name=sheet_name,
+                                  llm=self.llm)
+
+    def _base_loader(self) -> List[Document]:
+        try:
+            all_files = self.get_files_list()
+        except Exception as e:
+            raise ToolException(f"Unable to extract files: {e}")
+
+        docs: List[Document] = []
+        for file in all_files:
+            metadata = {
+                ("updated_at" if k == "Modified" else k): str(v)
+                for k, v in file.items()
+            }
+            docs.append(Document(page_content="", metadata=metadata))
+        return docs
+
+    def index_data(self,
+                   collection_suffix: str = '',
+                   progress_step: int = None,
+                   clean_index: bool = False):
+        docs = self._base_loader()
+        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) -> Document:
+        page_content = self.read_file(document.metadata['Path'], is_capture_image=True)
+
+        document.page_content = json.dumps(str(page_content))
+        return document
 
     def get_available_tools(self):
         return [
@@ -151,5 +209,11 @@ class SharepointApiWrapper(BaseToolApiWrapper):
                 "description": self.read_file.__doc__,
                 "args_schema": ReadDocument,
                 "ref": self.read_file
+            },
+            {
+                "name": "index_data",
+                "ref": self.index_data,
+                "description": self.index_data.__doc__,
+                "args_schema": indexData,
             }
         ]

alita_sdk/tools/testrail/__init__.py

@@ -16,7 +16,15 @@ def get_tools(tool):
         url=tool['settings']['url'],
         password=tool['settings'].get('password', None),
         email=tool['settings'].get('email', None),
-        toolkit_name=tool.get('toolkit_name')
+        toolkit_name=tool.get('toolkit_name'),
+        llm=tool['settings'].get('llm', None),
+
+        # 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()
 
 

alita_sdk/tools/testrail/api_wrapper.py

@@ -1,14 +1,21 @@
 import json
 import logging
-from typing import Dict, List, Optional, Union
+from typing import Dict, List, Optional, Union, Any, Generator
 
 import pandas as pd
 from langchain_core.tools import ToolException
 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 langchain_core.documents import Document
 
-from ..elitea_base import BaseToolApiWrapper
+from ...runtime.utils.utils import IndexerKeywords
+
+try:
+    from alita_sdk.runtime.langchain.interfaces.llm_processor import get_embeddings
+except ImportError:
+    from alita_sdk.langchain.interfaces.llm_processor import get_embeddings
 
 logger = logging.getLogger(__name__)
 
@@ -281,6 +288,19 @@ 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",
@@ -291,11 +311,19 @@ SUPPORTED_KEYS = {
 }
 
 
-class TestrailAPIWrapper(BaseToolApiWrapper):
+class TestrailAPIWrapper(BaseVectorStoreToolApiWrapper):
     url: str
     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
@@ -492,7 +520,7 @@ class TestrailAPIWrapper(BaseToolApiWrapper):
         you can submit and update specific fields only).
 
         :param case_id: T
-            he ID of the test case
+            He ID of the test case
         :param kwargs:
             :key title: str
                 The title of the test case
@@ -522,6 +550,118 @@ class TestrailAPIWrapper(BaseToolApiWrapper):
             f"Test case #{case_id} has been updated at '{updated_case['updated_on']}')"
         )
 
+    def _base_loader(self, project_id: str,
+                     suite_id: Optional[str] = None,
+                     section_id: Optional[int] = None,
+                     title_keyword: Optional[str] = None
+                     ) -> Generator[Document, None, None]:
+        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', [])
+        except StatusCodeError as e:
+            raise ToolException(f"Unable to extract test cases: {e}")
+            # 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={
+                'project_id': project_id,
+                '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,
+                'labels': [lbl['title'] for lbl in case.get('labels', [])],
+                'type': case.get('type_id') or -1,
+                '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',
+            })
+
+    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.
+        Used for late processing of documents after we ensure that the document has to be indexed to avoid
+        time-consuming operations for documents which might be useless.
+
+        Args:
+            document (Document): The base document to process.
+
+        Returns:
+            Generator[Document, None, None]: A generator yielding processed Document objects with metadata.
+        """
+        try:
+            # get base data from the document required to extract attachments and other metadata
+            base_data = json.loads(document.page_content)
+            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)
+
+            # process each attachment to extract its content
+            for attachment in attachments:
+                attachment_id = 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,
+                                                         'filename': attachment['filename'],
+                                                         'filetype': attachment['filetype'],
+                                                         'created_on': attachment['created_on'],
+                                                         'entity_type': 'test_case_attachment',
+                                                         'is_image': attachment['is_image'],
+                                                     })
+        except json.JSONDecodeError as e:
+            raise ToolException(f"Failed to decode JSON from document: {e}")
+
+    def _process_attachment(self, attachment: Dict[str, Any]) -> str:
+        """
+        Processes an attachment to extract its content.
+
+        Args:
+            attachment (Dict[str, Any]): The attachment data.
+
+        Returns:
+            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']}")
+        # TODO: add support for other file types
+        # use utility to handle different types (tools/utils)
+        return page_content
+
     def _to_markup(self, data: List[Dict], output_format: str) -> str:
         """
         Converts the given data into the specified format: 'json', 'csv', or 'markdown'.
@@ -550,7 +690,7 @@ class TestrailAPIWrapper(BaseToolApiWrapper):
             return df.to_markdown(index=False)
 
     def get_available_tools(self):
-        return [
+        tools = [
             {
                 "name": "get_case",
                 "ref": self.get_case,
@@ -587,4 +727,13 @@ class TestrailAPIWrapper(BaseToolApiWrapper):
                 "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