library-desk/src/services/rag_search_service.py

"""
RAG Search service for Library Desk.

Provides web, news, and image search with content extraction:
- Uses SearXNG for search queries
- Uses Trafilatura for content extraction
- Caches results in Redis
"""

import hashlib
import json
import logging
import time
from typing import List, Optional
from urllib.parse import urlparse

import redis.asyncio as aioredis

from src.clients.searxng_client import SearXNGClient
from src.clients.content_extractor import ContentExtractor
from src.config import Settings
from src.models.rag_search import (
    SearchType,
    RAGSearchResult,
    RAGSearchResponse,
)

logger = logging.getLogger(__name__)


def extract_domain(url: str) -> str:
    """Extract domain name from URL, removing 'www.' prefix."""
    try:
        parsed = urlparse(url)
        domain = parsed.netloc
        return domain.removeprefix("www.")
    except Exception:
        return url


class RAGSearchService:
    """
    Service for RAG-optimized web search with content extraction.

    Combines SearXNG search with Trafilatura content extraction
    and Redis caching for efficient RAG pipeline integration.
    """

    def __init__(
        self,
        searxng_client: SearXNGClient,
        content_extractor: ContentExtractor,
        redis_client: aioredis.Redis,
        settings: Settings
    ):
        """
        Initialize RAG search service.

        Args:
            searxng_client: SearXNG search client
            content_extractor: Trafilatura content extractor
            redis_client: Async Redis client for caching
            settings: Application settings
        """
        self.searxng = searxng_client
        self.extractor = content_extractor
        self.redis = redis_client
        self.settings = settings

        self.cache_ttl = settings.search_cache_ttl
        self.default_limit = settings.search_default_limit

        logger.info(
            f"Initialized RAGSearchService: cache_ttl={self.cache_ttl}s, "
            f"default_limit={self.default_limit}"
        )

    def _cache_key(self, query: str, search_type: str, limit: int) -> str:
        """Generate cache key from search parameters."""
        key_data = f"{query}:{search_type}:{limit}"
        key_hash = hashlib.md5(key_data.encode()).hexdigest()
        return f"rag_search:{key_hash}"

    async def _get_cached_result(self, cache_key: str) -> Optional[RAGSearchResponse]:
        """Try to get cached search result."""
        try:
            cached = await self.redis.get(cache_key)
            if cached:
                data = json.loads(cached)
                logger.debug(f"Cache hit: {cache_key}")
                return RAGSearchResponse(**data)
        except Exception as e:
            logger.warning(f"Cache read failed: {e}")
        return None

    async def _set_cached_result(self, cache_key: str, result: RAGSearchResponse):
        """Cache search result."""
        try:
            await self.redis.setex(
                cache_key,
                self.cache_ttl,
                result.model_dump_json()
            )
            logger.debug(f"Cached result: {cache_key} (TTL={self.cache_ttl}s)")
        except Exception as e:
            logger.warning(f"Cache write failed: {e}")

    async def _search_searxng(
        self,
        query: str,
        search_type: SearchType,
        limit: int
    ) -> List[dict]:
        """Execute search via SearXNG based on search type."""
        try:
            if search_type == SearchType.WEB:
                results = await self.searxng.search_general(
                    query=query,
                    limit=limit
                )
            elif search_type == SearchType.NEWS:
                results = await self.searxng.search_news(
                    query=query,
                    limit=limit
                )
            elif search_type == SearchType.IMAGES:
                results = await self.searxng.search_images(
                    query=query,
                    limit=limit
                )
            else:
                results = await self.searxng.search_general(
                    query=query,
                    limit=limit
                )

            return results
        except Exception as e:
            logger.error(f"SearXNG search failed: {e}")
            raise

    async def _extract_content_for_results(
        self,
        results: List[dict]
    ) -> List[RAGSearchResult]:
        """Extract full content from search result URLs."""
        # Get URLs for extraction
        urls = [r.get("url", "") for r in results if r.get("url")]

        # Extract content in parallel
        extraction_results = await self.extractor.extract_batch(urls)

        # Build result objects
        search_results = []
        for i, raw_result in enumerate(results):
            url = raw_result.get("url", "")

            # Find matching extraction result
            extracted_content = ""
            for ext_result in extraction_results:
                if ext_result.url == url and ext_result.success:
                    extracted_content = ext_result.content
                    break

            # Get original snippet
            snippet = raw_result.get("content", "")
            if len(snippet) > 300:
                snippet = snippet[:300] + "..."

            # Build result
            search_results.append(RAGSearchResult(
                title=raw_result.get("title", ""),
                url=url,
                content=extracted_content,
                snippet=snippet,
                source=extract_domain(url),
                published_date=raw_result.get("publishedDate")
            ))

        return search_results

    def _generate_sources_summary(self, results: List[RAGSearchResult]) -> str:
        """Generate markdown list of source URLs."""
        if not results:
            return ""

        lines = ["## Sources"]
        for i, r in enumerate(results, 1):
            lines.append(f"{i}. [{r.title}]({r.url})")

        return "\n".join(lines)

    async def search(
        self,
        query: str,
        search_type: SearchType = SearchType.WEB,
        limit: Optional[int] = None,
        user: str = "default"
    ) -> RAGSearchResponse:
        """
        Execute RAG-optimized search.

        Args:
            query: Search query string
            search_type: Type of search (web, news, images)
            limit: Maximum results to return (default from settings)
            user: User identifier for logging/rate limiting

        Returns:
            RAGSearchResponse with extracted content and sources

        Raises:
            ValueError: If query is empty
            Exception: If search fails
        """
        start_time = time.time()

        if not query or not query.strip():
            raise ValueError("Query cannot be empty")

        effective_limit = limit or self.default_limit

        # Check cache
        cache_key = self._cache_key(query, search_type.value, effective_limit)
        cached = await self._get_cached_result(cache_key)
        if cached:
            return cached

        logger.info(
            f"RAG search: '{query}' type={search_type.value} "
            f"limit={effective_limit} user={user}"
        )

        # Execute search
        raw_results = await self._search_searxng(query, search_type, effective_limit)

        # Extract content from results
        search_results = await self._extract_content_for_results(raw_results)

        # Generate sources summary
        sources_summary = self._generate_sources_summary(search_results)

        # Calculate timing
        search_time_ms = int((time.time() - start_time) * 1000)

        # Build response
        response = RAGSearchResponse(
            query=query,
            search_type=search_type,
            results=search_results,
            total_results=len(search_results),
            search_time_ms=search_time_ms,
            sources_summary=sources_summary
        )

        # Cache result
        await self._set_cached_result(cache_key, response)

        logger.info(
            f"RAG search completed: {len(search_results)} results "
            f"in {search_time_ms}ms"
        )

        return response