achilis1 commited on
Commit
a5d74b7
·
verified ·
1 Parent(s): 73f7203

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +51 -57
app.py CHANGED
@@ -1,39 +1,46 @@
1
  from smolagents import CodeAgent, DuckDuckGoSearchTool, HfApiModel, load_tool, tool
2
  import datetime
3
- import requests
4
  import pytz
5
  import yaml
 
6
  from tools.final_answer import FinalAnswerTool
7
  from Gradio_UI import GradioUI
8
- from PIL import Image
9
- import io
10
 
11
  @tool
12
- def web_researcher(query: str, max_results: int = 3) -> str:
13
- """A tool that performs web research using DuckDuckGo search.
14
  Args:
15
  query: The search query string
16
- max_results: Maximum number of results to return (default 3)
17
  Returns:
18
- A formatted string containing search results
19
  """
20
  try:
 
21
  search_results = DuckDuckGoSearchTool().search(query, max_results=max_results)
22
- if not search_results:
23
- return "No results found for your query."
24
 
 
 
 
 
25
  formatted_results = []
26
- for i, result in enumerate(search_results, 1):
 
 
 
 
27
  formatted_results.append(
28
- f"Result {i}:\n"
29
- f"Title: {result.get('title', 'No title')}\n"
30
- f"URL: {result.get('link', 'No URL')}\n"
31
- f"Snippet: {result.get('snippet', 'No description')}\n"
32
- f"{'-'*50}"
33
  )
 
34
  return "\n".join(formatted_results)
 
35
  except Exception as e:
36
- return f"Error performing search: {str(e)}"
37
 
38
  @tool
39
  def get_current_time_in_timezone(timezone: str) -> str:
@@ -48,36 +55,7 @@ def get_current_time_in_timezone(timezone: str) -> str:
48
  except Exception as e:
49
  return f"Error fetching time for timezone '{timezone}': {str(e)}"
50
 
51
- # Custom Final Answer Processor
52
- class CustomFinalAnswerTool:
53
- def __init__(self):
54
- self.tool = FinalAnswerTool()
55
-
56
- def __call__(self, answer):
57
- # Handle different answer types
58
- if hasattr(answer, '_image'): # If it's an image
59
- return self._process_image_answer(answer)
60
- elif isinstance(answer, str): # Regular text answer
61
- return answer
62
- else: # Fallback for other types
63
- return str(answer)
64
-
65
- def _process_image_answer(self, image_answer):
66
- """Process image answers to display properly"""
67
- try:
68
- # Convert the image to displayable format
69
- img = image_answer._image
70
- if isinstance(img, Image.Image):
71
- return img
72
- elif isinstance(img, bytes):
73
- return Image.open(io.BytesIO(img))
74
- return f"Generated image: {image_answer}"
75
- except Exception as e:
76
- return f"Image processing error: {str(e)}"
77
-
78
- # Initialize tools
79
- final_answer = CustomFinalAnswerTool()
80
- image_generation_tool = load_tool("agents-course/text-to-image", trust_remote_code=True)
81
 
82
  model = HfApiModel(
83
  max_tokens=2096,
@@ -86,27 +64,43 @@ model = HfApiModel(
86
  custom_role_conversions=None,
87
  )
88
 
 
 
89
  with open("prompts.yaml", 'r') as stream:
90
  prompt_templates = yaml.safe_load(stream)
91
 
92
  agent = CodeAgent(
93
  model=model,
94
- tools=[final_answer.tool, web_researcher, get_current_time_in_timezone, image_generation_tool],
95
- max_steps=10,
96
- verbosity_level=1,
97
  grammar=None,
98
  planning_interval=None,
99
- name=None,
100
- description=None,
101
  prompt_templates=prompt_templates
102
  )
103
 
104
- # Custom UI processor
105
- class CustomGradioUI(GradioUI):
106
  def process_final_answer(self, final_answer):
107
- """Override to handle different answer types"""
108
- if isinstance(final_answer, Image.Image):
109
- return final_answer
 
 
 
 
 
 
 
 
 
 
 
 
 
 
110
  return super().process_final_answer(final_answer)
111
 
112
- CustomGradioUI(agent).launch()
 
1
  from smolagents import CodeAgent, DuckDuckGoSearchTool, HfApiModel, load_tool, tool
2
  import datetime
 
3
  import pytz
4
  import yaml
5
+ from typing import List, Dict
6
  from tools.final_answer import FinalAnswerTool
7
  from Gradio_UI import GradioUI
 
 
8
 
9
  @tool
10
+ def enhanced_web_search(query: str, max_results: int = 5) -> str:
11
+ """An improved web search tool that provides reliable research capabilities.
12
  Args:
13
  query: The search query string
14
+ max_results: Maximum number of results to return (default 5)
15
  Returns:
16
+ A formatted string containing search results with titles, URLs, and snippets
17
  """
18
  try:
19
+ # Perform the search
20
  search_results = DuckDuckGoSearchTool().search(query, max_results=max_results)
 
 
21
 
22
+ if not search_results or len(search_results) == 0:
23
+ return "No relevant results found. Try refining your search query."
24
+
25
+ # Process and format results
26
  formatted_results = []
27
+ for idx, result in enumerate(search_results[:max_results], 1):
28
+ title = result.get('title', 'No title available').strip()
29
+ url = result.get('link', '').strip()
30
+ snippet = result.get('snippet', 'No description available').strip()
31
+
32
  formatted_results.append(
33
+ f"RESULT {idx}:\n"
34
+ f"Title: {title}\n"
35
+ f"URL: {url}\n"
36
+ f"Summary: {snippet}\n"
37
+ f"{'-'*60}"
38
  )
39
+
40
  return "\n".join(formatted_results)
41
+
42
  except Exception as e:
43
+ return f"Search failed. Error: {str(e)}. Please try again with a different query."
44
 
45
  @tool
46
  def get_current_time_in_timezone(timezone: str) -> str:
 
55
  except Exception as e:
56
  return f"Error fetching time for timezone '{timezone}': {str(e)}"
57
 
58
+ final_answer = FinalAnswerTool()
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
59
 
60
  model = HfApiModel(
61
  max_tokens=2096,
 
64
  custom_role_conversions=None,
65
  )
66
 
67
+ image_generation_tool = load_tool("agents-course/text-to-image", trust_remote_code=True)
68
+
69
  with open("prompts.yaml", 'r') as stream:
70
  prompt_templates = yaml.safe_load(stream)
71
 
72
  agent = CodeAgent(
73
  model=model,
74
+ tools=[final_answer, enhanced_web_search, get_current_time_in_timezone, image_generation_tool],
75
+ max_steps=12, # Increased for more complex research
76
+ verbosity_level=2, # More detailed logging
77
  grammar=None,
78
  planning_interval=None,
79
+ name="Research Assistant",
80
+ description="An AI assistant specialized in web research and information gathering",
81
  prompt_templates=prompt_templates
82
  )
83
 
84
+ # Custom UI with better display handling
85
+ class ResearchGradioUI(GradioUI):
86
  def process_final_answer(self, final_answer):
87
+ """Enhanced output processing for better readability"""
88
+ if isinstance(final_answer, str):
89
+ if final_answer.startswith("RESULT"):
90
+ # Format search results with HTML for better display
91
+ parts = final_answer.split("\n")
92
+ html_output = "<div style='font-family: sans-serif; margin: 10px;'>"
93
+ for part in parts:
94
+ if part.startswith("RESULT"):
95
+ html_output += f"<h3>{part}</h3>"
96
+ elif part.startswith(("Title:", "URL:", "Summary:")):
97
+ html_output += f"<p><b>{part}</b></p>"
98
+ elif part.startswith("-"):
99
+ html_output += "<hr>"
100
+ else:
101
+ html_output += f"<p>{part}</p>"
102
+ html_output += "</div>"
103
+ return html_output
104
  return super().process_final_answer(final_answer)
105
 
106
+ ResearchGradioUI(agent).launch()