File size: 5,426 Bytes
54eb2ce | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 | import asyncio
from fastapi import HTTPException
from sqlalchemy import select
from ..database.db import session_pool
from ..model.paper import Paper
from ..model.arxiv_catalog import ArxivCatalog
from ..schema.paper import RecommendationItem, RecommendationRequest, RecommendationsResponse
from ..core.recommendation_engine.recommender import (
get_similar_papers,
get_similar_on_topic,
get_papers_by_authors,
rerank_results,
)
from ..core.logger import SingletonLogger
logger = SingletonLogger().get_logger()
async def get_recommendations_by_metadata(
req: RecommendationRequest,
user_id: int,
limit: int = 10,
) -> RecommendationsResponse:
async with session_pool() as session:
all_papers_result = await session.execute(
select(Paper.arxiv_id).where(
Paper.user_id == user_id, Paper.arxiv_id.isnot(None)
)
)
saved_arxiv_ids = [row[0] for row in all_papers_result.all()]
user_saved_categories: list[str] = []
if saved_arxiv_ids:
catalog_result = await session.execute(
select(ArxivCatalog.primary_category)
.where(ArxivCatalog.arxiv_id.in_(saved_arxiv_ids))
.distinct()
)
user_saved_categories = [
row[0] for row in catalog_result.all() if row[0]
]
exclude_ids = saved_arxiv_ids
if req.arxiv_id:
exclude_ids = list(set(saved_arxiv_ids + [req.arxiv_id]))
primary_category = req.primary_category
async def _empty_list() -> list:
return []
similar_task = get_similar_papers(
req.title, req.abstract, exclude_ids, top_k=limit * 2
)
on_topic_task = (
get_similar_on_topic(
req.title, req.abstract, primary_category, exclude_ids, top_k=limit
)
if primary_category
else _empty_list()
)
by_authors_task = get_papers_by_authors(
req.authors, exclude_ids, top_k=limit
)
similar, on_topic, by_authors = await asyncio.gather(
similar_task, on_topic_task, by_authors_task
)
reranked_similar = rerank_results(similar, user_saved_categories)[:limit]
return RecommendationsResponse(
similar_papers=[RecommendationItem(**p) for p in reranked_similar],
on_this_topic=[RecommendationItem(**p) for p in on_topic[:limit]],
from_these_authors=[RecommendationItem(**p) for p in by_authors[:limit]],
)
async def get_recommendations_for_paper(
paper_id: int,
user_id: int,
limit: int = 10,
) -> RecommendationsResponse:
async with session_pool() as session:
result = await session.execute(
select(Paper).where(Paper.id == paper_id, Paper.user_id == user_id)
)
paper = result.scalar_one_or_none()
if not paper:
raise HTTPException(status_code=404, detail="Paper not found")
all_papers_result = await session.execute(
select(Paper.arxiv_id).where(
Paper.user_id == user_id, Paper.arxiv_id.isnot(None)
)
)
saved_arxiv_ids = [row[0] for row in all_papers_result.all()]
# Cross-reference with catalog for category affinity
user_saved_categories: list[str] = []
if saved_arxiv_ids:
catalog_result = await session.execute(
select(ArxivCatalog.primary_category)
.where(ArxivCatalog.arxiv_id.in_(saved_arxiv_ids))
.distinct()
)
user_saved_categories = [
row[0] for row in catalog_result.all() if row[0]
]
# Determine primary_category for on-topic query
primary_category = None
if paper.arxiv_id:
cat_result = await session.execute(
select(ArxivCatalog.primary_category).where(
ArxivCatalog.arxiv_id == paper.arxiv_id
)
)
row = cat_result.first()
if row:
primary_category = row[0]
# Fall back to first topic from Paper.topics
if not primary_category and paper.topics:
primary_category = paper.topics.split()[0] if paper.topics else None
async def _empty_list() -> list:
return []
# Run all three recommendation queries in parallel
similar_task = get_similar_papers(
paper.title, paper.abstract, saved_arxiv_ids, top_k=limit * 2
)
on_topic_task = (
get_similar_on_topic(
paper.title, paper.abstract, primary_category, saved_arxiv_ids, top_k=limit
)
if primary_category
else _empty_list()
)
by_authors_task = get_papers_by_authors(
paper.authors, saved_arxiv_ids, top_k=limit
)
similar, on_topic, by_authors = await asyncio.gather(
similar_task, on_topic_task, by_authors_task
)
reranked_similar = rerank_results(similar, user_saved_categories)[:limit]
return RecommendationsResponse(
similar_papers=[RecommendationItem(**p) for p in reranked_similar],
on_this_topic=[RecommendationItem(**p) for p in on_topic[:limit]],
from_these_authors=[RecommendationItem(**p) for p in by_authors[:limit]],
)
|