Skip to content

Commit 94bd6c9

Browse files
committed
Refactor email agent, add more tools, cleanup code
1 parent 8f5df2a commit 94bd6c9

File tree

6 files changed

+233
-150
lines changed

6 files changed

+233
-150
lines changed

mxtoai/agents/email_agent.py

Lines changed: 112 additions & 137 deletions
Large diffs are not rendered by default.

mxtoai/routed_litellm_model.py

Lines changed: 3 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -134,11 +134,11 @@ def _get_target_model(self) -> str:
134134
)
135135
return self.current_handle.target_model
136136

137-
return "gpt-4" # Default to gpt-4 model group
137+
return "gpt-4"
138138

139139
def __call__(
140140
self,
141-
messages: list[dict[str, Any]], # MODIFIED type hint for messages
141+
messages: list[dict[str, Any]], # MODIFIED type hint for messages
142142
stop_sequences: Optional[list[str]] = None,
143143
grammar: Optional[str] = None,
144144
tools_to_call_from: Optional[list[Tool]] = None,
@@ -165,9 +165,7 @@ def __call__(
165165
stop_sequences=stop_sequences,
166166
grammar=grammar,
167167
tools_to_call_from=tools_to_call_from,
168-
# Do not pass 'model' as an explicit argument here,
169-
# as self.model_id is now set to our target_model_group.
170-
**kwargs_for_super_generate,
168+
**kwargs_for_super_generate
171169
)
172170
finally:
173171
# Restore the original model_id for the instance.
@@ -176,7 +174,6 @@ def __call__(
176174
return chat_message
177175

178176
except Exception as e:
179-
# Log the error and re-raise with more context
180177
logger.error(f"Error in RoutedLiteLLMModel completion: {e!s}")
181178
msg = f"Failed to get completion from LiteLLM router: {e!s}"
182179
raise RuntimeError(msg) from e

mxtoai/tools/__init__.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,12 @@
11
# Tools package for email processing
22
from mxtoai.tools.attachment_processing_tool import AttachmentProcessingTool
33
from mxtoai.tools.deep_research_tool import DeepResearchTool
4-
from mxtoai.tools.fallback_search_tool import FallbackWebSearchTool
4+
from mxtoai.tools.search_with_fallback_tool import SearchWithFallbackTool
55
from mxtoai.tools.schedule_tool import ScheduleTool
66

77
__all__ = [
88
"AttachmentProcessingTool",
99
"DeepResearchTool",
10-
"FallbackWebSearchTool",
10+
"SearchWithFallbackTool",
1111
"ScheduleTool",
1212
]
Lines changed: 96 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,96 @@
1+
import logging
2+
from typing import Optional, List
3+
4+
from smolagents import Tool
5+
6+
logger = logging.getLogger(__name__)
7+
8+
class SearchWithFallbackTool(Tool):
9+
"""
10+
A web search tool that attempts a sequence of primary search tools
11+
and falls back to another tool if all primary attempts fail.
12+
"""
13+
14+
name = "web_search" # Consistent name for the agent to use
15+
description = (
16+
"Performs a web search using a sequence of search engines. "
17+
"It first attempts searches with primary engines (e.g., Bing, DuckDuckGo). "
18+
"If all primary searches fail or yield no results, it attempts a fallback engine (e.g., Google Search)."
19+
"If everything fails, rephrase the query and try again."
20+
)
21+
inputs = {
22+
"query": {"type": "string", "description": "The search query to perform."},
23+
}
24+
output_type = "string"
25+
26+
def __init__(
27+
self,
28+
primary_search_tools: List[Tool],
29+
fallback_search_tool: Optional[Tool] = None,
30+
):
31+
"""
32+
Initializes the SearchWithFallbackTool.
33+
34+
Args:
35+
primary_search_tools: A list of Tool instances to try in order as primary searchers.
36+
fallback_search_tool: An optional Tool instance to use if all primary tools fail.
37+
"""
38+
if not primary_search_tools and not fallback_search_tool:
39+
raise ValueError("SearchWithFallbackTool requires at least one primary or fallback search tool.")
40+
41+
self.primary_search_tools = primary_search_tools if primary_search_tools else []
42+
if not self.primary_search_tools:
43+
logger.warning(
44+
"SearchWithFallbackTool initialized without any primary search tools. "
45+
"It will only use the fallback tool if available and primary attempts are implicitly skipped."
46+
)
47+
48+
self.fallback_search_tool = fallback_search_tool
49+
super().__init__()
50+
51+
def _get_tool_identifier(self, tool_instance: Tool, default_name: str) -> str:
52+
"""Helper to get a descriptive name for a tool instance for logging."""
53+
base_name = getattr(tool_instance, 'name', default_name)
54+
if hasattr(tool_instance, 'engine'): # Specifically for WebSearchTool
55+
return f"{base_name} (engine: {tool_instance.engine})"
56+
return base_name
57+
58+
def forward(self, query: str) -> str:
59+
"""
60+
Execute the search, trying primary tools first, then the fallback tool.
61+
"""
62+
# Try primary search tools in order
63+
for i, tool_instance in enumerate(self.primary_search_tools):
64+
tool_identifier = self._get_tool_identifier(tool_instance, f"PrimaryTool_{i+1}")
65+
try:
66+
logger.debug(f"Attempting search with primary tool: {tool_identifier}")
67+
result = tool_instance.forward(query=query)
68+
# Underlying smolagents tools typically raise exceptions if no results are found.
69+
# So, a successful return here implies results were found.
70+
logger.info(f"Primary search tool {tool_identifier} succeeded.")
71+
return result
72+
except Exception as e:
73+
logger.warning(
74+
f"Primary search tool {tool_identifier} failed: {e!s}. "
75+
f"Trying next primary tool or fallback."
76+
)
77+
78+
# If all primary tools failed, try the fallback tool
79+
if self.fallback_search_tool:
80+
fallback_tool_instance = self.fallback_search_tool
81+
tool_identifier = self._get_tool_identifier(fallback_tool_instance, "FallbackTool")
82+
try:
83+
logger.debug(f"Attempting search with fallback tool: {tool_identifier}")
84+
result = fallback_tool_instance.forward(query=query)
85+
logger.info(f"Fallback search tool {tool_identifier} succeeded.")
86+
return result
87+
except Exception as e:
88+
logger.error(f"Fallback search tool ({tool_identifier}) also failed: {e!s}")
89+
# Ensure the original exception 'e' from the fallback tool is part of the new exception context
90+
raise SearchFailureException(
91+
f"All primary search tools failed, and the fallback search tool ({tool_identifier}) also failed. Last error: {e!s}"
92+
) from e
93+
else:
94+
logger.error("All primary search tools failed and no fallback tool is configured.")
95+
# It's important to raise an exception here if no tools succeeded and no fallback was available or fallback also failed.
96+
raise SearchFailureException("All primary search tools failed and no fallback tool is configured or the fallback also failed.")

poetry.lock

Lines changed: 19 additions & 5 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

pyproject.toml

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -45,6 +45,7 @@ dependencies = [
4545
"pydantic[email] (>=2.11.4,<3.0.0)",
4646
"python-multipart (>=0.0.20,<0.0.21)",
4747
"toml (>=0.10.2,<0.11.0)",
48+
"wikipedia-api (>=0.8.1,<0.9.0)",
4849
]
4950

5051
[tool.ruff]

0 commit comments

Comments
 (0)