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

alita_sdk/tools/ado/repos/repos_wrapper.py

@@ -250,6 +250,7 @@ class ReposApiWrapper(BaseCodeToolApiWrapper):
     token: Optional[SecretStr]
     _client: Optional[GitClient] = PrivateAttr()
 
+    llm: Optional[Any] = None
     # Vector store configuration
     connection_string: Optional[SecretStr] = None
     collection_name: Optional[str] = None
@@ -303,24 +304,30 @@ class ReposApiWrapper(BaseCodeToolApiWrapper):
 
     def _get_files(
             self,
-            directory_path: str = "",
-            branch_name: str = None,
+            path: str = "",
+            branch: str = None,
             recursion_level: str = "Full",
     ) -> str:
+        """Get list of files from a repository path and branch.
+
+        Args:
+            path (str): Path within the repository to list files from
+            branch (str): Branch to get files from. Defaults to base_branch if None.
+            recursion_level (str): OneLevel - includes immediate children, Full - includes all items, None - no recursion
+
+        Returns:
+            List[str]: List of file paths
         """
-        Params:
-            recursion_level: OneLevel - includes immediate children, Full - includes all items, None - no recursion
-        """
-        branch_name = branch_name if branch_name else self.base_branch
+        branch = branch if branch else self.base_branch
         files: List[str] = []
         try:
             version_descriptor = GitVersionDescriptor(
-                version=branch_name, version_type="branch"
+                version=branch, version_type="branch"
             )
             items = self._client.get_items(
                 repository_id=self.repository_id,
                 project=self.project,
-                scope_path=directory_path,
+                scope_path=path,
                 recursion_level=recursion_level,
                 version_descriptor=version_descriptor,
                 include_content_metadata=True,
@@ -334,7 +341,7 @@ class ReposApiWrapper(BaseCodeToolApiWrapper):
             item = items.pop(0)
             if item.git_object_type == "blob":
                 files.append(item.path)
-        return str(files)
+        return files # Changed to return list directly instead of str
 
     def set_active_branch(self, branch_name: str) -> str:
         """
@@ -389,7 +396,7 @@ class ReposApiWrapper(BaseCodeToolApiWrapper):
             logger.error(msg)
             return ToolException(msg)
 
-    def list_files(self, directory_path: str = "", branch_name: str = None) -> str:
+    def list_files(self, directory_path: str = "", branch_name: str = None) -> List[str]:
         """
         Recursively fetches files from a directory in the repo.
 
@@ -398,12 +405,12 @@ class ReposApiWrapper(BaseCodeToolApiWrapper):
             branch_name (str): The name of the branch where the files to be received.
 
         Returns:
-            str: List of file paths, or an error message.
+            List[str]: List of file paths, or an error message.
         """
         self.active_branch = branch_name if branch_name else self.active_branch
         return self._get_files(
-            directory_path=directory_path,
-            branch_name=self.active_branch if self.active_branch else self.base_branch,
+            path=directory_path,
+            branch=self.active_branch if self.active_branch else self.base_branch,
         )
 
     def parse_pull_request_comments(

alita_sdk/tools/aws/__init__.py

@@ -0,0 +1,7 @@
+from .delta_lake import DeltaLakeToolkit
+
+name = "aws"
+
+def get_tools(tool_type, tool):
+    if tool_type == 'delta_lake':
+        return DeltaLakeToolkit().get_toolkit().get_tools()

alita_sdk/tools/aws/delta_lake/__init__.py

@@ -0,0 +1,136 @@
+
+from functools import lru_cache
+from typing import List, Optional, Type
+
+from langchain_core.tools import BaseTool, BaseToolkit
+from pydantic import BaseModel, Field, SecretStr, computed_field, field_validator
+
+from ...utils import TOOLKIT_SPLITTER, clean_string, get_max_toolkit_length
+from .api_wrapper import DeltaLakeApiWrapper
+from .tool import DeltaLakeAction
+
+name = "delta_lake"
+
+@lru_cache(maxsize=1)
+def get_available_tools() -> dict[str, dict]:
+    api_wrapper = DeltaLakeApiWrapper.model_construct()
+    available_tools: dict = {
+        x["name"]: x["args_schema"].model_json_schema()
+        for x in api_wrapper.get_available_tools()
+    }
+    return available_tools
+
+toolkit_max_length = lru_cache(maxsize=1)(
+    lambda: get_max_toolkit_length(get_available_tools())
+)
+
+class DeltaLakeToolkitConfig(BaseModel):
+    class Config:
+        title = name
+        json_schema_extra = {
+            "metadata": {
+                "hidden": True,
+                "label": "AWS Delta Lake",
+                "icon_url": "delta-lake.svg",
+                "sections": {
+                    "auth": {
+                        "required": False,
+                        "subsections": [
+                            {"name": "AWS Access Key ID", "fields": ["aws_access_key_id"]},
+                            {"name": "AWS Secret Access Key", "fields": ["aws_secret_access_key"]},
+                            {"name": "AWS Session Token", "fields": ["aws_session_token"]},
+                            {"name": "AWS Region", "fields": ["aws_region"]},
+                        ],
+                    },
+                    "connection": {
+                        "required": False,
+                        "subsections": [
+                            {"name": "Delta Lake S3 Path", "fields": ["s3_path"]},
+                            {"name": "Delta Lake Table Path", "fields": ["table_path"]},
+                        ],
+                    },
+                },
+            }
+        }
+
+    aws_access_key_id: Optional[SecretStr] = Field(default=None, description="AWS access key ID", json_schema_extra={"secret": True, "configuration": True})
+    aws_secret_access_key: Optional[SecretStr] = Field(default=None, description="AWS secret access key", json_schema_extra={"secret": True, "configuration": True})
+    aws_session_token: Optional[SecretStr] = Field(default=None, description="AWS session token (optional)", json_schema_extra={"secret": True, "configuration": True})
+    aws_region: Optional[str] = Field(default=None, description="AWS region for Delta Lake storage", json_schema_extra={"configuration": True})
+    s3_path: Optional[str] = Field(default=None, description="S3 path to Delta Lake data (e.g., s3://bucket/path)", json_schema_extra={"configuration": True})
+    table_path: Optional[str] = Field(default=None, description="Delta Lake table path (if not using s3_path)", json_schema_extra={"configuration": True})
+    selected_tools: List[str] = Field(default=[], description="Selected tools", json_schema_extra={"args_schemas": get_available_tools()})
+
+    @field_validator("selected_tools", mode="before", check_fields=False)
+    @classmethod
+    def selected_tools_validator(cls, value: List[str]) -> list[str]:
+        return [i for i in value if i in get_available_tools()]
+
+def _get_toolkit(tool) -> BaseToolkit:
+    return DeltaLakeToolkit().get_toolkit(
+        selected_tools=tool["settings"].get("selected_tools", []),
+        aws_access_key_id=tool["settings"].get("aws_access_key_id", None),
+        aws_secret_access_key=tool["settings"].get("aws_secret_access_key", None),
+        aws_session_token=tool["settings"].get("aws_session_token", None),
+        aws_region=tool["settings"].get("aws_region", None),
+        s3_path=tool["settings"].get("s3_path", None),
+        table_path=tool["settings"].get("table_path", None),
+        toolkit_name=tool.get("toolkit_name"),
+    )
+
+def get_toolkit():
+    return DeltaLakeToolkit.toolkit_config_schema()
+
+def get_tools(tool):
+    return _get_toolkit(tool).get_tools()
+
+class DeltaLakeToolkit(BaseToolkit):
+    tools: List[BaseTool] = []
+    api_wrapper: Optional[DeltaLakeApiWrapper] = Field(default_factory=DeltaLakeApiWrapper.model_construct)
+    toolkit_name: Optional[str] = None
+
+    @computed_field
+    @property
+    def tool_prefix(self) -> str:
+        return (
+            clean_string(self.toolkit_name, toolkit_max_length()) + TOOLKIT_SPLITTER
+            if self.toolkit_name
+            else ""
+        )
+
+    @computed_field
+    @property
+    def available_tools(self) -> List[dict]:
+        return self.api_wrapper.get_available_tools()
+
+    @staticmethod
+    def toolkit_config_schema() -> Type[BaseModel]:
+        return DeltaLakeToolkitConfig
+
+    @classmethod
+    def get_toolkit(
+        cls,
+        selected_tools: list[str] | None = None,
+        toolkit_name: Optional[str] = None,
+        **kwargs,
+    ) -> "DeltaLakeToolkit":
+        delta_lake_api_wrapper = DeltaLakeApiWrapper(**kwargs)
+        instance = cls(
+            tools=[], api_wrapper=delta_lake_api_wrapper, toolkit_name=toolkit_name
+        )
+        if selected_tools:
+            selected_tools = set(selected_tools)
+            for t in instance.available_tools:
+                if t["name"] in selected_tools:
+                    instance.tools.append(
+                        DeltaLakeAction(
+                            api_wrapper=instance.api_wrapper,
+                            name=instance.tool_prefix + t["name"],
+                            description=f"S3 Path: {getattr(instance.api_wrapper, 's3_path', '')} Table Path: {getattr(instance.api_wrapper, 'table_path', '')}\n" + t["description"],
+                            args_schema=t["args_schema"],
+                        )
+                    )
+        return instance
+
+    def get_tools(self):
+        return self.tools

alita_sdk/tools/aws/delta_lake/api_wrapper.py

@@ -0,0 +1,220 @@
+import functools
+import json
+import logging
+from typing import Any, List, Optional
+
+from deltalake import DeltaTable
+from langchain_core.tools import ToolException
+from pydantic import (
+    ConfigDict,
+    Field,
+    PrivateAttr,
+    SecretStr,
+    field_validator,
+    model_validator,
+)
+from pydantic_core.core_schema import ValidationInfo
+from ...elitea_base import BaseToolApiWrapper
+from .schemas import ArgsSchema
+
+
+def process_output(func):
+    @functools.wraps(func)
+    def wrapper(self, *args, **kwargs):
+        try:
+            result = func(self, *args, **kwargs)
+            if isinstance(result, Exception):
+                return ToolException(str(result))
+            if isinstance(result, (dict, list)):
+                return json.dumps(result, default=str)
+            return str(result)
+        except Exception as e:
+            logging.error(f"Error in '{func.__name__}': {str(e)}")
+            return ToolException(str(e))
+    return wrapper
+
+
+class DeltaLakeApiWrapper(BaseToolApiWrapper):
+    """
+    API Wrapper for AWS Delta Lake. Handles authentication, querying, and utility methods.
+    """
+    model_config = ConfigDict(arbitrary_types_allowed=True, from_attributes=True)
+
+    aws_access_key_id: Optional[SecretStr] = Field(default=None, json_schema_extra={"env_key": "AWS_ACCESS_KEY_ID"})
+    aws_secret_access_key: Optional[SecretStr] = Field(default=None, json_schema_extra={"env_key": "AWS_SECRET_ACCESS_KEY"})
+    aws_session_token: Optional[SecretStr] = Field(default=None, json_schema_extra={"env_key": "AWS_SESSION_TOKEN"})
+    aws_region: Optional[str] = Field(default=None, json_schema_extra={"env_key": "AWS_REGION"})
+    s3_path: Optional[str] = Field(default=None, json_schema_extra={"env_key": "DELTA_LAKE_S3_PATH"})
+    table_path: Optional[str] = Field(default=None, json_schema_extra={"env_key": "DELTA_LAKE_TABLE_PATH"})
+    _delta_table: Optional[DeltaTable] = PrivateAttr(default=None)
+
+    @classmethod
+    def model_construct(cls, *args, **kwargs):
+        klass = super().model_construct(*args, **kwargs)
+        klass._delta_table = None
+        return klass
+
+    @field_validator(
+        "aws_access_key_id",
+        "aws_secret_access_key",
+        "aws_session_token",
+        "aws_region",
+        "s3_path",
+        "table_path",
+        mode="before",
+        check_fields=False,
+    )
+    @classmethod
+    def set_from_values_or_env(cls, value, info: ValidationInfo):
+        if value is None:
+            if json_schema_extra := cls.model_fields[info.field_name].json_schema_extra:
+                if env_key := json_schema_extra.get("env_key"):
+                    try:
+                        from langchain_core.utils import get_from_env
+                        return get_from_env(
+                            key=info.field_name,
+                            env_key=env_key,
+                            default=cls.model_fields[info.field_name].default,
+                        )
+                    except Exception:
+                        return None
+        return value
+
+    @model_validator(mode="after")
+    def validate_auth(self) -> "DeltaLakeApiWrapper":
+        if not (self.aws_access_key_id and self.aws_secret_access_key and self.aws_region):
+            raise ValueError("You must provide AWS credentials and region.")
+        if not (self.s3_path or self.table_path):
+            raise ValueError("You must provide either s3_path or table_path.")
+        return self
+
+    @property
+    def delta_table(self) -> DeltaTable:
+        if not self._delta_table:
+            path = self.table_path or self.s3_path
+            if not path:
+                raise ToolException("Delta Lake table path (table_path or s3_path) must be specified.")
+            try:
+                storage_options = {
+                    "AWS_ACCESS_KEY_ID": self.aws_access_key_id.get_secret_value() if self.aws_access_key_id else None,
+                    "AWS_SECRET_ACCESS_KEY": self.aws_secret_access_key.get_secret_value() if self.aws_secret_access_key else None,
+                    "AWS_REGION": self.aws_region,
+                }
+                if self.aws_session_token:
+                    storage_options["AWS_SESSION_TOKEN"] = self.aws_session_token.get_secret_value()
+                storage_options = {k: v for k, v in storage_options.items() if v is not None}
+                self._delta_table = DeltaTable(path, storage_options=storage_options)
+            except Exception as e:
+                raise ToolException(f"Error initializing DeltaTable: {e}")
+        return self._delta_table
+
+    @process_output
+    def query_table(self, query: Optional[str] = None, columns: Optional[List[str]] = None, filters: Optional[dict] = None) -> List[dict]:
+        """
+        Query Delta Lake table. Supports pandas-like filtering, column selection, and SQL-like queries (via pandas.DataFrame.query).
+        Args:
+            query: SQL-like query string (pandas.DataFrame.query syntax)
+            columns: List of columns to select
+            filters: Dict of column:value pairs for pandas-like filtering
+        Returns:
+            List of dicts representing rows
+        """
+        dt = self.delta_table
+        df = dt.to_pandas()
+        if filters:
+            for col, val in filters.items():
+                df = df[df[col] == val]
+        if query:
+            try:
+                df = df.query(query)
+            except Exception as e:
+                raise ToolException(f"Error in query param: {e}")
+        if columns:
+            df = df[columns]
+        return df.to_dict(orient="records")
+
+    @process_output
+    def vector_search(self, embedding: List[float], k: int = 5, embedding_column: str = "embedding") -> List[dict]:
+        """
+        Perform a vector similarity search on the Delta Lake table.
+        Args:
+            embedding: Query embedding vector.
+            k: Number of top results to return.
+            embedding_column: Name of the column containing embeddings.
+        Returns:
+            List of dicts for top k most similar rows.
+        """
+        import numpy as np
+
+        dt = self.delta_table
+        df = dt.to_pandas()
+        if embedding_column not in df.columns:
+            raise ToolException(f"Embedding column '{embedding_column}' not found in table.")
+
+        # Filter out rows with missing embeddings
+        df = df[df[embedding_column].notnull()]
+        if df.empty:
+            return []
+        # Convert embeddings to numpy arrays
+        emb_matrix = np.array(df[embedding_column].tolist())
+        query_vec = np.array(embedding)
+
+        # Normalize for cosine similarity
+        emb_matrix_norm = emb_matrix / np.linalg.norm(emb_matrix, axis=1, keepdims=True)
+        query_vec_norm = query_vec / np.linalg.norm(query_vec)
+        similarities = np.dot(emb_matrix_norm, query_vec_norm)
+
+        # Get top k indices
+        top_k_idx = np.argsort(similarities)[-k:][::-1]
+        top_rows = df.iloc[top_k_idx]
+        return top_rows.to_dict(orient="records")
+
+    @process_output
+    def get_table_schema(self) -> str:
+        dt = self.delta_table
+        return dt.schema().to_pyarrow().to_string()
+
+    def get_available_tools(self) -> List[dict]:
+        return [
+            {
+                "name": "query_table",
+                "description": self.query_table.__doc__,
+                "args_schema": ArgsSchema.QueryTableArgs.value,
+                "ref": self.query_table,
+            },
+            {
+                "name": "vector_search",
+                "description": self.vector_search.__doc__,
+                "args_schema": ArgsSchema.VectorSearchArgs.value,
+                "ref": self.vector_search,
+            },
+            {
+                "name": "get_table_schema",
+                "description": self.get_table_schema.__doc__,
+                "args_schema": ArgsSchema.NoInput.value,
+                "ref": self.get_table_schema,
+            },
+        ]
+
+    def run(self, name: str, *args: Any, **kwargs: Any):
+        for tool in self.get_available_tools():
+            if tool["name"] == name:
+                if len(args) == 1 and isinstance(args[0], dict) and not kwargs:
+                    kwargs = args[0]
+                    args = ()
+                try:
+                    return tool["ref"](*args, **kwargs)
+                except TypeError as e:
+                    if kwargs and not args:
+                        try:
+                            return tool["ref"](**kwargs)
+                        except TypeError:
+                            raise ValueError(
+                                f"Argument mismatch for tool '{name}'. Error: {e}"
+                            ) from e
+                    else:
+                        raise ValueError(
+                            f"Argument mismatch for tool '{name}'. Error: {e}"
+                        ) from e
+        else:
+            raise ValueError(f"Unknown tool name: {name}")

alita_sdk/tools/aws/delta_lake/schemas.py

@@ -0,0 +1,20 @@
+
+from enum import Enum
+from typing import List, Optional
+
+from pydantic import Field, create_model
+
+class ArgsSchema(Enum):
+    NoInput = create_model("NoInput")
+    QueryTableArgs = create_model(
+        "QueryTableArgs",
+        query=(Optional[str], Field(default=None, description="SQL query to execute on Delta Lake table. If None, returns all data.")),
+        columns=(Optional[List[str]], Field(default=None, description="List of columns to select.")),
+        filters=(Optional[dict], Field(default=None, description="Dict of column:value pairs for pandas-like filtering.")),
+    )
+    VectorSearchArgs = create_model(
+        "VectorSearchArgs",
+        embedding=(List[float], Field(description="Embedding vector for similarity search.")),
+        k=(int, Field(default=5, description="Number of top results to return.")),
+        embedding_column=(Optional[str], Field(default="embedding", description="Name of the column containing embeddings.")),
+    )

alita_sdk/tools/aws/delta_lake/tool.py

@@ -0,0 +1,35 @@
+
+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 DeltaLakeApiWrapper
+
+
+class DeltaLakeAction(BaseTool):
+    """Tool for interacting with the Delta Lake API on AWS."""
+
+    api_wrapper: DeltaLakeApiWrapper = Field(default_factory=DeltaLakeApiWrapper)
+    name: 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 Delta Lake API to run an operation."""
+        try:
+            # Use the tool name to dispatch to the correct API wrapper method
+            return self.api_wrapper.run(self.name, *args, **kwargs)
+        except Exception as e:
+            return f"Error: {format_exc()}"

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)