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

local_deep_research/advanced_search_system/strategies/entity_aware_source_strategy.py

@@ -0,0 +1,140 @@
+"""
+Entity-aware source-based search strategy for improved entity identification.
+"""
+
+import logging
+from typing import Dict
+
+from ..questions.entity_aware_question import EntityAwareQuestionGenerator
+from .source_based_strategy import SourceBasedSearchStrategy
+
+logger = logging.getLogger(__name__)
+
+
+class EntityAwareSourceStrategy(SourceBasedSearchStrategy):
+    """
+    Enhanced source-based strategy that better handles entity identification queries.
+
+    This strategy:
+    1. Detects when queries are seeking specific entities (names, places, etc.)
+    2. Generates more targeted search queries for entity identification
+    3. Focuses on extracting entity names from search results
+    """
+
+    def __init__(
+        self,
+        search=None,
+        model=None,
+        citation_handler=None,
+        include_text_content: bool = True,
+        use_cross_engine_filter: bool = True,
+        filter_reorder: bool = True,
+        filter_reindex: bool = True,
+        cross_engine_max_results: int = None,
+        all_links_of_system=None,
+        use_atomic_facts: bool = False,
+    ):
+        """Initialize with entity-aware components."""
+        # Initialize parent class
+        super().__init__(
+            search=search,
+            model=model,
+            citation_handler=citation_handler,
+            include_text_content=include_text_content,
+            use_cross_engine_filter=use_cross_engine_filter,
+            filter_reorder=filter_reorder,
+            filter_reindex=filter_reindex,
+            cross_engine_max_results=cross_engine_max_results,
+            all_links_of_system=all_links_of_system,
+            use_atomic_facts=use_atomic_facts,
+        )
+
+        # Replace the question generator with entity-aware version
+        self.question_generator = EntityAwareQuestionGenerator(self.model)
+
+    def _format_search_results_as_context(self, search_results):
+        """Format search results with emphasis on entity extraction."""
+        context_snippets = []
+        entity_mentions = []
+
+        for i, result in enumerate(search_results[:10]):
+            title = result.get("title", "Untitled")
+            snippet = result.get("snippet", "")
+            url = result.get("link", "")
+
+            if snippet:
+                context_snippets.append(
+                    f"Source {i + 1}: {title}\nURL: {url}\nSnippet: {snippet}"
+                )
+
+                # Extract potential entity names (basic approach)
+                # In production, this could use NER or more sophisticated extraction
+                words = snippet.split()
+                for j, word in enumerate(words):
+                    # Look for capitalized words that might be names
+                    if word and word[0].isupper() and len(word) > 2:
+                        # Get some context around the word
+                        context_start = max(0, j - 2)
+                        context_end = min(len(words), j + 3)
+                        entity_context = " ".join(
+                            words[context_start:context_end]
+                        )
+                        if entity_context not in entity_mentions:
+                            entity_mentions.append(entity_context)
+
+        # Add entity mentions to context if found
+        if entity_mentions:
+            context_snippets.append(
+                "\nPotential Entity Mentions:\n"
+                + "\n".join(f"- {mention}" for mention in entity_mentions[:10])
+            )
+
+        return "\n\n".join(context_snippets)
+
+    def analyze_topic(self, query: str) -> Dict:
+        """Analyze topic with enhanced entity identification."""
+        logger.info(
+            f"Starting entity-aware source-based research on topic: {query}"
+        )
+
+        # Detect if this is an entity identification query
+        entity_keywords = [
+            "who",
+            "what",
+            "which",
+            "identify",
+            "name",
+            "character",
+            "person",
+            "place",
+            "organization",
+            "company",
+            "author",
+            "scientist",
+            "inventor",
+            "city",
+            "country",
+            "book",
+            "movie",
+        ]
+
+        is_entity_query = any(
+            keyword in query.lower() for keyword in entity_keywords
+        )
+
+        if is_entity_query:
+            logger.info(
+                "Detected entity identification query - using enhanced search patterns"
+            )
+            self._update_progress(
+                "Detected entity identification query - optimizing search strategy",
+                3,
+                {
+                    "phase": "init",
+                    "query_type": "entity_identification",
+                    "strategy": "entity_aware_source",
+                },
+            )
+
+        # Use parent's analyze_topic with our enhanced components
+        return super().analyze_topic(query)