Add next_token to X Search tools (#169)

# PR Description Adds an optional `next_token` input parameter to the `X.SearchRecentTweetsByUsername` and `X.SearchRecentTweetsByKeywords` tools. This allows users to paginate through tweets. A `next_token` is provided in the tools's response. For example, to access the `next_token` when using the `tools.execute`, you can do `next_token = response.output.value["meta"].get("next_token", None)` and then pass it to the tool on your next call through the tools' `next_token` input parameter.
ArcadeAI · Dec 10, 2024 · 00d5bab · 00d5bab
1 parent 0344bc7
commit 00d5bab
Show file tree

Hide file tree

Showing 3 changed files with 95 additions and 7 deletions.
diff --git a/toolkits/x/arcade_x/tools/tweets.py b/toolkits/x/arcade_x/tools/tweets.py
@@ -1,4 +1,4 @@
-from typing import Annotated, Any
+from typing import Annotated, Any, Optional
 
 import httpx
 from arcade.sdk import ToolContext, tool
@@ -10,6 +10,7 @@
     get_headers_with_token,
     get_tweet_url,
     parse_search_recent_tweets_response,
+    remove_none_values,
 )
 
 TWEETS_URL = "https://api.x.com/2/tweets"
@@ -63,17 +64,25 @@ async def search_recent_tweets_by_username(
     context: ToolContext,
     username: Annotated[str, "The username of the X (Twitter) user to look up"],
     max_results: Annotated[
-        int, "The maximum number of results to return. Cannot be less than 10"
+        int, "The maximum number of results to return. Must be in range [1, 100] inclusive"
     ] = 10,
+    next_token: Annotated[
+        Optional[str], "The pagination token starting from which to return results"
+    ] = None,
 ) -> Annotated[dict[str, Any], "Dictionary containing the search results"]:
     """Search for recent tweets (last 7 days) on X (Twitter) by username.
     Includes replies and reposts."""
 
     headers = get_headers_with_token(context)
     params: dict[str, int | str] = {
         "query": f"from:{username}",
-        "max_results": max(max_results, 10),  # X API does not allow 'max_results' less than 10
+        "max_results": min(
+            max(max_results, 10), 100
+        ),  # X API does not allow 'max_results' less than 10 or greater than 100
+        "next_token": next_token,
     }
+    params = remove_none_values(params)
+
     url = (
         "https://api.x.com/2/tweets/search/recent?"
         "expansions=author_id&user.fields=id,name,username,entities&tweet.fields=entities"
@@ -106,8 +115,11 @@ async def search_recent_tweets_by_keywords(
         list[str] | None, "List of phrases that must be present in the tweet"
     ] = None,
     max_results: Annotated[
-        int, "The maximum number of results to return. Cannot be less than 10"
+        int, "The maximum number of results to return. Must be in range [1, 100] inclusive"
     ] = 10,
+    next_token: Annotated[
+        Optional[str], "The pagination token starting from which to return results"
+    ] = None,
 ) -> Annotated[dict[str, Any], "Dictionary containing the search results"]:
     """
     Search for recent tweets (last 7 days) on X (Twitter) by required keywords and phrases.
@@ -131,8 +143,13 @@ async def search_recent_tweets_by_keywords(
 
     params: dict[str, int | str] = {
         "query": query.strip(),
-        "max_results": max(max_results, 10),  # X API does not allow 'max_results' less than 10
+        "max_results": min(
+            max(max_results, 10), 100
+        ),  # X API does not allow 'max_results' less than 10 or greater than 100
+        "next_token": next_token,
     }
+    params = remove_none_values(params)
+
     url = (
         "https://api.x.com/2/tweets/search/recent?"
         "expansions=author_id&user.fields=id,name,username,entities&tweet.fields=entities"

diff --git a/toolkits/x/arcade_x/tools/utils.py b/toolkits/x/arcade_x/tools/utils.py
@@ -28,7 +28,7 @@ def parse_search_recent_tweets_response(response_data: dict[str, Any]) -> dict[s
     Returns the modified response data with added 'tweet_url', 'author_username', and 'author_name'.
     """
     if not sanity_check_tweets_data(response_data):
-        return {"data": []}
+        return {"data": [], "next_token": ""}
 
     # Add 'tweet_url' to each tweet
     for tweet in response_data["data"]:
@@ -110,3 +110,16 @@ def expand_urls_in_user_url(user_data: dict, delete_entities: bool = True) -> di
     if delete_entities:
         new_user_data.pop("entities", None)
     return new_user_data
+
+
+def remove_none_values(params: dict) -> dict:
+    """
+    Remove key/value pairs with None values from a dictionary.
+
+    Args:
+        params: The dictionary to clean
+
+    Returns:
+        A new dictionary with None values removed
+    """
+    return {k: v for k, v in params.items() if v is not None}
diff --git a/toolkits/x/evals/eval_x_tools.py b/toolkits/x/evals/eval_x_tools.py
@@ -26,6 +26,34 @@
 # Register the X tools
 catalog.add_module(arcade_x)
 
+search_recent_tweets_by_username_history = [
+    {"role": "user", "content": "list 1 tweet from elonmusk"},
+    {
+        "role": "assistant",
+        "content": "",
+        "tool_calls": [
+            {
+                "id": "call_kineaPbYCAof3n6qCwnYSKBb",
+                "type": "function",
+                "function": {
+                    "name": "X_SearchRecentTweetsByUsername",
+                    "arguments": '{"max_results":1,"username":"elonmusk"}',
+                },
+            }
+        ],
+    },
+    {
+        "role": "tool",
+        "content": '{"data":[{"author_id":"44196397","author_name":"Elon Musk","author_username":"elonmusk","edit_history_tweet_ids":["1866572304320466985"],"id":"1866572304320466985","text":"RT @chamath: Meanwhile the State of California is going to spend almost double this ($35B) to build a 171 mile stretch of rail between Merc…","tweet_url":"https://x.com/x/status/1866572304320466985"},{"author_id":"44196397","edit_history_tweet_ids":["1866571568266219998"],"id":"1866571568266219998","text":"This is awesome 🚀🇺🇸 https://twitter.com/cb_doge/status/1866565984502550905","tweet_url":"https://x.com/x/status/1866571568266219998"},{"author_id":"44196397","edit_history_tweet_ids":["1866571416969285954"],"id":"1866571416969285954","text":"@ajtourville @Tesla I’ve always felt that the climate predictions were too pessimistic and bound to backfire. \\n\\nExtreme environmentalists can’t say ridiculous things like the world is doomed in 5 years, because 5 years goes by, the world is ok and they lose credibility. \\n\\nIf we transition to… https://x.com/i/web/status/1866571416969285954","tweet_url":"https://x.com/x/status/1866571416969285954"},{"author_id":"44196397","edit_history_tweet_ids":["1866569957309603946"],"id":"1866569957309603946","text":"@shaunmmaguire Yes, please. This is gone on for too long. Enough.","tweet_url":"https://x.com/x/status/1866569957309603946"},{"author_id":"44196397","edit_history_tweet_ids":["1866569078539948491"],"id":"1866569078539948491","text":"@FatEmperor 😂","tweet_url":"https://x.com/x/status/1866569078539948491"},{"author_id":"44196397","edit_history_tweet_ids":["1866554579925577793"],"id":"1866554579925577793","text":"@cb_doge I’m not buying or building a house anywhere","tweet_url":"https://x.com/x/status/1866554579925577793"},{"author_id":"44196397","edit_history_tweet_ids":["1866536009833361915"],"id":"1866536009833361915","text":"RT @amuse: http://x.com/i/article/1866500805211123713","tweet_url":"https://x.com/x/status/1866536009833361915"},{"author_id":"44196397","edit_history_tweet_ids":["1866535704924483739"],"id":"1866535704924483739","text":"@benshapiro 😂","tweet_url":"https://x.com/x/status/1866535704924483739"},{"author_id":"44196397","edit_history_tweet_ids":["1866535550632550854"],"id":"1866535550632550854","text":"@AutismCapital 😂","tweet_url":"https://x.com/x/status/1866535550632550854"},{"author_id":"44196397","edit_history_tweet_ids":["1866535352024043804"],"id":"1866535352024043804","text":"@JDVance Yes","tweet_url":"https://x.com/x/status/1866535352024043804"}],"includes":{"users":[{"id":"44196397","name":"Elon Musk","username":"elonmusk"}]},"meta":{"newest_id":"1866572304320466985","next_token":"b26v89c19zqg8o3frr3tekall7a7ooom3sctaw30rz62l","oldest_id":"1866535352024043804","result_count":10}}',  # noqa: E501, RUF001
+        "tool_call_id": "call_kineaPbYCAof3n6qCwnYSKBb",
+        "name": "X_SearchRecentTweetsByUsername",
+    },
+    {
+        "role": "assistant",
+        "content": 'Here is a recent tweet from Elon Musk: \n\n"This is awesome 🚀🇺🇸" - [Tweet link](https://x.com/x/status/1866571568266219998)',
+    },
+]
+
 
 @tool_eval()
 def x_eval_suite() -> EvalSuite:
@@ -96,6 +124,36 @@ def x_eval_suite() -> EvalSuite:
         ],
     )
 
+    suite.add_case(
+        name="Search recent tweets by username with history",
+        user_message="Get the next 42",
+        additional_messages=search_recent_tweets_by_username_history,
+        expected_tool_calls=[
+            (
+                search_recent_tweets_by_username,
+                {
+                    "username": "elonmusk",
+                    "max_results": 42,
+                    "next_token": "b26v89c19zqg8o3frr3tekall7a7ooom3sctaw30rz62l",
+                },
+            )
+        ],
+        critics=[
+            BinaryCritic(
+                critic_field="username",
+                weight=0.2,
+            ),
+            BinaryCritic(
+                critic_field="max_results",
+                weight=0.2,
+            ),
+            BinaryCritic(
+                critic_field="next_token",
+                weight=0.6,
+            ),
+        ],
+    )
+
     suite.add_case(
         name="Lookup user by username",
         user_message="Can you get information about the user '@jack'?",
@@ -121,7 +179,7 @@ def x_eval_suite() -> EvalSuite:
             (
                 search_recent_tweets_by_keywords,
                 {
-                    "keywords": [],
+                    "keywords": None,
                     "phrases": ["Arcade AI"],
                     "max_results": 10,
                 },