123 lines
5.2 KiB
Python
123 lines
5.2 KiB
Python
# tools.py
|
|
from langchain_core.tools import tool, ToolException
|
|
from db_utils import add_url_or_get_id, update_document_content
|
|
from utils import format_bytes
|
|
import requests
|
|
import io
|
|
import pdfplumber
|
|
from playwright.async_api import async_playwright
|
|
import trafilatura
|
|
from trafilatura.settings import use_config
|
|
|
|
|
|
@tool
|
|
def human_approval(plan: str) -> str:
|
|
"""
|
|
Asks for human approval for a given plan.
|
|
The plan is a string that describes the actions to be taken.
|
|
Returns 'approved' or 'denied'.
|
|
"""
|
|
import logging
|
|
logger = logging.getLogger('KnowledgeAgent')
|
|
|
|
status = f"PROPOSED PLAN:\n{plan}"
|
|
print(f"\n[INFO] {status}")
|
|
logger.info(f"{status}")
|
|
try:
|
|
response = input("Do you approve this plan? (y/n): ").lower()
|
|
if response == 'y':
|
|
status = "Approved."
|
|
print(status)
|
|
logger.info(status)
|
|
return status
|
|
status = "Denied."
|
|
print(status)
|
|
logger.info(status)
|
|
return status
|
|
except Exception as e:
|
|
status = f"Error in human_approval: {e}"
|
|
print(status)
|
|
logger.error(status)
|
|
raise ToolException(status)
|
|
|
|
async def fetch_and_generate_markdown(url: str, logger):
|
|
"""Fetches raw content from a URL and generates markdown using a hybrid approach."""
|
|
raw_document = b''
|
|
markdown_content = ""
|
|
MIN_CONTENT_LENGTH = 200 # Minimum character length to be considered valid content
|
|
|
|
try:
|
|
# Use a HEAD request to check the content type first
|
|
head_response = requests.head(url, timeout=10, headers={'User-Agent': 'Mozilla/5.0'})
|
|
head_response.raise_for_status()
|
|
content_type = head_response.headers.get("Content-Type", "")
|
|
|
|
if "text/html" in content_type:
|
|
# 1. Try Trafilatura first
|
|
logger.info(f"Attempting to extract content with Trafilatura from: {url}")
|
|
config = use_config()
|
|
config.set("DEFAULT", "EXTRACTION_TIMEOUT", "0")
|
|
downloaded_html = trafilatura.fetch_url(url)
|
|
if downloaded_html:
|
|
raw_document = downloaded_html.encode('utf-8')
|
|
markdown_content = trafilatura.extract(
|
|
downloaded_html,
|
|
config=config,
|
|
include_comments=False,
|
|
include_tables=True,
|
|
)
|
|
|
|
# 2. Validate output and fallback to Playwright if necessary
|
|
if not markdown_content or len(markdown_content) < MIN_CONTENT_LENGTH:
|
|
logger.warning(f"Trafilatura extraction failed or content too short. Falling back to Playwright for: {url}")
|
|
async with async_playwright() as p:
|
|
browser = await p.chromium.launch()
|
|
page = await browser.new_page()
|
|
await page.goto(url, wait_until="networkidle", timeout=15000)
|
|
raw_document = (await page.content()).encode('utf-8')
|
|
# Use a robust JS evaluation to get main content text
|
|
markdown_content = await page.evaluate('''() => {
|
|
const main = document.querySelector('main, #main, #content, [role="main"]');
|
|
return main ? main.innerText : document.body.innerText;
|
|
}''')
|
|
await browser.close()
|
|
logger.info(f"Successfully fetched content with Playwright for url: {url}")
|
|
else:
|
|
logger.info(f"Successfully extracted content with Trafilatura for url: {url}")
|
|
|
|
elif "application/pdf" in content_type:
|
|
logger.info(f"Downloading PDF content from: {url}")
|
|
response = requests.get(url, timeout=10, headers={'User-Agent': 'Mozilla/5.0'})
|
|
response.raise_for_status()
|
|
raw_document = response.content
|
|
with pdfplumber.open(io.BytesIO(raw_document)) as pdf:
|
|
markdown_content = "\n".join(page.extract_text() for page in pdf.pages if page.extract_text())
|
|
logger.info(f"Successfully processed PDF for url: {url}")
|
|
|
|
else:
|
|
markdown_content = f"[MARKDOWN_GENERATION_FAILED: Unsupported content type '{content_type}']"
|
|
|
|
except Exception as e:
|
|
logger.error(f"An unexpected error occurred while processing {url}: {e}", exc_info=True)
|
|
markdown_content = f"[MARKDOWN_GENERATION_FAILED: {e}]"
|
|
|
|
return raw_document, markdown_content
|
|
|
|
async def process_url(url: str, logger):
|
|
"""Downloads, processes, and stores content from a URL."""
|
|
url_id, url_status = add_url_or_get_id(url)
|
|
if url_status == "existing":
|
|
# Optionally, we could check here if the content is missing and re-process if needed
|
|
return url_id, url_status
|
|
|
|
raw_document, markdown_content = await fetch_and_generate_markdown(url, logger)
|
|
|
|
if raw_document or markdown_content:
|
|
logger.info(f"Updating document content for url_id: {url_id}")
|
|
try:
|
|
update_document_content(url_id, raw_document, markdown_content)
|
|
logger.info(f"Successfully updated document content for url_id: {url_id}")
|
|
except Exception as e:
|
|
logger.error(f"Failed to update document content for url_id {url_id}: {e}", exc_info=True)
|
|
|
|
return url_id, url_status |