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

alita_sdk/runtime/langchain/document_loaders/utils.py

@@ -2,7 +2,8 @@ import re
 import string
 from gensim.parsing import remove_stopwords
 
-from ..tools.log import print_log
+from ..tools.utils import bytes_to_base64
+from langchain_core.messages import HumanMessage
 
 
 def cleanse_data(document: str) -> str:
@@ -32,3 +33,31 @@ def cleanse_data(document: str) -> str:
     #     document = document.replace(kw, "")
 
     return document
+
+def perform_llm_prediction_for_image_bytes(image_bytes: bytes, llm, prompt: str) -> str:
+    """Performs LLM prediction for image content."""
+    base64_string = bytes_to_base64(image_bytes)
+    result = llm.invoke([
+        HumanMessage(
+            content=[
+                {"type": "text", "text": prompt},
+                {
+                    "type": "image_url",
+                    "image_url": {"url": f"data:image/png;base64,{base64_string}"},
+                },
+            ]
+        )
+    ])
+    return result.content
+
+def create_temp_file(file_content: bytes):
+    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
+        return temp_file.name

alita_sdk/runtime/toolkits/artifact.py

@@ -26,6 +26,11 @@ class ArtifactToolkit(BaseToolkit):
             connection_string = (Optional[SecretStr], Field(description="Connection string for vectorstore",
                                                             default=None,
                                                             json_schema_extra={'secret': True})),
+            # embedding model settings
+            embedding_model=(str, Field(default="HuggingFaceEmbeddings", description="Embedding model to use")),
+            embedding_model_params=(dict, Field(default={"model_name": "sentence-transformers/all-MiniLM-L6-v2"},
+                                                description="Parameters for embedding model")),
+
             __config__=ConfigDict(json_schema_extra={'metadata': {"label": "Artifact", "icon_url": None}})
         )
     

alita_sdk/runtime/tools/artifact.py

@@ -61,21 +61,19 @@ class ArtifactWrapper(BaseVectorStoreToolApiWrapper):
     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 _base_loader(self, **kwargs) -> List[Document]:
+    def _base_loader(self, **kwargs) -> Generator[Document, None, None]:
         try:
             all_files = self.list_files(self.bucket, False)
         except Exception as e:
             raise ToolException(f"Unable to extract files: {e}")
 
-        docs: List[Document] = []
         for file in all_files['rows']:
             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'])
-            docs.append(Document(page_content="", metadata=metadata))
-        return docs
+            yield Document(page_content="", metadata=metadata)
 
     def get_hash_from_bucket_and_file_name(self, bucket, file_name):
         hasher = hashlib.sha256()

alita_sdk/runtime/tools/vectorstore.py

@@ -197,7 +197,8 @@ class VectorStoreWrapper(BaseToolApiWrapper):
             tool_name="_clean_collection"
         )
         data = self.vectoradapter.vectorstore.get(include=['metadatas'])
-        self.vectoradapter.vectorstore.delete(ids=data['ids'])
+        if data['ids']:
+            self.vectoradapter.vectorstore.delete(ids=data['ids'])
         self._log_data(
             f"Collection '{self.dataset}' has been cleaned. ",
             tool_name="_clean_collection"

alita_sdk/tools/ado/test_plan/test_plan_wrapper.py

@@ -1,22 +1,21 @@
 import json
 import logging
-from typing import Any, Dict, Generator, List, Optional
-
-from langchain_core.documents import Document
+import xml.etree.ElementTree as ET
+from typing import Generator, Optional
 
-from alita_sdk.tools.elitea_base import BaseIndexParams
 from azure.devops.connection import Connection
 from azure.devops.v7_0.test_plan.models import TestPlanCreateParams, TestSuiteCreateParams, \
     SuiteTestCaseCreateUpdateParameters
 from azure.devops.v7_0.test_plan.test_plan_client import TestPlanClient
+from langchain_core.documents import Document
 from langchain_core.tools import ToolException
 from msrest.authentication import BasicAuthentication
 from pydantic import create_model, PrivateAttr, model_validator, SecretStr
 from pydantic.fields import FieldInfo as Field
-import xml.etree.ElementTree as ET
 
 from ..work_item import AzureDevOpsApiWrapper
 from ...elitea_base import BaseVectorStoreToolApiWrapper, extend_with_vector_tools
+
 try:
     from alita_sdk.runtime.langchain.interfaces.llm_processor import get_embeddings
 except ImportError:
@@ -164,18 +163,6 @@ TestCasesGetModel = create_model(
     suite_id=(int, Field(description="ID of the test suite for which test cases are requested"))
 )
 
-# Schema for indexing ADO Wiki pages into vector store
-indexData = create_model(
-    "indexData",
-    __base__=BaseIndexParams,
-    plan_id=(int, Field(description="ID of the test plan for which test cases are requested")),
-    suite_ids=(list[int], Field(description="List of test suite IDs for which test cases are requested (can be empty)")),
-    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 TestPlanApiWrapper(BaseVectorStoreToolApiWrapper):
     __test__ = False
     organization_url: str
@@ -377,19 +364,6 @@ class TestPlanApiWrapper(BaseVectorStoreToolApiWrapper):
             logger.error(f"Error getting test cases: {e}")
             return ToolException(f"Error getting test cases: {e}")
 
-    def index_data(self,
-            plan_id: str,
-            suite_ids: list[str] = [],
-            collection_suffix: str = '',
-            progress_step: int = None,
-            clean_index: bool = False
-    ):
-        """Load ADO TestCases into the vector store."""
-        docs = self._base_loader(plan_id, suite_ids)
-        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 _base_loader(self, plan_id: str, suite_ids: Optional[list[str]] = []) -> Generator[Document, None, None]:
         cases = []
         for sid in suite_ids:
@@ -410,7 +384,15 @@ class TestPlanApiWrapper(BaseVectorStoreToolApiWrapper):
                 })
 
     def _process_document(self, document: Document) -> Generator[Document, None, None]:
-        yield document
+        if False:
+            yield  # Unreachable, but keeps the function a generator
+
+    def _index_tool_params(self):
+        """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)"))
+        }
 
     @extend_with_vector_tools
     def get_available_tools(self):
@@ -481,11 +463,5 @@ class TestPlanApiWrapper(BaseVectorStoreToolApiWrapper):
                 "description": self.get_test_cases.__doc__,
                 "args_schema": TestCasesGetModel,
                 "ref": self.get_test_cases,
-            },
-            {
-                "name": "index_data",
-                "ref": self.index_data,
-                "description": self.index_data.__doc__,
-                "args_schema": indexData,
             }
         ]

alita_sdk/tools/ado/wiki/ado_wrapper.py

@@ -1,28 +1,26 @@
 import logging
-from typing import Any, Dict, Generator, List, Optional
+from typing import Any, Dict, Generator, Optional
 
-from alita_sdk.tools.elitea_base import BaseIndexParams
-from langchain_core.documents import Document
-
-from ...elitea_base import BaseVectorStoreToolApiWrapper, extend_with_vector_tools
 from azure.devops.connection import Connection
 from azure.devops.exceptions import AzureDevOpsServiceError
 from azure.devops.v7_0.core import CoreClient
 from azure.devops.v7_0.wiki import WikiClient, WikiPageCreateOrUpdateParameters, WikiCreateParametersV2, \
     WikiPageMoveParameters
 from azure.devops.v7_0.wiki.models import GitVersionDescriptor
+from langchain_core.documents import Document
 from langchain_core.tools import ToolException
 from msrest.authentication import BasicAuthentication
 from pydantic import create_model, PrivateAttr, SecretStr
 from pydantic import model_validator
 from pydantic.fields import Field
+
+from ...elitea_base import BaseVectorStoreToolApiWrapper, extend_with_vector_tools
+
 try:
     from alita_sdk.runtime.langchain.interfaces.llm_processor import get_embeddings
 except ImportError:
     from alita_sdk.langchain.interfaces.llm_processor import get_embeddings
 
-from ...elitea_base import BaseToolApiWrapper
-
 logger = logging.getLogger(__name__)
 
 GetWikiInput = create_model(
@@ -60,17 +58,6 @@ RenamePageInput = create_model(
     version_type=(Optional[str], Field(description="Version type (branch, tag, or commit). Determines how Id is interpreted", default="branch"))
 )
 
-# Schema for indexing ADO Wiki pages into vector store
-indexData = create_model(
-    "indexData",
-    __base__=BaseIndexParams,
-    wiki_identifier=(str, Field(description="Wiki identifier to index, e.g., 'ABCProject.wiki'")),
-    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 AzureDevOpsApiWrapper(BaseVectorStoreToolApiWrapper):
     organization_url: str
@@ -242,19 +229,6 @@ class AzureDevOpsApiWrapper(BaseVectorStoreToolApiWrapper):
             logger.error(f"Unable to modify wiki page: {str(e)}")
             return ToolException(f"Unable to modify wiki page: {str(e)}")
 
-    def index_data(
-            self,
-            wiki_identifier: str,
-            collection_suffix: str = '',
-            progress_step: int = None,
-            clean_index: bool = False
-    ):
-        """Load ADO Wiki pages into the vector store."""
-        docs = self._base_loader(wiki_identifier)
-        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 _base_loader(self, wiki_identifier: str) -> Generator[Document, None, None]:
         pages = self._client.get_pages_batch(pages_batch_request={}, project=self.project, wiki_identifier=wiki_identifier)
         #
@@ -266,8 +240,11 @@ class AzureDevOpsApiWrapper(BaseVectorStoreToolApiWrapper):
                 'updated_on': ''
             })
 
-    def _process_document(self, document: Document) -> Generator[Document, None, None]:
-        yield document
+    def _index_tool_params(self):
+        """Return the parameters for indexing data."""
+        return {
+            "wiki_identifier": (str, Field(description="Wiki identifier to index, e.g., 'ABCProject.wiki'"))
+        }
 
     @extend_with_vector_tools
     def get_available_tools(self):
@@ -314,11 +291,5 @@ class AzureDevOpsApiWrapper(BaseVectorStoreToolApiWrapper):
                 "description": self.rename_wiki_page.__doc__,
                 "args_schema": RenamePageInput,
                 "ref": self.rename_wiki_page,
-            },
-            {
-                "name": "index_data",
-                "ref": self.index_data,
-                "description": self.index_data.__doc__,
-                "args_schema": indexData,
             }
         ]

alita_sdk/tools/confluence/api_wrapper.py

@@ -839,6 +839,8 @@ class ConfluenceAPIWrapper(BaseVectorStoreToolApiWrapper):
         loader = AlitaConfluenceLoader(self.client, self.llm, bins_with_llm, **confluence_loader_params)
 
         for document in loader._lazy_load(kwargs={}):
+            if 'updated_on' not in document.metadata and 'when' in document.metadata:
+                document.metadata['updated_on'] = document.metadata['when']
             yield document
 
     def _process_document(self, document: Document) -> Generator[Document, None, None]:

alita_sdk/tools/elitea_base.py

@@ -41,6 +41,11 @@ BaseCodeIndexParams = create_model(
     blacklist=(Optional[List[str]], Field(description="File extensions or paths to exclude. Defaults to no exclusions if None.", default=None)),
 )
 
+RemoveIndexParams = create_model(
+    "RemoveIndexParams",
+    collection_suffix=(Optional[str], Field(description="Optional suffix for collection name (max 7 characters)", default="", max_length=7)),
+)
+
 BaseSearchParams = create_model(
     "BaseSearchParams",
     query=(str, Field(description="Query text to search in the index")),
@@ -109,7 +114,7 @@ BaseIndexDataParams = create_model(
                          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")),
-    chunking_tool=(Literal['markdown', 'statistical', 'proposal'], Field(description="Name of chunking tool", default="markdown")),
+    chunking_tool=(Literal['markdown', 'statistical', 'proposal'], Field(description="Name of chunking tool", default=None)),
     chunking_config=(Optional[dict], Field(description="Chunking tool configuration", default_factory=dict)),
 )
 
@@ -345,7 +350,9 @@ class BaseVectorStoreToolApiWrapper(BaseToolApiWrapper):
                 "alita_sdk_options": {
                     "target_schema": collection_name,
                 },
-                "connection_string": self.connection_string.get_secret_value()
+                # "connection_string": self.connection_string.get_secret_value()
+                # 'postgresql+psycopg://project_23_user:Rxu4QtM2InLVNnm62GX7@pgvector:5432/project_23'
+                "connection_string": 'postgresql+psycopg://postgres:yourpassword@localhost:5432/postgres'
             }
         elif self.vectorstore_type == 'Chroma':
             vectorstore_params = {
@@ -363,6 +370,13 @@ class BaseVectorStoreToolApiWrapper(BaseToolApiWrapper):
             process_document_func=self._process_documents,
         )
 
+    def remove_index(self, collection_suffix: str = ""):
+        """
+            Cleans the indexed data in the collection
+        """
+
+        self._init_vector_store(collection_suffix)._clean_collection()
+
     def search_index(self,
                      query: str,
                      collection_suffix: str = "",
@@ -463,7 +477,14 @@ class BaseVectorStoreToolApiWrapper(BaseToolApiWrapper):
                 "ref": self.stepback_summary_index,
                 "description": self.stepback_summary_index.__doc__,
                 "args_schema": BaseStepbackSearchParams
-            }
+            },
+            {
+                "name": "remove_index",
+                "mode": "remove_index",
+                "ref": self.remove_index,
+                "description": self.remove_index.__doc__,
+                "args_schema": RemoveIndexParams
+            },
         ]
 
 

alita_sdk/tools/gitlab/__init__.py

@@ -85,13 +85,14 @@ class AlitaGitlabToolkit(BaseToolkit):
             if selected_tools:
                 if tool["name"] not in selected_tools:
                     continue
+
             tools.append(BaseAction(
                 api_wrapper=gitlab_api_wrapper,
                 name=prefix + tool["name"],
-                description=tool["description"] + "\nrepo: " + gitlab_api_wrapper.repository,
+                description=tool["description"] +  f"\nrepo: {gitlab_api_wrapper.repository}",
                 args_schema=tool["args_schema"]
             ))
         return cls(tools=tools)
 
-    def get_tools(self):
+    def get_tools(self)-> List[BaseTool]:
         return self.tools

alita_sdk/tools/gitlab/api_wrapper.py

@@ -1,6 +1,7 @@
 # api_wrapper.py
 from typing import Any, Dict, List, Optional
-
+import fnmatch
+from alita_sdk.tools.elitea_base import extend_with_vector_tools
 from alita_sdk.tools.elitea_base import BaseCodeToolApiWrapper
 from pydantic import create_model, Field, model_validator, SecretStr, PrivateAttr
 
@@ -53,7 +54,11 @@ CreateBranchModel = create_model(
 )
 ListBranchesInRepoModel = create_model(
     "ListBranchesInRepoModel",
+    limit=(Optional[int], Field(default=20, description="Maximum number of branches to return. If not provided, all branches will be returned.")),
+    branch_wildcard=(Optional[str], Field(default=None, description="Wildcard pattern to filter branches by name. If not provided, all branches will be returned."))
+
 )
+
 ListFilesModel = create_model(
     "ListFilesModel",
     path=(Optional[str], Field(description="The path to list files from")),
@@ -142,9 +147,30 @@ class GitLabAPIWrapper(BaseCodeToolApiWrapper):
         self._repo_instance.default_branch = branch_name
         return f"Active branch set to {branch_name}"
 
-    def list_branches_in_repo(self) -> List[str]:
-        branches = self._repo_instance.branches.list()
-        return [branch.name for branch in branches]
+    def list_branches_in_repo(self, limit: Optional[int] = 20, branch_wildcard: Optional[str] = None) -> List[str]:
+        """
+        Lists branches in the repository with optional limit and wildcard filtering.
+
+        Parameters:
+            limit (Optional[int]): Maximum number of branches to return
+            branch_wildcard (Optional[str]): Wildcard pattern to filter branches (e.g., '*dev')
+
+        Returns:
+            List[str]: List containing names of branches
+        """
+        try:
+            branches = self._repo_instance.branches.list(get_all=True)
+            
+            if branch_wildcard:
+                branches = [branch for branch in branches if fnmatch.fnmatch(branch.name, branch_wildcard)]
+            
+            if limit is not None:
+                branches = branches[:limit]
+            
+            branch_names = [branch.name for branch in branches]
+            return branch_names
+        except Exception as e:
+            return f"Failed to list branches: {str(e)}"
 
     def list_files(self, path: str = None, recursive: bool = True, branch: str = None) -> List[str]:
         branch = branch if branch else self._active_branch
@@ -404,84 +430,85 @@ class GitLabAPIWrapper(BaseCodeToolApiWrapper):
             for commit in commits
         ]
 
+    @extend_with_vector_tools
     def get_available_tools(self):
         return [
             {
                 "name": "create_branch",
                 "ref": self.create_branch,
-                "description": self.create_branch.__doc__,
+                "description": self.create_branch.__doc__ or "Create a new branch in the repository.",
                 "args_schema": CreateBranchModel,
             },
             {
                 "name": "list_branches_in_repo",
                 "ref": self.list_branches_in_repo,
-                "description": self.list_branches_in_repo.__doc__,
+                "description": self.list_branches_in_repo.__doc__ or "List branches in the repository with optional limit and wildcard filtering.",
                 "args_schema": ListBranchesInRepoModel,
             },
             {
                 "name": "list_files",
                 "ref": self.list_files,
-                "description": self.list_files.__doc__,
+                "description": self.list_files.__doc__ or "List files in the repository with optional path, recursive search, and branch.",
                 "args_schema": ListFilesModel,
             },
             {
                 "name": "list_folders",
                 "ref": self.list_folders,
-                "description": self.list_folders.__doc__,
+                "description": self.list_folders.__doc__ or "List folders in the repository with optional path, recursive search, and branch.",
                 "args_schema": ListFoldersModel,
             },
             {
                 "name": "get_issues",
                 "ref": self.get_issues,
-                "description": self.get_issues.__doc__,
+                "description": self.get_issues.__doc__ or "Get all open issues in the repository.",
                 "args_schema": GetIssuesModel,
             },
             {
                 "name": "get_issue",
                 "ref": self.get_issue,
-                "description": self.get_issue.__doc__,
+                "description": self.get_issue.__doc__ or "Get details of a specific issue by its number.",
                 "args_schema": GetIssueModel,
             },
             {
                 "name": "create_pull_request",
                 "ref": self.create_pull_request,
-                "description": self.create_pull_request.__doc__,
+                "description": self.create_pull_request.__doc__ or "Create a pull request (merge request) in the repository.",
                 "args_schema": CreatePullRequestModel,
             },
             {
                 "name": "comment_on_issue",
                 "ref": self.comment_on_issue,
-                "description": self.comment_on_issue.__doc__,
+                "description": self.comment_on_issue.__doc__ or "Comment on an issue in the repository.",
                 "args_schema": CommentOnIssueModel,
             },
             {
                 "name": "create_file",
                 "ref": self.create_file,
-                "description": self.create_file.__doc__,
+                "description": self.create_file.__doc__ or "Create a new file in the repository.",
                 "args_schema": CreateFileModel,
             },
             {
                 "name": "read_file",
                 "ref": self.read_file,
-                "description": self.read_file.__doc__,
+                "description": self.read_file.__doc__ or "Read the contents of a file in the repository.",
                 "args_schema": ReadFileModel,
             },
             {
                 "name": "update_file",
                 "ref": self.update_file,
-                "description": self.update_file.__doc__,
+                "description": self.update_file.__doc__ or "Update the contents of a file in the repository.",
                 "args_schema": UpdateFileModel,
             },
             {
                 "name": "append_file",
                 "ref": self.append_file,
-                "description": self.append_file.__doc__,
+                "description": self.append_file.__doc__ or "Append content to a file in the repository.",
                 "args_schema": AppendFileModel,
             },
             {
                 "name": "delete_file",
                 "ref": self.delete_file,
-                "description": self.delete_file.__doc__,
+                "description": self.delete_file.__doc__ or "Delete a file from the repository.",
                 "args_schema": DeleteFileModel,
             },
             {
@@ -507,5 +534,5 @@ class GitLabAPIWrapper(BaseCodeToolApiWrapper):
                 "ref": self.get_commits,
                 "description": "Retrieve a list of commits from the repository.",
                 "args_schema": GetCommitsModel,
-            },
+            }
         ] + self._get_vector_search_tools()