Skip to content

Commit d772453

Browse files
committed
Refactor model_name attribute access in llm_model in robots_node.py
- Changed the access of model_name from dictionary-style to attribute-style in llm_model to comply with langchain BaseChatModel. - Updated the conditional and split operations accordingly.
1 parent 9ef73d7 commit d772453

File tree

1 file changed

+7
-11
lines changed

1 file changed

+7
-11
lines changed

scrapegraphai/nodes/robots_node.py

Lines changed: 7 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,6 @@
99
from langchain.prompts import PromptTemplate
1010
from langchain.output_parsers import CommaSeparatedListOutputParser
1111

12-
from .base_node import BaseNode
1312
from langchain.output_parsers import CommaSeparatedListOutputParser
1413
from langchain.prompts import PromptTemplate
1514
from langchain_community.document_loaders import AsyncChromiumLoader
@@ -18,7 +17,6 @@
1817
from ..utils.logging import get_logger
1918
from .base_node import BaseNode
2019

21-
2220
class RobotsNode(BaseNode):
2321
"""
2422
A node responsible for checking if a website is scrapeable or not based on the robots.txt file.
@@ -48,13 +46,14 @@ def __init__(
4846
output: List[str],
4947
node_config: Optional[dict] = None,
5048
node_name: str = "Robots",
51-
5249
):
5350
super().__init__(node_name, "node", input, output, 1)
5451

5552
self.llm_model = node_config["llm_model"]
5653

57-
self.force_scraping = False if node_config is None else node_config.get("force_scraping", False)
54+
self.force_scraping = (
55+
False if node_config is None else node_config.get("force_scraping", False)
56+
)
5857
self.verbose = (
5958
True if node_config is None else node_config.get("verbose", False)
6059
)
@@ -111,14 +110,11 @@ def execute(self, state: dict) -> dict:
111110
base_url = f"{parsed_url.scheme}://{parsed_url.netloc}"
112111
loader = AsyncChromiumLoader(f"{base_url}/robots.txt")
113112
document = loader.load()
114-
if "ollama" in self.llm_model["model_name"]:
115-
self.llm_model["model_name"] = self.llm_model["model_name"].split("/")[
116-
-1
117-
]
118-
model = self.llm_model["model_name"].split("/")[-1]
119-
113+
if "ollama" in self.llm_model.model_name:
114+
self.llm_model.model_name = self.llm_model.model_name.split("/")[-1]
115+
model = self.llm_model.model_name.split("/")[-1]
120116
else:
121-
model = self.llm_model["model_name"]
117+
model = self.llm_model.model_name
122118
try:
123119
agent = robots_dictionary[model]
124120

0 commit comments

Comments
 (0)