AI/ontology_platform/ont_platform/api/routes/extraction.py

"""Phase 1 URL/HTML ingestion routes."""

from __future__ import annotations

import time
from typing import Annotated

from fastapi import APIRouter, Body, HTTPException, Query
from pydantic import BaseModel, Field, model_validator

from ont_platform.core.extraction.lightweight_extractor import LightweightExtractor
from ont_platform.core.extractors.web_extractor import extract_web_content
from ont_platform.models.content_unit import PlatformContentUnit
from ont_platform.storage.dedup_cache import get_default_dedup_cache

router = APIRouter(tags=["extraction"])


class UrlIngestRequest(BaseModel):
    """URL/HTML request accepted by Phase 1 ingestion endpoints."""

    url: str | None = None
    html: str | None = None
    project_id: str = "default"
    language: str | None = None
    skip_if_duplicate: bool = True
    ontology_user_instruction: str = ""
    facts_user_instruction: str = ""

    @model_validator(mode="after")
    def _require_url_or_html(self) -> UrlIngestRequest:
        if not self.url and not self.html:
            raise ValueError("url or html is required")
        return self


class UrlIngestResponse(BaseModel):
    status: str
    url: str | None
    title: str | None
    author: str | None
    published_date: str | None
    language: str | None
    text_length: int
    source_document: dict = Field(default_factory=dict)
    evidence_spans: list[dict] = Field(default_factory=list)
    content_unit: dict = Field(default_factory=dict)
    dedup: dict = Field(default_factory=dict)
    entities: list = Field(default_factory=list)
    relations: list = Field(default_factory=list)
    extraction_time_sec: float
    entity_count: int
    relation_count: int
    warnings: list[str] = Field(default_factory=list)


def _build_payload(
    payload: UrlIngestRequest | None,
    query_url: str | None,
    query_project_id: str,
) -> UrlIngestRequest:
    if payload is None:
        try:
            return UrlIngestRequest(url=query_url, project_id=query_project_id)
        except ValueError as exc:
            raise HTTPException(status_code=400, detail=str(exc)) from exc

    data = payload.model_dump()
    if query_url:
        data["url"] = query_url
    if query_project_id != "default" and payload.project_id == "default":
        data["project_id"] = query_project_id
    try:
        return UrlIngestRequest(**data)
    except ValueError as exc:
        raise HTTPException(status_code=400, detail=str(exc)) from exc


@router.post("/api/v1/extract/url", response_model=UrlIngestResponse)
@router.post("/api/v1/process/url", response_model=UrlIngestResponse)
@router.post("/process/url", response_model=UrlIngestResponse)
async def extract_url(
    payload: Annotated[UrlIngestRequest | None, Body()] = None,
    url: Annotated[str | None, Query()] = None,
    project_id: Annotated[str, Query()] = "default",
):
    """Extract a SourceDocument-ready payload from URL or supplied HTML."""

    request = _build_payload(payload, query_url=url, query_project_id=project_id)
    start_time = time.time()

    try:
        extracted = extract_web_content(
            html=request.html,
            url=request.url,
            lang=request.language,
        )
        source_document = extracted.to_source_document_dict(project_id=request.project_id)
        evidence_spans = [
            span.to_dict()
            for span in extracted.evidence_spans(
                project_id=request.project_id,
                document_id=source_document["id"],
            )
        ]
        content_unit = PlatformContentUnit.from_extracted(extracted).to_dict()

        dedup_result = get_default_dedup_cache().check_and_remember(
            project_id=request.project_id,
            document_id=source_document["id"],
            content_hash=extracted.content_hash,
            fingerprint=extracted.fingerprint,
        )

        entities: list = []
        relations: list = []
        warnings: list[str] = []
        if dedup_result.is_duplicate and request.skip_if_duplicate:
            warnings.append("Duplicate source document skipped by fingerprint.")
        else:
            lightweight = LightweightExtractor(use_llm=False)
            candidates = lightweight.extract(
                text=extracted.text,
                project_id=request.project_id,
                document_id=source_document["id"],
            )
            entities = candidates.entities
            relations = candidates.relations
            warnings = candidates.warnings

        extraction_time = time.time() - start_time

        return {
            "status": "success",
            "url": request.url,
            "title": extracted.title,
            "author": extracted.author,
            "published_date": extracted.publish_date,
            "language": extracted.language,
            "text_length": len(extracted.text),
            "source_document": source_document,
            "evidence_spans": evidence_spans,
            "content_unit": content_unit,
            "dedup": dedup_result.to_dict(),
            "entities": entities,
            "relations": relations,
            "extraction_time_sec": round(extraction_time, 2),
            "entity_count": len(entities),
            "relation_count": len(relations),
            "warnings": warnings,
        }

    except HTTPException:
        raise
    except Exception as exc:
        raise HTTPException(status_code=500, detail=f"Extraction failed: {exc}") from exc


__all__ = ["router"]