saksham
1a4c80958f
FastAPI + Celery + Next.js + Postgres/Redis app with company monitoring, source collection, LLM-based change analysis, enrichment, and account security (Turnstile, escalating lockout, email verification).
199 lines
6.8 KiB
Python
199 lines
6.8 KiB
Python
"""Orchestrates collectors against the database: turns `DiscoveredSource`s
|
|
into `Source` rows, and a collector's `CollectionResult` into persisted
|
|
`SourceDocument` + `Snapshot` rows. Collectors themselves stay
|
|
database-free (see collectors/base.py) so they're trivially unit-testable;
|
|
this module is the seam where that plain-dataclass world meets the ORM.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import hashlib
|
|
import uuid
|
|
from datetime import UTC, datetime
|
|
|
|
from sqlalchemy.ext.asyncio import AsyncSession
|
|
|
|
from app.collectors.base import CollectionResult, CompanyContext, SourceConfig
|
|
from app.collectors.registry import get_collector
|
|
from app.core.config import Settings
|
|
from app.core.logging import get_logger
|
|
from app.models.company import Company
|
|
from app.models.enums import EnrichmentStatus, SourceStatus, SourceType
|
|
from app.models.source import Source
|
|
from app.repositories.source_repository import (
|
|
SnapshotRepository,
|
|
SourceDocumentRepository,
|
|
SourceRepository,
|
|
)
|
|
|
|
logger = get_logger(__name__)
|
|
|
|
# Collector types with real (non-fixture) auto-discovery.
|
|
_DISCOVERABLE_TYPES = (
|
|
SourceType.WEBSITE,
|
|
SourceType.GITHUB,
|
|
SourceType.SEC_EDGAR,
|
|
SourceType.JOB_POSTING,
|
|
SourceType.RSS,
|
|
SourceType.GOV_CONTRACT,
|
|
SourceType.PATENT,
|
|
)
|
|
|
|
|
|
def _leadership_names(company: Company) -> list[str]:
|
|
enrichment = company.enrichment
|
|
if enrichment is None or enrichment.status == EnrichmentStatus.FAILED:
|
|
return []
|
|
return [
|
|
member["name"]
|
|
for member in enrichment.data.get("leadership_team", []) or []
|
|
if member.get("name")
|
|
]
|
|
|
|
|
|
def to_company_context(company: Company, settings: Settings | None = None) -> CompanyContext:
|
|
return CompanyContext(
|
|
id=str(company.id),
|
|
name=company.name,
|
|
official_website=company.official_website,
|
|
monitoring_focus=company.monitoring_focus,
|
|
aliases=[a.alias for a in company.aliases],
|
|
competitors=[c.name for c in company.competitors],
|
|
leadership_names=_leadership_names(company),
|
|
uspto_api_key=settings.uspto_api_key if settings is not None else None,
|
|
)
|
|
|
|
|
|
def _to_source_config(source: Source) -> SourceConfig:
|
|
return SourceConfig(
|
|
id=str(source.id),
|
|
source_type=source.source_type,
|
|
name=source.name,
|
|
base_url=source.base_url,
|
|
configuration_metadata=source.configuration_metadata or {},
|
|
)
|
|
|
|
|
|
async def discover_sources_for_company(
|
|
db: AsyncSession, company: Company, settings: Settings | None = None
|
|
) -> list[Source]:
|
|
"""Runs discovery for every collector type capable of it and creates a
|
|
`Source` row per suggestion, skipping ones that already exist for this
|
|
company (same type + base_url)."""
|
|
repo = SourceRepository(db)
|
|
existing = await repo.list_for_company(company.id)
|
|
existing_keys = {(s.source_type, s.base_url) for s in existing}
|
|
|
|
context = to_company_context(company, settings)
|
|
created: list[Source] = []
|
|
|
|
for source_type in _DISCOVERABLE_TYPES:
|
|
collector = get_collector(source_type)
|
|
try:
|
|
discovered = await collector.discover(context)
|
|
except Exception as exc: # pragma: no cover - defensive, discovery is best-effort
|
|
logger.warning("source_discovery_failed", source_type=source_type, error=str(exc))
|
|
continue
|
|
|
|
for candidate in discovered:
|
|
key = (candidate.source_type, candidate.base_url)
|
|
if key in existing_keys:
|
|
continue
|
|
source = await repo.create(
|
|
company_id=company.id,
|
|
source_type=candidate.source_type,
|
|
name=candidate.name,
|
|
base_url=candidate.base_url,
|
|
configuration_metadata=candidate.configuration_metadata,
|
|
)
|
|
existing_keys.add(key)
|
|
created.append(source)
|
|
|
|
return created
|
|
|
|
|
|
def _summarize_documents(documents) -> dict:
|
|
return {
|
|
"document_count": len(documents),
|
|
"titles": [d.title for d in documents if d.title][:50],
|
|
"urls": [d.url for d in documents][:50],
|
|
"content_hashes": [d.content_hash for d in documents][:50],
|
|
}
|
|
|
|
|
|
def _build_text_summary(documents) -> str:
|
|
"""Concatenated per-document excerpts used for bounded text diffing
|
|
(change_detection's text-diff layer) - a title-only summary is too thin
|
|
to catch wording-level changes within a page."""
|
|
parts = []
|
|
for doc in documents[:8]:
|
|
title = doc.title or doc.url
|
|
excerpt = doc.content_text[:600]
|
|
parts.append(f"### {title}\n{excerpt}")
|
|
return "\n\n".join(parts)
|
|
|
|
|
|
async def collect_source(
|
|
db: AsyncSession,
|
|
settings: Settings,
|
|
source: Source,
|
|
company: Company,
|
|
*,
|
|
monitoring_run_id: uuid.UUID | None = None,
|
|
) -> CollectionResult:
|
|
"""Runs one source's collector, persists new documents (deduped by
|
|
content hash within the source), writes a Snapshot summarizing the
|
|
batch, and updates the Source's health/status fields."""
|
|
collector = get_collector(source.source_type)
|
|
context = to_company_context(company, settings)
|
|
config = _to_source_config(source)
|
|
|
|
result = await collector.collect(config, context)
|
|
|
|
doc_repo = SourceDocumentRepository(db)
|
|
persisted_hashes: list[str] = []
|
|
for doc in result.documents:
|
|
if await doc_repo.exists_with_hash(source.id, doc.content_hash):
|
|
continue
|
|
await doc_repo.create(
|
|
source_id=source.id,
|
|
company_id=company.id,
|
|
url=doc.url,
|
|
canonical_url=doc.canonical_url,
|
|
title=doc.title,
|
|
author=doc.author,
|
|
publication_date=doc.publication_date,
|
|
retrieved_date=doc.retrieved_date,
|
|
content_text=doc.content_text,
|
|
content_hash=doc.content_hash,
|
|
metadata_json=doc.metadata,
|
|
language=doc.language,
|
|
http_status=doc.http_status,
|
|
extraction_method=doc.extraction_method,
|
|
trust_score=doc.trust_score,
|
|
)
|
|
persisted_hashes.append(doc.content_hash)
|
|
|
|
now = datetime.now(UTC)
|
|
if result.documents:
|
|
batch_hash = hashlib.sha256(
|
|
"|".join(sorted(d.content_hash for d in result.documents)).encode("utf-8")
|
|
).hexdigest()
|
|
snapshot_repo = SnapshotRepository(db)
|
|
await snapshot_repo.create(
|
|
company_id=company.id,
|
|
source_id=source.id,
|
|
snapshot_type=source.source_type.value,
|
|
hash=batch_hash,
|
|
structured_summary=_summarize_documents(result.documents),
|
|
text_summary=_build_text_summary(result.documents),
|
|
monitoring_run_id=monitoring_run_id,
|
|
)
|
|
|
|
source_repo = SourceRepository(db)
|
|
success = result.status in (SourceStatus.ACTIVE,)
|
|
await source_repo.mark_checked(source, status=result.status, checked_at=now, success=success)
|
|
|
|
await db.commit()
|
|
return result
|