Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat: update file upload and download mechanism on toolset #1124

Open
wants to merge 9 commits into
base: master
Choose a base branch
from
111 changes: 37 additions & 74 deletions python/composio/client/collections.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,10 +2,8 @@
Composio server object collections
"""

import base64
import difflib
import json
import os
import time
import traceback
import typing as t
Expand Down Expand Up @@ -988,39 +986,20 @@ def get( # type: ignore
return super().get(queries=queries)


def _check_file_uploadable(param_field: dict) -> bool:
return (
isinstance(param_field, dict)
and (param_field.get("title") in ["File", "FileType"])
and all(
field_name in param_field.get("properties", {})
for field_name in ["name", "content"]
)
)


def _check_file_downloadable(param_field: dict) -> bool:
return set(param_field.keys()) == {"name", "content"}


class ActionParametersModel(BaseModel):
"""Action parameter data models."""

class OpenAPISchema(BaseModel):
properties: t.Dict[str, t.Any]
title: str
type: str

required: t.Optional[t.List[str]] = None
examples: t.Optional[t.List[t.Any]] = None


class ActionResponseModel(BaseModel):
"""Action response data model."""
class ActionParametersModel(OpenAPISchema):
"""Action parameter data models."""

properties: t.Dict[str, t.Any]
title: str
type: str

required: t.Optional[t.List[str]] = None
class ActionResponseModel(OpenAPISchema):
"""Action response data model."""


class ActionModel(BaseModel):
Expand Down Expand Up @@ -1077,6 +1056,12 @@ class SearchResultTask(BaseModel):
)


class CreateUploadURLResponse(BaseModel):
id: str = Field(..., description="ID of the file")
url: str = Field(..., description="Onetime upload URL")
key: str = Field(..., description="S3 upload location")


class Actions(Collection[ActionModel]):
"""Collection of composio actions.."""

Expand Down Expand Up @@ -1316,58 +1301,13 @@ def execute(
:param session_id: ID of the current workspace session
:return: A dictionary containing the response from the executed action.
"""
# TOFIX: Remvoe this
if action.is_local:
return self.client.local.execute_action(action=action, request_data=params)

actions = self.get(actions=[action])
if len(actions) == 0:
raise ComposioClientError(f"Action {action} not found")

(action_model,) = actions
action_req_schema = action_model.parameters.properties
modified_params: t.Dict[str, t.Union[str, t.Dict[str, str]]] = {}
for param, value in params.items():
request_param_schema = action_req_schema.get(param)
if request_param_schema is None:
# User has sent a parameter that is not used by this action,
# so we can ignore it.
continue

file_readable = request_param_schema.get("file_readable", False)
file_uploadable = _check_file_uploadable(request_param_schema)

if file_readable and isinstance(value, str) and os.path.isfile(value):
with open(value, "rb") as file:
file_content = file.read()
try:
modified_params[param] = file_content.decode("utf-8")
except UnicodeDecodeError:
# If decoding fails, treat as binary and encode in base64
modified_params[param] = base64.b64encode(file_content).decode(
"utf-8"
)
elif file_uploadable and isinstance(value, str):
if not os.path.isfile(value):
raise ValueError(f"Attachment File with path `{value}` not found.")

with open(value, "rb") as file:
file_content = file.read()

modified_params[param] = {
"name": os.path.basename(value),
"content": base64.b64encode(file_content).decode("utf-8"),
}
else:
modified_params[param] = value

if action.no_auth:
return self._raise_if_required(
self.client.long_timeout_http.post(
url=str(self.endpoint / action.slug / "execute"),
json={
"appName": action.app,
"input": modified_params,
"input": params,
"text": text,
"version": action.version,
"sessionInfo": {
Expand All @@ -1390,7 +1330,7 @@ def execute(
"connectedAccountId": connected_account,
"entityId": entity_id,
"appName": action.app,
"input": modified_params,
"input": params,
"text": text,
"version": action.version,
"authConfig": self._serialize_auth(auth=auth),
Expand Down Expand Up @@ -1460,6 +1400,29 @@ def search_for_a_task(
for task in response.json().get("items", [])
]

def create_file_upload(
self,
app: str,
action: str,
filename: str,
mimetype: str,
md5: str,
) -> CreateUploadURLResponse:
return CreateUploadURLResponse(
**self._raise_if_required(
response=self.client.http.post(
url=str(self.endpoint / "files" / "upload" / "request"),
json={
"md5": md5,
"app": app,
"action": action,
"filename": filename,
"mimetype": mimetype,
},
)
).json()
)


class ExpectedFieldInput(BaseModel):
name: str
Expand Down
94 changes: 94 additions & 0 deletions python/composio/client/files.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,94 @@
from __future__ import annotations

import hashlib
import typing as t
from pathlib import Path

import requests
import typing_extensions as te
from pydantic import BaseModel, ConfigDict, Field

from composio.exceptions import ComposioSDKError
from composio.utils import mimetypes


# pylint: disable=missing-timeout

if te.TYPE_CHECKING:
from composio.client import Composio


_DEFAULT_CHUNK_SIZE = 1024 * 1024 * 100


def get_md5(file: Path):
obj = hashlib.md5()
with file.open("rb") as fp:
while True:
line = fp.read(_DEFAULT_CHUNK_SIZE)
if not line:
break
obj.update(line)
return obj.hexdigest()


def upload(url: str, file: Path) -> bool:
with file.open("rb") as data:
return requests.put(url=url, data=data).status_code == 200


class FileUploadable(BaseModel):
model_config = ConfigDict(json_schema_extra={"file_uploadable": True})

name: str
mimetype: str
s3key: str

@classmethod
def from_path(
cls,
file: t.Union[str, Path],
client: Composio,
action: str,
app: str,
) -> te.Self:
file = Path(file)
if not file.exists():
raise ComposioSDKError(f"File not found: {file}")

mimetype = mimetypes.guess(file=file)
s3meta = client.actions.create_file_upload(
app=app,
action=action,
filename=file.name,
mimetype=mimetype,
md5=get_md5(file=file),
)
if not upload(url=s3meta.url, file=file):
raise ComposioSDKError(f"Error uploading file: {file}")

return cls(
name=file.name,
mimetype=mimetype,
s3key=s3meta.key,
)


class FileDownloadable(BaseModel):
model_config = ConfigDict(json_schema_extra={"file_downloadable": True})

name: str = Field(..., description="Name of the file")
mimetype: str = Field(..., description="Mime type of the file.")
s3url: str = Field(..., description="URL of the file.")

def download(self, outdir: Path, chunk_size: int = _DEFAULT_CHUNK_SIZE) -> Path:
outfile = outdir / self.name
outdir.mkdir(exist_ok=True, parents=True)
response = requests.get(url=self.s3url, stream=True)
if response.status_code != 200:
raise ComposioSDKError(f"Error downloading file: {self.s3url}")

with outfile.open("wb") as fd:
for chunk in response.iter_content(chunk_size=chunk_size):
fd.write(chunk)
return outfile
22 changes: 20 additions & 2 deletions python/composio/tools/base/abs.py
Original file line number Diff line number Diff line change
Expand Up @@ -154,7 +154,25 @@ def __init__(self, model: t.Type[ModelType]) -> None:

@classmethod
def wrap(cls, model: t.Type[ModelType]) -> t.Type[BaseModel]:
class wrapper(model): # type: ignore
if "data" not in model.__annotations__:

class wrapper(BaseModel): # type: ignore
data: model = Field( # type: ignore
...,
description="Data from the action execution",
)
successful: bool = Field(
...,
description="Whether or not the action execution was successful or not",
)
error: t.Optional[str] = Field(
None,
description="Error if any occurred during the execution of the action",
)

return t.cast(t.Type[BaseModel], wrapper)

class wrapper(model): # type: ignore # pylint: disable=function-redefined
successful: bool = Field(
...,
description="Whether or not the action execution was successful or not",
Expand Down Expand Up @@ -199,7 +217,7 @@ def schema(self) -> t.Dict:
] += f" Note: choose value only from following options - {prop['enum']}"

schema["properties"] = properties
schema["title"] = self.model.__name__
schema["title"] = f"{self.model.__name__}Wrapper"
return remove_json_ref(schema)


Expand Down
2 changes: 1 addition & 1 deletion python/composio/tools/base/runtime.py
Original file line number Diff line number Diff line change
Expand Up @@ -179,7 +179,7 @@ def execute(self, request: t.Any, metadata: dict) -> t.Any:
type(inflection.camelize(f.__name__), (WrappedAction,), {}),
)
cls.__doc__ = f.__doc__
cls.description = f.__doc__ # type: ignore
cls.description = f.__doc__ or f.__name__ # type: ignore

# Normalize app name
toolname = toolname.upper()
Expand Down
Loading
Loading