|
| 1 | +from __future__ import annotations |
| 2 | + |
| 3 | +from typing import Any, Dict, List |
| 4 | + |
| 5 | +from scoring_policy import LLM_RELEVANCE_MODE_DEFAULT, normalize_relevance_mode |
| 6 | +from search_models import DigestStats, Paper, SearchRequest |
| 7 | +from text_utils import clean_text |
| 8 | + |
| 9 | +AGENT_SEARCH_SCHEMA_VERSION = "1.1" |
| 10 | + |
| 11 | + |
| 12 | +def map_agent_status(stats: DigestStats, papers: List[Paper]) -> str: |
| 13 | + if papers: |
| 14 | + if stats.source_warning_messages: |
| 15 | + return "partial_source_failure" |
| 16 | + return "ok" |
| 17 | + if stats.no_results_reason == "outside_horizon": |
| 18 | + return "outside_horizon" |
| 19 | + if stats.no_results_reason == "below_threshold": |
| 20 | + return "below_threshold" |
| 21 | + if stats.no_results_reason == "source_failure": |
| 22 | + return "error" |
| 23 | + if stats.no_results_reason in {"none_retrieved", "no_candidates"}: |
| 24 | + return "no_candidates" |
| 25 | + return "error" |
| 26 | + |
| 27 | + |
| 28 | +def describe_agent_llm_backend(config: Any) -> Dict[str, str]: |
| 29 | + if getattr(config, "gemini_api_key", ""): |
| 30 | + return {"provider": "gemini", "model": str(getattr(config, "gemini_model", "") or "")} |
| 31 | + if ( |
| 32 | + getattr(config, "enable_openai_compat_fallback", False) |
| 33 | + and clean_text(str(getattr(config, "openai_compat_api_base", "") or "")) |
| 34 | + and clean_text(str(getattr(config, "openai_compat_model", "") or "")) |
| 35 | + ): |
| 36 | + return { |
| 37 | + "provider": "openai_compatible", |
| 38 | + "model": str(getattr(config, "openai_compat_model", "") or ""), |
| 39 | + } |
| 40 | + if getattr(config, "enable_cerebras_fallback", False) and getattr(config, "cerebras_api_key", ""): |
| 41 | + return {"provider": "cerebras", "model": str(getattr(config, "cerebras_model", "") or "")} |
| 42 | + return {"provider": "none", "model": ""} |
| 43 | + |
| 44 | + |
| 45 | +def build_agent_search_response( |
| 46 | + *, |
| 47 | + project_name: str, |
| 48 | + llm_projects: List[Dict[str, str]], |
| 49 | + generated_topics: List[Dict[str, Any]], |
| 50 | + request_config: Any, |
| 51 | + search_request: SearchRequest, |
| 52 | + stats: DigestStats, |
| 53 | + papers: List[Paper], |
| 54 | + diagnostics: List[str], |
| 55 | + include_abstracts: bool = False, |
| 56 | + include_diagnostics: bool = False, |
| 57 | +) -> Dict[str, Any]: |
| 58 | + primary_topic = generated_topics[0] if generated_topics else {} |
| 59 | + backend = describe_agent_llm_backend(request_config) |
| 60 | + return { |
| 61 | + "schema_version": AGENT_SEARCH_SCHEMA_VERSION, |
| 62 | + "status": map_agent_status(stats, papers), |
| 63 | + "request": { |
| 64 | + "project_name": clean_text(project_name) or llm_projects[0]["name"], |
| 65 | + "search_intent": search_request.intent, |
| 66 | + "time_horizon": search_request.time_horizon_key, |
| 67 | + "top_k": request_config.max_papers, |
| 68 | + "output_language": request_config.output_language, |
| 69 | + "include_diagnostics": include_diagnostics, |
| 70 | + "include_abstracts": include_abstracts, |
| 71 | + }, |
| 72 | + "meta": { |
| 73 | + "intent_label": search_request.intent_label, |
| 74 | + "requested_horizon_label": search_request.time_horizon_label, |
| 75 | + "window_used_label": stats.window_used_label or search_request.time_horizon_label, |
| 76 | + "query_plan_label": stats.query_plan_label or "generated topic queries", |
| 77 | + "used_provider": backend["provider"], |
| 78 | + "used_model": backend["model"], |
| 79 | + "sources_queried": [ |
| 80 | + label |
| 81 | + for enabled, label in [ |
| 82 | + (bool(request_config.arxiv_queries), "arXiv"), |
| 83 | + (bool(request_config.pubmed_queries), "PubMed"), |
| 84 | + ( |
| 85 | + request_config.enable_semantic_scholar |
| 86 | + and bool(request_config.semantic_scholar_queries), |
| 87 | + "Semantic Scholar", |
| 88 | + ), |
| 89 | + ( |
| 90 | + request_config.enable_google_scholar |
| 91 | + and bool(request_config.google_scholar_queries), |
| 92 | + "Google Scholar", |
| 93 | + ), |
| 94 | + ] |
| 95 | + if enabled |
| 96 | + ], |
| 97 | + "scanned_count": stats.post_time_filter_candidates or stats.total_candidates, |
| 98 | + "selected_count": len(papers), |
| 99 | + "threshold_used": stats.ranking_threshold, |
| 100 | + "notice": stats.search_notice, |
| 101 | + "partial_source_failure": bool(stats.source_warning_messages), |
| 102 | + "source_warning_count": len(stats.source_warning_messages), |
| 103 | + "source_warnings": list(stats.source_warning_messages), |
| 104 | + }, |
| 105 | + "topic": { |
| 106 | + "name": clean_text(str(primary_topic.get("name", ""))), |
| 107 | + "keywords": [ |
| 108 | + clean_text(str(item)) |
| 109 | + for item in primary_topic.get("keywords", []) |
| 110 | + if clean_text(str(item)) |
| 111 | + ], |
| 112 | + "relevance_mode": normalize_relevance_mode( |
| 113 | + primary_topic.get("relevance_mode", LLM_RELEVANCE_MODE_DEFAULT) |
| 114 | + ), |
| 115 | + "arxiv_query": clean_text(str(primary_topic.get("arxiv_query", ""))), |
| 116 | + "pubmed_query": clean_text(str(primary_topic.get("pubmed_query", ""))), |
| 117 | + "semantic_scholar_query": clean_text( |
| 118 | + str(primary_topic.get("semantic_scholar_query", "")) |
| 119 | + ), |
| 120 | + "google_scholar_query": clean_text( |
| 121 | + str(primary_topic.get("google_scholar_query", "")) |
| 122 | + ), |
| 123 | + }, |
| 124 | + "papers": [ |
| 125 | + { |
| 126 | + "rank": index, |
| 127 | + "id": paper.paper_id, |
| 128 | + "title": paper.title, |
| 129 | + "authors": ", ".join(paper.authors), |
| 130 | + "source": paper.source, |
| 131 | + "url": paper.url, |
| 132 | + "published_at": paper.published_at_utc.isoformat(), |
| 133 | + "relevance_score": paper.score, |
| 134 | + "relevance_reason": paper.llm_relevance_text, |
| 135 | + "core_point": paper.llm_core_point_text, |
| 136 | + "usefulness": paper.llm_usefulness_text, |
| 137 | + "evidence_spans": list(paper.llm_evidence_spans or []), |
| 138 | + "topic": paper.topic, |
| 139 | + "project_name": paper.project_name, |
| 140 | + "relevance_mode": paper.relevance_mode, |
| 141 | + **({"abstract": paper.abstract} if include_abstracts else {}), |
| 142 | + } |
| 143 | + for index, paper in enumerate(papers, start=1) |
| 144 | + ], |
| 145 | + "diagnostics": diagnostics, |
| 146 | + } |
0 commit comments