feat : managing search credits display.

agent.py

@@ -25,7 +25,7 @@ import os
 from smolagents import CodeAgent, InferenceClientModel
 
 from other_tools import ImageQueryTool
-from web_tools import TavilySearchTool, TavilyExtractTool, TavilyImageURLSearchTool
+from web_tools import TavilyBaseClient, TavilySearchTool, TavilyExtractTool, TavilyImageURLSearchTool
 
 class SmolAlbert(CodeAgent):
     """
@@ -81,3 +81,10 @@ class SmolAlbert(CodeAgent):
         Reset the agent's internal state.
         """
         self.agent.memory.reset()
+
+    @staticmethod
+    def get_search_credits() -> str:
+        """
+        Get the current search credits of the Tavily API.
+        """
+        return TavilyBaseClient.get_usage()

agent_ui.py

@@ -42,6 +42,8 @@ from smolagents.agents import MultiStepAgent, PlanningStep
 from smolagents.memory import ActionStep, FinalAnswerStep
 from smolagents.models import ChatMessageStreamDelta, MessageRole, agglomerate_stream_deltas
 
+FINAL_ANSWER_TAG = "Final answer:"
+
 def get_step_footnote_content(step_log: ActionStep | PlanningStep, step_name: str) -> str:
     """Get a footnote string for a step log with duration and token information"""
     step_footnote = f"**{step_name}**"
@@ -218,7 +220,7 @@ def _process_final_answer_step(step_log: FinalAnswerStep) -> Generator:
     if isinstance(final_answer, AgentText):
         yield gr.ChatMessage(
             role=MessageRole.ASSISTANT,
-            content=f"**Final answer:**\n{final_answer.to_string()}\n",
+            content=f"**{FINAL_ANSWER_TAG}**\n{final_answer.to_string()}\n",
             metadata={"status": "done"},
         )
     elif isinstance(final_answer, AgentImage):
@@ -235,7 +237,7 @@ def _process_final_answer_step(step_log: FinalAnswerStep) -> Generator:
         )
     else:
         yield gr.ChatMessage(
-            role=MessageRole.ASSISTANT, content=f"**Final answer:** {str(final_answer)}", metadata={"status": "done"}
+            role=MessageRole.ASSISTANT, content=f"**{FINAL_ANSWER_TAG}** {str(final_answer)}", metadata={"status": "done"}
         )
 
 
@@ -331,11 +333,11 @@ class AgentUI:
                     content_text = msg.content if isinstance(msg.content, str) else ""
 
                     # Detect final answer messages and append to quiet
-                    # HACK : FinalAnswerStep messages are produced by _process_final_answer_step and use "**Final answer:**" text
-                    if "Final answer:" in content_text:
-                        # Remove everything before and including the "Final answer:" label (and any leading/trailing whitespace/newlines)
+                    # HACK : FinalAnswerStep messages are produced by _process_final_answer_step and use FINAL_ANSWER_TAG
+                    if FINAL_ANSWER_TAG in content_text:
+                        # Remove everything before and including the FINAL_ANSWER_TAG label (and any leading/trailing whitespace/newlines)
                         answer_only = re.sub(
-                            r"(?s)^.*?\*\*Final answer:\*\*\s*[\n]*",  # (?s) allows . to match newlines
+                            rf"(?s)^.*?\*\*{FINAL_ANSWER_TAG}\*\*\s*[\n]*",  # (?s) allows . to match newlines
                             "",
                             content_text,
                             flags=re.IGNORECASE,
@@ -429,10 +431,17 @@ class AgentUI:
         """
         self.create_app().launch(debug=True, share=share, **kwargs)
 
+    def get_tavily_credits(self):
+        """
+        Fetch the Tavily credits.
+        """
+        return self.agent.get_search_credits()
+
     def create_app(self):
         import gradio as gr
 
-        with gr.Blocks(theme="glass", fill_height=True) as agent:
+        # some nice thmes available here: https://huggingface.co/spaces/gradio/theme-gallery
+        with gr.Blocks(theme="JohnSmith9982/small_and_pretty", fill_height=True) as agent:
 
             # Set up states to hold the session information
             stored_query = gr.State("")             # current user query
@@ -454,31 +463,38 @@ class AgentUI:
                     )
                     submit_btn = gr.Button("Submit", variant="primary")
 
+                tavily_credits = gr.Textbox(
+                    label="Tavily Credits",
+                    value=self.get_tavily_credits(),
+                    interactive=False,
+                    container=True,
+                )
+
                 gr.HTML(
                     "<br><br><h4><center>Powered by <a target='_blank' href='https://github.com/huggingface/smolagents'><b>smolagents</b></a></center></h4>"
                 )
 
             with gr.Tab("Quiet", scale=1):
+
                 quiet_chatbot = gr.Chatbot(
-                        label="Agent",
-                        type="messages",
-                        avatar_images=(
-                            None,
-                            "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/smolagents/mascot_smol.png",
-                        ),
-                        resizeable=True,
-                        scale=1,
-                        latex_delimiters=[
-                            {"left": r"$$", "right": r"$$", "display": True},
-                            {"left": r"$", "right": r"$", "display": False},
-                            {"left": r"\[", "right": r"\]", "display": True},
-                            {"left": r"\(", "right": r"\)", "display": False},
-                        ],
-                    )
+                    label="Agent",
+                    type="messages",
+                    avatar_images=(
+                        None,
+                        "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/smolagents/mascot_smol.png",
+                    ),
+                    resizeable=True,
+                    scale=1,
+                    latex_delimiters=[
+                        {"left": r"$$", "right": r"$$", "display": True},
+                        {"left": r"$", "right": r"$", "display": False},
+                        {"left": r"\[", "right": r"\]", "display": True},
+                        {"left": r"\(", "right": r"\)", "display": False},
+                    ],
+                )
 
             with gr.Tab("Chatterbox", scale=1):
 
-                # Main chat interface
                 verbose_chatbot = gr.Chatbot(
                     label="Agent",
                     type="messages",
@@ -505,6 +521,10 @@ class AgentUI:
                 self.interact_with_agent,
                 [stored_query, stored_messages_verbose, stored_messages_quiet],
                 [verbose_chatbot, quiet_chatbot],
+            ).then(
+                self.get_tavily_credits,
+                None,
+                tavily_credits,
             ).then(
                 self.enable_query,
                 None,
@@ -519,6 +539,10 @@ class AgentUI:
                 self.interact_with_agent,
                 [stored_query, stored_messages_verbose, stored_messages_quiet],
                 [verbose_chatbot, quiet_chatbot],
+            ).then(
+                self.get_tavily_credits,
+                None,
+                tavily_credits,
             ).then(
                 self.enable_query,
                 None,

app.py

@@ -23,6 +23,7 @@
 from agent import SmolAlbert
 from agent_ui import AgentUI
 
-agent = SmolAlbert()
-agent_ui = AgentUI(agent)
-agent_ui.launch(share=False)
+if __name__ == "__main__":
+    agent = SmolAlbert()
+    agent_ui = AgentUI(agent)
+    agent_ui.launch(share=False)

web_tools.py

@@ -21,13 +21,32 @@
 # THE SOFTWARE.
 
 import os
+import requests
 
 from smolagents import Tool
 from tavily import TavilyClient
 
-tavily_client = TavilyClient(api_key=os.getenv("TAVILY_API_KEY"))
+class TavilyBaseClient:
+    __api_key = os.getenv("TAVILY_API_KEY")
+    _tavily_client = TavilyClient(api_key=__api_key)
 
-class TavilySearchTool(Tool):
+    @staticmethod
+    def get_usage() -> str:
+        url = "https://api.tavily.com/usage"
+        headers = {
+            "Authorization": f"Bearer {TavilyBaseClient.__api_key}",
+            "Content-Type": "application/json",
+        }
+        res = requests.get(url, headers=headers)
+        res.raise_for_status()
+        
+        account = res.json().get("account", {})
+        plan_usage = account.get("plan_usage")
+        plan_limit = account.get("plan_limit")
+
+        return f"{plan_usage}/{plan_limit}"
+
+class TavilySearchTool(TavilyBaseClient, Tool):
     """
     A tool to perform web searches using the Tavily API.
     """
@@ -42,10 +61,10 @@ class TavilySearchTool(Tool):
     output_type = "string"
 
     def forward(self, query: str):
-        response = tavily_client.search(query)
+        response = self._tavily_client.search(query)
         return response
-    
-class TavilyExtractTool(Tool):
+
+class TavilyExtractTool(TavilyBaseClient, Tool):
     """
     A tool to extract information from web pages using the Tavily API.
     """
@@ -60,10 +79,10 @@ class TavilyExtractTool(Tool):
     output_type = "string"
 
     def forward(self, url: str):
-        response = tavily_client.extract(url)
+        response = self._tavily_client.extract(url)
         return response
 
-class TavilyImageURLSearchTool(Tool):
+class TavilyImageURLSearchTool(TavilyBaseClient, Tool):
     """
     A tool to search for image URLs using the Tavily API.
     """
@@ -78,7 +97,7 @@ class TavilyImageURLSearchTool(Tool):
     output_type = "string"
 
     def forward(self, query: str):
-        response = tavily_client.search(query, include_images=True)
+        response = self._tavily_client.search(query, include_images=True)
         
         images = response.get("images", [])