local-deep-research 0.4.4__py3-none-any.whl → 0.5.2__py3-none-any.whl

local_deep_research/utilities/log_utils.py

@@ -2,10 +2,30 @@
 Utilities for logging.
 """
 
+# Needed for loguru annotations
+from __future__ import annotations
+
 import inspect
 import logging
+import sys
+from functools import wraps
+from pathlib import Path
+from typing import Any, Callable
 
+import loguru
+from flask import g, has_app_context
 from loguru import logger
+from sqlalchemy.exc import OperationalError
+
+from ..web.database.models import ResearchLog
+from ..web.services.socket_service import SocketIOService
+from .db_utils import get_db_session
+
+_LOG_DIR = Path(__file__).parents[2] / "data" / "logs"
+_LOG_DIR.mkdir(exist_ok=True)
+"""
+Default log directory to use.
+"""
 
 
 class InterceptHandler(logging.Handler):
@@ -34,3 +54,134 @@ class InterceptHandler(logging.Handler):
         logger.opt(depth=depth, exception=record.exc_info).log(
             level, record.getMessage()
         )
+
+
+def log_for_research(
+    to_wrap: Callable[[int, ...], Any],
+) -> Callable[[int, ...], Any]:
+    """
+    Decorator for a function that's part of the research process. It expects the function to
+    take the research ID as the first parameter, and configures all log
+    messages made during this request to include the research ID.
+
+    Args:
+        to_wrap: The function to wrap. Should take the research ID as the first parameter.
+
+    Returns:
+        The wrapped function.
+
+    """
+
+    @wraps(to_wrap)
+    def wrapped(research_id: int, *args: Any, **kwargs: Any) -> Any:
+        g.research_id = research_id
+        to_wrap(research_id, *args, **kwargs)
+        g.pop("research_id")
+
+    return wrapped
+
+
+def _get_research_id() -> int | None:
+    """
+    Gets the current research ID, if present.
+
+    Returns:
+        The current research ID, or None if it does not exist.
+
+    """
+    research_id = None
+    if has_app_context():
+        research_id = g.get("research_id")
+    return research_id
+
+
+def database_sink(message: loguru.Message) -> None:
+    """
+    Sink that saves messages to the database.
+
+    Args:
+        message: The log message to save.
+
+    """
+    record = message.record
+    research_id = _get_research_id()
+
+    # Create a new database entry.
+    db_log = ResearchLog(
+        timestamp=record["time"],
+        message=str(message),
+        module=record["name"],
+        function=record["function"],
+        line_no=int(record["line"]),
+        level=record["level"].name,
+        research_id=research_id,
+    )
+
+    # Save the entry to the database.
+    db_session = get_db_session()
+    try:
+        db_session.add(db_log)
+        db_session.commit()
+    except OperationalError:
+        # Something else is probably using the DB and we can't write to it
+        # right now. Ignore this.
+        db_session.rollback()
+        return
+
+
+def frontend_progress_sink(message: loguru.Message) -> None:
+    """
+    Sink that sends messages to the frontend.
+
+    Args:
+        message: The log message to send.
+
+    """
+    research_id = _get_research_id()
+    if research_id is None:
+        # If we don't have a research ID, don't send anything.
+        return
+
+    record = message.record
+    frontend_log = dict(
+        log_entry=dict(
+            message=record["message"],
+            type=record["level"].name,
+            time=record["time"].isoformat(),
+        ),
+    )
+    SocketIOService().emit_to_subscribers(
+        "progress", research_id, frontend_log, enable_logging=False
+    )
+
+
+def config_logger(name: str) -> None:
+    """
+    Configures the default logger.
+
+    Args:
+        name: The name to use for the log file.
+
+    """
+    logger.enable("local_deep_research")
+    logger.remove()
+
+    # Log more important stuff to the console.
+    logger.add(sys.stderr, level="INFO")
+    logger.add(
+        _LOG_DIR / f"{name}.log",
+        level="DEBUG",
+        enqueue=True,
+        rotation="00:00",
+        retention="30 days",
+        compression="zip",
+    )
+    logger.add(database_sink)
+    logger.add(frontend_progress_sink)
+
+    # Add a special log level for milestones.
+    try:
+        logger.level("milestone", no=26, color="<magenta><bold>")
+    except ValueError:
+        # Level already exists, that's fine
+        pass

local_deep_research/utilities/search_cache.py

@@ -0,0 +1,387 @@
+"""
+Search Cache Utility
+Provides intelligent caching for search results to avoid repeated queries.
+Includes TTL, LRU eviction, and query normalization.
+"""
+
+import hashlib
+import json
+import os
+import sqlite3
+import time
+from functools import lru_cache
+from typing import Any, Dict, List, Optional
+
+from loguru import logger
+
+
+class SearchCache:
+    """
+    Persistent cache for search results with TTL and LRU eviction.
+    Stores results in SQLite for persistence across sessions.
+    """
+
+    def __init__(
+        self,
+        cache_dir: str = None,
+        max_memory_items: int = 1000,
+        default_ttl: int = 3600,
+    ):
+        """
+        Initialize search cache.
+
+        Args:
+            cache_dir: Directory for cache database. Defaults to data/__CACHE_DIR__
+            max_memory_items: Maximum items in memory cache
+            default_ttl: Default time-to-live in seconds (1 hour default)
+        """
+        self.max_memory_items = max_memory_items
+        self.default_ttl = default_ttl
+
+        # Setup cache directory
+        if cache_dir is None:
+            cache_dir = os.path.join(
+                os.getcwd(), "data", "__CACHE_DIR__", "search_cache"
+            )
+
+        os.makedirs(cache_dir, exist_ok=True)
+        self.db_path = os.path.join(cache_dir, "search_cache.db")
+
+        # Initialize database
+        self._init_db()
+
+        # In-memory cache for frequently accessed items
+        self._memory_cache = {}
+        self._access_times = {}
+
+    def _init_db(self):
+        """Initialize SQLite database for persistent cache."""
+        try:
+            with sqlite3.connect(self.db_path) as conn:
+                conn.execute(
+                    """
+                    CREATE TABLE IF NOT EXISTS search_cache (
+                        query_hash TEXT PRIMARY KEY,
+                        query_text TEXT NOT NULL,
+                        results TEXT NOT NULL,
+                        created_at INTEGER NOT NULL,
+                        expires_at INTEGER NOT NULL,
+                        access_count INTEGER DEFAULT 1,
+                        last_accessed INTEGER NOT NULL
+                    )
+                """
+                )
+                conn.execute(
+                    """
+                    CREATE INDEX IF NOT EXISTS idx_expires_at ON search_cache(expires_at)
+                """
+                )
+                conn.execute(
+                    """
+                    CREATE INDEX IF NOT EXISTS idx_last_accessed ON search_cache(last_accessed)
+                """
+                )
+                conn.commit()
+        except Exception as e:
+            logger.error(f"Failed to initialize search cache database: {e}")
+
+    def _normalize_query(self, query: str) -> str:
+        """Normalize query for consistent caching."""
+        # Convert to lowercase and remove extra whitespace
+        normalized = " ".join(query.lower().strip().split())
+
+        # Remove common punctuation that doesn't affect search
+        normalized = normalized.replace('"', "").replace("'", "")
+
+        return normalized
+
+    def _get_query_hash(
+        self, query: str, search_engine: str = "default"
+    ) -> str:
+        """Generate hash for query + search engine combination."""
+        normalized_query = self._normalize_query(query)
+        cache_key = f"{search_engine}:{normalized_query}"
+        return hashlib.md5(cache_key.encode()).hexdigest()
+
+    def _cleanup_expired(self):
+        """Remove expired entries from database."""
+        try:
+            current_time = int(time.time())
+            with sqlite3.connect(self.db_path) as conn:
+                cursor = conn.cursor()
+                cursor.execute(
+                    "DELETE FROM search_cache WHERE expires_at < ?",
+                    (current_time,),
+                )
+                deleted = cursor.rowcount
+                conn.commit()
+                if deleted > 0:
+                    logger.debug(f"Cleaned up {deleted} expired cache entries")
+        except Exception as e:
+            logger.error(f"Failed to cleanup expired cache entries: {e}")
+
+    def _evict_lru_memory(self):
+        """Evict least recently used items from memory cache."""
+        if len(self._memory_cache) <= self.max_memory_items:
+            return
+
+        # Sort by access time and remove oldest
+        sorted_items = sorted(self._access_times.items(), key=lambda x: x[1])
+        items_to_remove = (
+            len(self._memory_cache) - self.max_memory_items + 100
+        )  # Remove extra for efficiency
+
+        for query_hash, _ in sorted_items[:items_to_remove]:
+            self._memory_cache.pop(query_hash, None)
+            self._access_times.pop(query_hash, None)
+
+    def get(
+        self, query: str, search_engine: str = "default"
+    ) -> Optional[List[Dict[str, Any]]]:
+        """
+        Get cached search results for a query.
+
+        Args:
+            query: Search query
+            search_engine: Search engine identifier for cache partitioning
+
+        Returns:
+            Cached results or None if not found/expired
+        """
+        query_hash = self._get_query_hash(query, search_engine)
+        current_time = int(time.time())
+
+        # Check memory cache first
+        if query_hash in self._memory_cache:
+            entry = self._memory_cache[query_hash]
+            if entry["expires_at"] > current_time:
+                self._access_times[query_hash] = current_time
+                logger.debug(f"Cache hit (memory) for query: {query[:50]}...")
+                return entry["results"]
+            else:
+                # Expired, remove from memory
+                self._memory_cache.pop(query_hash, None)
+                self._access_times.pop(query_hash, None)
+
+        # Check database cache
+        try:
+            with sqlite3.connect(self.db_path) as conn:
+                cursor = conn.cursor()
+                cursor.execute(
+                    """
+                    SELECT results, expires_at FROM search_cache
+                    WHERE query_hash = ? AND expires_at > ?
+                """,
+                    (query_hash, current_time),
+                )
+
+                row = cursor.fetchone()
+                if row:
+                    results_json, expires_at = row
+                    results = json.loads(results_json)
+
+                    # Update access statistics
+                    cursor.execute(
+                        """
+                        UPDATE search_cache
+                        SET access_count = access_count + 1, last_accessed = ?
+                        WHERE query_hash = ?
+                    """,
+                        (current_time, query_hash),
+                    )
+                    conn.commit()
+
+                    # Add to memory cache
+                    self._memory_cache[query_hash] = {
+                        "results": results,
+                        "expires_at": expires_at,
+                    }
+                    self._access_times[query_hash] = current_time
+                    self._evict_lru_memory()
+
+                    logger.debug(
+                        f"Cache hit (database) for query: {query[:50]}..."
+                    )
+                    return results
+
+        except Exception as e:
+            logger.error(f"Failed to retrieve from search cache: {e}")
+
+        logger.debug(f"Cache miss for query: {query[:50]}...")
+        return None
+
+    def put(
+        self,
+        query: str,
+        results: List[Dict[str, Any]],
+        search_engine: str = "default",
+        ttl: Optional[int] = None,
+    ) -> bool:
+        """
+        Store search results in cache.
+
+        Args:
+            query: Search query
+            results: Search results to cache
+            search_engine: Search engine identifier
+            ttl: Time-to-live in seconds (uses default if None)
+
+        Returns:
+            True if successfully cached
+        """
+        if not results:  # Don't cache empty results
+            return False
+
+        query_hash = self._get_query_hash(query, search_engine)
+        current_time = int(time.time())
+        expires_at = current_time + (ttl or self.default_ttl)
+
+        try:
+            results_json = json.dumps(results)
+
+            # Store in database
+            with sqlite3.connect(self.db_path) as conn:
+                conn.execute(
+                    """
+                    INSERT OR REPLACE INTO search_cache
+                    (query_hash, query_text, results, created_at, expires_at, access_count, last_accessed)
+                    VALUES (?, ?, ?, ?, ?, 1, ?)
+                """,
+                    (
+                        query_hash,
+                        self._normalize_query(query),
+                        results_json,
+                        current_time,
+                        expires_at,
+                        current_time,
+                    ),
+                )
+                conn.commit()
+
+            # Store in memory cache
+            self._memory_cache[query_hash] = {
+                "results": results,
+                "expires_at": expires_at,
+            }
+            self._access_times[query_hash] = current_time
+            self._evict_lru_memory()
+
+            logger.debug(f"Cached results for query: {query[:50]}...")
+            return True
+
+        except Exception as e:
+            logger.error(f"Failed to store in search cache: {e}")
+            return False
+
+    def invalidate(self, query: str, search_engine: str = "default") -> bool:
+        """Invalidate cached results for a specific query."""
+        query_hash = self._get_query_hash(query, search_engine)
+
+        try:
+            # Remove from memory
+            self._memory_cache.pop(query_hash, None)
+            self._access_times.pop(query_hash, None)
+
+            # Remove from database
+            with sqlite3.connect(self.db_path) as conn:
+                cursor = conn.cursor()
+                cursor.execute(
+                    "DELETE FROM search_cache WHERE query_hash = ?",
+                    (query_hash,),
+                )
+                deleted = cursor.rowcount
+                conn.commit()
+
+            logger.debug(f"Invalidated cache for query: {query[:50]}...")
+            return deleted > 0
+
+        except Exception as e:
+            logger.error(f"Failed to invalidate cache: {e}")
+            return False
+
+    def clear_all(self) -> bool:
+        """Clear all cached results."""
+        try:
+            self._memory_cache.clear()
+            self._access_times.clear()
+
+            with sqlite3.connect(self.db_path) as conn:
+                conn.execute("DELETE FROM search_cache")
+                conn.commit()
+
+            logger.info("Cleared all search cache")
+            return True
+
+        except Exception as e:
+            logger.error(f"Failed to clear search cache: {e}")
+            return False
+
+    def get_stats(self) -> Dict[str, Any]:
+        """Get cache statistics."""
+        try:
+            current_time = int(time.time())
+            with sqlite3.connect(self.db_path) as conn:
+                cursor = conn.cursor()
+
+                # Total entries
+                cursor.execute(
+                    "SELECT COUNT(*) FROM search_cache WHERE expires_at > ?",
+                    (current_time,),
+                )
+                total_entries = cursor.fetchone()[0]
+
+                # Total expired entries
+                cursor.execute(
+                    "SELECT COUNT(*) FROM search_cache WHERE expires_at <= ?",
+                    (current_time,),
+                )
+                expired_entries = cursor.fetchone()[0]
+
+                # Average access count
+                cursor.execute(
+                    "SELECT AVG(access_count) FROM search_cache WHERE expires_at > ?",
+                    (current_time,),
+                )
+                avg_access = cursor.fetchone()[0] or 0
+
+                return {
+                    "total_valid_entries": total_entries,
+                    "expired_entries": expired_entries,
+                    "memory_cache_size": len(self._memory_cache),
+                    "average_access_count": round(avg_access, 2),
+                    "cache_hit_potential": (
+                        f"{(total_entries / (total_entries + 1)) * 100:.1f}%"
+                        if total_entries > 0
+                        else "0%"
+                    ),
+                }
+
+        except Exception as e:
+            logger.error(f"Failed to get cache stats: {e}")
+            return {"error": str(e)}
+
+
+# Global cache instance
+_global_cache = None
+
+
+def get_search_cache() -> SearchCache:
+    """Get global search cache instance."""
+    global _global_cache
+    if _global_cache is None:
+        _global_cache = SearchCache()
+    return _global_cache
+
+
+@lru_cache(maxsize=100)
+def normalize_entity_query(entity: str, constraint: str) -> str:
+    """
+    Normalize entity + constraint combination for consistent caching.
+    Uses LRU cache for frequent normalizations.
+    """
+    # Remove quotes and normalize whitespace
+    entity_clean = " ".join(entity.strip().lower().split())
+    constraint_clean = " ".join(constraint.strip().lower().split())
+
+    # Create canonical form
+    return f"{entity_clean} {constraint_clean}"

local_deep_research/utilities/search_utilities.py

@@ -47,7 +47,6 @@ def format_links_to_markdown(all_links: List[Dict]) -> str:
     logger.info(f"Formatting {len(all_links)} links to markdown...")
 
     if all_links:
-
         # Group links by URL and collect all their indices
         url_to_indices = {}
         for link in all_links:
@@ -57,7 +56,6 @@ def format_links_to_markdown(all_links: List[Dict]) -> str:
             index = link.get("index", "")
             # logger.info(f"URL \n {str(url)} ")
             if url:
-
                 if url not in url_to_indices:
                     url_to_indices[url] = []
                 url_to_indices[url].append(index)
@@ -139,7 +137,9 @@ def format_findings(
         logger.info(f"Formatting {len(findings_list)} detailed finding items.")
 
         for idx, finding in enumerate(findings_list):
-            logger.debug(f"Formatting finding item {idx}. Keys: {list(finding.keys())}")
+            logger.debug(
+                f"Formatting finding item {idx}. Keys: {list(finding.keys())}"
+            )
             # Use .get() for safety
             phase = finding.get("phase", "Unknown Phase")
             content = finding.get("content", "No content available.")
@@ -201,8 +201,14 @@ def format_findings(
                     )
 
             # If the question is in the finding itself, display it
-            if not question_displayed and "question" in finding and finding["question"]:
-                formatted_text += f"### SEARCH QUESTION:\n{finding['question']}\n\n"
+            if (
+                not question_displayed
+                and "question" in finding
+                and finding["question"]
+            ):
+                formatted_text += (
+                    f"### SEARCH QUESTION:\n{finding['question']}\n\n"
+                )
 
             # Content
             formatted_text += f"\n\n{content}\n\n"
@@ -213,7 +219,9 @@ def format_findings(
                     links = extract_links_from_search_results(search_results)
                     if links:
                         formatted_text += "### SOURCES USED IN THIS SECTION:\n"
-                        formatted_text += format_links_to_markdown(links) + "\n\n"
+                        formatted_text += (
+                            format_links_to_markdown(links) + "\n\n"
+                        )
                 except Exception:
                     logger.exception(
                         f"Error processing search results/links for finding {idx}"

local_deep_research/utilities/threading_utils.py

@@ -0,0 +1,92 @@
+import threading
+from functools import wraps
+from typing import Any, Callable, Tuple
+from loguru import logger
+
+from cachetools import cached, keys
+from flask import current_app, g
+from flask.ctx import AppContext
+
+
+def thread_specific_cache(*args: Any, **kwargs: Any) -> Callable:
+    """
+    A version of `cached()` that is local to a single thread. In other words,
+    cache entries will only be valid in the thread where they were created.
+
+    Args:
+        *args: Will be forwarded to `cached()`.
+        **kwargs: Will be forwarded to `cached()`.
+
+    Returns:
+        The wrapped function.
+
+    """
+
+    def _key_func(*args_: Any, **kwargs_: Any) -> Tuple[int, ...]:
+        base_hash = keys.hashkey(*args_, **kwargs_)
+        return (threading.get_ident(),) + base_hash
+
+    return cached(*args, **kwargs, key=_key_func)
+
+
+def thread_with_app_context(to_wrap: Callable) -> Callable:
+    """
+    Decorator that wraps the entry point to a thread and injects the current
+    app context from Flask. This is useful when we want to use multiple
+    threads to handle a single request.
+
+    When using this wrapped function, `current_app.app_context()` should be
+    passed as the first argument when initializing the thread.
+
+    Args:
+        to_wrap: The function to wrap.
+
+    Returns:
+        The wrapped function.
+
+    """
+
+    @wraps(to_wrap)
+    def _run_with_context(
+        app_context: AppContext | None, *args: Any, **kwargs: Any
+    ) -> Any:
+        if app_context is None:
+            # Do nothing.
+            return to_wrap(*args, **kwargs)
+
+        with app_context:
+            return to_wrap(*args, **kwargs)
+
+    return _run_with_context
+
+
+def thread_context() -> AppContext | None:
+    """
+    Pushes a new app context for a thread that is being spawned to handle the
+    current request. Will copy all the global data from the current context.
+
+    Returns:
+        The new context, or None if no context is active.
+
+    """
+    # Copy global data.
+    global_data = {}
+    try:
+        for key in g:
+            global_data[key] = g.get(key)
+    except TypeError:
+        # Context is not initialized. Don't change anything.
+        pass
+
+    try:
+        context = current_app.app_context()
+    except RuntimeError:
+        # Context is not initialized.
+        logger.debug("No current app context, not passing to thread.")
+        return None
+
+    with context:
+        for key, value in global_data.items():
+            setattr(g, key, value)
+
+    return context