diff --git a/lib/crewai-tools/src/crewai_tools/tools/arxiv_paper_tool/arxiv_paper_tool.py b/lib/crewai-tools/src/crewai_tools/tools/arxiv_paper_tool/arxiv_paper_tool.py index 56f0e4828..bb3bc25ca 100644 --- a/lib/crewai-tools/src/crewai_tools/tools/arxiv_paper_tool/arxiv_paper_tool.py +++ b/lib/crewai-tools/src/crewai_tools/tools/arxiv_paper_tool/arxiv_paper_tool.py @@ -6,9 +6,10 @@ from typing import ClassVar import urllib.error import urllib.parse import urllib.request +from xml.etree.ElementTree import Element from crewai.tools import BaseTool, EnvVar -import defusedxml.ElementTree as ET +from defusedxml.ElementTree import fromstring from pydantic import BaseModel, ConfigDict, Field @@ -90,7 +91,7 @@ class ArxivPaperTool(BaseTool): logger.error(f"Error fetching data from Arxiv: {e}") raise - root = ET.fromstring(data) + root = fromstring(data) papers = [] for entry in root.findall(self.ATOM_NAMESPACE + "entry"): @@ -121,11 +122,11 @@ class ArxivPaperTool(BaseTool): return papers @staticmethod - def _get_element_text(entry: ET.Element, element_name: str) -> str | None: + def _get_element_text(entry: Element, element_name: str) -> str | None: elem = entry.find(f"{ArxivPaperTool.ATOM_NAMESPACE}{element_name}") return elem.text.strip() if elem is not None and elem.text else None - def _extract_pdf_url(self, entry: ET.Element) -> str | None: + def _extract_pdf_url(self, entry: Element) -> str | None: for link in entry.findall(self.ATOM_NAMESPACE + "link"): if link.attrib.get("title", "").lower() == "pdf": return link.attrib.get("href") diff --git a/lib/crewai-tools/tests/rag/test_xml_loader.py b/lib/crewai-tools/tests/rag/test_xml_loader.py index 1576e84b4..5e2e75426 100644 --- a/lib/crewai-tools/tests/rag/test_xml_loader.py +++ b/lib/crewai-tools/tests/rag/test_xml_loader.py @@ -196,5 +196,4 @@ class TestXMLLoader: """Verify that ArxivPaperTool imports from defusedxml, not xml.etree.ElementTree.""" import crewai_tools.tools.arxiv_paper_tool.arxiv_paper_tool as arxiv_module - ET = arxiv_module.ET - assert "defusedxml" in ET.__name__ + assert "defusedxml" in arxiv_module.fromstring.__module__ diff --git a/lib/crewai-tools/tests/tools/arxiv_paper_tool_test.py b/lib/crewai-tools/tests/tools/arxiv_paper_tool_test.py index f9134b421..36f8210f6 100644 --- a/lib/crewai-tools/tests/tools/arxiv_paper_tool_test.py +++ b/lib/crewai-tools/tests/tools/arxiv_paper_tool_test.py @@ -1,7 +1,7 @@ from pathlib import Path from unittest.mock import MagicMock, patch import urllib.error -import defusedxml.ElementTree as ET +from defusedxml.ElementTree import ParseError from crewai_tools import ArxivPaperTool import pytest @@ -108,7 +108,7 @@ def test_invalid_xml_response(mock_urlopen, tool): mock_response.status = 200 mock_urlopen.return_value.__enter__.return_value = mock_response - with pytest.raises(ET.ParseError): + with pytest.raises(ParseError): tool.fetch_arxiv_data("quantum", 1)