stonks-oracle/services/extractor/inference_adapter.py

"""Inference adapter bridging the document extractor to the InferenceGateway.

Replaces direct use of llm_factory / VLLMClient / OllamaClient in the
extraction pipeline. Uses the shared InferenceGateway with extraction-
specific prompt construction and records actual endpoint, deployment,
model, and protocol lineage in the result.

Requirements: 2.12, 13.6
"""
from __future__ import annotations

import logging
import time
from dataclasses import dataclass, field

from services.extractor.client import (
    ExtractionAttempt,
    ExtractionResponse,
    _repair_json,
    _strip_markdown_fences,
)
from services.extractor.prompts import (
    build_extraction_prompt,
    get_json_schema,
    get_prompt_metadata,
)
from services.extractor.schemas import validate_extraction
from services.shared.inference.gateway import InferenceGateway
from services.shared.inference.lineage import ModelLineage, build_lineage_from_result
from services.shared.inference.models import (
    ChatMessage,
    InferenceResult,
    InferenceTarget,
    StructuredGenerationRequest,
)

logger = logging.getLogger("extractor.inference_adapter")


@dataclass
class ExtractionWithLineage:
    """Extraction result bundled with inference lineage metadata.

    The lineage records the actual endpoint, deployment, model, and
    protocol used — fixing the hardcoded ``model_provider = 'ollama'``.
    """

    response: ExtractionResponse
    lineage: ModelLineage
    raw_inference_results: list[InferenceResult] = field(default_factory=list)


async def extract_document(
    gateway: InferenceGateway,
    target: InferenceTarget,
    document_text: str,
    document_type: str = "article",
    document_id: str = "",
    known_tickers: list[str] | None = None,
    max_retries: int = 3,
    retry_base_delay: float = 2.0,
    retry_max_delay: float = 30.0,
    retry_backoff_multiplier: float = 2.0,
) -> ExtractionWithLineage:
    """Extract structured intelligence from a document via the InferenceGateway.

    This adapter:
    1. Builds extraction-specific prompts (same as current pipeline)
    2. Constructs a StructuredGenerationRequest
    3. Routes through the InferenceGateway (correct client per protocol)
    4. Parses / repairs JSON, validates against the extraction schema
    5. Records actual lineage (endpoint_id, deployment_id, model, protocol)

    Args:
        gateway: The shared InferenceGateway instance.
        target: Resolved inference target for extraction.
        document_text: The document to extract from.
        document_type: Type of document (article, filing, transcript, etc.).
        document_id: UUID of the source document.
        known_tickers: Optional list of tracked tickers for context.
        max_retries: Maximum number of retry attempts.
        retry_base_delay: Initial retry delay in seconds.
        retry_max_delay: Maximum retry delay in seconds.
        retry_backoff_multiplier: Backoff multiplier for retries.

    Returns:
        ExtractionWithLineage containing the extraction response and lineage.
    """
    import asyncio

    prompts = build_extraction_prompt(
        document_text=document_text,
        document_type=document_type,
        document_id=document_id,
        known_tickers=known_tickers,
    )
    json_schema = get_json_schema()
    prompt_meta = get_prompt_metadata()

    response = ExtractionResponse(
        prompt_metadata=prompt_meta,
        model=target.model,
    )
    inference_results: list[InferenceResult] = []
    last_lineage: ModelLineage | None = None

    total_start = time.monotonic()

    for attempt_num in range(max_retries + 1):
        # Build request
        request = StructuredGenerationRequest(
            messages=[
                ChatMessage(role="system", content=prompts["system"]),
                ChatMessage(role="user", content=prompts["user"]),
            ],
            json_schema=json_schema,
            max_output_tokens=target.extra_body.get("max_tokens", 4096),
            temperature=0.0,
            seed=0,
            timeout_seconds=target.timeout_seconds,
            trace_id=document_id,
        )

        # Call via gateway
        result = await gateway.generate(target, request)
        inference_results.append(result)
        last_lineage = build_lineage_from_result(result, trace_id=document_id)

        # Convert to ExtractionAttempt for compatibility
        attempt = _inference_result_to_attempt(result, target.model, document_text)
        response.attempts.append(attempt)

        if attempt.error is None and attempt.validation and attempt.validation.valid:
            response.success = True
            response.result = attempt.validation.parsed
            break

        # Determine if retryable
        retryable = _is_result_retryable(result)
        attempt.retryable = retryable

        if not retryable:
            logger.warning(
                "Non-retryable error for doc %s: %s — stopping retries",
                document_id or "unknown",
                attempt.error,
            )
            break

        if attempt_num < max_retries:
            delay = retry_base_delay * (retry_backoff_multiplier ** attempt_num)
            delay = min(delay, retry_max_delay)
            logger.warning(
                "Extraction attempt %d/%d failed for doc %s: %s — retrying in %.1fs",
                attempt_num + 1,
                max_retries + 1,
                document_id or "unknown",
                attempt.error or "validation failed",
                delay,
            )
            await asyncio.sleep(delay)

    response.total_duration_ms = int((time.monotonic() - total_start) * 1000)

    # Use actual lineage from last inference call
    lineage = last_lineage or ModelLineage(model=target.model, protocol=target.protocol)

    return ExtractionWithLineage(
        response=response,
        lineage=lineage,
        raw_inference_results=inference_results,
    )


def _inference_result_to_attempt(
    result: InferenceResult,
    model: str,
    document_text: str,
) -> ExtractionAttempt:
    """Convert an InferenceResult to the legacy ExtractionAttempt format.

    Applies the same markdown-fence stripping, JSON repair, and schema
    validation as the existing VLLMClient and OllamaClient.
    """
    attempt = ExtractionAttempt(model=model)
    attempt.duration_ms = result.latency_ms
    attempt.raw_output = result.content

    # Check for gateway-level errors
    if result.error:
        attempt.error = result.error
        attempt.retryable = _is_result_retryable(result)
        return attempt

    content = result.content
    if not content:
        attempt.error = "empty_model_response"
        return attempt

    # Strip markdown fences if present
    content = _strip_markdown_fences(content)

    # Repair malformed JSON
    content = _repair_json(content)

    # Validate against extraction schema
    attempt.validation = validate_extraction(content, document_text=document_text)
    if not attempt.validation.valid:
        attempt.error = "; ".join(attempt.validation.errors)

    return attempt


def _is_result_retryable(result: InferenceResult) -> bool:
    """Determine if an inference result error is retryable."""
    if result.error_category in (
        "timeout",
        "rate_limit",
        "server_error",
        "connection_error",
    ):
        return True
    if result.error and "empty" in result.error.lower():
        return True
    return False