CodeCow/Scrapegraph-ai
1
0
Fork 0
mirror of https://github.com/VinciGit00/Scrapegraph-ai.git synced 2026-07-12 21:01:56 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

feat(telemetry): add telemetry module

Browse Source
This commit is contained in:
Marco Perini 2024-06-17 13:00:33 +02:00
parent 93342b4708
commit 080a318ff6
22 changed files with 277 additions and 30 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
examples/openai/smart_scraper_openai.py
View File

@ -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

4
scrapegraphai/graphs/abstract_graph.py
View File

@ -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()

60
scrapegraphai/graphs/base_graph.py
View File

@ -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})
self.edges = self._create_edges({e for e in self.raw_edges})

3
scrapegraphai/graphs/csv_scraper_graph.py
View File

@ -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:

3
scrapegraphai/graphs/csv_scraper_multi_graph.py
View File

@ -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:

3
scrapegraphai/graphs/deep_scraper_graph.py
View File

@ -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__
)

3
scrapegraphai/graphs/json_scraper_graph.py
View File

@ -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:

3
scrapegraphai/graphs/json_scraper_multi_graph.py
View File

@ -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:

3
scrapegraphai/graphs/omni_scraper_graph.py
View File

@ -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:

3
scrapegraphai/graphs/omni_search_graph.py
View File

@ -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:

3
scrapegraphai/graphs/pdf_scraper_graph.py
View File

@ -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:

3
scrapegraphai/graphs/pdf_scraper_multi_graph.py
View File

@ -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:

3
scrapegraphai/graphs/script_creator_graph.py
View File

@ -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:

3
scrapegraphai/graphs/script_creator_multi_graph.py
View File

@ -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:

3
scrapegraphai/graphs/search_graph.py
View File

@ -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:

3
scrapegraphai/graphs/smart_scraper_graph.py
View File

@ -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:

3
scrapegraphai/graphs/smart_scraper_multi_graph.py
View File

@ -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:

3
scrapegraphai/graphs/speech_graph.py
View File

@ -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:

3
scrapegraphai/graphs/xml_scraper_graph.py
View File

@ -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:

3
scrapegraphai/graphs/xml_scraper_multi_graph.py
View File

@ -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:

5
scrapegraphai/telemetry/__init__.py Normal file
View File

@ -0,0 +1,5 @@
"""
This module contains the telemetry module for the scrapegraphai package.
"""
from .telemetry import log_graph_execution, log_event, disable_telemetry

183
scrapegraphai/telemetry/telemetry.py Normal file
View File

@ -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