handling the agent model

Author: VinciGit00

examples/benchmarks/SmartScraper/benchmark_docker.py

@@ -2,7 +2,6 @@
 Basic example of scraping pipeline using SmartScraper from text
 """
 
-import os
 from scrapegraphai.graphs import SmartScraperGraph
 from scrapegraphai.utils import prettify_exec_info
 

scrapegraphai/helpers/robots.py

@@ -5,7 +5,7 @@
 robots_dictionary = {
     "gpt-3.5-turbo": ["GPTBot", "ChatGPT-user"],
     "gpt-4-turbo": ["GPTBot", "ChatGPT-user"],
-    "claude": "Claude-Web",
+    "claude": ["Claude-Web", "ClaudeBot"],
     "perplexity": "PerplexityBot",
     "cohere": "cohere-ai",
     "anthropic": "anthropic-ai"

scrapegraphai/nodes/rag_node.py

@@ -9,9 +9,9 @@
 from langchain_community.document_transformers import EmbeddingsRedundantFilter
 from langchain_community.embeddings import HuggingFaceHubEmbeddings
 from langchain_community.vectorstores import FAISS
+from langchain_community.embeddings import OllamaEmbeddings
 from langchain_openai import OpenAIEmbeddings, AzureOpenAIEmbeddings
 from ..models import OpenAI, Ollama, AzureOpenAI, HuggingFace
-from langchain_community.embeddings import OllamaEmbeddings
 from .base_node import BaseNode
 
 
@@ -97,7 +97,7 @@ def execute(self, state):
             # remove streaming and temperature
             params.pop("streaming", None)
             params.pop("temperature", None)
-            
+
             embeddings = OllamaEmbeddings(**params)
         elif isinstance(embedding_model, HuggingFace):
             embeddings = HuggingFaceHubEmbeddings(model=embedding_model.model)

scrapegraphai/nodes/robots_node.py

@@ -102,11 +102,17 @@ def execute(self, state):
 
             document = loader.load()
 
-            # TODO: look at the agent
-            agent = "TODO"
-            # mandare la richiesta
-            # if errore -> manda l'eccezione
-            # poi faccio un return
+            model = self.llm_model["model"]
+
+            if "ollama" in model:
+                model = model.split("/", maxsplit=1)[-1]
+
+            try:
+                agent = robots_dictionary[model]
+
+            except KeyError:
+                agent = model
+
             prompt = PromptTemplate(
                 template=template,
                 partial_variables={"context": document,