cyberyan_linkedin_search/backend/apps/search/backends/elastic.py

"""
Elasticsearch search backend.

The query builder is a pure function (:func:`build_query`) that turns a
:class:`~apps.search.types.SearchRequest` into an Elasticsearch request body.
Keeping it free of I/O means the ranking logic is unit-testable without a
running cluster - see ``tests/test_query_builder.py``.

Ranking model
-------------
``must``   the keyword must match somewhere, with per-field boosts that encode
           "a name hit is worth more than a summary hit", plus ``fuzziness:
           AUTO`` so ``pyhton`` still finds ``python``.
``should`` an exact phrase match on the strong fields adds a large bonus, which
           pushes "product manager" above documents that merely contain both
           words far apart.
``filter`` structured constraints (skills, role, seniority, country, industry,
           employer, years of experience).  They live in the filter context, so
           they are cached by Elasticsearch and contribute nothing to the score.
"""

from __future__ import annotations

import logging
import time
from typing import Any

from django.conf import settings
from elasticsearch import ApiError, Elasticsearch, TransportError

from apps.core.exceptions import SearchBackendError
from apps.search.backends.base import SearchBackend
from apps.search.indexing import alias_name, get_client
from apps.search.types import (
    FACET_FIELDS,
    SORT_EXPERIENCE_ASC,
    SORT_EXPERIENCE_DESC,
    SORT_NAME,
    FacetBucket,
    SearchHit,
    SearchRequest,
    SearchResult,
)

logger = logging.getLogger(__name__)

#: Field boosts for the free-text part of the query.
SEARCH_FIELDS = [
    "full_name^6",
    "job_title^4",
    "skills^3.5",
    "company^2.5",
    "past_companies^1.5",
    "industry^1.5",
    "majors^1.5",
    "education",
    "location_name",
    "headline^2",
    "summary",
    "interests",
]

#: Narrower field set for the phrase bonus - phrases in a summary are noise.
PHRASE_FIELDS = ["full_name^8", "job_title^6", "skills^5", "company^4", "headline^3"]

HIGHLIGHT_FIELDS = {
    "full_name": {"number_of_fragments": 0},
    "job_title": {"number_of_fragments": 0},
    "skills": {"number_of_fragments": 3},
    "summary": {"fragment_size": 160, "number_of_fragments": 1},
}

#: Maps an API filter to the indexed keyword field it constrains.
FILTER_FIELDS = {
    "skills": "skills.keyword",
    "roles": "job_title_role",
    "seniority": "seniority_levels",
    "countries": "location_country",
    "industries": "industry.keyword",
    "companies": "company.keyword",
}


def build_query(request: SearchRequest) -> dict[str, Any]:
    """Translate a search request into an Elasticsearch request body."""
    must: list[dict[str, Any]] = []
    should: list[dict[str, Any]] = []

    if request.q:
        must.append(
            {
                "multi_match": {
                    "query": request.q,
                    "type": "best_fields",
                    "fields": SEARCH_FIELDS,
                    "operator": "or",
                    # One-word queries must match; longer queries tolerate a
                    # missing term so "senior python django" still returns
                    # people who only list two of the three.
                    "minimum_should_match": "2<70%",
                    "fuzziness": "AUTO",
                    "prefix_length": 1,
                    "max_expansions": 30,
                }
            }
        )
        should.append(
            {
                "multi_match": {
                    "query": request.q,
                    "type": "phrase",
                    "fields": PHRASE_FIELDS,
                    "slop": 1,
                    "boost": 4,
                }
            }
        )
    else:
        must.append({"match_all": {}})

    body: dict[str, Any] = {
        "query": {
            "bool": {
                "must": must,
                "should": should,
                "filter": _build_filters(request),
            }
        },
        "from": request.offset,
        "size": request.page_size,
        "track_total_hits": True,
        "sort": _build_sort(request),
    }

    if request.q:
        body["highlight"] = {
            "pre_tags": ["<mark>"],
            "post_tags": ["</mark>"],
            "fields": HIGHLIGHT_FIELDS,
        }

    if request.include_facets:
        body["aggs"] = _build_aggregations()

    return body


def _build_filters(request: SearchRequest) -> list[dict[str, Any]]:
    filters: list[dict[str, Any]] = []

    for attribute, field_name in FILTER_FIELDS.items():
        values = getattr(request, attribute)
        if values:
            # ``terms`` is OR within one filter; separate filters are ANDed,
            # which matches how users read the UI: "python OR django" skills,
            # AND located in Germany.
            filters.append({"terms": {field_name: list(values)}})

    if request.min_experience is not None or request.max_experience is not None:
        bounds: dict[str, int] = {}
        if request.min_experience is not None:
            bounds["gte"] = request.min_experience
        if request.max_experience is not None:
            bounds["lte"] = request.max_experience
        filters.append({"range": {"years_experience": bounds}})

    return filters


def _build_sort(request: SearchRequest) -> list[Any]:
    if request.sort == SORT_NAME:
        return [{"full_name.keyword": "asc"}, {"id": "asc"}]
    if request.sort == SORT_EXPERIENCE_DESC:
        return [{"years_experience": {"order": "desc", "missing": "_last"}}, "_score"]
    if request.sort == SORT_EXPERIENCE_ASC:
        return [{"years_experience": {"order": "asc", "missing": "_last"}}, "_score"]
    return ["_score", {"id": "asc"}]  # deterministic tie-break


def _build_aggregations() -> dict[str, Any]:
    size = settings.SEARCH["FACET_SIZE"]
    return {
        name: {"terms": {"field": field_name, "size": size, "order": {"_count": "desc"}}}
        for name, field_name in FACET_FIELDS.items()
    }


class ElasticsearchBackend(SearchBackend):
    name = "elasticsearch"

    def __init__(self, client: Elasticsearch | None = None) -> None:
        self._client = client or get_client()
        self._index = alias_name()

    # -- SearchBackend ------------------------------------------------------
    def search(self, request: SearchRequest) -> SearchResult:
        body = build_query(request)
        started = time.perf_counter()
        try:
            response = self._client.search(index=self._index, **body)
        except (ApiError, TransportError) as exc:
            logger.exception("search_backend_error", extra={"event": "search_backend_error"})
            raise SearchBackendError() from exc

        took_ms = int((time.perf_counter() - started) * 1000)
        hits = [self._to_hit(raw) for raw in response["hits"]["hits"]]
        return SearchResult(
            total=response["hits"]["total"]["value"],
            hits=hits,
            facets=self._to_facets(response.get("aggregations", {})),
            took_ms=took_ms,
            backend=self.name,
        )

    def suggest(self, prefix: str, limit: int = 8) -> list[str]:
        if not prefix:
            return []
        try:
            response = self._client.search(
                index=self._index,
                size=limit,
                query={
                    "bool": {
                        "should": [
                            {"match": {"full_name.autocomplete": {"query": prefix, "boost": 2}}},
                            {"match": {"skills.autocomplete": prefix}},
                        ],
                        "minimum_should_match": 1,
                    }
                },
                _source=["full_name", "job_title"],
            )
        except (ApiError, TransportError) as exc:
            raise SearchBackendError() from exc

        seen: list[str] = []
        for raw in response["hits"]["hits"]:
            value = raw["_source"].get("full_name", "")
            if value and value not in seen:
                seen.append(value)
        return seen[:limit]

    def ping(self) -> bool:
        if not self._client.ping():
            raise SearchBackendError("Elasticsearch did not answer a ping.")
        return True

    # -- Mapping helpers ----------------------------------------------------
    @staticmethod
    def _to_hit(raw: dict[str, Any]) -> SearchHit:
        source = raw["_source"]
        return SearchHit(
            id=source["id"],
            score=raw.get("_score") or 0.0,
            full_name=source.get("full_name", ""),
            headline=source.get("headline", ""),
            job_title=source.get("job_title", ""),
            job_title_role=source.get("job_title_role", ""),
            seniority_levels=source.get("seniority_levels", []),
            company=source.get("company", ""),
            industry=source.get("industry", ""),
            location_name=source.get("location_name", ""),
            location_country=source.get("location_country", ""),
            years_experience=source.get("years_experience"),
            linkedin_url=source.get("linkedin_url", ""),
            skills=source.get("skills", [])[:12],
            highlight=raw.get("highlight", {}),
        )

    @staticmethod
    def _to_facets(aggregations: dict[str, Any]) -> dict[str, list[FacetBucket]]:
        return {
            name: [
                FacetBucket(key=str(bucket["key"]), count=bucket["doc_count"])
                for bucket in payload.get("buckets", [])
                if bucket["key"] != ""
            ]
            for name, payload in aggregations.items()
        }