From 8e3d5deaaa369dc127bdeddbc0193a559a71d35d Mon Sep 17 00:00:00 2001
From: Marco Vinciguerra <mvincig11@gmail.com>
Date: Sun, 18 Aug 2024 19:39:49 +0200
Subject: [PATCH 1/9] add screenshot scraper

---
 examples/openai/screenshot_scraper.py         | 38 ++++++++++
 examples/openai/smart_scraper_openai.py       |  8 +-
 scrapegraphai/graphs/__init__.py              |  1 +
 .../graphs/screenshot_scraper_graph.py        | 71 ++++++++++++++++++
 scrapegraphai/nodes/__init__.py               |  4 +-
 scrapegraphai/nodes/fetch_screen_node.py      | 56 ++++++++++++++
 .../nodes/generate_answer_from_image_node.py  | 74 +++++++++++++++++++
 7 files changed, 248 insertions(+), 4 deletions(-)
 create mode 100644 examples/openai/screenshot_scraper.py
 create mode 100644 scrapegraphai/graphs/screenshot_scraper_graph.py
 create mode 100644 scrapegraphai/nodes/fetch_screen_node.py
 create mode 100644 scrapegraphai/nodes/generate_answer_from_image_node.py

diff --git a/examples/openai/screenshot_scraper.py b/examples/openai/screenshot_scraper.py
new file mode 100644
index 00000000..795dea9d
--- /dev/null
+++ b/examples/openai/screenshot_scraper.py
@@ -0,0 +1,38 @@
+""" 
+Basic example of scraping pipeline using SmartScraper
+"""
+
+import os
+import json
+from dotenv import load_dotenv
+from scrapegraphai.graphs import ScreenshotScraperGraph
+from scrapegraphai.utils import prettify_exec_info
+
+load_dotenv()
+
+# ************************************************
+# Define the configuration for the graph
+# ************************************************
+
+
+graph_config = {
+    "llm": {
+        "api_key": os.getenv("OPENAI_API_KEY"),
+        "model": "gpt-4o",
+    },
+    "verbose": True,
+    "headless": False,
+}
+
+# ************************************************
+# Create the ScreenshotScraperGraph instance and run it
+# ************************************************
+
+smart_scraper_graph = ScreenshotScraperGraph(
+    prompt="List me the email of the company",
+    source="https://scrapegraphai.com/",
+    config=graph_config
+)
+
+result = smart_scraper_graph.run()
+print(json.dumps(result, indent=4))
diff --git a/examples/openai/smart_scraper_openai.py b/examples/openai/smart_scraper_openai.py
index 6771b817..119f67e5 100644
--- a/examples/openai/smart_scraper_openai.py
+++ b/examples/openai/smart_scraper_openai.py
@@ -2,10 +2,12 @@
 Basic example of scraping pipeline using SmartScraper
 """
 
-import os, json
+import os
+import json
+from dotenv import load_dotenv
 from scrapegraphai.graphs import SmartScraperGraph
 from scrapegraphai.utils import prettify_exec_info
-from dotenv import load_dotenv
+
 load_dotenv()
 
 # ************************************************
@@ -16,7 +18,7 @@ load_dotenv()
 graph_config = {
     "llm": {
         "api_key": os.getenv("OPENAI_API_KEY"),
-        "model": "gpt-3.5-turbo",
+        "model": "gpt-4o",
     },
     "verbose": True,
     "headless": False,
diff --git a/scrapegraphai/graphs/__init__.py b/scrapegraphai/graphs/__init__.py
index 26a0b9e1..6dda222d 100644
--- a/scrapegraphai/graphs/__init__.py
+++ b/scrapegraphai/graphs/__init__.py
@@ -24,3 +24,4 @@ from .script_creator_multi_graph import ScriptCreatorMultiGraph
 from .markdown_scraper_graph import MDScraperGraph
 from .markdown_scraper_multi_graph import MDScraperMultiGraph
 from .search_link_graph import SearchLinkGraph
+from .screenshot_scraper_graph import ScreenshotScraperGraph
diff --git a/scrapegraphai/graphs/screenshot_scraper_graph.py b/scrapegraphai/graphs/screenshot_scraper_graph.py
new file mode 100644
index 00000000..fb37c03a
--- /dev/null
+++ b/scrapegraphai/graphs/screenshot_scraper_graph.py
@@ -0,0 +1,71 @@
+"""
+ScreenshotScraperGraph Module
+"""
+
+from typing import Optional
+import logging
+from pydantic import BaseModel
+from .base_graph import BaseGraph
+from .abstract_graph import AbstractGraph
+
+from ..nodes import (
+    FetchScreenNode,
+    GenerateAnswerFromImageNode,
+)
+
+class ScreenshotScraperGraph(AbstractGraph):
+    """
+   smart_scraper.run()
+        )
+    """
+
+    def __init__(self, prompt: str, source: str, config: dict, schema: Optional[BaseModel] = None):
+        super().__init__(prompt, config, source, schema)
+
+
+    def _create_graph(self) -> BaseGraph:
+        """
+        Creates the graph of nodes representing the workflow for web scraping.
+
+        Returns:
+            BaseGraph: A graph instance representing the web scraping workflow.
+        """
+        fetch_screen_node = FetchScreenNode(
+            input="url",
+            output=["imgs"],
+            node_config={
+                "link": self.source
+            }
+        )
+        generate_answer_from_image_node = GenerateAnswerFromImageNode(
+            input="doc",
+            output=["parsed_doc"],
+            node_config={
+                "config": self.config
+            }
+        )
+
+        return BaseGraph(
+            nodes=[
+                fetch_screen_node,
+                generate_answer_from_image_node,
+            ],
+            edges=[
+                (fetch_screen_node, generate_answer_from_image_node),
+            ],
+            entry_point=fetch_screen_node,
+            graph_name=self.__class__.__name__
+        )
+
+    def run(self) -> str:
+        """
+        Executes the scraping process and returns the answer to the prompt.
+
+        Returns:
+            str: The answer to the prompt.
+        """
+
+        inputs = {"user_prompt": self.prompt}
+        self.final_state, self.execution_info = self.graph.execute(inputs)
+
+        return self.final_state.get("answer", "No answer found.")
diff --git a/scrapegraphai/nodes/__init__.py b/scrapegraphai/nodes/__init__.py
index 856438cd..dd1c3fcc 100644
--- a/scrapegraphai/nodes/__init__.py
+++ b/scrapegraphai/nodes/__init__.py
@@ -19,4 +19,6 @@ from .generate_answer_pdf_node import GenerateAnswerPDFNode
 from .graph_iterator_node import GraphIteratorNode
 from .merge_answers_node import MergeAnswersNode
 from .generate_answer_omni_node import GenerateAnswerOmniNode
-from .merge_generated_scripts import MergeGeneratedScriptsNode 
+from .merge_generated_scripts import MergeGeneratedScriptsNode
+from .fetch_screen_node import FetchScreenNode
+from .generate_answer_from_image_node import GenerateAnswerFromImageNode
diff --git a/scrapegraphai/nodes/fetch_screen_node.py b/scrapegraphai/nodes/fetch_screen_node.py
new file mode 100644
index 00000000..c869966b
--- /dev/null
+++ b/scrapegraphai/nodes/fetch_screen_node.py
@@ -0,0 +1,56 @@
+from typing import List, Optional
+from playwright.sync_api import sync_playwright
+from .base_node import BaseNode
+
+class FetchScreenNode(BaseNode):
+    """
+    FetchScreenNode captures screenshots from a given URL and stores the image data as bytes.
+    """
+
+    def __init__(
+        self,
+        input: str,
+        output: List[str],
+        node_config: Optional[dict] = None,
+        node_name: str = "FetchScreenNode",
+    ):
+        super().__init__(node_name, "node", input, output, 2, node_config)
+        self.url = node_config.get("link")
+
+    def execute(self, state: dict) -> dict:
+        """Captures screenshots from the input URL and stores them in the state dictionary as bytes."""
+        
+        screenshots = []
+
+        with sync_playwright() as p:
+            browser = p.chromium.launch()
+            page = browser.new_page()
+            page.goto(self.url)
+
+            viewport_height = page.viewport_size["height"]
+
+            # Initialize screenshot counter
+            screenshot_counter = 1
+
+            # List to keep track of screenshot data
+            screenshot_data_list = []
+
+            # Function to capture screenshots
+            def capture_screenshot(scroll_position, counter):
+                page.evaluate(f"window.scrollTo(0, {scroll_position});")
+                screenshot_data = page.screenshot()
+                screenshot_data_list.append(screenshot_data)
+
+            # Capture screenshots
+            capture_screenshot(0, screenshot_counter)  # First screenshot
+            screenshot_counter += 1
+            capture_screenshot(viewport_height, screenshot_counter)  # Second screenshot
+
+            browser.close()
+
+        # Store screenshot data as bytes in the state dictionary
+        for screenshot_data in screenshot_data_list:
+            screenshots.append(screenshot_data)
+        state["link"] = self.url
+        state['screenshots'] = screenshots
+        return state
diff --git a/scrapegraphai/nodes/generate_answer_from_image_node.py b/scrapegraphai/nodes/generate_answer_from_image_node.py
new file mode 100644
index 00000000..8844990b
--- /dev/null
+++ b/scrapegraphai/nodes/generate_answer_from_image_node.py
@@ -0,0 +1,74 @@
+from typing import List, Optional
+from .base_node import BaseNode
+import base64
+import requests
+
+class GenerateAnswerFromImageNode(BaseNode):
+    """
+    GenerateAnswerFromImageNode analyzes images from the state dictionary using the OpenAI API
+    and updates the state with the generated answers.
+    """
+
+    def __init__(
+        self,
+        input: str,
+        output: List[str],
+        node_config: Optional[dict] = None,
+        node_name: str = "GenerateAnswerFromImageNode",
+    ):
+        super().__init__(node_name, "node", input, output, 2, node_config)
+
+    def execute(self, state: dict) -> dict:
+        """Processes images from the state, generates answers, and updates the state."""
+        # Retrieve the image data from the state dictionary
+        images = state.get('screenshots', [])
+        results = []
+
+        # OpenAI API Key
+        for image_data in images:
+            # Encode the image data to base64
+            base64_image = base64.b64encode(image_data).decode('utf-8')
+
+            # Prepare API request
+            headers = {
+                "Content-Type": "application/json",
+                "Authorization": f"Bearer {self.node_config.get("config").get("llm").get("api_key")}"
+            }
+
+            payload = {
+                "model": "gpt-4o-mini",
+                "messages": [
+                    {
+                        "role": "user",
+                        "content": [
+                            {
+                                "type": "text",
+                                "text": state.get("user_prompt", "Extract information from the image")
+                            },
+                            {
+                                "type": "image_url",
+                                "image_url": {
+                                    "url": f"data:image/jpeg;base64,{base64_image}"
+                                }
+                            }
+                        ]
+                    }
+                ],
+                "max_tokens": 300
+            }
+
+            # Make the API request
+            response = requests.post("https://api.openai.com/v1/chat/completions", headers=headers, json=payload)
+            result = response.json()
+
+            # Extract the response text
+            response_text = result.get('choices', [{}])[0].get('message', {}).get('content', 'No response')
+
+            # Append the result to the results list
+            results.append({
+                "analysis": response_text
+            })
+
+        # Update the state dictionary with the results
+        state['answer'] = results
+        return state

From 5eb3cff64f5becf7e107325117364b67b5fe7348 Mon Sep 17 00:00:00 2001
From: Marco Vinciguerra <mvincig11@gmail.com>
Date: Sun, 18 Aug 2024 20:53:35 +0200
Subject: [PATCH 2/9] feat: refactoring of the code

---
 examples/openai/screenshot_scraper.py         |  4 +-
 scrapegraphai/nodes/fetch_screen_node.py      | 20 ++++----
 .../nodes/generate_answer_from_image_node.py  | 49 ++++++++++++-------
 ...n_ollama.py => scrape_json_ollama_test.py} |  0
 tests/graphs/screenshot_scraper_test.py       | 39 +++++++++++++++
 5 files changed, 82 insertions(+), 30 deletions(-)
 rename tests/graphs/{scrape_json_ollama.py => scrape_json_ollama_test.py} (100%)
 create mode 100644 tests/graphs/screenshot_scraper_test.py

diff --git a/examples/openai/screenshot_scraper.py b/examples/openai/screenshot_scraper.py
index 795dea9d..826dcc50 100644
--- a/examples/openai/screenshot_scraper.py
+++ b/examples/openai/screenshot_scraper.py
@@ -29,8 +29,8 @@ graph_config = {
 # ************************************************
 
 smart_scraper_graph = ScreenshotScraperGraph(
-    prompt="List me the email of the company",
-    source="https://scrapegraphai.com/",
+    prompt="List me all the projects",
+    source="https://perinim.github.io/projects/",
     config=graph_config
 )
 
diff --git a/scrapegraphai/nodes/fetch_screen_node.py b/scrapegraphai/nodes/fetch_screen_node.py
index c869966b..1477f4e4 100644
--- a/scrapegraphai/nodes/fetch_screen_node.py
+++ b/scrapegraphai/nodes/fetch_screen_node.py
@@ -1,3 +1,6 @@
+"""
+fetch_screen_node module
+"""
 from typing import List, Optional
 from playwright.sync_api import sync_playwright
 from .base_node import BaseNode
@@ -18,8 +21,10 @@ class FetchScreenNode(BaseNode):
         self.url = node_config.get("link")
 
     def execute(self, state: dict) -> dict:
-        """Captures screenshots from the input URL and stores them in the state dictionary as bytes."""
-        
+        """
+        Captures screenshots from the input URL and stores them in the state dictionary as bytes.
+        """
+
         screenshots = []
 
         with sync_playwright() as p:
@@ -29,28 +34,25 @@ class FetchScreenNode(BaseNode):
 
             viewport_height = page.viewport_size["height"]
 
-            # Initialize screenshot counter
             screenshot_counter = 1
 
-            # List to keep track of screenshot data
             screenshot_data_list = []
 
-            # Function to capture screenshots
             def capture_screenshot(scroll_position, counter):
                 page.evaluate(f"window.scrollTo(0, {scroll_position});")
                 screenshot_data = page.screenshot()
                 screenshot_data_list.append(screenshot_data)
 
-            # Capture screenshots
-            capture_screenshot(0, screenshot_counter)  # First screenshot
+            capture_screenshot(0, screenshot_counter)
             screenshot_counter += 1
-            capture_screenshot(viewport_height, screenshot_counter)  # Second screenshot
+            capture_screenshot(viewport_height, screenshot_counter)
 
             browser.close()
 
-        # Store screenshot data as bytes in the state dictionary
         for screenshot_data in screenshot_data_list:
             screenshots.append(screenshot_data)
+
         state["link"] = self.url
         state['screenshots'] = screenshots
+
         return state
diff --git a/scrapegraphai/nodes/generate_answer_from_image_node.py b/scrapegraphai/nodes/generate_answer_from_image_node.py
index 8844990b..a1d2769b 100644
--- a/scrapegraphai/nodes/generate_answer_from_image_node.py
+++ b/scrapegraphai/nodes/generate_answer_from_image_node.py
@@ -6,7 +6,7 @@ import requests
 class GenerateAnswerFromImageNode(BaseNode):
     """
     GenerateAnswerFromImageNode analyzes images from the state dictionary using the OpenAI API
-    and updates the state with the generated answers.
+    and updates the state with the consolidated answers.
     """
 
     def __init__(
@@ -19,20 +19,28 @@ class GenerateAnswerFromImageNode(BaseNode):
         super().__init__(node_name, "node", input, output, 2, node_config)
 
     def execute(self, state: dict) -> dict:
-        """Processes images from the state, generates answers, and updates the state."""
-        # Retrieve the image data from the state dictionary
+        """
+        Processes images from the state, generates answers, 
+        consolidates the results, and updates the state.
+        """
         images = state.get('screenshots', [])
-        results = []
+        analyses = []
+
+        api_key = self.node_config.get("config", {}).get("llm", {}).get("api_key", "")
+
+        supported_models = ["gpt-4o", "gpt-4o-mini", "gpt-4-turbo"]
+
+        if self.node_config["config"]["llm"]["model"] not in supported_models:
+            raise ValueError(f"""Model '{self.node_config['config']['llm']['model']}' 
+                             is not supported. Supported models are: 
+                             {', '.join(supported_models)}.""")
 
-        # OpenAI API Key
         for image_data in images:
-            # Encode the image data to base64
             base64_image = base64.b64encode(image_data).decode('utf-8')
 
-            # Prepare API request
             headers = {
                 "Content-Type": "application/json",
-                "Authorization": f"Bearer {self.node_config.get("config").get("llm").get("api_key")}"
+                "Authorization": f"Bearer {api_key}"
             }
 
             payload = {
@@ -43,7 +51,8 @@ class GenerateAnswerFromImageNode(BaseNode):
                         "content": [
                             {
                                 "type": "text",
-                                "text": state.get("user_prompt", "Extract information from the image")
+                                "text": state.get("user_prompt", 
+                                                  "Extract information from the image")
                             },
                             {
                                 "type": "image_url",
@@ -57,18 +66,20 @@ class GenerateAnswerFromImageNode(BaseNode):
                 "max_tokens": 300
             }
 
-            # Make the API request
-            response = requests.post("https://api.openai.com/v1/chat/completions", headers=headers, json=payload)
+            response = requests.post("https://api.openai.com/v1/chat/completions",
+                                     headers=headers,
+                                     json=payload,
+                                     timeout=10 )
             result = response.json()
 
-            # Extract the response text
-            response_text = result.get('choices', [{}])[0].get('message', {}).get('content', 'No response')
+            response_text = result.get('choices',
+                                       [{}])[0].get('message', {}).get('content', 'No response')
+            analyses.append(response_text)
 
-            # Append the result to the results list
-            results.append({
-                "analysis": response_text
-            })
+        consolidated_analysis = " ".join(analyses)
+
+        state['answer'] = {
+            "consolidated_analysis": consolidated_analysis
+        }
 
-        # Update the state dictionary with the results
-        state['answer'] = results
         return state
diff --git a/tests/graphs/scrape_json_ollama.py b/tests/graphs/scrape_json_ollama_test.py
similarity index 100%
rename from tests/graphs/scrape_json_ollama.py
rename to tests/graphs/scrape_json_ollama_test.py
diff --git a/tests/graphs/screenshot_scraper_test.py b/tests/graphs/screenshot_scraper_test.py
new file mode 100644
index 00000000..c4f436d2
--- /dev/null
+++ b/tests/graphs/screenshot_scraper_test.py
@@ -0,0 +1,39 @@
+import os
+import pytest
+import json
+from scrapegraphai.graphs import ScreenshotScraperGraph
+from dotenv import load_dotenv
+
+# Load environment variables
+load_dotenv()
+
+# Define a fixture for the graph configuration
+@pytest.fixture
+def graph_config():
+    """
+    Creation of the graph
+    """
+    return {
+        "llm": {
+            "api_key": os.getenv("OPENAI_API_KEY"),
+            "model": "gpt-4o",
+        },
+        "verbose": True,
+        "headless": False,
+    }
+
+def test_screenshot_scraper_graph(graph_config):
+    """
+    test
+    """
+    smart_scraper_graph = ScreenshotScraperGraph(
+        prompt="List me all the projects",
+        source="https://perinim.github.io/projects/",
+        config=graph_config
+    )
+
+    result = smart_scraper_graph.run()
+
+    assert result is not None, "The result should not be None"
+
+    print(json.dumps(result, indent=4))

From 103c21c86ebf0909a968799f3f0b8edb303338fe Mon Sep 17 00:00:00 2001
From: Marco Vinciguerra <mvincig11@gmail.com>
Date: Sun, 18 Aug 2024 20:54:35 +0200
Subject: [PATCH 3/9] Update generate_answer_from_image_node.py

---
 scrapegraphai/nodes/generate_answer_from_image_node.py | 7 +++++--
 1 file changed, 5 insertions(+), 2 deletions(-)

diff --git a/scrapegraphai/nodes/generate_answer_from_image_node.py b/scrapegraphai/nodes/generate_answer_from_image_node.py
index a1d2769b..10cca551 100644
--- a/scrapegraphai/nodes/generate_answer_from_image_node.py
+++ b/scrapegraphai/nodes/generate_answer_from_image_node.py
@@ -1,7 +1,10 @@
-from typing import List, Optional
-from .base_node import BaseNode
+"""
+generate answer from image module
+"""
 import base64
+from typing import List, Optional
 import requests
+from .base_node import BaseNode
 
 class GenerateAnswerFromImageNode(BaseNode):
     """

From c72c077eb6bf78bdc48e45535c96277abe0092bf Mon Sep 17 00:00:00 2001
From: Marco Vinciguerra <mvincig11@gmail.com>
Date: Mon, 19 Aug 2024 01:15:45 +0200
Subject: [PATCH 4/9] refactoring of the nodes

---
 .../graphs/screenshot_scraper_graph.py        | 40 ++++++++++++-------
 scrapegraphai/nodes/fetch_screen_node.py      |  7 +---
 .../nodes/generate_answer_from_image_node.py  |  8 ++--
 3 files changed, 30 insertions(+), 25 deletions(-)

diff --git a/scrapegraphai/graphs/screenshot_scraper_graph.py b/scrapegraphai/graphs/screenshot_scraper_graph.py
index fb37c03a..613c1e3a 100644
--- a/scrapegraphai/graphs/screenshot_scraper_graph.py
+++ b/scrapegraphai/graphs/screenshot_scraper_graph.py
@@ -1,22 +1,32 @@
+""" 
+ScreenshotScraperGraph Module 
 """
-ScreenshotScraperGraph Module
-"""
-
 from typing import Optional
 import logging
 from pydantic import BaseModel
 from .base_graph import BaseGraph
 from .abstract_graph import AbstractGraph
+from ..nodes import ( FetchScreenNode, GenerateAnswerFromImageNode, )
 
-from ..nodes import (
-    FetchScreenNode,
-    GenerateAnswerFromImageNode,
-)
+class ScreenshotScraperGraph(AbstractGraph): 
+    """ 
+    A graph instance representing the web scraping workflow for images.
 
-class ScreenshotScraperGraph(AbstractGraph):
-    """
-   smart_scraper.run()
-        )
+    Attributes:
+        prompt (str): The input text to be scraped.
+        config (dict): Configuration parameters for the graph.
+        source (str): The source URL or image link to scrape from.
+
+    Methods:
+        __init__(prompt: str, source: str, config: dict, schema: Optional[BaseModel] = None)
+            Initializes the ScreenshotScraperGraph instance with the given prompt, 
+            source, and configuration parameters.
+
+        _create_graph()
+            Creates a graph of nodes representing the web scraping workflow for images.
+
+        run()
+            Executes the scraping process and returns the answer to the prompt.
     """
 
     def __init__(self, prompt: str, source: str, config: dict, schema: Optional[BaseModel] = None):
@@ -25,10 +35,10 @@ class ScreenshotScraperGraph(AbstractGraph):
 
     def _create_graph(self) -> BaseGraph:
         """
-        Creates the graph of nodes representing the workflow for web scraping.
+        Creates the graph of nodes representing the workflow for web scraping with images.
 
         Returns:
-            BaseGraph: A graph instance representing the web scraping workflow.
+            BaseGraph: A graph instance representing the web scraping workflow for images.
         """
         fetch_screen_node = FetchScreenNode(
             input="url",
@@ -38,8 +48,8 @@ class ScreenshotScraperGraph(AbstractGraph):
             }
         )
         generate_answer_from_image_node = GenerateAnswerFromImageNode(
-            input="doc",
-            output=["parsed_doc"],
+            input="imgs",
+            output=["answer"],
             node_config={
                 "config": self.config
             }
diff --git a/scrapegraphai/nodes/fetch_screen_node.py b/scrapegraphai/nodes/fetch_screen_node.py
index 1477f4e4..1534e8a3 100644
--- a/scrapegraphai/nodes/fetch_screen_node.py
+++ b/scrapegraphai/nodes/fetch_screen_node.py
@@ -25,8 +25,6 @@ class FetchScreenNode(BaseNode):
         Captures screenshots from the input URL and stores them in the state dictionary as bytes.
         """
 
-        screenshots = []
-
         with sync_playwright() as p:
             browser = p.chromium.launch()
             page = browser.new_page()
@@ -49,10 +47,7 @@ class FetchScreenNode(BaseNode):
 
             browser.close()
 
-        for screenshot_data in screenshot_data_list:
-            screenshots.append(screenshot_data)
-
         state["link"] = self.url
-        state['screenshots'] = screenshots
+        state['screenshots'] = screenshot_data_list
 
         return state
diff --git a/scrapegraphai/nodes/generate_answer_from_image_node.py b/scrapegraphai/nodes/generate_answer_from_image_node.py
index 10cca551..7e17a6f7 100644
--- a/scrapegraphai/nodes/generate_answer_from_image_node.py
+++ b/scrapegraphai/nodes/generate_answer_from_image_node.py
@@ -31,10 +31,10 @@ class GenerateAnswerFromImageNode(BaseNode):
 
         api_key = self.node_config.get("config", {}).get("llm", {}).get("api_key", "")
 
-        supported_models = ["gpt-4o", "gpt-4o-mini", "gpt-4-turbo"]
+        supported_models = ("gpt-4o", "gpt-4o-mini", "gpt-4-turbo")
 
         if self.node_config["config"]["llm"]["model"] not in supported_models:
-            raise ValueError(f"""Model '{self.node_config['config']['llm']['model']}' 
+            raise ValueError(f"""Model '{self.node_config['config']['llm']['model']}'
                              is not supported. Supported models are: 
                              {', '.join(supported_models)}.""")
 
@@ -47,7 +47,7 @@ class GenerateAnswerFromImageNode(BaseNode):
             }
 
             payload = {
-                "model": "gpt-4o-mini",
+                "model": self.node_config["config"]["llm"]["model"],
                 "messages": [
                     {
                         "role": "user",
@@ -72,7 +72,7 @@ class GenerateAnswerFromImageNode(BaseNode):
             response = requests.post("https://api.openai.com/v1/chat/completions",
                                      headers=headers,
                                      json=payload,
-                                     timeout=10 )
+                                     timeout=10)
             result = response.json()
 
             response_text = result.get('choices',

From 79fa3f6bd60ac0d40576466a0f568b4f15d0221e Mon Sep 17 00:00:00 2001
From: Marco Vinciguerra <mvincig11@gmail.com>
Date: Mon, 19 Aug 2024 02:01:41 +0200
Subject: [PATCH 5/9] add if and cool stuff

---
 .../graphs/screenshot_scraper_graph.py        |  1 +
 scrapegraphai/nodes/fetch_screen_node.py      |  2 +
 .../nodes/generate_answer_from_image_node.py  | 90 ++++++++++---------
 3 files changed, 50 insertions(+), 43 deletions(-)

diff --git a/scrapegraphai/graphs/screenshot_scraper_graph.py b/scrapegraphai/graphs/screenshot_scraper_graph.py
index 613c1e3a..e9df554a 100644
--- a/scrapegraphai/graphs/screenshot_scraper_graph.py
+++ b/scrapegraphai/graphs/screenshot_scraper_graph.py
@@ -79,3 +79,4 @@ class ScreenshotScraperGraph(AbstractGraph):
         self.final_state, self.execution_info = self.graph.execute(inputs)
 
         return self.final_state.get("answer", "No answer found.")
+        
\ No newline at end of file
diff --git a/scrapegraphai/nodes/fetch_screen_node.py b/scrapegraphai/nodes/fetch_screen_node.py
index 1534e8a3..0bb71c37 100644
--- a/scrapegraphai/nodes/fetch_screen_node.py
+++ b/scrapegraphai/nodes/fetch_screen_node.py
@@ -4,6 +4,7 @@ fetch_screen_node module
 from typing import List, Optional
 from playwright.sync_api import sync_playwright
 from .base_node import BaseNode
+from ..utils.logging import get_logger
 
 class FetchScreenNode(BaseNode):
     """
@@ -24,6 +25,7 @@ class FetchScreenNode(BaseNode):
         """
         Captures screenshots from the input URL and stores them in the state dictionary as bytes.
         """
+        self.logger.info(f"--- Executing {self.node_name} Node ---")
 
         with sync_playwright() as p:
             browser = p.chromium.launch()
diff --git a/scrapegraphai/nodes/generate_answer_from_image_node.py b/scrapegraphai/nodes/generate_answer_from_image_node.py
index 7e17a6f7..e2dcc617 100644
--- a/scrapegraphai/nodes/generate_answer_from_image_node.py
+++ b/scrapegraphai/nodes/generate_answer_from_image_node.py
@@ -5,6 +5,7 @@ import base64
 from typing import List, Optional
 import requests
 from .base_node import BaseNode
+from ..utils.logging import get_logger
 
 class GenerateAnswerFromImageNode(BaseNode):
     """
@@ -26,6 +27,8 @@ class GenerateAnswerFromImageNode(BaseNode):
         Processes images from the state, generates answers, 
         consolidates the results, and updates the state.
         """
+        self.logger.info(f"--- Executing {self.node_name} Node ---")
+
         images = state.get('screenshots', [])
         analyses = []
 
@@ -38,51 +41,52 @@ class GenerateAnswerFromImageNode(BaseNode):
                              is not supported. Supported models are: 
                              {', '.join(supported_models)}.""")
 
-        for image_data in images:
-            base64_image = base64.b64encode(image_data).decode('utf-8')
+        if self.node_config["config"]["llm"]["model"].startswith("gpt"):
+            for image_data in images:
+                base64_image = base64.b64encode(image_data).decode('utf-8')
 
-            headers = {
-                "Content-Type": "application/json",
-                "Authorization": f"Bearer {api_key}"
-            }
+                headers = {
+                    "Content-Type": "application/json",
+                    "Authorization": f"Bearer {api_key}"
+                }
 
-            payload = {
-                "model": self.node_config["config"]["llm"]["model"],
-                "messages": [
-                    {
-                        "role": "user",
-                        "content": [
-                            {
-                                "type": "text",
-                                "text": state.get("user_prompt", 
-                                                  "Extract information from the image")
-                            },
-                            {
-                                "type": "image_url",
-                                "image_url": {
-                                    "url": f"data:image/jpeg;base64,{base64_image}"
+                payload = {
+                    "model": self.node_config["config"]["llm"]["model"],
+                    "messages": [
+                        {
+                            "role": "user",
+                            "content": [
+                                {
+                                    "type": "text",
+                                    "text": state.get("user_prompt", 
+                                                    "Extract information from the image")
+                                },
+                                {
+                                    "type": "image_url",
+                                    "image_url": {
+                                        "url": f"data:image/jpeg;base64,{base64_image}"
+                                    }
                                 }
-                            }
-                        ]
-                    }
-                ],
-                "max_tokens": 300
+                            ]
+                        }
+                    ],
+                    "max_tokens": 300
+                }
+
+                response = requests.post("https://api.openai.com/v1/chat/completions",
+                                        headers=headers,
+                                        json=payload,
+                                        timeout=10)
+                result = response.json()
+
+                response_text = result.get('choices',
+                                        [{}])[0].get('message', {}).get('content', 'No response')
+                analyses.append(response_text)
+
+            consolidated_analysis = " ".join(analyses)
+
+            state['answer'] = {
+                "consolidated_analysis": consolidated_analysis
             }
 
-            response = requests.post("https://api.openai.com/v1/chat/completions",
-                                     headers=headers,
-                                     json=payload,
-                                     timeout=10)
-            result = response.json()
-
-            response_text = result.get('choices',
-                                       [{}])[0].get('message', {}).get('content', 'No response')
-            analyses.append(response_text)
-
-        consolidated_analysis = " ".join(analyses)
-
-        state['answer'] = {
-            "consolidated_analysis": consolidated_analysis
-        }
-
-        return state
+            return state

From 0bf79b5926e073fa2bd6723bf595b0be292de766 Mon Sep 17 00:00:00 2001
From: Marco Vinciguerra <mvincig11@gmail.com>
Date: Mon, 19 Aug 2024 02:21:45 +0200
Subject: [PATCH 6/9] Update generate_answer_from_image_node.py

---
 scrapegraphai/nodes/generate_answer_from_image_node.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/scrapegraphai/nodes/generate_answer_from_image_node.py b/scrapegraphai/nodes/generate_answer_from_image_node.py
index e2dcc617..7f4aa687 100644
--- a/scrapegraphai/nodes/generate_answer_from_image_node.py
+++ b/scrapegraphai/nodes/generate_answer_from_image_node.py
@@ -32,8 +32,6 @@ class GenerateAnswerFromImageNode(BaseNode):
         images = state.get('screenshots', [])
         analyses = []
 
-        api_key = self.node_config.get("config", {}).get("llm", {}).get("api_key", "")
-
         supported_models = ("gpt-4o", "gpt-4o-mini", "gpt-4-turbo")
 
         if self.node_config["config"]["llm"]["model"] not in supported_models:
@@ -42,6 +40,8 @@ class GenerateAnswerFromImageNode(BaseNode):
                              {', '.join(supported_models)}.""")
 
         if self.node_config["config"]["llm"]["model"].startswith("gpt"):
+            api_key = self.node_config.get("config", {}).get("llm", {}).get("api_key", "")
+
             for image_data in images:
                 base64_image = base64.b64encode(image_data).decode('utf-8')
 

From f60aa3acde3c9bead2250e81eb8fc77d2e1e450c Mon Sep 17 00:00:00 2001
From: Marco Vinciguerra <mvincig11@gmail.com>
Date: Mon, 19 Aug 2024 11:22:40 +0200
Subject: [PATCH 7/9] feat: add async call

---
 .../nodes/generate_answer_from_image_node.py  | 109 ++++++++++--------
 1 file changed, 58 insertions(+), 51 deletions(-)

diff --git a/scrapegraphai/nodes/generate_answer_from_image_node.py b/scrapegraphai/nodes/generate_answer_from_image_node.py
index 7f4aa687..7d145f0e 100644
--- a/scrapegraphai/nodes/generate_answer_from_image_node.py
+++ b/scrapegraphai/nodes/generate_answer_from_image_node.py
@@ -1,9 +1,7 @@
-"""
-generate answer from image module
-"""
 import base64
+import asyncio
 from typing import List, Optional
-import requests
+import aiohttp
 from .base_node import BaseNode
 from ..utils.logging import get_logger
 
@@ -22,10 +20,46 @@ class GenerateAnswerFromImageNode(BaseNode):
     ):
         super().__init__(node_name, "node", input, output, 2, node_config)
 
-    def execute(self, state: dict) -> dict:
+    async def process_image(self, session, api_key, image_data, user_prompt):
+        # Convert image data to base64
+        base64_image = base64.b64encode(image_data).decode('utf-8')
+
+        headers = {
+            "Content-Type": "application/json",
+            "Authorization": f"Bearer {api_key}"
+        }
+
+        payload = {
+            "model": self.node_config["config"]["llm"]["model"],
+            "messages": [
+                {
+                    "role": "user",
+                    "content": [
+                        {
+                            "type": "text",
+                            "text": user_prompt
+                        },
+                        {
+                            "type": "image_url",
+                            "image_url": {
+                                "url": f"data:image/jpeg;base64,{base64_image}"
+                            }
+                        }
+                    ]
+                }
+            ],
+            "max_tokens": 300
+        }
+
+        async with session.post("https://api.openai.com/v1/chat/completions",
+                                headers=headers, json=payload) as response:
+            result = await response.json()
+            return result.get('choices', [{}])[0].get('message', {}).get('content', 'No response')
+
+    async def execute_async(self, state: dict) -> dict:
         """
         Processes images from the state, generates answers, 
-        consolidates the results, and updates the state.
+        consolidates the results, and updates the state asynchronously.
         """
         self.logger.info(f"--- Executing {self.node_name} Node ---")
 
@@ -39,54 +73,27 @@ class GenerateAnswerFromImageNode(BaseNode):
                              is not supported. Supported models are: 
                              {', '.join(supported_models)}.""")
 
-        if self.node_config["config"]["llm"]["model"].startswith("gpt"):
-            api_key = self.node_config.get("config", {}).get("llm", {}).get("api_key", "")
+        api_key = self.node_config.get("config", {}).get("llm", {}).get("api_key", "")
 
-            for image_data in images:
-                base64_image = base64.b64encode(image_data).decode('utf-8')
+        async with aiohttp.ClientSession() as session:
+            tasks = [
+                self.process_image(session, api_key, image_data,
+                                   state.get("user_prompt", "Extract information from the image"))
+                for image_data in images
+            ]
 
-                headers = {
-                    "Content-Type": "application/json",
-                    "Authorization": f"Bearer {api_key}"
-                }
+            analyses = await asyncio.gather(*tasks)
 
-                payload = {
-                    "model": self.node_config["config"]["llm"]["model"],
-                    "messages": [
-                        {
-                            "role": "user",
-                            "content": [
-                                {
-                                    "type": "text",
-                                    "text": state.get("user_prompt", 
-                                                    "Extract information from the image")
-                                },
-                                {
-                                    "type": "image_url",
-                                    "image_url": {
-                                        "url": f"data:image/jpeg;base64,{base64_image}"
-                                    }
-                                }
-                            ]
-                        }
-                    ],
-                    "max_tokens": 300
-                }
+        consolidated_analysis = " ".join(analyses)
 
-                response = requests.post("https://api.openai.com/v1/chat/completions",
-                                        headers=headers,
-                                        json=payload,
-                                        timeout=10)
-                result = response.json()
+        state['answer'] = {
+            "consolidated_analysis": consolidated_analysis
+        }
 
-                response_text = result.get('choices',
-                                        [{}])[0].get('message', {}).get('content', 'No response')
-                analyses.append(response_text)
+        return state
 
-            consolidated_analysis = " ".join(analyses)
-
-            state['answer'] = {
-                "consolidated_analysis": consolidated_analysis
-            }
-
-            return state
+    def execute(self, state: dict) -> dict:
+        """
+        Wrapper to run the asynchronous execute_async function in a synchronous context.
+        """
+        return asyncio.run(self.execute_async(state))

From f774fe40e5ff83b3cd88ef0b6ba64ed99239b8d3 Mon Sep 17 00:00:00 2001
From: Marco Vinciguerra <mvincig11@gmail.com>
Date: Mon, 19 Aug 2024 14:24:27 +0200
Subject: [PATCH 8/9] add try catch and robust integration

---
 .../graphs/screenshot_scraper_graph.py        |  4 ++--
 .../nodes/generate_answer_from_image_node.py  | 20 +++++++++++++++++--
 2 files changed, 20 insertions(+), 4 deletions(-)

diff --git a/scrapegraphai/graphs/screenshot_scraper_graph.py b/scrapegraphai/graphs/screenshot_scraper_graph.py
index e9df554a..13046d93 100644
--- a/scrapegraphai/graphs/screenshot_scraper_graph.py
+++ b/scrapegraphai/graphs/screenshot_scraper_graph.py
@@ -42,14 +42,14 @@ class ScreenshotScraperGraph(AbstractGraph):
         """
         fetch_screen_node = FetchScreenNode(
             input="url",
-            output=["imgs"],
+            output=["screenshots"],
             node_config={
                 "link": self.source
             }
         )
         generate_answer_from_image_node = GenerateAnswerFromImageNode(
             input="imgs",
-            output=["answer"],
+            output=["screenshots"],
             node_config={
                 "config": self.config
             }
diff --git a/scrapegraphai/nodes/generate_answer_from_image_node.py b/scrapegraphai/nodes/generate_answer_from_image_node.py
index 7d145f0e..4cc93d18 100644
--- a/scrapegraphai/nodes/generate_answer_from_image_node.py
+++ b/scrapegraphai/nodes/generate_answer_from_image_node.py
@@ -1,3 +1,6 @@
+"""
+GenerateAnswerFromImageNode Module
+"""
 import base64
 import asyncio
 from typing import List, Optional
@@ -21,7 +24,9 @@ class GenerateAnswerFromImageNode(BaseNode):
         super().__init__(node_name, "node", input, output, 2, node_config)
 
     async def process_image(self, session, api_key, image_data, user_prompt):
-        # Convert image data to base64
+        """
+        async process image
+        """
         base64_image = base64.b64encode(image_data).decode('utf-8')
 
         headers = {
@@ -96,4 +101,15 @@ class GenerateAnswerFromImageNode(BaseNode):
         """
         Wrapper to run the asynchronous execute_async function in a synchronous context.
         """
-        return asyncio.run(self.execute_async(state))
+        try:
+            eventloop = asyncio.get_event_loop()
+        except RuntimeError:
+            eventloop = None
+
+        if eventloop and eventloop.is_running():
+            task = eventloop.create_task(self.execute_async(state))
+            state = eventloop.run_until_complete(asyncio.gather(task))[0]
+        else:
+            state = asyncio.run(self.execute_async(state))
+
+        return state

From fee77d1fe6e9f55d2fe0082274a6d01b952e552f Mon Sep 17 00:00:00 2001
From: Marco Vinciguerra <mvincig11@gmail.com>
Date: Mon, 19 Aug 2024 14:35:13 +0200
Subject: [PATCH 9/9] Update screenshot_scraper_graph.py

---
 scrapegraphai/graphs/screenshot_scraper_graph.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/scrapegraphai/graphs/screenshot_scraper_graph.py b/scrapegraphai/graphs/screenshot_scraper_graph.py
index 13046d93..f3ce608d 100644
--- a/scrapegraphai/graphs/screenshot_scraper_graph.py
+++ b/scrapegraphai/graphs/screenshot_scraper_graph.py
@@ -48,8 +48,8 @@ class ScreenshotScraperGraph(AbstractGraph):
             }
         )
         generate_answer_from_image_node = GenerateAnswerFromImageNode(
-            input="imgs",
-            output=["screenshots"],
+            input="screenshots",
+            output=["answer"],
             node_config={
                 "config": self.config
             }