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

alita_sdk/tools/bitbucket/api_wrapper.py

@@ -5,7 +5,7 @@ import logging
 from typing import TYPE_CHECKING, Any, Dict, List, Optional
 
 from langchain_core.tools import ToolException
-from pydantic import BaseModel, model_validator, SecretStr
+from pydantic import model_validator, SecretStr
 from .bitbucket_constants import create_pr_data
 from .cloud_api_wrapper import BitbucketCloudApi, BitbucketServerApi
 from pydantic.fields import PrivateAttr
@@ -172,26 +172,26 @@ class BitbucketAPIWrapper(BaseCodeToolApiWrapper):
         """
         return self._bitbucket.get_pull_requests()
     
-    def get_pull_request(self, pr_id: str) -> Any:
+    def get_pull_request(self, pr_id: str) -> Dict[str, Any]:
         """
         Get details of a pull request
         Parameters:
             pr_id(str): the pull request ID
         Returns:
-            Any: Details of the pull request
+            dict: Details of the pull request as a dictionary
         """
         try:
             return self._bitbucket.get_pull_request(pr_id=pr_id)
         except Exception as e:
             return ToolException(f"Can't get pull request `{pr_id}` due to error:\n{str(e)}")
         
-    def get_pull_requests_changes(self, pr_id: str) -> Any:
+    def get_pull_requests_changes(self, pr_id: str) -> Dict[str, Any]:
         """
         Get changes of a pull request
         Parameters:
             pr_id(str): the pull request ID
         Returns:
-            Any: Changes of the pull request
+            dict: Changes of the pull request as a dictionary
         """
         try:
             return self._bitbucket.get_pull_requests_changes(pr_id=pr_id)

alita_sdk/tools/bitbucket/cloud_api_wrapper.py

@@ -17,6 +17,21 @@ if TYPE_CHECKING:
     pass
 
 
+def normalize_response(response) -> Dict[str, Any]:
+    """
+    Normalize API response to dictionary format.
+    Handles different response types from Bitbucket APIs.
+    """
+    if isinstance(response, dict):
+        return response
+    if hasattr(response, 'to_dict'):
+        return response.to_dict()
+    if hasattr(response, '__dict__'):
+        return {k: v for k, v in response.__dict__.items()
+                if isinstance(v, (str, int, float, bool, list, dict, type(None)))}
+    return {"raw_response": str(response)}
+
+
 class BitbucketApiAbstract(ABC):
 
     @abstractmethod
@@ -48,11 +63,11 @@ class BitbucketApiAbstract(ABC):
         pass
 
     @abstractmethod
-    def get_pull_request(self, pr_id: str) -> Any:
+    def get_pull_request(self, pr_id: str) -> Dict[str, Any]:
         pass
 
     @abstractmethod
-    def get_pull_requests_changes(self, pr_id: str) -> Any:
+    def get_pull_requests_changes(self, pr_id: str) -> Dict[str, Any]:
         pass
 
     @abstractmethod
@@ -95,7 +110,7 @@ class BitbucketServerApi(BitbucketApiAbstract):
 
     def get_files_list(self, file_path: str, branch: str) -> list:
         files = self.api_client.get_file_list(project_key=self.project, repository_slug=self.repository, query=branch,
-                                      sub_folder=file_path)
+                                              sub_folder=file_path)
         files_list = []
         for file in files:
             files_list.append(file['path'])
@@ -112,7 +127,7 @@ class BitbucketServerApi(BitbucketApiAbstract):
         )
 
     def update_file(self, file_path: str, update_query: str, branch: str) -> str:
-        file_content = self.get_file(file_path = file_path, branch=branch)
+        file_content = self.get_file(file_path=file_path, branch=branch)
         updated_file_content = file_content
         for old, new in extract_old_new_pairs(update_query):
             if not old.strip():
@@ -126,8 +141,8 @@ class BitbucketServerApi(BitbucketApiAbstract):
                 "the current file contents."
             )
 
-
-        source_commit_generator = self.api_client.get_commits(project_key=self.project, repository_slug=self.repository, hash_newest=branch, limit=1)
+        source_commit_generator = self.api_client.get_commits(project_key=self.project, repository_slug=self.repository,
+                                                              hash_newest=branch, limit=1)
         source_commit = next(source_commit_generator)
         return self.api_client.update_file(
             project_key=self.project,
@@ -138,7 +153,7 @@ class BitbucketServerApi(BitbucketApiAbstract):
             filename=file_path,
             source_commit_id=source_commit['id']
         )
-    
+
     def get_pull_request_commits(self, pr_id: str) -> List[Dict[str, Any]]:
         """
         Get commits from a pull request
@@ -147,27 +162,32 @@ class BitbucketServerApi(BitbucketApiAbstract):
         Returns:
             List[Dict[str, Any]]: List of commits in the pull request
         """
-        return self.api_client.get_pull_requests_commits(project_key=self.project, repository_slug=self.repository, pull_request_id=pr_id)
-    
+        return self.api_client.get_pull_requests_commits(project_key=self.project, repository_slug=self.repository,
+                                                         pull_request_id=pr_id)
+
     def get_pull_request(self, pr_id):
         """ Get details of a pull request
         Parameters:
             pr_id(str): the pull request ID
         Returns:
-            Any: Details of the pull request
+            Dict[str, Any]: Details of the pull request
         """
-        return self.api_client.get_pull_request(project_key=self.project, repository_slug=self.repository, pull_request_id=pr_id)
-    
-    def get_pull_requests_changes(self, pr_id: str) -> Any:
+        response = self.api_client.get_pull_request(project_key=self.project, repository_slug=self.repository,
+                                                    pull_request_id=pr_id)
+        return normalize_response(response)
+
+    def get_pull_requests_changes(self, pr_id: str) -> Dict[str, Any]:
         """
         Get changes of a pull request
         Parameters:
             pr_id(str): the pull request ID
         Returns:
-            Any: Changes of the pull request
+            Dict[str, Any]: Changes of the pull request
         """
-        return self.api_client.get_pull_requests_changes(project_key=self.project, repository_slug=self.repository, pull_request_id=pr_id)
-    
+        response = self.api_client.get_pull_requests_changes(project_key=self.project, repository_slug=self.repository,
+                                                             pull_request_id=pr_id)
+        return normalize_response(response)
+
     def add_pull_request_comment(self, pr_id, content, inline=None):
         """
         Add a comment to a pull request. Supports multiple content types and inline comments.
@@ -204,6 +224,7 @@ class BitbucketServerApi(BitbucketApiAbstract):
             **comment_data
         )
 
+
 class BitbucketCloudApi(BitbucketApiAbstract):
     api_client: Cloud
 
@@ -249,8 +270,9 @@ class BitbucketCloudApi(BitbucketApiAbstract):
         index = 0
         # TODO: add pagination
         for item in \
-        self.repository.get(path=f'src/{branch}/{file_path}?max_depth=100&pagelen=100&fields=values.path&q=type="commit_file"')[
-            'values']:
+                self.repository.get(
+                    path=f'src/{branch}/{file_path}?max_depth=100&pagelen=100&fields=values.path&q=type="commit_file"')[
+                    'values']:
             files_list.append(item['path'])
         return files_list
 
@@ -259,7 +281,8 @@ class BitbucketCloudApi(BitbucketApiAbstract):
             'branch': f'{branch}',
             f'{file_path}': f'{file_contents}',
         }
-        return self.repository.post(path='src', data=form_data, files={}, headers={'Content-Type': 'application/x-www-form-urlencoded'})
+        return self.repository.post(path='src', data=form_data, files={},
+                                    headers={'Content-Type': 'application/x-www-form-urlencoded'})
 
     def update_file(self, file_path: str, update_query: str, branch: str) -> ToolException | str:
 
@@ -277,7 +300,7 @@ class BitbucketCloudApi(BitbucketApiAbstract):
                 "the current file contents."
             )
         return self.create_file(file_path, updated_file_content, branch)
-    
+
     def get_pull_request_commits(self, pr_id: str) -> List[Dict[str, Any]]:
         """
         Get commits from a pull request
@@ -287,26 +310,28 @@ class BitbucketCloudApi(BitbucketApiAbstract):
             List[Dict[str, Any]]: List of commits in the pull request
         """
         return self.repository.pullrequests.get(pr_id).get('commits', {}).get('values', [])
-    
-    def get_pull_request(self, pr_id: str) -> Any:
+
+    def get_pull_request(self, pr_id: str) -> Dict[str, Any]:
         """ Get details of a pull request
         Parameters:
             pr_id(str): the pull request ID
         Returns:
-            Any: Details of the pull request
+            Dict[str, Any]: Details of the pull request
         """
-        return self.repository.pullrequests.get(pr_id)
-    
-    def get_pull_requests_changes(self, pr_id: str) -> Any:
+        response = self.repository.pullrequests.get(pr_id)
+        return normalize_response(response)
+
+    def get_pull_requests_changes(self, pr_id: str) -> Dict[str, Any]:
         """
         Get changes of a pull request
         Parameters:
             pr_id(str): the pull request ID
         Returns:
-            Any: Changes of the pull request
+            Dict[str, Any]: Changes of the pull request
         """
-        return self.repository.pullrequests.get(pr_id).get('diff', {})
-    
+        response = self.repository.pullrequests.get(pr_id).get('diff', {})
+        return normalize_response(response)
+
     def add_pull_request_comment(self, pr_id: str, content, inline=None) -> str:
         """
         Add a comment to a pull request. Supports multiple content types and inline comments.

alita_sdk/tools/elitea_base.py

@@ -2,7 +2,7 @@ import ast
 import fnmatch
 import logging
 import traceback
-from typing import Any, Optional, List, Dict
+from typing import Any, Optional, List, Dict, Generator
 
 from langchain_core.documents import Document
 from langchain_core.tools import ToolException
@@ -208,7 +208,7 @@ class BaseVectorStoreToolApiWrapper(BaseToolApiWrapper):
             Document: The processed document with metadata."""
         pass
 
-    def _process_documents(self, documents: List[Document]) -> List[Document]:
+    def _process_documents(self, documents: List[Document]) -> Generator[Document, None, None]:
         """
         Process a list of base documents to extract relevant metadata for full document preparation.
         Used for late processing of documents after we ensure that the documents have to be indexed to avoid
@@ -219,9 +219,14 @@ class BaseVectorStoreToolApiWrapper(BaseToolApiWrapper):
             documents (List[Document]): The base documents to process.
 
         Returns:
-            List[Document]: The processed documents with metadata.
+            Generator[Document, None, None]: A generator yielding processed documents with metadata.
         """
-        return [self._process_document(doc) for doc in documents]
+        for doc in documents:
+            processed_docs = self._process_document(doc)
+            if processed_docs:  # Only proceed if the list is not empty
+                for processed_doc in processed_docs:
+                    yield processed_doc
+
 
     def _init_vector_store(self, collection_suffix: str = "", embeddings: Optional[Any] = None):
         """ Initializes the vector store wrapper with the provided parameters."""

alita_sdk/tools/gitlab/__init__.py

@@ -1,16 +1,18 @@
 from typing import Dict, List, Literal, Optional
 
+from alita_sdk.tools.base.tool import BaseAction
 from langchain_core.tools import BaseTool
 from langchain_core.tools import BaseToolkit
 from pydantic import create_model, BaseModel, ConfigDict, SecretStr
 from pydantic.fields import Field
 
 from .api_wrapper import GitLabAPIWrapper
-from .tools import __all__
+# from .tools import __all__
 from ..utils import clean_string, TOOLKIT_SPLITTER, get_max_toolkit_length
 
 name = "gitlab"
 
+
 def get_tools(tool):
     return AlitaGitlabToolkit().get_toolkit(
         selected_tools=tool['settings'].get('selected_tools', []),
@@ -18,6 +20,15 @@ def get_tools(tool):
         repository=tool['settings']['repository'],
         branch=tool['settings']['branch'],
         private_token=tool['settings']['private_token'],
+
+        llm=tool['settings'].get('llm', None),
+        alita=tool['settings'].get('alita', None),
+        connection_string=tool['settings'].get('connection_string', None),
+        collection_name=f"{tool.get('toolkit_name')}_{str(tool['id'])}",
+        doctype='code',
+        embedding_model="HuggingFaceEmbeddings",
+        embedding_model_params={"model_name": "sentence-transformers/all-MiniLM-L6-v2"},
+        vectorstore_type="PGVector",
         toolkit_name=tool.get('toolkit_name')
     ).get_tools()
 
@@ -27,10 +38,8 @@ class AlitaGitlabToolkit(BaseToolkit):
 
     @staticmethod
     def toolkit_config_schema() -> BaseModel:
-        selected_tools = {}
-        for t in __all__:
-            default = t['tool'].__pydantic_fields__['args_schema'].default
-            selected_tools[t['name']] = default.schema() if default else default
+        selected_tools = {x['name']: x['args_schema'].schema() for x in
+                          GitLabAPIWrapper.model_construct().get_available_tools()}
         AlitaGitlabToolkit.toolkit_max_length = get_max_toolkit_length(selected_tools)
         return create_model(
             name,
@@ -66,15 +75,18 @@ class AlitaGitlabToolkit(BaseToolkit):
             selected_tools = []
         gitlab_api_wrapper = GitLabAPIWrapper(**kwargs)
         prefix = clean_string(toolkit_name, cls.toolkit_max_length) + TOOLKIT_SPLITTER if toolkit_name else ''
-        available_tools: List[Dict] = __all__
+        available_tools: List[Dict] = gitlab_api_wrapper.get_available_tools()
         tools = []
         for tool in available_tools:
             if selected_tools:
-                if tool['name'] not in selected_tools:
+                if tool["name"] not in selected_tools:
                     continue
-                initiated_tool = tool['tool'](api_wrapper=gitlab_api_wrapper)
-                initiated_tool.name = prefix + tool['name']
-                tools.append(initiated_tool)
+            tools.append(BaseAction(
+                api_wrapper=gitlab_api_wrapper,
+                name=prefix + tool["name"],
+                description=tool["description"] + "\nrepo: " + gitlab_api_wrapper.repository,
+                args_schema=tool["args_schema"]
+            ))
         return cls(tools=tools)
 
     def get_tools(self):