alita-sdk 0.3.554__py3-none-any.whl → 0.3.602__py3-none-any.whl

alita_sdk/tools/gitlab_org/api_wrapper.py

@@ -8,7 +8,7 @@ from langchain_core.tools import ToolException
 from pydantic import model_validator, PrivateAttr, create_model, SecretStr
 from pydantic.fields import Field
 
-from ..elitea_base import BaseToolApiWrapper
+from ..elitea_base import BaseToolApiWrapper, BaseCodeToolApiWrapper
 from ..gitlab.utils import get_diff_w_position, get_position
 
 logger = logging.getLogger(__name__)
@@ -24,7 +24,7 @@ GitLabCreateBranch = create_model(
 GitLabListBranches = create_model(
     "GitLabListBranchesModel",
     repository=(Optional[str], Field(description="Name of the repository", default=None)),
-    limit=(Optional[int], Field(description="Maximum number of branches to return. If not provided, all branches will be returned.", default=20)),
+    limit=(Optional[int], Field(description="Maximum number of branches to return. If not provided, all branches will be returned.", default=20, gt=0)),
     branch_wildcard=(Optional[str], Field(description="Wildcard pattern to filter branches by name. If not provided, all branches will be returned.", default=None))
 )
 
@@ -159,6 +159,9 @@ class GitLabWorkspaceAPIWrapper(BaseToolApiWrapper):
     repo_instances: Dict[str, Any] = {}
     _active_branch: Optional[str] = PrivateAttr(default='main')
 
+    # Reuse common file helpers from BaseCodeToolApiWrapper where applicable
+    edit_file = BaseCodeToolApiWrapper.edit_file
+
     class Config:
         arbitrary_types_allowed = True
 
@@ -371,51 +374,76 @@ class GitLabWorkspaceAPIWrapper(BaseToolApiWrapper):
         except Exception as e:
             return ToolException(e)
 
-    def update_file(self, file_path: str, update_query: str, branch: str, repository: Optional[str] = None) -> str:
-        """Updates a file with new content.
-        Parameters:
-            branch (str): The name of the branch where update the file.
-            update_query(str): Contains file contents.
-                The old file contents is wrapped in OLD <<<< and >>>> OLD
-                The new file contents is wrapped in NEW <<<< and >>>> NEW
-                For example:
-                /test/hello.txt
-                OLD <<<<
-                Hello Earth!
-                >>>> OLD
-                NEW <<<<
-                Hello Mars!
-                >>>> NEW
+    def _read_file(self, file_path: str, branch: str, **kwargs) -> str:
+        """
+        Internal read_file used by BaseCodeToolApiWrapper.edit_file.
+        Delegates to the public `read_file` implementation which supports an optional repository argument.
+        The repository may be passed via kwargs or provided earlier through `update_file` which sets
+        a temporary attribute `_tmp_repository_for_edit`.
+        """
+        # Repository from temporary context, then None
+        repository = getattr(self, "_tmp_repository_for_edit", None)
+        try:
+            # Public read_file signature: read_file(file_path, branch, repository=None)
+            return self.read_file(file_path, branch, repository)
+        except Exception as e:
+            raise ToolException(f"Can't extract file content (`{file_path}`) due to error:\n{str(e)}")
+
+    def _write_file(self, file_path: str, content: str, branch: str = None, commit_message: str = None) -> str:
+        """
+        Write content to a file (update only) in the specified GitLab repository.
+
+        This implementation follows the same commit flow as the previous `update_file`:
+        it does not attempt to create the file when it is missing — it will always
+        create a commit with a single `update` action. If the file does not exist on
+        the target branch, the underlying GitLab API will typically return an error.
         """
         try:
+            branch = branch if branch else (self._active_branch if self._active_branch else self.branch)
+            # pick repository from temporary edit context
+            repository = getattr(self, "_tmp_repository_for_edit", None)
             repo_instance = self._get_repo(repository)
-            file_content = self.read_file(file_path, branch, repository)
-            updated_file_content = file_content
-            for old, new in self.extract_old_new_pairs(update_query):
-                if not old.strip():
-                    continue
-                updated_file_content = updated_file_content.replace(old, new)
-            if file_content == updated_file_content:
-                return (
-                    "File content was not updated because old content was not found or empty."
-                    "It may be helpful to use the read_file action to get "
-                    "the current file contents."
-                )
+
+            # Always perform an 'update' action commit (do not create file when missing)
             commit = {
                 "branch": branch,
-                "commit_message": "Update " + file_path,
+                "commit_message": commit_message or f"Update {file_path}",
                 "actions": [
                     {
                         "action": "update",
                         "file_path": file_path,
-                        "content": updated_file_content,
+                        "content": content,
                     }
                 ],
             }
             repo_instance.commits.create(commit)
-            return "Updated file " + file_path
+            return f"Updated file {file_path}"
+        except ToolException:
+            raise
+        except Exception as e:
+            return ToolException(f"Unable to write file due to error: {str(e)}")
+
+    def update_file(self, file_path: str, update_query: str, branch: str, repository: Optional[str] = None) -> str:
+        """Updates a file with new content using OLD/NEW markers by delegating to `edit_file`.
+
+        The method sets a temporary repository context so that `edit_file`'s internal
+        calls to `_read_file` and `_write_file` operate on the requested repository.
+        """
+        # Set temporary repository context used by _read_file/_write_file
+        self._tmp_repository_for_edit = repository
+        try:
+            commit_message = f"Update {file_path}"
+            return self.edit_file(file_path=file_path, file_query=update_query, branch=branch, commit_message=commit_message)
+        except ToolException as e:
+            return str(e)
         except Exception as e:
             return ToolException(f"Unable to update file due to error: {str(e)}")
+        finally:
+            # Clear temporary context
+            try:
+                delattr(self, "_tmp_repository_for_edit")
+            except Exception:
+                self._tmp_repository_for_edit = None
 
     def delete_file(self, file_path: str, branch: str, repository: Optional[str] = None) -> str:
         """Deletes a file from the repo."""
@@ -428,36 +456,6 @@ class GitLabWorkspaceAPIWrapper(BaseToolApiWrapper):
         except Exception as e:
             return ToolException(f"Unable to delete file due to error: {str(e)}")
 
-    def extract_old_new_pairs(self, file_query):
-        """Extract old and new content pairs from the file query."""
-        code_lines = file_query.split("\n")
-        old_contents = []
-        new_contents = []
-        in_old_section = False
-        in_new_section = False
-        current_section_content = []
-        for line in code_lines:
-            if "OLD <<<" in line:
-                in_old_section = True
-                current_section_content = []
-                continue
-            if ">>>> OLD" in line:
-                in_old_section = False
-                old_contents.append("\n".join(current_section_content).strip())
-                current_section_content = []
-                continue
-            if "NEW <<<" in line:
-                in_new_section = True
-                current_section_content = []
-                continue
-            if ">>>> NEW" in line:
-                in_new_section = False
-                new_contents.append("\n".join(current_section_content).strip())
-                current_section_content = []
-                continue
-            if in_old_section or in_new_section:
-                current_section_content.append(line)
-        return list(zip(old_contents, new_contents))
 
     def append_file(self, file_path: str, content: str, branch: str, repository: Optional[str] = None) -> str:
         """

alita_sdk/tools/google/bigquery/__init__.py

@@ -8,6 +8,7 @@ from ....configurations.bigquery import BigQueryConfiguration
 from ...utils import clean_string, get_max_toolkit_length
 from .api_wrapper import BigQueryApiWrapper
 from .tool import BigQueryAction
+from ....runtime.utils.constants import TOOLKIT_NAME_META, TOOL_NAME_META, TOOLKIT_TYPE_META
 
 name = "bigquery"
 
@@ -129,7 +130,7 @@ class BigQueryToolkit(BaseToolkit):
                             name=t["name"],
                             description=description,
                             args_schema=t["args_schema"],
-                            metadata={"toolkit_name": toolkit_name} if toolkit_name else {}
+                            metadata={TOOLKIT_NAME_META: toolkit_name, TOOLKIT_TYPE_META: name, TOOL_NAME_META: t["name"]} if toolkit_name else {TOOL_NAME_META: t["name"]}
                         )
                     )
         return instance

alita_sdk/tools/google_places/__init__.py

@@ -8,6 +8,7 @@ from ..base.tool import BaseAction
 from ..elitea_base import filter_missconfigured_index_tools
 from ..utils import clean_string, get_max_toolkit_length
 from ...configurations.google_places import GooglePlacesConfiguration
+from ...runtime.utils.constants import TOOLKIT_NAME_META, TOOL_NAME_META, TOOLKIT_TYPE_META
 
 name = "google_places"
 
@@ -67,7 +68,7 @@ class GooglePlacesToolkit(BaseToolkit):
                 name=tool["name"],
                 description=description,
                 args_schema=tool["args_schema"],
-                metadata={"toolkit_name": toolkit_name} if toolkit_name else {}
+                metadata={TOOLKIT_NAME_META: toolkit_name, TOOLKIT_TYPE_META: name, TOOL_NAME_META: tool["name"]} if toolkit_name else {TOOL_NAME_META: tool["name"]}
             ))
         return cls(tools=tools)
 

alita_sdk/tools/jira/__init__.py

@@ -9,6 +9,7 @@ from ..elitea_base import filter_missconfigured_index_tools
 from ..utils import clean_string, get_max_toolkit_length, parse_list, check_connection_response
 from ...configurations.jira import JiraConfiguration
 from ...configurations.pgvector import PgVectorConfiguration
+from ...runtime.utils.constants import TOOLKIT_NAME_META, TOOLKIT_TYPE_META, TOOL_NAME_META
 
 name = "jira"
 
@@ -126,7 +127,7 @@ class JiraToolkit(BaseToolkit):
                 name=tool["name"],
                 description=description,
                 args_schema=tool["args_schema"],
-                metadata={"toolkit_name": toolkit_name} if toolkit_name else {}
+                metadata={TOOLKIT_NAME_META: toolkit_name, TOOLKIT_TYPE_META: name, TOOL_NAME_META: tool["name"]} if toolkit_name else {TOOL_NAME_META: tool["name"]}
             ))
         return cls(tools=tools)
 

alita_sdk/tools/keycloak/__init__.py

@@ -6,6 +6,7 @@ from pydantic import BaseModel, ConfigDict, create_model, Field, SecretStr
 from .api_wrapper import KeycloakApiWrapper
 from ..base.tool import BaseAction
 from ..utils import clean_string, get_max_toolkit_length
+from ...runtime.utils.constants import TOOLKIT_NAME_META, TOOL_NAME_META, TOOLKIT_TYPE_META
 
 name = "keycloak"
 
@@ -54,7 +55,7 @@ class KeycloakToolkit(BaseToolkit):
                 name=tool["name"],
                 description=description,
                 args_schema=tool["args_schema"],
-                metadata={"toolkit_name": toolkit_name} if toolkit_name else {}
+                metadata={TOOLKIT_NAME_META: toolkit_name, TOOLKIT_TYPE_META: name, TOOL_NAME_META: tool["name"]} if toolkit_name else {TOOL_NAME_META: tool["name"]}
             ))
         return cls(tools=tools)
 

alita_sdk/tools/localgit/__init__.py

@@ -5,6 +5,7 @@ from pydantic import BaseModel, ConfigDict, create_model, Field
 
 from .local_git import LocalGit
 from .tool import LocalGitAction
+from ...runtime.utils.constants import TOOLKIT_NAME_META, TOOL_NAME_META, TOOLKIT_TYPE_META
 
 name = "localgit"
 
@@ -55,7 +56,7 @@ class AlitaLocalGitToolkit(BaseToolkit):
                 mode=tool["mode"],
                 description=description,
                 args_schema=tool["args_schema"],
-                metadata={"toolkit_name": toolkit_name} if toolkit_name else {}
+                metadata={TOOLKIT_NAME_META: toolkit_name, TOOLKIT_TYPE_META: name, TOOL_NAME_META: tool["name"]} if toolkit_name else {TOOL_NAME_META: tool["name"]}
             ))
         return cls(tools=tools)
 

alita_sdk/tools/memory/__init__.py

@@ -118,7 +118,7 @@ class MemoryToolkit(BaseToolkit):
         # Add metadata to tools if toolkit_name is provided
         if toolkit_name:
             for tool in tools:
-                tool.metadata = {"toolkit_name": toolkit_name}
+                tool.metadata = {"toolkit_name": toolkit_name, "toolkit_type": name}
         
         return cls(tools=tools)
 

alita_sdk/tools/ocr/__init__.py

@@ -6,6 +6,7 @@ from pydantic import create_model, BaseModel, ConfigDict, Field
 from .api_wrapper import OCRApiWrapper
 from ..base.tool import BaseAction
 from ..utils import clean_string, get_max_toolkit_length
+from ...runtime.utils.constants import TOOLKIT_NAME_META, TOOL_NAME_META, TOOLKIT_TYPE_META
 
 name = "ocr"
 
@@ -59,7 +60,7 @@ class OCRToolkit(BaseToolkit):
                 name=tool["name"],
                 description=description,
                 args_schema=tool["args_schema"],
-                metadata={"toolkit_name": toolkit_name} if toolkit_name else {}
+                metadata={TOOLKIT_NAME_META: toolkit_name, TOOLKIT_TYPE_META: name, TOOL_NAME_META: tool["name"]} if toolkit_name else {TOOL_NAME_META: tool["name"]}
             ))
         return cls(tools=tools)
 

alita_sdk/tools/openapi/__init__.py

@@ -1,19 +1,218 @@
 from __future__ import annotations
 
+import base64
 import json
-from typing import Any, Dict, List, Optional
+import logging
+import threading
+import time
+from typing import Any, Dict, List, Optional, Tuple
+from urllib.parse import urlparse
 
 from langchain_core.tools import BaseTool, BaseToolkit
 from pydantic import BaseModel, ConfigDict, Field, create_model
+import requests
 import yaml
 
 from .api_wrapper import _get_base_url_from_spec, build_wrapper
 from .tool import OpenApiAction
 from ..elitea_base import filter_missconfigured_index_tools
 from ...configurations.openapi import OpenApiConfiguration
+from ...runtime.utils.constants import TOOLKIT_NAME_META, TOOL_NAME_META, TOOLKIT_TYPE_META
+
+logger = logging.getLogger(__name__)
 
 name = 'openapi'
 
+# Module-level token cache: {cache_key: (access_token, expires_at_timestamp)}
+# Protected by _oauth_token_cache_lock for thread-safe access
+_oauth_token_cache: Dict[str, Tuple[str, float]] = {}
+_oauth_token_cache_lock = threading.Lock()
+
+# Token expiry buffer in seconds (refresh 60 seconds before actual expiry)
+_TOKEN_EXPIRY_BUFFER = 60
+
+
+def _get_oauth_cache_key(client_id: str, token_url: str, scope: Optional[str]) -> str:
+    """Generate a cache key for OAuth tokens."""
+    return f"{client_id}:{token_url}:{scope or ''}"
+
+
+def _get_cached_token(cache_key: str) -> Optional[str]:
+    """Get a cached token if it exists and is not expired. Thread-safe."""
+    with _oauth_token_cache_lock:
+        if cache_key not in _oauth_token_cache:
+            return None
+        token, expires_at = _oauth_token_cache[cache_key]
+        if time.time() >= expires_at - _TOKEN_EXPIRY_BUFFER:
+            # Token expired or about to expire
+            del _oauth_token_cache[cache_key]
+            return None
+        return token
+
+
+def _cache_token(cache_key: str, token: str, expires_in: Optional[int]) -> None:
+    """Cache a token with its expiry time. Thread-safe."""
+    # Default to 1 hour if expires_in not provided
+    expires_in = expires_in or 3600
+    expires_at = time.time() + expires_in
+    with _oauth_token_cache_lock:
+        _oauth_token_cache[cache_key] = (token, expires_at)
+
+
+def _obtain_oauth_token(
+    client_id: str,
+    client_secret: str,
+    token_url: str,
+    scope: Optional[str] = None,
+    method: str = 'default',
+    timeout: int = 30,
+) -> Tuple[str, Optional[str]]:
+    """
+    Obtain an OAuth2 access token using client credentials flow.
+    
+    Args:
+        client_id: OAuth client ID
+        client_secret: OAuth client secret
+        token_url: OAuth token endpoint URL
+        scope: Optional OAuth scope(s), space-separated if multiple
+        method: Token exchange method - 'default' (POST body) or 'Basic' (Basic auth header)
+        timeout: Request timeout in seconds
+    
+    Returns:
+        Tuple of (access_token, error_message)
+        On success: (token, None)
+        On failure: (None, error_message)
+    """
+    try:
+        headers = {
+            'Content-Type': 'application/x-www-form-urlencoded',
+            'Accept': 'application/json',
+        }
+        
+        # Build form data
+        data: Dict[str, str] = {
+            'grant_type': 'client_credentials',
+        }
+        
+        if method == 'Basic':
+            # Use Basic auth header for client credentials
+            credentials = f"{client_id}:{client_secret}"
+            encoded_credentials = base64.b64encode(credentials.encode('utf-8')).decode('utf-8')
+            headers['Authorization'] = f'Basic {encoded_credentials}'
+        else:
+            # Default: include credentials in POST body
+            data['client_id'] = client_id
+            data['client_secret'] = client_secret
+        
+        if scope:
+            data['scope'] = scope
+        
+        # Log only the domain to avoid exposing sensitive path parameters (e.g., tenant IDs)
+        token_domain = urlparse(token_url).netloc or 'unknown'
+        logger.debug(f"OAuth token request to {token_domain} using method '{method}'")
+        
+        response = requests.post(
+            token_url,
+            headers=headers,
+            data=data,
+            timeout=timeout,
+        )
+        
+        if response.status_code == 200:
+            try:
+                token_data = response.json()
+                access_token = token_data.get('access_token')
+                if not access_token:
+                    return None, "OAuth response did not contain 'access_token'"
+                
+                # Cache the token
+                cache_key = _get_oauth_cache_key(client_id, token_url, scope)
+                expires_in = token_data.get('expires_in')
+                _cache_token(cache_key, access_token, expires_in)
+                
+                logger.debug(f"OAuth token obtained successfully (expires_in: {expires_in})")
+                return access_token, None
+            except json.JSONDecodeError as e:
+                return None, f"Failed to parse OAuth token response as JSON: {e}"
+        
+        # Handle error responses
+        error_msg = f"OAuth token request failed with status {response.status_code}"
+        try:
+            error_data = response.json()
+            if 'error' in error_data:
+                error_msg = f"{error_msg}: {error_data.get('error')}"
+                if 'error_description' in error_data:
+                    error_msg = f"{error_msg} - {error_data.get('error_description')}"
+        except Exception:
+            if response.text:
+                error_msg = f"{error_msg}: {response.text[:500]}"
+        
+        return None, error_msg
+        
+    except requests.exceptions.Timeout:
+        return None, f"OAuth token request to {token_url} timed out"
+    except requests.exceptions.ConnectionError as e:
+        return None, f"Failed to connect to OAuth token endpoint {token_url}: {e}"
+    except requests.exceptions.RequestException as e:
+        return None, f"OAuth token request failed: {e}"
+    except Exception as e:
+        return None, f"Unexpected error during OAuth token exchange: {e}"
+
+
+def _secret_to_str(value: Any) -> Optional[str]:
+    """Convert a secret value to string, handling SecretStr and other types."""
+    if value is None:
+        return None
+    if hasattr(value, 'get_secret_value'):
+        try:
+            value = value.get_secret_value()
+        except Exception:
+            pass
+    if isinstance(value, str):
+        return value
+    return str(value)
+
+
+def _get_oauth_access_token(settings: Dict[str, Any]) -> Tuple[Optional[str], Optional[str]]:
+    """
+    Get an OAuth access token from settings, using cache if available.
+    
+    Args:
+        settings: Dictionary containing OAuth configuration
+    
+    Returns:
+        Tuple of (access_token, error_message)
+        On success: (token, None)
+        On failure: (None, error_message)
+        If OAuth not configured: (None, None)
+    """
+    client_id = settings.get('client_id')
+    client_secret = _secret_to_str(settings.get('client_secret'))
+    token_url = settings.get('token_url')
+    
+    # Check if OAuth is configured
+    if not client_id or not client_secret or not token_url:
+        return None, None  # OAuth not configured
+    
+    scope = settings.get('scope')
+    method = settings.get('method', 'default') or 'default'
+    
+    # Try to get cached token
+    cache_key = _get_oauth_cache_key(client_id, token_url, scope)
+    cached_token = _get_cached_token(cache_key)
+    if cached_token:
+        logger.debug("Using cached OAuth token")
+        return cached_token, None
+    
+    # Obtain new token
+    return _obtain_oauth_token(
+        client_id=client_id,
+        client_secret=client_secret,
+        token_url=token_url,
+        scope=scope,
+        method=method,
+    )
+
 
 def get_toolkit(tool) -> BaseToolkit:
     settings = tool.get('settings', {}) or {}
@@ -221,7 +420,7 @@ class AlitaOpenAPIToolkit(BaseToolkit):
                     name=tool_def['name'],
                     description=description,
                     args_schema=tool_def.get('args_schema'),
-                    metadata={"toolkit_name": toolkit_name} if toolkit_name else {},
+                    metadata={TOOLKIT_NAME_META: toolkit_name, TOOLKIT_TYPE_META: name, TOOL_NAME_META: tool_def["name"]} if toolkit_name else {TOOL_NAME_META: tool_def["name"]},
                 )
             )
 
@@ -249,22 +448,35 @@ def _coerce_selected_tool_names(selected_tools: Any) -> list[str]:
     return []
 
 
-def _secret_to_str(value: Any) -> Optional[str]:
-    if value is None:
-        return None
-    if hasattr(value, 'get_secret_value'):
-        try:
-            value = value.get_secret_value()
-        except Exception:
-            pass
-    if isinstance(value, str):
-        return value
-    return str(value)
-
-
 def _build_headers_from_settings(settings: Dict[str, Any]) -> Dict[str, str]:
+    """
+    Build HTTP headers from settings, supporting API key and OAuth authentication.
+    
+    Authentication priority:
+    1. OAuth (client credentials flow) - if client_id, client_secret, and token_url are provided
+    2. API Key - if api_key is provided
+    3. Legacy authentication structure (for backward compatibility)
+    
+    Args:
+        settings: Dictionary containing authentication settings
+        
+    Returns:
+        Dictionary of HTTP headers to include in requests
+    """
     headers: Dict[str, str] = {}
 
+    # First, try OAuth authentication (client credentials flow)
+    # This takes priority because it's more secure and commonly used with modern APIs
+    oauth_token, oauth_error = _get_oauth_access_token(settings)
+    if oauth_token:
+        headers['Authorization'] = f'Bearer {oauth_token}'
+        logger.debug("Using OAuth Bearer token for authentication")
+        return headers
+    elif oauth_error:
+        # OAuth was configured but failed - log the error
+        # We'll still try API key auth as fallback
+        logger.warning(f"OAuth token exchange failed: {oauth_error}")
+
     # Legacy structure used by the custom OpenAPI UI
     auth = settings.get('authentication')
     if isinstance(auth, dict) and auth.get('type') == 'api_key':