Update tools/web_search.py

tools/web_search.py

@@ -1,27 +1,63 @@
 from typing import Any, Optional
 from smolagents.tools import Tool
-import duckduckgo_search
+# import duckduckgo_search
+import requests
+from xml.etree import ElementTree
 
-class DuckDuckGoSearchTool(Tool):
-    name = "web_search"
-    description = "Performs a duckduckgo web search based on your query (think a Google search) then returns the top search results."
+class PubmedSearchTool(Tool):
+    name = "pubmed_search"
+    description = "Performs a pubmed search based on your query then returns the top search results."
     inputs = {'query': {'type': 'string', 'description': 'The search query to perform.'}}
     output_type = "string"
 
-    def __init__(self, max_results=10, **kwargs):
-        super().__init__()
-        self.max_results = max_results
-        try:
-            from duckduckgo_search import DDGS
-        except ImportError as e:
-            raise ImportError(
-                "You must install package `duckduckgo_search` to run this tool: for instance run `pip install duckduckgo-search`."
-            ) from e
-        self.ddgs = DDGS(**kwargs)
+    def pubmed_search_tool(query: str):
+        """Outil pour rechercher des articles médicaux sur PubMed en fonction d'une requête."""
+        
+        base_url = "https://eutils.ncbi.nlm.nih.gov/entrez/eutils/esearch.fcgi"
+        params = {
+            "db": "pubmed",
+            "term": query,  # Terme de recherche (ex: "diabetes treatment")
+            "retmax": 5,  # Nombre maximum de résultats
+            "retmode": "xml"  # Format XML pour récupérer les résultats
+        }
+        
+        response = requests.get(base_url, params=params)
+        
+        if response.status_code != 200:
+            return "Erreur lors de la recherche sur PubMed."
+    
+        # Parser la réponse XML pour extraire les IDs des articles
+        root = ElementTree.fromstring(response.content)
+        article_ids = [id_elem.text for id_elem in root.findall(".//Id")]
+    
+        if not article_ids:
+            return "Aucun article trouvé pour cette requête."
+    
+        # Construire les liens vers PubMed
+        pubmed_links = [f"https://pubmed.ncbi.nlm.nih.gov/{article_id}/" for article_id in article_ids]
+        
+        return "\n".join(pubmed_links)
 
-    def forward(self, query: str) -> str:
-        results = self.ddgs.text(query, max_results=self.max_results)
-        if len(results) == 0:
-            raise Exception("No results found! Try a less restrictive/shorter query.")
-        postprocessed_results = [f"[{result['title']}]({result['href']})\n{result['body']}" for result in results]
-        return "## Search Results\n\n" + "\n\n".join(postprocessed_results)
+# class DuckDuckGoSearchTool(Tool):
+#     name = "web_search"
+#     description = "Performs a duckduckgo web search based on your query (think a Google search) then returns the top search results."
+#     inputs = {'query': {'type': 'string', 'description': 'The search query to perform.'}}
+#     output_type = "string"
+
+#     def __init__(self, max_results=10, **kwargs):
+#         super().__init__()
+#         self.max_results = max_results
+#         try:
+#             from duckduckgo_search import DDGS
+#         except ImportError as e:
+#             raise ImportError(
+#                 "You must install package `duckduckgo_search` to run this tool: for instance run `pip install duckduckgo-search`."
+#             ) from e
+#         self.ddgs = DDGS(**kwargs)
+
+#     def forward(self, query: str) -> str:
+#         results = self.ddgs.text(query, max_results=self.max_results)
+#         if len(results) == 0:
+#             raise Exception("No results found! Try a less restrictive/shorter query.")
+#         postprocessed_results = [f"[{result['title']}]({result['href']})\n{result['body']}" for result in results]
+#         return "## Search Results\n\n" + "\n\n".join(postprocessed_results)