src/core/recommendation_engine/recommender.py

import asyncio
from datetime import datetime, timedelta
from typing import Optional

from qdrant_client import models
from qdrant_client.http.models import Document

from qdrant_client.http.exceptions import ResponseHandlingException

from ...lib.qdrant import get_qdrant_client, reset_qdrant_client, CATALOG_COLLECTION, CATALOG_EMBED_MODEL
from ...core.logger import SingletonLogger

logger = SingletonLogger().get_logger()


async def _qdrant_query(method_name: str, *args, **kwargs):
    """Call a QdrantClient method by name, retrying once on stale connection."""
    client = get_qdrant_client()
    try:
        return await asyncio.to_thread(getattr(client, method_name), *args, **kwargs)
    except ResponseHandlingException:
        logger.warning("Qdrant connection error, resetting client and retrying")
        await asyncio.sleep(2)
        client = reset_qdrant_client()
        return await asyncio.to_thread(getattr(client, method_name), *args, **kwargs)


def _point_to_dict(point) -> dict:
    payload = point.payload or {}
    score = getattr(point, "score", None)
    return {
        "arxiv_id": payload.get("arxiv_id"),
        "title": payload.get("title"),
        "abstract": payload.get("abstract"),
        "authors": payload.get("authors"),
        "categories": payload.get("categories"),
        "primary_category": payload.get("primary_category"),
        "published_date": payload.get("published_date"),
        "paper_url": payload.get("paper_url"),
        "pdf_url": payload.get("pdf_url"),
        "score": score,
        "final_score": None,
    }


def _exclude(results: list[dict], exclude_ids: set[str]) -> list[dict]:
    if not exclude_ids:
        return results
    return [r for r in results if r.get("arxiv_id") not in exclude_ids]


async def get_similar_papers(
    title: str,
    abstract: str,
    exclude_ids: list[str],
    top_k: int = 20,
) -> list[dict]:
    results = await _qdrant_query(
        "query_points",
        collection_name=CATALOG_COLLECTION,
        query=Document(
            text=f"{title} {abstract}",
            model=CATALOG_EMBED_MODEL,
        ),
        using=CATALOG_EMBED_MODEL,
        limit=top_k + len(exclude_ids),
        with_payload=True,
    )
    items = [_point_to_dict(p) for p in results.points]
    return _exclude(items, set(exclude_ids))[:top_k]


async def get_similar_on_topic(
    title: str,
    abstract: str,
    primary_category: str,
    exclude_ids: list[str],
    top_k: int = 10,
) -> list[dict]:
    results = await _qdrant_query(
        "query_points",
        collection_name=CATALOG_COLLECTION,
        query=Document(
            text=f"{title} {abstract}",
            model=CATALOG_EMBED_MODEL,
        ),
        using=CATALOG_EMBED_MODEL,
        query_filter=models.Filter(
            must=[
                models.FieldCondition(
                    key="primary_category",
                    match=models.MatchValue(value=primary_category),
                )
            ],
        ),
        limit=top_k + len(exclude_ids),
        with_payload=True,
    )
    items = [_point_to_dict(p) for p in results.points]
    return _exclude(items, set(exclude_ids))[:top_k]


async def get_papers_by_authors(
    authors_str: str,
    exclude_ids: list[str],
    top_k: int = 10,
) -> list[dict]:
    author_list = [a.strip() for a in authors_str.split(";") if a.strip()][:5]

    results = []
    seen = set(exclude_ids)

    for author in author_list:
        scroll_filter = models.Filter(
            must=[
                models.FieldCondition(
                    key="authors",
                    match=models.MatchText(text=author),
                )
            ]
        )

        scroll_result, _ = await _qdrant_query(
            "scroll",
            collection_name=CATALOG_COLLECTION,
            scroll_filter=scroll_filter,
            limit=top_k + len(seen),
            with_payload=True,
        )

        for p in scroll_result:
            arxiv_id = p.payload.get("arxiv_id")
            if arxiv_id and arxiv_id not in seen:
                seen.add(arxiv_id)
                results.append(_point_to_dict(p))

        if len(results) >= top_k:
            break

    return results[:top_k]


async def get_feed_recommendations(
    query_texts: list[str],
    category_filter: list[str],
    exclude_ids: list[str],
    offset: int = 0,
    limit: int = 24,
) -> list[dict]:
    """Aggregate vector search across multiple saved papers, filtered by categories."""
    combined = " ".join(query_texts)[:2000]

    fetch_count = offset + limit + len(exclude_ids)

    category_must = []
    if category_filter:
        category_must.append(
            models.FieldCondition(
                key="primary_category",
                match=models.MatchAny(any=category_filter),
            )
        )

    results = await _qdrant_query(
        "query_points",
        collection_name=CATALOG_COLLECTION,
        query=Document(
            text=combined,
            model=CATALOG_EMBED_MODEL,
        ),
        using=CATALOG_EMBED_MODEL,
        query_filter=models.Filter(must=category_must) if category_must else None,
        limit=fetch_count,
        with_payload=True,
    )

    items = [_point_to_dict(p) for p in results.points]
    items = _exclude(items, set(exclude_ids))
    return items[offset : offset + limit]


async def search_catalog(
    query: str,
    offset: int = 0,
    limit: int = 20,
) -> list[dict]:
    """Free-text vector search over the catalog collection."""
    fetch_count = offset + limit

    results = await _qdrant_query(
        "query_points",
        collection_name=CATALOG_COLLECTION,
        query=Document(
            text=query,
            model=CATALOG_EMBED_MODEL,
        ),
        using=CATALOG_EMBED_MODEL,
        limit=fetch_count,
        with_payload=True,
    )

    items = [_point_to_dict(p) for p in results.points]
    return items[offset : offset + limit]


def rerank_results(
    candidates: list[dict],
    user_saved_categories: list[str],
) -> list[dict]:
    """Re-rank recommendation candidates by similarity, category affinity, and recency."""
    now = datetime.utcnow()
    thirty_days_ago = now - timedelta(days=30)
    six_months_ago = now - timedelta(days=180)

    for candidate in candidates:
        similarity = candidate.get("score") or 0.0

        category_affinity = 0.0
        if candidate.get("primary_category") in user_saved_categories:
            category_affinity = 1.0

        recency = 0.0
        pub_date_str = candidate.get("published_date")
        if pub_date_str:
            try:
                pub_date = datetime.strptime(pub_date_str, "%Y-%m-%d")
                if pub_date >= thirty_days_ago:
                    recency = 1.0
                elif pub_date >= six_months_ago:
                    recency = 0.5
            except ValueError:
                pass

        candidate["final_score"] = (
            0.60 * similarity
            + 0.15 * category_affinity
            + 0.15 * recency
            + 0.10 * 0.5
        )

    candidates.sort(key=lambda c: c.get("final_score", 0), reverse=True)
    return candidates