From 1534f9aa73332dc764c0d241fcf42c1379b7c096 Mon Sep 17 00:00:00 2001
From: Viraj <35092918+angrybayblade@users.noreply.github.com>
Date: Tue, 26 Nov 2024 22:40:17 +0530
Subject: [PATCH] Advanced use case search (#916)

<!-- ELLIPSIS_HIDDEN -->



> [!IMPORTANT]
> Adds advanced task search functionality based on use cases with new
model and methods in `collections.py` and `toolset.py`.
>
>   - **Behavior**:
> - Adds `SearchResultTask` model in `collections.py` to represent
search results for tasks.
> - Implements `search_for_a_task()` in `Actions` class in
`collections.py` to perform advanced task searches based on use cases.
> - Updates `find_actions_by_use_case()` in `ComposioToolSet` in
`toolset.py` to support advanced search with a new `advanced` parameter.
>   - **Misc**:
>     - Minor changes to docstrings and comments for clarity.
>
> <sup>This description was created by </sup>[<img alt="Ellipsis"
src="https://img.shields.io/badge/Ellipsis-blue?color=175173">](https://www.ellipsis.dev?ref=ComposioHQ%2Fcomposio&utm_source=github&utm_medium=referral)<sup>
for 6201d4f28ba9c627e79569bc7aaa31f45d807219. It will automatically
update as commits are pushed.</sup>

<!-- ELLIPSIS_HIDDEN -->
---
 python/composio/client/collections.py | 53 +++++++++++++++++++++++++++
 python/composio/tools/toolset.py      | 26 +++++++++----
 2 files changed, 71 insertions(+), 8 deletions(-)

diff --git a/python/composio/client/collections.py b/python/composio/client/collections.py
index 249994f831a..9ccdc81608b 100644
--- a/python/composio/client/collections.py
+++ b/python/composio/client/collections.py
@@ -1045,6 +1045,26 @@ class CustomAuthObject(BaseModel):
     parameters: t.List[CustomAuthParameter] = Field(default_factory=lambda: [])
 
 
+class SearchResultTask(BaseModel):
+
+    app: str = Field(
+        description="Name of the app required to perform the subtask.",
+    )
+    actions: list[str] = Field(
+        description=(
+            "List of possible actions in order of relevance that can be used to "
+            "perform the task, provide minimum of {-min_actions-} and maximum of "
+            "{-max_actions-} actions."
+        ),
+    )
+    description: str = Field(
+        description="Descrption of the subtask.",
+    )
+    order: int = Field(
+        description="Order of the subtask, SHOULD START FROM 0",
+    )
+
+
 class Actions(Collection[ActionModel]):
     """Collection of composio actions.."""
 
@@ -1316,6 +1336,39 @@ def _serialize_auth(auth: t.Optional[CustomAuthObject]) -> t.Optional[t.Dict]:
         ]
         return data
 
+    def search_for_a_task(
+        self,
+        use_case: str,
+        limit: t.Optional[int] = None,
+        min_actions_per_task: t.Optional[int] = None,
+        max_actions_per_task: t.Optional[int] = None,
+        apps: t.Optional[t.List[str]] = None,
+    ) -> t.List[SearchResultTask]:
+        params: t.Dict[str, t.Any] = {"useCase": use_case}
+        if limit is not None:
+            params["limit"] = limit
+
+        if min_actions_per_task is not None:
+            params["minActionsPerTask"] = min_actions_per_task
+
+        if max_actions_per_task is not None:
+            params["maxActionsPerTask"] = max_actions_per_task
+
+        if apps is not None:
+            params["apps"] = ",".join(apps)
+
+        response = self._raise_if_required(
+            response=self.client.http.get(
+                str(self.endpoint / "search" / "advanced"),
+                params=params,
+            )
+        )
+
+        return [
+            SearchResultTask.model_validate(task)
+            for task in response.json().get("items", [])
+        ]
+
 
 class ExpectedFieldInput(BaseModel):
     name: str
diff --git a/python/composio/tools/toolset.py b/python/composio/tools/toolset.py
index 774cbe8bc0c..c8cb70e4cbf 100644
--- a/python/composio/tools/toolset.py
+++ b/python/composio/tools/toolset.py
@@ -1009,22 +1009,32 @@ def find_actions_by_use_case(
         self,
         *apps: AppType,
         use_case: str,
+        advanced: bool = False,
     ) -> t.List[Action]:
         """
         Find actions by specified use case.
 
         :param apps: List of apps to search.
         :param use_case: String describing the use case.
+        :param advanced: Use advanced search (will be slower than the normal search)
         :return: A list of actions matching the relevant use case.
         """
-        actions = self.client.actions.get(
-            apps=[App(app) for app in apps],
-            use_case=use_case,
-            allow_all=True,
-        )
-        return [
-            Action(value=get_enum_key(name=action.name).lower()) for action in actions
-        ]
+        if advanced:
+            actions = []
+            for task in self.client.actions.search_for_a_task(use_case=use_case):
+                actions += task.actions
+        else:
+            actions = list(
+                map(
+                    lambda x: get_enum_key(x.name),
+                    self.client.actions.get(
+                        apps=[App(app) for app in apps],
+                        use_case=use_case,
+                        allow_all=True,
+                    ),
+                )
+            )
+        return [Action(action) for action in actions]
 
     def find_actions_by_tags(
         self,