CIAgent/apps/api/app/collectors/website.py

"""Official website collector: sitemap.xml + heuristic page discovery,
robots.txt-respecting, capped crawl depth/page count.
"""

from __future__ import annotations

from datetime import UTC, datetime
from urllib.parse import urljoin
from xml.etree import ElementTree

import httpx

from app.collectors.base import (
    CollectedDocument,
    CollectionResult,
    CompanyContext,
    DiscoveredSource,
    SourceConfig,
)
from app.collectors.extraction import (
    canonicalize_url,
    compute_content_hash,
    extract_readable_text,
    extract_title,
)
from app.collectors.robots import is_allowed
from app.core.config import Settings, get_settings
from app.core.http import FetchError, SsrfBlockedError, fetch_with_retries
from app.core.logging import get_logger
from app.models.enums import SourceStatus, SourceType

logger = get_logger(__name__)

HEURISTIC_PATHS = [
    "",
    "/about",
    "/about-us",
    "/products",
    "/services",
    "/news",
    "/press",
    "/press-releases",
    "/careers",
    "/jobs",
    "/leadership",
    "/team",
    "/investors",
    "/investor-relations",
    "/sustainability",
    "/contact",
]


class WebsiteCollector:
    source_type = SourceType.WEBSITE

    async def discover(self, company: CompanyContext) -> list[DiscoveredSource]:
        if not company.official_website:
            return []
        settings = get_settings()
        pages = await self._discover_pages(company.official_website, settings)
        return [
            DiscoveredSource(
                source_type=SourceType.WEBSITE,
                name=f"{company.name} — Official Website",
                base_url=company.official_website,
                configuration_metadata={"pages": pages},
            )
        ]

    async def _discover_pages(self, base_url: str, settings: Settings) -> list[str]:
        pages: list[str] = []

        sitemap_urls = await self._read_sitemap(base_url, settings)
        pages.extend(sitemap_urls[: settings.max_pages_per_domain])

        if len(pages) < settings.max_pages_per_domain:
            for path in HEURISTIC_PATHS:
                candidate = urljoin(base_url, path)
                if candidate not in pages:
                    pages.append(candidate)
                if len(pages) >= settings.max_pages_per_domain:
                    break

        return pages[: settings.max_pages_per_domain]

    async def _read_sitemap(self, base_url: str, settings: Settings) -> list[str]:
        sitemap_url = urljoin(base_url, "/sitemap.xml")
        try:
            result = await fetch_with_retries(sitemap_url, settings=settings, max_attempts=1)
        except (SsrfBlockedError, FetchError, httpx.HTTPError):
            return []
        if result.status_code != 200:
            return []
        try:
            root = ElementTree.fromstring(result.content)
        except ElementTree.ParseError:
            return []

        ns = {"sm": "http://www.sitemaps.org/schemas/sitemap/0.9"}
        urls = [loc.text.strip() for loc in root.findall(".//sm:url/sm:loc", ns) if loc.text]
        return urls

    async def collect(self, source: SourceConfig, company: CompanyContext) -> CollectionResult:
        settings = get_settings()
        pages: list[str] = source.configuration_metadata.get("pages") or (
            [source.base_url] if source.base_url else []
        )
        pages = pages[: settings.max_pages_per_domain]

        documents: list[CollectedDocument] = []
        seen_hashes: set[str] = set()
        attempted = 0
        any_success = False
        last_error: str | None = None

        for page_url in pages:
            attempted += 1
            try:
                if not await is_allowed(page_url, settings=settings):
                    logger.info("website_collector_robots_disallowed", url=page_url)
                    continue

                result = await fetch_with_retries(page_url, settings=settings)
                if result.status_code == 401 or result.status_code == 403:
                    last_error = f"HTTP {result.status_code} (auth required) for {page_url}"
                    continue
                if result.status_code >= 400:
                    last_error = f"HTTP {result.status_code} for {page_url}"
                    continue

                text, method = extract_readable_text(result.text, page_url)
                if not text:
                    continue
                content_hash = compute_content_hash(text)
                if content_hash in seen_hashes:
                    continue
                seen_hashes.add(content_hash)

                documents.append(
                    CollectedDocument(
                        url=page_url,
                        canonical_url=canonicalize_url(result.final_url),
                        title=extract_title(result.text),
                        author=None,
                        publication_date=None,
                        retrieved_date=datetime.now(UTC),
                        content_text=text,
                        content_hash=content_hash,
                        metadata={"http_status": result.status_code},
                        extraction_method=method,
                        http_status=result.status_code,
                        trust_score=0.85,
                    )
                )
                any_success = True
            except SsrfBlockedError as exc:
                last_error = str(exc)
                logger.warning("website_collector_ssrf_blocked", url=page_url, error=str(exc))
            except (FetchError, httpx.HTTPError) as exc:
                last_error = str(exc)
                logger.warning("website_collector_fetch_failed", url=page_url, error=str(exc))

        if not documents:
            status = SourceStatus.FAILED if attempted > 0 else SourceStatus.ACTIVE
            return CollectionResult(
                status=status, documents=[], error=last_error, pages_attempted=attempted
            )

        status = SourceStatus.ACTIVE if any_success else SourceStatus.FAILED
        return CollectionResult(
            status=status,
            documents=documents,
            error=last_error if not any_success else None,
            pages_attempted=attempted,
        )