wehub-resource-sync
217 lines
7.5 KiB
Python
217 lines
7.5 KiB
Python
"""
|
|
SmartScraperMultiBatchGraph Module
|
|
|
|
A scraping pipeline that uses the OpenAI Batch API for LLM calls,
|
|
providing 50% cost savings compared to real-time API calls.
|
|
"""
|
|
|
|
import asyncio
|
|
from copy import deepcopy
|
|
from typing import Dict, List, Optional, Type
|
|
|
|
from pydantic import BaseModel
|
|
|
|
from ..nodes import FetchNode, GraphIteratorNode, ParseNode
|
|
from ..nodes.batch_generate_answer_node import BatchGenerateAnswerNode
|
|
from ..nodes.merge_answers_node import MergeAnswersNode
|
|
from ..utils.copy import safe_deepcopy
|
|
from .abstract_graph import AbstractGraph
|
|
from .base_graph import BaseGraph
|
|
from .smart_scraper_graph import SmartScraperGraph
|
|
|
|
|
|
class _FetchParseOnlyGraph(AbstractGraph):
|
|
"""Internal graph that only fetches and parses a URL (no LLM generation).
|
|
|
|
This is used to separate the fetch/parse phase from the LLM generation
|
|
phase, allowing all LLM calls to be batched together.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
prompt: str,
|
|
source: str,
|
|
config: dict,
|
|
schema: Optional[Type[BaseModel]] = None,
|
|
):
|
|
super().__init__(prompt, config, source, schema)
|
|
self.input_key = "url" if source.startswith("http") else "local_dir"
|
|
|
|
def _create_graph(self) -> BaseGraph:
|
|
fetch_node = FetchNode(
|
|
input="url | local_dir",
|
|
output=["doc"],
|
|
node_config={
|
|
"llm_model": self.llm_model,
|
|
"force": self.config.get("force", False),
|
|
"cut": self.config.get("cut", True),
|
|
"loader_kwargs": self.config.get("loader_kwargs", {}),
|
|
"browser_base": self.config.get("browser_base"),
|
|
"scrape_do": self.config.get("scrape_do"),
|
|
"storage_state": self.config.get("storage_state"),
|
|
},
|
|
)
|
|
parse_node = ParseNode(
|
|
input="doc",
|
|
output=["parsed_doc"],
|
|
node_config={
|
|
"llm_model": self.llm_model,
|
|
"chunk_size": self.model_token,
|
|
},
|
|
)
|
|
|
|
return BaseGraph(
|
|
nodes=[fetch_node, parse_node],
|
|
edges=[(fetch_node, parse_node)],
|
|
entry_point=fetch_node,
|
|
graph_name=self.__class__.__name__,
|
|
)
|
|
|
|
def run(self) -> str:
|
|
inputs = {"user_prompt": self.prompt, self.input_key: self.source}
|
|
self.final_state, self.execution_info = self.graph.execute(inputs)
|
|
return self.final_state.get("parsed_doc", "")
|
|
|
|
|
|
class SmartScraperMultiBatchGraph(AbstractGraph):
|
|
"""A scraping pipeline that uses OpenAI Batch API for cost savings.
|
|
|
|
Similar to SmartScraperMultiGraph, but instead of making individual
|
|
LLM calls per URL, it:
|
|
1. Fetches and parses all URLs concurrently (Phase 1)
|
|
2. Collects all prompts and submits them as a single OpenAI Batch (Phase 2)
|
|
3. Polls for batch completion (Phase 3)
|
|
4. Merges all results into a final answer (Phase 4)
|
|
|
|
This provides ~50% cost savings on OpenAI API calls at the expense
|
|
of higher latency (up to 24 hours for batch completion).
|
|
|
|
Attributes:
|
|
prompt (str): The user prompt for scraping.
|
|
source (List[str]): List of URLs to scrape.
|
|
config (dict): Configuration including 'llm' and optional 'batch_api' settings.
|
|
schema (Optional[BaseModel]): Optional Pydantic schema for structured output.
|
|
|
|
Config options under 'batch_api':
|
|
poll_interval (int): Seconds between batch status checks (default: 30).
|
|
max_wait_time (int): Maximum wait time in seconds (default: 86400 = 24h).
|
|
model (str): Override model for batch requests (optional).
|
|
temperature (float): Temperature for batch requests (default: 0.0).
|
|
|
|
Example:
|
|
>>> graph = SmartScraperMultiBatchGraph(
|
|
... prompt="Extract the main topic and key points",
|
|
... source=[
|
|
... "https://example.com/page1",
|
|
... "https://example.com/page2",
|
|
... ],
|
|
... config={
|
|
... "llm": {"model": "openai/gpt-4o-mini"},
|
|
... "batch_api": {
|
|
... "poll_interval": 30,
|
|
... "max_wait_time": 3600,
|
|
... },
|
|
... }
|
|
... )
|
|
>>> result = graph.run()
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
prompt: str,
|
|
source: List[str],
|
|
config: dict,
|
|
schema: Optional[Type[BaseModel]] = None,
|
|
):
|
|
self.copy_config = safe_deepcopy(config)
|
|
self.copy_schema = deepcopy(schema)
|
|
self.batch_config = config.get("batch_api", {})
|
|
|
|
# Validate that the model is OpenAI-based
|
|
model_str = config.get("llm", {}).get("model", "")
|
|
if "/" in model_str:
|
|
provider = model_str.split("/")[0]
|
|
else:
|
|
provider = ""
|
|
if provider and provider != "openai":
|
|
raise ValueError(
|
|
f"SmartScraperMultiBatchGraph only supports OpenAI models. "
|
|
f"Got provider '{provider}'. "
|
|
f"Use SmartScraperMultiGraph for other providers."
|
|
)
|
|
|
|
super().__init__(prompt, config, source, schema)
|
|
|
|
def _create_graph(self) -> BaseGraph:
|
|
"""Creates the graph of nodes for the batch scraping pipeline.
|
|
|
|
The graph has two phases:
|
|
1. GraphIteratorNode runs _FetchParseOnlyGraph per URL (concurrent)
|
|
2. BatchGenerateAnswerNode submits all prompts via Batch API
|
|
3. MergeAnswersNode combines the results
|
|
|
|
Returns:
|
|
BaseGraph: A graph instance representing the batch scraping workflow.
|
|
"""
|
|
# Phase 1: Fetch and parse all URLs concurrently
|
|
graph_iterator_node = GraphIteratorNode(
|
|
input="user_prompt & urls",
|
|
output=["parsed_docs"],
|
|
node_config={
|
|
"graph_instance": _FetchParseOnlyGraph,
|
|
"scraper_config": self.copy_config,
|
|
},
|
|
schema=self.copy_schema,
|
|
)
|
|
|
|
# Phase 2: Submit all prompts to OpenAI Batch API
|
|
batch_generate_node = BatchGenerateAnswerNode(
|
|
input="user_prompt & parsed_docs",
|
|
output=["results"],
|
|
node_config={
|
|
"llm_model": self.llm_model,
|
|
"schema": self.copy_schema,
|
|
"batch_config": self.batch_config,
|
|
},
|
|
)
|
|
|
|
# Phase 3: Merge all results
|
|
merge_answers_node = MergeAnswersNode(
|
|
input="user_prompt & results",
|
|
output=["answer"],
|
|
node_config={
|
|
"llm_model": self.llm_model,
|
|
"schema": self.copy_schema,
|
|
},
|
|
)
|
|
|
|
return BaseGraph(
|
|
nodes=[
|
|
graph_iterator_node,
|
|
batch_generate_node,
|
|
merge_answers_node,
|
|
],
|
|
edges=[
|
|
(graph_iterator_node, batch_generate_node),
|
|
(batch_generate_node, merge_answers_node),
|
|
],
|
|
entry_point=graph_iterator_node,
|
|
graph_name=self.__class__.__name__,
|
|
)
|
|
|
|
def run(self) -> str:
|
|
"""Executes the full batch scraping pipeline.
|
|
|
|
This will:
|
|
1. Fetch and parse all URLs concurrently
|
|
2. Submit all LLM prompts as an OpenAI Batch
|
|
3. Poll until the batch completes (may take minutes to hours)
|
|
4. Merge results into a final answer
|
|
|
|
Returns:
|
|
str: The merged answer from all scraped URLs.
|
|
"""
|
|
inputs = {"user_prompt": self.prompt, "urls": self.source}
|
|
self.final_state, self.execution_info = self.graph.execute(inputs)
|
|
return self.final_state.get("answer", "No answer found.")
|