refactor: Enhance logging with loguru and standardize code documentation

- Refactor token counting logic in Gradio_UI.py for improved maintainability
- Improve error handling messages in API call retry mechanism
- Update .gitignore

.gitignore

@@ -1,3 +1,7 @@
 .env
 .venv
-uv.lock
+uv.lock
+.python-version
+.gradio
+__pycache__/
+*.pyc

Gradio_UI.py

@@ -124,14 +124,35 @@ def stream_to_gradio(
     total_input_tokens = 0
     total_output_tokens = 0
 
-    for step_log in agent.run(task, stream=True, reset=reset_agent_memory, additional_args=additional_args):
-        # Track tokens if model provides them
-        if hasattr(agent.model, "last_input_token_count"):
+    def has_input_token_count(model):
+        """Checks if the model has an input token counter"""
+        return hasattr(model, "last_input_token_count") and model.last_input_token_count is not None
+    
+    def has_output_token_count(model):
+        """Checks if the model has an output token counter"""
+        return hasattr(model, "last_output_token_count") and model.last_output_token_count is not None
+    
+    def update_token_counters(step):
+        """Updates token counters for the step if possible"""
+        nonlocal total_input_tokens, total_output_tokens
+        
+        if has_input_token_count(agent.model):
+            # Update the total input token counter
             total_input_tokens += agent.model.last_input_token_count
-            total_output_tokens += agent.model.last_output_token_count
-            if isinstance(step_log, ActionStep):
-                step_log.input_token_count = agent.model.last_input_token_count
-                step_log.output_token_count = agent.model.last_output_token_count
+            
+            # Update the output token counter if available
+            if has_output_token_count(agent.model):
+                total_output_tokens += agent.model.last_output_token_count
+            
+            # If the step is an ActionStep, save the counters in it
+            if isinstance(step, ActionStep):
+                step.input_token_count = agent.model.last_input_token_count
+                if has_output_token_count(agent.model):
+                    step.output_token_count = agent.model.last_output_token_count
+    
+    for step_log in agent.run(task, stream=True, reset=reset_agent_memory, additional_args=additional_args):
+        # Update token counters if the model provides them
+        update_token_counters(step_log)
 
         for message in pull_messages_from_step(
             step_log,
@@ -284,7 +305,7 @@ class GradioUI:
                 [stored_messages, text_input],
             ).then(self.interact_with_agent, [stored_messages, chatbot], [chatbot])
 
-        demo.launch(debug=False, share=True, **kwargs)
+        demo.launch(**kwargs)
 
 
 __all__ = ["stream_to_gradio", "GradioUI"]

app.py

@@ -1,5 +1,8 @@
 import os
-import yaml
+import time
+import random
+import traceback
+from loguru import logger
 from Gradio_UI import GradioUI
 from litellm import completion
 from smolagents import (
@@ -9,69 +12,208 @@ from smolagents import (
     LiteLLMModel,
     VisitWebpageTool,
     tool,
+    Tool,
 )
 
+# Setting up logging with loguru - only terminal output
+logger.remove()  # Remove default handlers
+logger.add(
+    lambda msg: print(msg, end=""), 
+    level="INFO", 
+    format="<green>{time:YYYY-MM-DD HH:mm:ss}</green> | <level>{level: <8}</level> | <cyan>{message}</cyan>")
+
+# API key configuration
 os.environ["GEMINI_API_KEY"] = os.getenv("GEMINI_API_KEY")
 
+
+class LyricsSearchTool(Tool):
+    """
+    Uses web search to find song lyrics based on song title and artist name
+
+    The search query should include the song title and artist name. The tool
+    will return the lyrics of the song if found.
+
+    Parameters
+    ----------
+    query : str
+        The search query for finding song lyrics. Should include song title and artist name.
+
+    Returns
+    -------
+    str
+        The lyrics of the song if found, otherwise an empty string.
+    """
+    name = "lyrics_search_tool"
+    description = "Uses web search to find song lyrics based on song title and artist name"
+    inputs = {
+        "query": {
+            "type": "string",
+            "description": "The search query for finding song lyrics. Should include song title and artist name.",
+        }
+    }
+    output_type = "string"
+
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+
+    def forward(self, query: str) -> str:
+        assert isinstance(query, str), "Your search query must be a string"
+        # TODO: Implement lyrics search functionality
+        return "Lyrics search not implemented yet"
+
+
 @tool
 def analyze_lyrics_tool(song_title: str, artist: str, lyrics: str) -> str:
     """
     Performs a deep analysis of the musical track, given its metadata.
-    
+
     Args:
-        song_title: title of the song or music trach.
+        song_title: Title of the song or music track.
         artist: The name of the artist.
         lyrics: The lyrics of the song.
-    
+
     Returns:
         A summary of the song's meaning in English.
     """
 
-    prompt = f'''You are an expert in songs and their meanings. 
+    prompt = f"""You are an expert in songs and their meanings. 
     Summarize the meaning of {song_title} by {artist} and identify 
     key themes based on the lyrics: 
     {lyrics}. 
 
-    Includs deep idea and vibes analysis with explainations 
-    based on references to the exact lines
-    '''
+    Include deep idea and vibes analysis with explanations 
+    based on references to the exact lines.
+    """
+
+    # If the USE_ANTHROPIC environment variable is defined, use the Claude model
+    if os.getenv("USE_ANTHROPIC", "false").lower() == "true":
+        model_to_use = "claude-3-haiku-20240307"
+        logger.info("Using Anthropic model: {} for lyrics analysis", model_to_use)
+    else:
+        model_to_use = "gemini/gemini-2.0-flash"
+        logger.info("Using Gemini model: {} for lyrics analysis", model_to_use)
+
+    # Use the function with retry mechanism
+    logger.info("Analyzing lyrics for song: '{}' by '{}'", song_title, artist)
+    return _make_api_call_with_retry(model_to_use, prompt)
+
+
+# Function with manual implementation of retry mechanism
+def _make_api_call_with_retry(model: str, prompt: str) -> str:
+    """
+    Makes an API call with a retry mechanism for error handling.
 
-    response = completion(
-        model="gemini/gemini-2.0-flash", 
-        messages=[
-            {"role": "user", "content": prompt}
-            ])
+    Args:
+        model: The model identifier to use.
+        prompt: The prompt text to send to the model.
+
+    Returns:
+        The response from the model as a string.
+    """
+    max_attempts = 20
+    base_delay = 10
+    max_delay = 60
+    attempt = 0
+    last_exception = None
     
-    try:
-        lyrics = response.choices[0].message.content.strip()
-        return lyrics
-    except (AttributeError, KeyError, IndexError):
+    while attempt < max_attempts:
         try:
-            lyrics = response['choices'][0]['message']['content'].strip()
-            return lyrics
-        except (AttributeError, KeyError, IndexError):
-            pass
-
-final_answer = FinalAnswerTool()
-model = LiteLLMModel(model_id="gemini/gemini-2.0-flash")
-with open("prompts.yaml", 'r') as stream:
-    prompt_templates = yaml.safe_load(stream)
-
-# Example usage within the agent
-agent = CodeAgent(
-    tools=[
-      FinalAnswerTool(), 
-      DuckDuckGoSearchTool(),
-      VisitWebpageTool(),
-      analyze_lyrics_tool
-    ], 
+            # Add a small random delay to prevent simultaneous requests
+            jitter = random.uniform(0.1, 1.0)
+            time.sleep(jitter)
+            
+            # If this is a retry attempt, add exponential backoff delay
+            if attempt > 0:
+                delay = min(base_delay * (2 ** (attempt - 1)), max_delay)
+                time.sleep(delay)
+            
+            response = completion(
+                model=model,
+                messages=[{"role": "user", "content": prompt}],
+                num_retries=2,  # Built-in retry mechanism of LiteLLM
+            )
+
+            # Try to extract the content from the response
+            try:
+                analysis_result = response.choices[0].message.content.strip()
+                return analysis_result
+            except (AttributeError, KeyError, IndexError):
+                try:
+                    analysis_result = response["choices"][0]["message"]["content"].strip()
+                    return analysis_result
+                except (AttributeError, KeyError, IndexError):
+                    # If we couldn't extract the content, return an error
+                    raise ValueError("Failed to extract content from response")
+                    
+        except (ConnectionError, TimeoutError) as e:
+            last_exception = e
+            logger.warning("API call failed (attempt {}/{}) for model {}: {}. Retrying...", attempt+1, max_attempts, model, str(e))
+            attempt += 1
+            continue
+        except Exception as e:
+            logger.error("Unexpected error: {}", str(e))
+            logger.error(traceback.format_exc())
+            raise  # For other exceptions, we don't retry
+    
+    # If all attempts failed, re-raise the last exception
+    if last_exception:
+        logger.error("All {} attempts failed. Last error: {}", max_attempts, str(last_exception))
+        raise last_exception
+
+
+# TODO: use DuckDuckGoSearchTool to find related information
+# for explanation in case the LLM itself is not confident or doesn't know
+#
+
+# Check if we need to use Anthropic for local testing
+use_anthropic = os.getenv("USE_ANTHROPIC", "false").lower() == "true"
+
+# Configure Anthropic API key if needed
+if use_anthropic:
+    os.environ["ANTHROPIC_API_KEY"] = os.getenv("ANTHROPIC_API_KEY")
+    model = LiteLLMModel(model_id="claude-3-haiku-20240307")
+    logger.info("Using Anthropic Claude model for local testing")
+else:
+    model = LiteLLMModel(model_id="gemini/gemini-2.0-flash")
+    logger.info("Using Gemini model as default")
+
+web_agent = CodeAgent(
     model=model,
-    additional_authorized_imports=['numpy', 'bs4'],
+    tools=[DuckDuckGoSearchTool(), VisitWebpageTool()],
+    name="lyrics_search_agent",
+    description="Browses the web to find original full lyrics and scrape them. Excels at building effective search queries",
+    additional_authorized_imports=["numpy", "bs4"],
     max_steps=22,
-    verbosity_level=0,
-    name="Song Lyrics Analyzer",
-    description="Analyze the meaning of song lyrics and identify key themes based on the lyrics using web search tools and deep lyrics analysis tool.",
-    prompt_templates=prompt_templates
+    verbosity_level=2,
+)
+
+
+analysis_agent = CodeAgent(
+    model=model,
+    tools=[DuckDuckGoSearchTool(), VisitWebpageTool(), analyze_lyrics_tool],
+    name="lyrics_analysis_agent",
+    description="You are a Song Analysis Expert with deep knowledge of music theory, lyrical interpretation, cultural contexts, and music history. Your role is to analyze song lyrics to uncover their deeper meaning, artistic significance, and historical context.",
+    additional_authorized_imports=["numpy", "bs4"],
+    max_steps=50,
+    verbosity_level=2,
 )
 
-GradioUI(agent).launch()
+
+# When using the DuckDuckGoSearchTool, clearly indicate when information comes from external research versus your own knowledge base.
+manager_agent = CodeAgent(
+    model=model,
+    tools=[FinalAnswerTool()],
+    name="manager_agent",
+    description="Manages the search process and coordinates the search and analysis of song lyrics.",
+    managed_agents=[web_agent, analysis_agent],
+    additional_authorized_imports=["json"],
+    planning_interval=5,
+    verbosity_level=2,
+    max_steps=15,
+)
+
+logger.info("Initializing Gradio UI and launching server")
+GradioUI(manager_agent).launch(
+    debug=True, share=False, server_name="127.0.0.1", server_port=3000
+)
+logger.success("Server started successfully")

pyproject.toml

@@ -8,5 +8,6 @@ dependencies = [
     "gradio>=5.20.0",
     "huggingface-hub>=0.29.1",
     "litellm>=1.61.20",
+    "loguru>=0.7.3",
     "smolagents>=1.9.2",
 ]