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

alita_sdk/runtime/clients/client.py

@@ -58,8 +58,8 @@ class AlitaClient:
         self.list_apps_url = f"{self.base_url}{self.api_path}/applications/applications/prompt_lib/{self.project_id}"
         self.integration_details = f"{self.base_url}{self.api_path}/integrations/integration/{self.project_id}"
         self.secrets_url = f"{self.base_url}{self.api_path}/secrets/secret/{self.project_id}"
-        self.artifacts_url = f"{self.base_url}{self.api_path}/artifacts/artifacts/{self.project_id}"
-        self.artifact_url = f"{self.base_url}{self.api_path}/artifacts/artifact/{self.project_id}"
+        self.artifacts_url = f"{self.base_url}{self.api_path}/artifacts/artifacts/default/{self.project_id}"
+        self.artifact_url = f"{self.base_url}{self.api_path}/artifacts/artifact/default/{self.project_id}"
         self.bucket_url = f"{self.base_url}{self.api_path}/artifacts/buckets/{self.project_id}"
         self.configurations_url = f'{self.base_url}{self.api_path}/integrations/integrations/default/{self.project_id}?section=configurations&unsecret=true'
         self.ai_section_url = f'{self.base_url}{self.api_path}/integrations/integrations/default/{self.project_id}?section=ai'
@@ -291,7 +291,7 @@ class AlitaClient:
         return self._process_requst(data)
 
     def download_artifact(self, bucket_name, artifact_name):
-        url = f'{self.artifact_url}/{bucket_name}/{artifact_name}'
+        url = f'{self.artifact_url}/{bucket_name.lower()}/{artifact_name}'
         data = requests.get(url, headers=self.headers, verify=False)
         if data.status_code == 403:
             return {"error": "You are not authorized to access this resource"}

alita_sdk/runtime/tools/vectorstore.py

@@ -1,13 +1,14 @@
 import json
-from json import dumps
-from typing import Any, Optional, List, Dict
+import math
+from typing import Any, Optional, List, Dict, Callable
 from pydantic import BaseModel, model_validator, Field
-from langchain_core.tools import ToolException
 from ..langchain.tools.vector import VectorAdapter
 from langchain_core.messages import HumanMessage
 from alita_sdk.tools.elitea_base import BaseToolApiWrapper
 from logging import getLogger
 
+from ..utils.logging import dispatch_custom_event
+
 logger = getLogger(__name__)
 
 class IndexDocumentsModel(BaseModel):
@@ -139,6 +140,7 @@ class VectorStoreWrapper(BaseToolApiWrapper):
     vectoradapter: Any = None
     pg_helper: Any = None
     embeddings: Any = None
+    process_document_func: Optional[Callable] = None
     
     @model_validator(mode='before')
     @classmethod
@@ -182,18 +184,122 @@ class VectorStoreWrapper(BaseToolApiWrapper):
             except Exception as e:
                 logger.error(f"Failed to initialize PGVectorSearch: {str(e)}")
 
-    def index_documents(self, documents):
+    def _get_indexed_data(self, store):
+        """ Get all indexed data from vectorstore """
+
+        # get already indexed data
+        result = {}
+        try:
+            self._log_data("Retrieving already indexed data from vectorstore",
+                           tool_name="index_documents")
+            data = store.get(include=['documents', 'metadatas'])
+            # re-structure data to be more usable
+            for doc_str, meta, db_id in zip(data['documents'], data['metadatas'], data['ids']):
+                doc = json.loads(doc_str)
+                doc_id = str(meta['id'])
+                result[doc_id] = {
+                    'metadata': meta,
+                    'document': doc,
+                    'id': db_id
+                }
+        except Exception as e:
+            logger.error(f"Failed to get indexed data from vectorstore: {str(e)}. Continuing with empty index.")
+        return result
+
+    def _reduce_duplicates(self, documents, store) -> List[Any]:
+        """Remove documents already indexed in the vectorstore based on metadata 'id' and 'updated_on' fields."""
+
+        self._log_data("Verification of documents to index started", tool_name="index_documents")
+
+        data = self._get_indexed_data(store)
+        indexed_ids = set(data.keys())
+        if not indexed_ids:
+            self._log_data("Vectorstore is empty, indexing all incoming documents", tool_name="index_documents")
+            return documents
+
+        final_docs = []
+        docs_to_remove = []
+
+        for document in documents:
+            doc_id = document.metadata.get('id')
+            # get document's metadata and id and check if already indexed
+            if doc_id in indexed_ids:
+                # document has been indexed already, then verify `updated_on`
+                to_index_updated_on = document.metadata.get('updated_on')
+                indexed_meta = data[doc_id]['metadata']
+                indexed_updated_on = indexed_meta.get('updated_on')
+                if to_index_updated_on and indexed_updated_on and to_index_updated_on == indexed_updated_on:
+                    # same updated_on, skip indexing
+                    continue
+                # if updated_on is missing or different, we will re-index the document and remove old one
+                docs_to_remove.append(data[doc_id]['id'])
+            else:
+                final_docs.append(document)
+
+        if docs_to_remove:
+            self._log_data(
+                f"Removing {len(docs_to_remove)} documents from vectorstore that are already indexed with different updated_on.",
+                tool_name="index_documents"
+            )
+            store.delete(ids=docs_to_remove)
+
+        return final_docs
+
+    def index_documents(self, documents, progress_step: int = 20, clean_index: bool = True):
+        """ Index documents in the vectorstore.
+
+        Args:
+            documents (Any): Generator or list of documents to index.
+            document_processing_func (Optional[Callable]): Function to process documents after duplicates removal and before indexing.
+            progress_step (int): Step for progress reporting, default is 20.
+            clean_index (bool): If True, clean the index before re-indexing all documents.
+        """
+
         from ..langchain.interfaces.llm_processor import add_documents
+
+        # pre-process documents if needed (find duplicates, etc.)
+        if clean_index:
+            logger.info("Cleaning index before re-indexing all documents.")
+            self._log_data("Cleaning index before re-indexing all documents. Previous index will be removed", tool_name="index_documents")
+            try:
+                self.vectoradapter.delete_dataset(self.dataset)
+                self.vectoradapter.persist()
+                self.vectoradapter.vacuum()
+                self._log_data("Previous index has been removed",
+                               tool_name="index_documents")
+            except Exception as e:
+                logger.warning(f"Failed to clean index: {str(e)}. Continuing with re-indexing.")
+        else:
+            # remove duplicates based on metadata 'id' and 'updated_on' fields
+            documents = self._reduce_duplicates(documents, self.vectoradapter.vectorstore)
+
+
+        if not documents or len(documents) == 0:
+            logger.info("No new documents to index after duplicate check.")
+            return {"status": "ok", "message": "No new documents to index."}
+
+        # if func is provided, apply it to documents
+        # used for processing of documents before indexing,
+        # e.g. to avoid time-consuming operations for documents that are already indexed
+        self.process_document_func(documents) if self.process_document_func else None
+
+        # notify user about missed required metadata fields: id, updated_on
+        # it is not required to have them, but it is recommended to have them for proper re-indexing and duplicate detection
+        for doc in documents:
+            if 'id' not in doc.metadata or 'updated_on' not in doc.metadata:
+                logger.warning(f"Document is missing required metadata field 'id' or 'updated_on': {doc.metadata}")
+
         logger.debug(f"Indexing documents: {documents}")
         logger.debug(self.vectoradapter)
-        self.vectoradapter.delete_dataset(self.dataset)
-        self.vectoradapter.persist()
-        logger.debug(f"Deleted Dataset")
-        #
-        self.vectoradapter.vacuum()
-        #
+
+        documents = list(documents)
+        total_docs = len(documents)
         documents_count = 0
         _documents = []
+
+        # set default progress step to 20 if out of 0...100 or None
+        progress_step = 20 if progress_step not in range(0, 100) else progress_step
+        next_progress_point = progress_step
         for document in documents:
             documents_count += 1
             # logger.debug(f"Indexing document: {document}")
@@ -203,7 +309,14 @@ class VectorStoreWrapper(BaseToolApiWrapper):
                     add_documents(vectorstore=self.vectoradapter.vectorstore, documents=_documents)
                     self.vectoradapter.persist()
                     _documents = []
-            except Exception as e:
+
+                percent = math.floor((documents_count / total_docs) * 100)
+                if percent >= next_progress_point:
+                    msg = f"Indexing progress: {percent}%. Processed {documents_count} of {total_docs} documents."
+                    logger.debug(msg)
+                    self._log_data(msg)
+                    next_progress_point += progress_step
+            except Exception:
                 from traceback import format_exc
                 logger.error(f"Error: {format_exc()}")
                 return {"status": "error", "message": f"Error: {format_exc()}"}
@@ -383,9 +496,11 @@ class VectorStoreWrapper(BaseToolApiWrapper):
             combined_items = [item for item in combined_items if abs(item[1]) >= cut_off]
         
         # Sort by score and limit results
-        combined_items.sort(key=lambda x: x[1], reverse=True)
+
+        # for chroma we want ascending order (lower score is better), for others descending
+        combined_items.sort(key=lambda x: x[1], reverse= self.vectorstore_type.lower() != 'chroma')
         combined_items = combined_items[:search_top]
-        
+
         # Format output based on doctype
         if doctype == 'code':
             return code_format(combined_items)
@@ -498,6 +613,21 @@ class VectorStoreWrapper(BaseToolApiWrapper):
         ])
         return result.content
 
+    def _log_data(self, message: str, tool_name: str = "index_data"):
+        """Log data and dispatch custom event for indexing progress"""
+
+        try:
+            dispatch_custom_event(
+                name="thinking_step",
+                data={
+                    "message": message,
+                    "tool_name": tool_name,
+                    "toolkit": "vectorstore",
+                },
+            )
+        except Exception as e:
+            logger.warning(f"Failed to dispatch progress event: {str(e)}")
+
     def get_available_tools(self):
         return [
             {

alita_sdk/tools/__init__.py

@@ -80,6 +80,8 @@ _safe_import_tool('postman', 'postman', 'get_tools', 'PostmanToolkit')
 _safe_import_tool('memory', 'memory', 'get_tools', 'MemoryToolkit')
 _safe_import_tool('zephyr_squad', 'zephyr_squad', 'get_tools', 'ZephyrSquadToolkit')
 _safe_import_tool('slack', 'slack', 'get_tools', 'SlackToolkit')
+_safe_import_tool('bigquery', 'google.bigquery', 'get_tools', 'BigQueryToolkit')
+_safe_import_tool('delta_lake', 'aws.delta_lake', 'get_tools', 'DeltaLakeToolkit')
 
 # Log import summary
 available_count = len(AVAILABLE_TOOLS)

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()}"