cite-agent 1.0.4__py3-none-any.whl → 1.0.5__py3-none-any.whl

src/services/search_service/__init__.py

@@ -0,0 +1,5 @@
+"""Search service package exposing the high-level search engine."""
+
+from .search_engine import SearchEngine
+
+__all__ = ["SearchEngine"]

src/services/search_service/indexer.py

@@ -0,0 +1,186 @@
+#indexer.py
+from typing import List, Dict, Optional
+import asyncio
+from datetime import datetime, timezone
+import redis.asyncio as redis
+import json
+
+from ...utils.logger import logger, log_operation
+from ...storage.db.operations import DatabaseOperations
+from .vector_search import VectorSearchEngine
+
+
+def _utc_timestamp() -> str:
+    return datetime.now(timezone.utc).isoformat()
+
+class DocumentIndexer:
+    def __init__(self, db_ops: DatabaseOperations, vector_search: VectorSearchEngine, redis_url: str):
+        #logger.info("Initializing DocumentIndexer")
+        self.db = db_ops
+        self.vector_search = vector_search
+        self.redis_client = redis.from_url(redis_url)
+        self.indexing_queue = asyncio.Queue()
+        self.active_sessions: Dict[str, Dict] = {}
+        self._running = False
+        #logger.info("DocumentIndexer initialized")
+
+    @log_operation("start_indexing")
+    async def start(self):
+        """Start the indexing process."""
+        #logger.info("Starting indexing service")
+        self._running = True
+        try:
+            await asyncio.gather(
+                self._process_queue(),
+                self._monitor_research_sessions()
+            )
+        except Exception as e:
+            logger.error(f"Error in indexing service: {str(e)}")
+            self._running = False
+            raise
+
+    @log_operation("monitor_sessions")
+    async def _monitor_research_sessions(self):
+        """Monitor active research sessions for new documents."""
+        while self._running:
+            try:
+                # Subscribe to research session updates
+                async for message in self.redis_client.subscribe("research_updates"):
+                    update = json.loads(message["data"])
+                    session_id = update.get("session_id")
+                    
+                    if update.get("type") == "new_papers":
+                        await self._handle_new_papers(session_id, update.get("papers", []))
+                        
+            except Exception as e:
+                logger.error(f"Error monitoring sessions: {str(e)}")
+                await asyncio.sleep(1)
+
+    async def _handle_new_papers(self, session_id: str, papers: List[str]):
+        """Handle new papers added to a research session."""
+        #logger.info(f"Processing new papers for session {session_id}")
+        for paper_id in papers:
+            await self.queue_document(paper_id, session_id)
+
+    @log_operation("queue_document")
+    async def queue_document(self, doc_id: str, session_id: Optional[str] = None):
+        """Queue a document for indexing."""
+        #logger.info(f"Queuing document for indexing: {doc_id}")
+        await self.indexing_queue.put({
+            "doc_id": doc_id,
+            "session_id": session_id,
+            "queued_at": _utc_timestamp()
+        })
+        
+        if session_id:
+            await self._update_session_progress(session_id, "queued", doc_id)
+
+    async def _process_queue(self):
+        """Process documents in the indexing queue."""
+        #logger.info("Starting queue processing")
+        while self._running:
+            try:
+                item = await self.indexing_queue.get()
+                doc_id = item["doc_id"]
+                session_id = item.get("session_id")
+
+                #logger.info(f"Processing document: {doc_id}")
+                
+                if session_id:
+                    await self._update_session_progress(session_id, "processing", doc_id)
+
+                processed_doc = await self.db.get_processed_paper(doc_id)
+                if not processed_doc:
+                    #logger.warning(f"Processed content not found: {doc_id}")
+                    continue
+
+                success = await self._index_document(doc_id, processed_doc.content)
+                
+                if success:
+                    #logger.info(f"Successfully indexed: {doc_id}")
+                    await self.db.update_paper_status(doc_id, "indexed")
+                    if session_id:
+                        await self._update_session_progress(session_id, "completed", doc_id)
+                else:
+                    #logger.error(f"Failed to index: {doc_id}")
+                    if session_id:
+                        await self._update_session_progress(session_id, "failed", doc_id)
+
+            except asyncio.CancelledError:
+                break
+            except Exception as e:
+                logger.error(f"Error processing document: {str(e)}")
+            finally:
+                self.indexing_queue.task_done()
+
+    async def _update_session_progress(self, session_id: str, status: str, doc_id: str):
+        """Update indexing progress for research session."""
+        try:
+            progress_key = f"indexing_progress:{session_id}"
+            await self.redis_client.hset(
+                progress_key,
+                mapping={
+                    doc_id: json.dumps({
+                        "status": status,
+                        "updated_at": _utc_timestamp()
+                    })
+                }
+            )
+            
+            # Publish update
+            await self.redis_client.publish(
+                "indexing_updates",
+                json.dumps({
+                    "session_id": session_id,
+                    "doc_id": doc_id,
+                    "status": status
+                })
+            )
+        except Exception as e:
+            logger.error(f"Error updating session progress: {str(e)}")
+
+    @log_operation("batch_index")
+    async def batch_index(self, doc_ids: List[str], session_id: Optional[str] = None) -> Dict[str, bool]:
+        """Batch index multiple documents with session tracking."""
+        #logger.info(f"Starting batch indexing: {len(doc_ids)} documents")
+        results = {}
+        
+        for doc_id in doc_ids:
+            await self.queue_document(doc_id, session_id)
+            results[doc_id] = True
+        
+        await self.indexing_queue.join()
+        return results
+
+    @log_operation("reindex_all")
+    async def reindex_all(self):
+        """Reindex all documents with progress tracking."""
+        #logger.info("Starting full reindexing")
+        try:
+            papers = await self.db.search_papers({"status": "processed"})
+            doc_ids = [paper.id for paper in papers]
+            
+            if not doc_ids:
+                logger.info("No documents found for reindexing")
+                return
+            
+            total = len(doc_ids)
+            progress_key = "reindex_progress"
+            
+            await self.redis_client.set(progress_key, "0")
+            results = await self.batch_index(doc_ids)
+            
+            success_count = sum(1 for success in results.values() if success)
+            await self.redis_client.set(progress_key, str(success_count / total * 100))
+            
+            #logger.info(f"Reindexing completed: {success_count}/{total} successful")
+            
+        except Exception as e:
+            logger.error(f"Error during reindexing: {str(e)}")
+            raise
+
+    async def cleanup(self):
+        """Cleanup indexing resources."""
+        #logger.info("Cleaning up indexer resources")
+        self._running = False
+        await self.redis_client.close()

src/services/search_service/search_engine.py

@@ -0,0 +1,342 @@
+"""High-level scholarly search orchestration service.
+
+The engine aggregates OpenAlex, optional PubMed, and lightweight web results to
+provide the advanced research surface required by the beta release. Emphasis is
+placed on resilience: network failures or missing API keys degrade gracefully
+instead of crashing the request path.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+import os
+from dataclasses import dataclass
+from datetime import datetime
+from typing import Any, Dict, Iterable, List, Optional, Sequence, Tuple
+
+import httpx
+
+from ..paper_service import OpenAlexClient
+from ..performance_service.rust_performance import HighPerformanceService
+
+logger = logging.getLogger(__name__)
+
+_PUBMED_BASE = "https://eutils.ncbi.nlm.nih.gov/entrez/eutils"
+_DDG_PROXY = os.getenv("DDG_PROXY_URL", "https://ddg-webapp-aagd.vercel.app/search")
+
+
+@dataclass(slots=True)
+class SearchResult:
+    """Canonical representation of a scholarly item."""
+
+    id: str
+    title: str
+    abstract: str
+    source: str
+    authors: List[str]
+    year: Optional[int]
+    doi: str
+    url: str
+    relevance: float
+    citations: int
+    keywords: List[str]
+    metadata: Dict[str, Any]
+
+
+class SearchEngine:
+    """Advanced scholarly search with optional enrichment."""
+
+    def __init__(
+        self,
+        *,
+        openalex_client: Optional[OpenAlexClient] = None,
+        performance_service: Optional[HighPerformanceService] = None,
+        redis_url: str = os.getenv("REDIS_URL", "redis://localhost:6379"),
+        timeout: float = float(os.getenv("SEARCH_TIMEOUT", "12.0")),
+    ) -> None:
+        self.openalex = openalex_client or OpenAlexClient()
+        self.performance = performance_service or HighPerformanceService()
+        self.redis_url = redis_url
+        self.timeout = timeout
+        self._session = httpx.AsyncClient(timeout=timeout)
+        self._lock = asyncio.Lock()
+
+    # ------------------------------------------------------------------
+    async def search_papers(
+        self,
+        query: str,
+        *,
+        limit: int = 10,
+        sources: Optional[Sequence[str]] = None,
+        include_metadata: bool = True,
+        include_abstracts: bool = True,
+        include_citations: bool = True,
+    ) -> Dict[str, Any]:
+        """Search across configured scholarly sources."""
+
+        if not query or not query.strip():
+            raise ValueError("Query must be a non-empty string")
+
+        limit = max(1, min(limit, 100))
+        sources = sources or ("openalex",)
+        gathered: List[SearchResult] = []
+
+        if "openalex" in sources:
+            gathered.extend(await self._search_openalex(query, limit))
+
+        if "pubmed" in sources:
+            try:
+                gathered.extend(await self._search_pubmed(query, limit))
+            except Exception as exc:  # pragma: no cover - optional remote dependency
+                logger.info("PubMed search failed", extra={"error": str(exc)})
+
+        deduped = self._deduplicate_results(gathered)
+        sorted_results = sorted(deduped, key=lambda item: item.relevance, reverse=True)[:limit]
+
+        payload = {
+            "query": query,
+            "count": len(sorted_results),
+            "sources_used": list(dict.fromkeys([res.source for res in sorted_results])),
+            "papers": [self._result_to_payload(res, include_metadata, include_abstracts, include_citations) for res in sorted_results],
+            "timestamp": datetime.utcnow().isoformat() + "Z",
+        }
+        return payload
+
+    async def web_search(self, query: str, *, num_results: int = 5) -> List[Dict[str, Any]]:
+        """Perform a lightweight DuckDuckGo-backed web search."""
+
+        params = {"q": query, "max_results": max(1, min(num_results, 10))}
+        try:
+            response = await self._session.get(_DDG_PROXY, params=params)
+            response.raise_for_status()
+            data = response.json()
+            results = data.get("results", []) if isinstance(data, dict) else []
+            formatted = [
+                {
+                    "title": item.get("title", ""),
+                    "url": item.get("href") or item.get("link") or "",
+                    "snippet": item.get("body") or item.get("snippet") or "",
+                    "source": item.get("source", "duckduckgo"),
+                }
+                for item in results[:num_results]
+            ]
+            return formatted
+        except Exception as exc:  # pragma: no cover - network optional
+            logger.info("Web search fallback", extra={"error": str(exc)})
+            return []
+
+    async def fetch_paper_bundle(self, paper_ids: Iterable[str]) -> List[Dict[str, Any]]:
+        """Convenience helper for fetching OpenAlex metadata for multiple IDs."""
+
+        papers = await self.openalex.get_papers_bulk(paper_ids)
+        formatted: List[Dict[str, Any]] = []
+        for paper in papers:
+            formatted.append(self._format_openalex_work(paper))
+        return formatted
+
+    async def close(self) -> None:
+        await self.openalex.close()
+        try:
+            await self._session.aclose()
+        except Exception:
+            pass
+
+    # ------------------------------------------------------------------
+    async def _search_openalex(self, query: str, limit: int) -> List[SearchResult]:
+        payload = await self.openalex.search_works(
+            query,
+            limit=limit,
+            filters={"type": "journal-article"},
+        )
+        results = payload.get("results", []) if isinstance(payload, dict) else []
+        formatted = [self._format_openalex_work(item) for item in results]
+
+        if formatted:
+            try:
+                combined = "\n".join(res.abstract for res in formatted if res.abstract)
+                if combined:
+                    keywords = await self.performance.extract_keywords(combined, max_keywords=10)
+                    for res in formatted:
+                        res.metadata.setdefault("query_keywords", keywords)
+            except Exception:
+                # Keyword enrichment is best-effort
+                pass
+
+        return formatted
+
+    async def _search_pubmed(self, query: str, limit: int) -> List[SearchResult]:
+        params = {
+            "db": "pubmed",
+            "term": query,
+            "retmax": max(1, min(limit, 50)),
+            "retmode": "json",
+            "sort": "relevance",
+        }
+        search_url = f"{_PUBMED_BASE}/esearch.fcgi"
+        response = await self._session.get(search_url, params=params)
+        response.raise_for_status()
+        id_list = response.json().get("esearchresult", {}).get("idlist", [])
+        if not id_list:
+            return []
+
+        summary_params = {
+            "db": "pubmed",
+            "id": ",".join(id_list[:limit]),
+            "retmode": "json",
+        }
+        summary_url = f"{_PUBMED_BASE}/esummary.fcgi"
+        summary_resp = await self._session.get(summary_url, params=summary_params)
+        summary_resp.raise_for_status()
+        summaries = summary_resp.json().get("result", {})
+
+        results: List[SearchResult] = []
+        for pmid in id_list[:limit]:
+            raw = summaries.get(pmid)
+            if not isinstance(raw, dict):
+                continue
+            title = raw.get("title", "")
+            abstract = raw.get("elocationid", "") or raw.get("source", "")
+            authors = [author.get("name", "") for author in raw.get("authors", []) if isinstance(author, dict)]
+            results.append(
+                SearchResult(
+                    id=f"PMID:{pmid}",
+                    title=title,
+                    abstract=abstract,
+                    source="pubmed",
+                    authors=[a for a in authors if a],
+                    year=self._safe_int(raw.get("pubdate", "")[:4]),
+                    doi=raw.get("elocationid", ""),
+                    url=f"https://pubmed.ncbi.nlm.nih.gov/{pmid}/",
+                    relevance=0.6,
+                    citations=raw.get("pmc", 0) or 0,
+                    keywords=[],
+                    metadata={"pmid": pmid},
+                )
+            )
+        return results
+
+    def _deduplicate_results(self, results: List[SearchResult]) -> List[SearchResult]:
+        seen: Dict[str, SearchResult] = {}
+        for result in results:
+            key = result.doi.lower() if result.doi else result.title.lower()
+            if key in seen:
+                existing = seen[key]
+                if result.relevance > existing.relevance:
+                    seen[key] = result
+            else:
+                seen[key] = result
+        return list(seen.values())
+
+    def _format_openalex_work(self, work: Dict[str, Any]) -> SearchResult:
+        title = work.get("title", "")
+        abstract = self._extract_openalex_abstract(work)
+        authors = [
+            auth.get("author", {}).get("display_name", "")
+            for auth in work.get("authorships", [])
+            if isinstance(auth, dict)
+        ]
+        doi = work.get("doi", "") or ""
+        url = work.get("id", "")
+        concepts = [concept.get("display_name", "") for concept in work.get("concepts", []) if isinstance(concept, dict)]
+        citations = work.get("cited_by_count", 0) or 0
+        relevance = work.get("relevance_score", 0.0) or 0.5 + min(0.4, citations / 1000)
+
+        keywords = concepts[:5] or self._quick_keywords(abstract, limit=5)
+
+        metadata = {
+            "openalex_id": url,
+            "open_access": work.get("open_access", {}).get("is_oa", False),
+            "primary_location": work.get("primary_location"),
+        }
+
+        return SearchResult(
+            id=url.split("/")[-1] if url else doi or title,
+            title=title,
+            abstract=abstract,
+            source="openalex",
+            authors=[a for a in authors if a],
+            year=work.get("publication_year"),
+            doi=doi.replace("https://doi.org/", ""),
+            url=url or f"https://openalex.org/{title[:50].replace(' ', '_')}",
+            relevance=float(relevance),
+            citations=citations,
+            keywords=keywords,
+            metadata=metadata,
+        )
+
+    def _result_to_payload(
+        self,
+        result: SearchResult,
+        include_metadata: bool,
+        include_abstracts: bool,
+        include_citations: bool,
+    ) -> Dict[str, Any]:
+        payload = {
+            "id": result.id,
+            "title": result.title,
+            "source": result.source,
+            "authors": result.authors,
+            "year": result.year,
+            "doi": result.doi,
+            "url": result.url,
+            "keywords": result.keywords,
+            "relevance": result.relevance,
+        }
+        if include_abstracts:
+            payload["abstract"] = result.abstract
+        if include_citations:
+            payload["citations_count"] = result.citations
+        if include_metadata:
+            payload["metadata"] = result.metadata
+        return payload
+
+    def _extract_openalex_abstract(self, work: Dict[str, Any]) -> str:
+        inverted = work.get("abstract_inverted_index")
+        if isinstance(inverted, dict) and inverted:
+            # Convert inverted index back to human-readable abstract
+            index_map: Dict[int, str] = {}
+            for token, positions in inverted.items():
+                for position in positions:
+                    index_map[position] = token
+            abstract_tokens = [token for _, token in sorted(index_map.items())]
+            return " ".join(abstract_tokens)
+        return work.get("abstract", "") or ""
+
+    def _safe_int(self, value: Any) -> Optional[int]:
+        try:
+            return int(value)
+        except Exception:
+            return None
+
+    def _quick_keywords(self, text: str, limit: int = 5) -> List[str]:
+        import re
+        from collections import Counter
+
+        if not text:
+            return []
+
+        words = re.findall(r"[a-zA-Z]{3,}", text.lower())
+        stop_words = {
+            "the",
+            "and",
+            "for",
+            "with",
+            "from",
+            "that",
+            "have",
+            "this",
+            "were",
+            "also",
+            "into",
+            "which",
+            "their",
+            "between",
+            "within",
+        }
+        filtered = [word for word in words if word not in stop_words]
+        most_common = Counter(filtered).most_common(limit)
+        return [word for word, _ in most_common]
+
+
+__all__ = ["SearchEngine"]