add if and cool stuff

Author: VinciGit00

scrapegraphai/graphs/screenshot_scraper_graph.py

@@ -79,3 +79,4 @@ def run(self) -> str:
         self.final_state, self.execution_info = self.graph.execute(inputs)
 
         return self.final_state.get("answer", "No answer found.")
+        

scrapegraphai/nodes/fetch_screen_node.py

@@ -4,6 +4,7 @@
 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 @@ def execute(self, state: dict) -> dict:
         """
         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()

scrapegraphai/nodes/generate_answer_from_image_node.py

@@ -5,6 +5,7 @@
 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 @@ def execute(self, state: dict) -> dict:
         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 @@ def execute(self, state: dict) -> dict:
                              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 = 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)
+                response_text = result.get('choices',
+                                        [{}])[0].get('message', {}).get('content', 'No response')
+                analyses.append(response_text)
 
-        consolidated_analysis = " ".join(analyses)
+            consolidated_analysis = " ".join(analyses)
 
-        state['answer'] = {
-            "consolidated_analysis": consolidated_analysis
-        }
+            state['answer'] = {
+                "consolidated_analysis": consolidated_analysis
+            }
 
-        return state
+            return state