From 080a318ff68652a3c81a6890cd40fd20c48ac6d0 Mon Sep 17 00:00:00 2001 From: Marco Perini Date: Mon, 17 Jun 2024 13:00:33 +0200 Subject: [PATCH] feat(telemetry): add telemetry module --- examples/openai/smart_scraper_openai.py | 4 +- scrapegraphai/graphs/abstract_graph.py | 4 +- scrapegraphai/graphs/base_graph.py | 60 +++++- scrapegraphai/graphs/csv_scraper_graph.py | 3 +- .../graphs/csv_scraper_multi_graph.py | 3 +- scrapegraphai/graphs/deep_scraper_graph.py | 3 +- scrapegraphai/graphs/json_scraper_graph.py | 3 +- .../graphs/json_scraper_multi_graph.py | 3 +- scrapegraphai/graphs/omni_scraper_graph.py | 3 +- scrapegraphai/graphs/omni_search_graph.py | 3 +- scrapegraphai/graphs/pdf_scraper_graph.py | 3 +- .../graphs/pdf_scraper_multi_graph.py | 3 +- scrapegraphai/graphs/script_creator_graph.py | 3 +- .../graphs/script_creator_multi_graph.py | 3 +- scrapegraphai/graphs/search_graph.py | 3 +- scrapegraphai/graphs/smart_scraper_graph.py | 3 +- .../graphs/smart_scraper_multi_graph.py | 3 +- scrapegraphai/graphs/speech_graph.py | 3 +- scrapegraphai/graphs/xml_scraper_graph.py | 3 +- .../graphs/xml_scraper_multi_graph.py | 3 +- scrapegraphai/telemetry/__init__.py | 5 + scrapegraphai/telemetry/telemetry.py | 183 ++++++++++++++++++ 22 files changed, 277 insertions(+), 30 deletions(-) create mode 100644 scrapegraphai/telemetry/__init__.py create mode 100644 scrapegraphai/telemetry/telemetry.py diff --git a/examples/openai/smart_scraper_openai.py b/examples/openai/smart_scraper_openai.py index e353fd9b..bae4f688 100644 --- a/examples/openai/smart_scraper_openai.py +++ b/examples/openai/smart_scraper_openai.py @@ -2,7 +2,7 @@ Basic example of scraping pipeline using SmartScraper """ -import os +import os, json from dotenv import load_dotenv from scrapegraphai.graphs import SmartScraperGraph from scrapegraphai.utils import prettify_exec_info @@ -37,7 +37,7 @@ smart_scraper_graph = SmartScraperGraph( ) result = smart_scraper_graph.run() -print(result) +print(json.dumps(result, indent=4)) # ************************************************ # Get graph execution info diff --git a/scrapegraphai/graphs/abstract_graph.py b/scrapegraphai/graphs/abstract_graph.py index b5e15e8f..6cd4ac45 100644 --- a/scrapegraphai/graphs/abstract_graph.py +++ b/scrapegraphai/graphs/abstract_graph.py @@ -26,7 +26,7 @@ from ..models import ( OneApi ) from ..models.ernie import Ernie -from ..utils.logging import set_verbosity_debug, set_verbosity_warning +from ..utils.logging import set_verbosity_debug, set_verbosity_warning, set_verbosity_info from ..helpers import models_tokens from ..models import AzureOpenAI, Bedrock, Gemini, Groq, HuggingFace, Ollama, OpenAI, Anthropic, DeepSeek @@ -90,7 +90,7 @@ class AbstractGraph(ABC): verbose = bool(config and config.get("verbose")) if verbose: - set_verbosity_debug() + set_verbosity_info() else: set_verbosity_warning() diff --git a/scrapegraphai/graphs/base_graph.py b/scrapegraphai/graphs/base_graph.py index 1b2cb4da..90585e6a 100644 --- a/scrapegraphai/graphs/base_graph.py +++ b/scrapegraphai/graphs/base_graph.py @@ -1,12 +1,10 @@ -""" -BaseGraph Module -""" - import time import warnings from langchain_community.callbacks import get_openai_callback from typing import Tuple +# Import telemetry functions +from ..telemetry import log_graph_execution, log_event class BaseGraph: """ @@ -46,12 +44,12 @@ class BaseGraph: ... ) """ - def __init__(self, nodes: list, edges: list, entry_point: str, use_burr: bool = False, burr_config: dict = None): - + def __init__(self, nodes: list, edges: list, entry_point: str, use_burr: bool = False, burr_config: dict = None, graph_name: str = "Custom"): self.nodes = nodes self.raw_edges = edges self.edges = self._create_edges({e for e in edges}) self.entry_point = entry_point.node_name + self.graph_name = graph_name self.initial_state = {} if nodes[0].node_name != entry_point.node_name: @@ -103,12 +101,46 @@ class BaseGraph: "total_cost_USD": 0.0, } + start_time = time.time() + error_node = None + source_type = None + llm_model = None + embedder_model = None + while current_node_name: curr_time = time.time() current_node = next(node for node in self.nodes if node.node_name == current_node_name) + # check if there is a "source" key in the node config + if current_node.__class__.__name__ == "FetchNode": + # get the second key name of the state dictionary + source_type = list(state.keys())[1] + # quick fix for local_dir source type + if source_type == "local_dir": + source_type = "html_dir" + + # check if there is an "llm_model" variable in the class + if hasattr(current_node, "llm_model") and llm_model is None: + llm_model = current_node.llm_model + if hasattr(llm_model, "model_name"): + llm_model = llm_model.model_name + elif hasattr(llm_model, "model"): + llm_model = llm_model.model + + # check if there is an "embedder_model" variable in the class + if hasattr(current_node, "embedder_model") and embedder_model is None: + embedder_model = current_node.embedder_model + if hasattr(embedder_model, "model_name"): + embedder_model = embedder_model.model_name + elif hasattr(embedder_model, "model"): + embedder_model = embedder_model.model + with get_openai_callback() as cb: - result = current_node.execute(state) + try: + result = current_node.execute(state) + except Exception as e: + error_node = current_node.node_name + raise e node_exec_time = time.time() - curr_time total_exec_time += node_exec_time @@ -147,6 +179,17 @@ class BaseGraph: "exec_time": total_exec_time, }) + # Log the graph execution telemetry + graph_execution_time = time.time() - start_time + log_graph_execution( + graph_name=self.graph_name, + llm_model=llm_model, + embedder_model=embedder_model, + source_type=source_type, + execution_time=graph_execution_time, + error_node=error_node + ) + return state, exec_info def execute(self, initial_state: dict) -> Tuple[dict, list]: @@ -162,7 +205,6 @@ class BaseGraph: self.initial_state = initial_state if self.use_burr: - from ..integrations import BurrBridge bridge = BurrBridge(self, self.burr_config) @@ -190,4 +232,4 @@ class BaseGraph: # add the node to the list of nodes self.nodes.append(node) # update the edges connecting the last node to the new node - self.edges = self._create_edges({e for e in self.raw_edges}) \ No newline at end of file + self.edges = self._create_edges({e for e in self.raw_edges}) diff --git a/scrapegraphai/graphs/csv_scraper_graph.py b/scrapegraphai/graphs/csv_scraper_graph.py index d8d25b4a..48fb5bdb 100644 --- a/scrapegraphai/graphs/csv_scraper_graph.py +++ b/scrapegraphai/graphs/csv_scraper_graph.py @@ -64,7 +64,8 @@ class CSVScraperGraph(AbstractGraph): (fetch_node, rag_node), (rag_node, generate_answer_node) ], - entry_point=fetch_node + entry_point=fetch_node, + graph_name=self.__class__.__name__ ) def run(self) -> str: diff --git a/scrapegraphai/graphs/csv_scraper_multi_graph.py b/scrapegraphai/graphs/csv_scraper_multi_graph.py index 85ed1727..fd15f49a 100644 --- a/scrapegraphai/graphs/csv_scraper_multi_graph.py +++ b/scrapegraphai/graphs/csv_scraper_multi_graph.py @@ -100,7 +100,8 @@ class CSVScraperMultiGraph(AbstractGraph): edges=[ (graph_iterator_node, merge_answers_node), ], - entry_point=graph_iterator_node + entry_point=graph_iterator_node, + graph_name=self.__class__.__name__ ) def run(self) -> str: diff --git a/scrapegraphai/graphs/deep_scraper_graph.py b/scrapegraphai/graphs/deep_scraper_graph.py index d8d5525f..e9e41771 100644 --- a/scrapegraphai/graphs/deep_scraper_graph.py +++ b/scrapegraphai/graphs/deep_scraper_graph.py @@ -141,7 +141,8 @@ class DeepScraperGraph(AbstractGraph): (search_node, graph_iterator_node), (graph_iterator_node, merge_answers_node) ], - entry_point=fetch_node + entry_point=fetch_node, + graph_name=self.__class__.__name__ ) diff --git a/scrapegraphai/graphs/json_scraper_graph.py b/scrapegraphai/graphs/json_scraper_graph.py index 2dbee471..09a5f02e 100644 --- a/scrapegraphai/graphs/json_scraper_graph.py +++ b/scrapegraphai/graphs/json_scraper_graph.py @@ -89,7 +89,8 @@ class JSONScraperGraph(AbstractGraph): (fetch_node, rag_node), (rag_node, generate_answer_node) ], - entry_point=fetch_node + entry_point=fetch_node, + graph_name=self.__class__.__name__ ) def run(self) -> str: diff --git a/scrapegraphai/graphs/json_scraper_multi_graph.py b/scrapegraphai/graphs/json_scraper_multi_graph.py index f86fdc67..2824c416 100644 --- a/scrapegraphai/graphs/json_scraper_multi_graph.py +++ b/scrapegraphai/graphs/json_scraper_multi_graph.py @@ -104,7 +104,8 @@ class JSONScraperMultiGraph(AbstractGraph): edges=[ (graph_iterator_node, merge_answers_node), ], - entry_point=graph_iterator_node + entry_point=graph_iterator_node, + graph_name=self.__class__.__name__ ) def run(self) -> str: diff --git a/scrapegraphai/graphs/omni_scraper_graph.py b/scrapegraphai/graphs/omni_scraper_graph.py index 3234dd02..a5eefad2 100644 --- a/scrapegraphai/graphs/omni_scraper_graph.py +++ b/scrapegraphai/graphs/omni_scraper_graph.py @@ -122,7 +122,8 @@ class OmniScraperGraph(AbstractGraph): (image_to_text_node, rag_node), (rag_node, generate_answer_omni_node) ], - entry_point=fetch_node + entry_point=fetch_node, + graph_name=self.__class__.__name__ ) def run(self) -> str: diff --git a/scrapegraphai/graphs/omni_search_graph.py b/scrapegraphai/graphs/omni_search_graph.py index d5783729..df525949 100644 --- a/scrapegraphai/graphs/omni_search_graph.py +++ b/scrapegraphai/graphs/omni_search_graph.py @@ -115,7 +115,8 @@ class OmniSearchGraph(AbstractGraph): (search_internet_node, graph_iterator_node), (graph_iterator_node, merge_answers_node) ], - entry_point=search_internet_node + entry_point=search_internet_node, + graph_name=self.__class__.__name__ ) def run(self) -> str: diff --git a/scrapegraphai/graphs/pdf_scraper_graph.py b/scrapegraphai/graphs/pdf_scraper_graph.py index c476e629..41099d8b 100644 --- a/scrapegraphai/graphs/pdf_scraper_graph.py +++ b/scrapegraphai/graphs/pdf_scraper_graph.py @@ -105,7 +105,8 @@ class PDFScraperGraph(AbstractGraph): (parse_node, rag_node), (rag_node, generate_answer_node_pdf) ], - entry_point=fetch_node + entry_point=fetch_node, + graph_name=self.__class__.__name__ ) def run(self) -> str: diff --git a/scrapegraphai/graphs/pdf_scraper_multi_graph.py b/scrapegraphai/graphs/pdf_scraper_multi_graph.py index 60e81bf7..e9b5660b 100644 --- a/scrapegraphai/graphs/pdf_scraper_multi_graph.py +++ b/scrapegraphai/graphs/pdf_scraper_multi_graph.py @@ -105,7 +105,8 @@ class PdfScraperMultiGraph(AbstractGraph): edges=[ (graph_iterator_node, merge_answers_node), ], - entry_point=graph_iterator_node + entry_point=graph_iterator_node, + graph_name=self.__class__.__name__ ) def run(self) -> str: diff --git a/scrapegraphai/graphs/script_creator_graph.py b/scrapegraphai/graphs/script_creator_graph.py index 0697db0b..ce3fa319 100644 --- a/scrapegraphai/graphs/script_creator_graph.py +++ b/scrapegraphai/graphs/script_creator_graph.py @@ -95,7 +95,8 @@ class ScriptCreatorGraph(AbstractGraph): (fetch_node, parse_node), (parse_node, generate_scraper_node), ], - entry_point=fetch_node + entry_point=fetch_node, + graph_name=self.__class__.__name__ ) def run(self) -> str: diff --git a/scrapegraphai/graphs/script_creator_multi_graph.py b/scrapegraphai/graphs/script_creator_multi_graph.py index 1660fd83..2b36f4ed 100644 --- a/scrapegraphai/graphs/script_creator_multi_graph.py +++ b/scrapegraphai/graphs/script_creator_multi_graph.py @@ -99,7 +99,8 @@ class ScriptCreatorMultiGraph(AbstractGraph): edges=[ (graph_iterator_node, merge_scripts_node), ], - entry_point=graph_iterator_node + entry_point=graph_iterator_node, + graph_name=self.__class__.__name__ ) def run(self) -> str: diff --git a/scrapegraphai/graphs/search_graph.py b/scrapegraphai/graphs/search_graph.py index 23d08854..6bece062 100644 --- a/scrapegraphai/graphs/search_graph.py +++ b/scrapegraphai/graphs/search_graph.py @@ -114,7 +114,8 @@ class SearchGraph(AbstractGraph): (search_internet_node, graph_iterator_node), (graph_iterator_node, merge_answers_node) ], - entry_point=search_internet_node + entry_point=search_internet_node, + graph_name=self.__class__.__name__ ) def run(self) -> str: diff --git a/scrapegraphai/graphs/smart_scraper_graph.py b/scrapegraphai/graphs/smart_scraper_graph.py index 0cc6a701..9ee0c3cc 100644 --- a/scrapegraphai/graphs/smart_scraper_graph.py +++ b/scrapegraphai/graphs/smart_scraper_graph.py @@ -104,7 +104,8 @@ class SmartScraperGraph(AbstractGraph): (parse_node, rag_node), (rag_node, generate_answer_node) ], - entry_point=fetch_node + entry_point=fetch_node, + graph_name=self.__class__.__name__ ) def run(self) -> str: diff --git a/scrapegraphai/graphs/smart_scraper_multi_graph.py b/scrapegraphai/graphs/smart_scraper_multi_graph.py index 70fd570a..996beff1 100644 --- a/scrapegraphai/graphs/smart_scraper_multi_graph.py +++ b/scrapegraphai/graphs/smart_scraper_multi_graph.py @@ -104,7 +104,8 @@ class SmartScraperMultiGraph(AbstractGraph): edges=[ (graph_iterator_node, merge_answers_node), ], - entry_point=graph_iterator_node + entry_point=graph_iterator_node, + graph_name=self.__class__.__name__ ) def run(self) -> str: diff --git a/scrapegraphai/graphs/speech_graph.py b/scrapegraphai/graphs/speech_graph.py index 9eb9b44a..1058d127 100644 --- a/scrapegraphai/graphs/speech_graph.py +++ b/scrapegraphai/graphs/speech_graph.py @@ -109,7 +109,8 @@ class SpeechGraph(AbstractGraph): (rag_node, generate_answer_node), (generate_answer_node, text_to_speech_node) ], - entry_point=fetch_node + entry_point=fetch_node, + graph_name=self.__class__.__name__ ) def run(self) -> str: diff --git a/scrapegraphai/graphs/xml_scraper_graph.py b/scrapegraphai/graphs/xml_scraper_graph.py index 2ef5a1c4..dbab0b73 100644 --- a/scrapegraphai/graphs/xml_scraper_graph.py +++ b/scrapegraphai/graphs/xml_scraper_graph.py @@ -91,7 +91,8 @@ class XMLScraperGraph(AbstractGraph): (fetch_node, rag_node), (rag_node, generate_answer_node) ], - entry_point=fetch_node + entry_point=fetch_node, + graph_name=self.__class__.__name__ ) def run(self) -> str: diff --git a/scrapegraphai/graphs/xml_scraper_multi_graph.py b/scrapegraphai/graphs/xml_scraper_multi_graph.py index a9127d5b..e1f4423c 100644 --- a/scrapegraphai/graphs/xml_scraper_multi_graph.py +++ b/scrapegraphai/graphs/xml_scraper_multi_graph.py @@ -105,7 +105,8 @@ class XMLScraperMultiGraph(AbstractGraph): edges=[ (graph_iterator_node, merge_answers_node), ], - entry_point=graph_iterator_node + entry_point=graph_iterator_node, + graph_name=self.__class__.__name__ ) def run(self) -> str: diff --git a/scrapegraphai/telemetry/__init__.py b/scrapegraphai/telemetry/__init__.py new file mode 100644 index 00000000..9586734d --- /dev/null +++ b/scrapegraphai/telemetry/__init__.py @@ -0,0 +1,5 @@ +""" +This module contains the telemetry module for the scrapegraphai package. +""" + +from .telemetry import log_graph_execution, log_event, disable_telemetry \ No newline at end of file diff --git a/scrapegraphai/telemetry/telemetry.py b/scrapegraphai/telemetry/telemetry.py new file mode 100644 index 00000000..73e7c9cb --- /dev/null +++ b/scrapegraphai/telemetry/telemetry.py @@ -0,0 +1,183 @@ +""" +This module contains code that relates to sending ScrapeGraphAI usage telemetry. + +To disable sending telemetry there are three ways: + +1. Set it to false programmatically in your driver: + >>> from scrapegraphai import telemetry + >>> telemetry.disable_telemetry() +2. Set it to `false` in ~/.scrapegraphai.conf under `DEFAULT` + [DEFAULT] + telemetry_enabled = False +3. Set SCRAPEGRAPHAI_TELEMETRY_ENABLED=false as an environment variable: + SCRAPEGRAPHAI_TELEMETRY_ENABLED=false python run.py + or: + export SCRAPEGRAPHAI_TELEMETRY_ENABLED=false +""" + +import configparser +import functools +import importlib.metadata +import json +import os +import platform +import threading +import logging +import uuid +from typing import Callable, Dict +from urllib import request + +VERSION = importlib.metadata.version("scrapegraphai") +STR_VERSION = ".".join([str(i) for i in VERSION]) +HOST = "https://eu.i.posthog.com" +TRACK_URL = f"{HOST}/capture/" # https://posthog.com/docs/api/post-only-endpoints +API_KEY = "phc_orsfU4aHhtpTSLVcUE2hdUkQDLM4OEQZndKGFBKMEtn" +TIMEOUT = 2 +DEFAULT_CONFIG_LOCATION = os.path.expanduser("~/.scrapegraphai.conf") + + +logger = logging.getLogger(__name__) + + +def _load_config(config_location: str) -> configparser.ConfigParser: + config = configparser.ConfigParser() + try: + with open(config_location) as f: + config.read_file(f) + except Exception: + config["DEFAULT"] = {} + else: + if "DEFAULT" not in config: + config["DEFAULT"] = {} + + if "anonymous_id" not in config["DEFAULT"]: + config["DEFAULT"]["anonymous_id"] = str(uuid.uuid4()) + try: + with open(config_location, "w") as f: + config.write(f) + except Exception: + pass + return config + + +def _check_config_and_environ_for_telemetry_flag( + telemetry_default: bool, config_obj: configparser.ConfigParser +) -> bool: + telemetry_enabled = telemetry_default + if "telemetry_enabled" in config_obj["DEFAULT"]: + try: + telemetry_enabled = config_obj.getboolean("DEFAULT", "telemetry_enabled") + except ValueError as e: + logger.debug(f"Unable to parse value for `telemetry_enabled` from config. Encountered {e}") + if os.environ.get("SCRAPEGRAPHAI_TELEMETRY_ENABLED") is not None: + env_value = os.environ.get("SCRAPEGRAPHAI_TELEMETRY_ENABLED") + config_obj["DEFAULT"]["telemetry_enabled"] = env_value + try: + telemetry_enabled = config_obj.getboolean("DEFAULT", "telemetry_enabled") + except ValueError as e: + logger.debug(f"Unable to parse value for `SCRAPEGRAPHAI_TELEMETRY_ENABLED` from environment. Encountered {e}") + return telemetry_enabled + + +config = _load_config(DEFAULT_CONFIG_LOCATION) +g_telemetry_enabled = _check_config_and_environ_for_telemetry_flag(True, config) +g_anonymous_id = config["DEFAULT"]["anonymous_id"] +call_counter = 0 +MAX_COUNT_SESSION = 1000 + +BASE_PROPERTIES = { + "os_type": os.name, + "os_version": platform.platform(), + "python_version": f"{platform.python_version()}/{platform.python_implementation()}", + "distinct_id": g_anonymous_id, + "scrapegraphai_version": VERSION, + "telemetry_version": "0.0.1", +} + + +def disable_telemetry(): + global g_telemetry_enabled + g_telemetry_enabled = False + + +def is_telemetry_enabled() -> bool: + if g_telemetry_enabled: + global call_counter + if call_counter == 0: + logger.debug( + "Note: ScrapeGraphAI collects anonymous usage data to improve the library. " + "You can disable telemetry by setting SCRAPEGRAPHAI_TELEMETRY_ENABLED=false or " + "by editing ~/.scrapegraphai.conf." + ) + call_counter += 1 + if call_counter > MAX_COUNT_SESSION: + return False + return True + else: + return False + + +def _send_event_json(event_json: dict): + headers = { + "Content-Type": "application/json", + "Authorization": f"Bearer {API_KEY}", + "User-Agent": f"scrapegraphai/{STR_VERSION}", + } + try: + data = json.dumps(event_json).encode() + req = request.Request(TRACK_URL, data=data, headers=headers) + with request.urlopen(req, timeout=TIMEOUT) as f: + res = f.read() + if f.code != 200: + raise RuntimeError(res) + except Exception as e: + logger.debug(f"Failed to send telemetry data: {e}") + else: + logger.debug(f"Telemetry data sent: {data}") + + +def send_event_json(event_json: dict): + if not g_telemetry_enabled: + raise RuntimeError("Telemetry tracking is disabled!") + try: + th = threading.Thread(target=_send_event_json, args=(event_json,)) + th.start() + except Exception as e: + logger.debug(f"Failed to send telemetry data in a thread: {e}") + + +def log_event(event: str, properties: Dict[str, any]): + if is_telemetry_enabled(): + event_json = { + "api_key": API_KEY, + "event": event, + "properties": {**BASE_PROPERTIES, **properties}, + } + send_event_json(event_json) + + +def log_graph_execution(graph_name: str, llm_model: str, embedder_model: str, source_type: str, execution_time: float, error_node: str = None): + properties = { + "graph_name": graph_name, + "llm_model": llm_model, + "embedder_model": embedder_model, + "source_type": source_type, + "execution_time": execution_time, + "error_node": error_node, + } + log_event("graph_execution", properties) + + +def capture_function_usage(call_fn: Callable) -> Callable: + @functools.wraps(call_fn) + def wrapped_fn(*args, **kwargs): + try: + return call_fn(*args, **kwargs) + finally: + if is_telemetry_enabled(): + try: + function_name = call_fn.__name__ + log_event("function_usage", {"function_name": function_name}) + except Exception as e: + logger.debug(f"Failed to send telemetry for function usage. Encountered: {e}") + return wrapped_fn \ No newline at end of file