emdash-core 0.1.37__py3-none-any.whl → 0.1.60__py3-none-any.whl

emdash_core/context/providers/explored_areas.py

@@ -4,6 +4,14 @@ from dataclasses import asdict
 from typing import Optional, Union
 
 from ..models import ContextItem, ContextProviderSpec
+from ..tool_relevance import (
+    TOOL_RELEVANCE,
+    SEARCH_TOOLS,
+    TOP_RESULTS_LIMIT,
+    NON_TOP_RESULT_MULTIPLIER,
+    get_tool_relevance,
+    is_search_tool,
+)
 from .base import ContextProvider
 from ..registry import ContextProviderRegistry
 from ...graph.connection import KuzuConnection
@@ -16,44 +24,16 @@ class ExploredAreasProvider(ContextProvider):
     Analyzes the steps recorded during an agent session and assigns
     relevance scores based on the tool type used to discover each entity.
 
-    High relevance: deliberate investigation (expand_node, get_callers, etc.)
-    Medium relevance: targeted search (semantic_search, text_search)
-    Low relevance: broad search (grep, get_top_pagerank)
+    Scoring is defined in tool_relevance.py:
+    - Highest: Code modifications (write_to_file, apply_diff)
+    - High: Deliberate investigation (expand_node, get_callers, read_file)
+    - Medium: Targeted search (semantic_search, text_search, grep)
+    - Low: Broad discovery (list_files, graph algorithms)
     """
 
-    # Tool-based relevance scores
-    TOOL_RELEVANCE = {
-        # High relevance - deliberate investigation
-        "expand_node": 1.0,
-        "get_callers": 0.9,
-        "get_callees": 0.9,
-        "get_class_hierarchy": 0.9,
-        "get_neighbors": 0.85,
-        "get_impact_analysis": 0.85,
-        "read_file": 0.8,  # Reading a file is deliberate investigation
-        # Medium relevance - targeted search
-        "semantic_search": 0.7,
-        "text_search": 0.6,
-        "get_file_dependencies": 0.6,
-        "find_entity": 0.6,
-        # Lower relevance - broad search/modification
-        "grep": 0.4,
-        "write_to_file": 0.4,
-        "apply_diff": 0.4,
-        "get_top_pagerank": 0.3,
-        "get_communities": 0.3,
-        "list_files": 0.2,
-        "execute_command": 0.1,
-    }
-
-    # Only top N results from search tools are considered highly relevant
-    TOP_RESULTS_LIMIT = 3
-
-    # Tools where we limit to top results
-    SEARCH_TOOLS = {"semantic_search", "text_search", "grep", "find_entity"}
-
     def __init__(self, connection: KuzuConnection, config: Optional[dict] = None):
         super().__init__(connection, config)
+        self._neighbor_cache: dict[str, list[str]] = {}
 
     @property
     def spec(self) -> ContextProviderSpec:
@@ -88,10 +68,10 @@ class ExploredAreasProvider(ContextProvider):
                 entities = step.get("entities_discovered", [])
 
             # Get base relevance score for this tool
-            base_score = self.TOOL_RELEVANCE.get(tool_name, 0.2)
+            base_score = get_tool_relevance(tool_name)
 
             # For search tools, only top results are highly relevant
-            if tool_name in self.SEARCH_TOOLS:
+            if is_search_tool(tool_name):
                 # Process top results with full score, others with reduced score
                 for i, entity in enumerate(entities):
                     qname = self._extract_qualified_name(entity)
@@ -99,10 +79,10 @@ class ExploredAreasProvider(ContextProvider):
                         continue
 
                     # Top results get full score, others get reduced
-                    if i < self.TOP_RESULTS_LIMIT:
+                    if i < TOP_RESULTS_LIMIT:
                         score = base_score
                     else:
-                        score = base_score * 0.5  # Reduced score for non-top results
+                        score = base_score * NON_TOP_RESULT_MULTIPLIER
 
                     self._update_entity_score(entity_scores, qname, score, entity)
             else:
@@ -120,13 +100,17 @@ class ExploredAreasProvider(ContextProvider):
             display_name = qname
             if qname.startswith("file:"):
                 display_name = qname[5:]  # Remove "file:" prefix
+
+            # Fetch neighbors from graph
+            neighbors = self._fetch_neighbors(display_name, entity_type)
+
             items.append(
                 ContextItem(
                     qualified_name=display_name,
                     entity_type=entity_type or "Unknown",
                     file_path=file_path,
                     score=score,
-                    neighbors=[],  # Could fetch from graph if needed
+                    neighbors=neighbors,
                 )
             )
 
@@ -178,6 +162,66 @@ class ExploredAreasProvider(ContextProvider):
             return entity.get("file_path") or entity.get("path")
         return None
 
+    def _fetch_neighbors(
+        self, qualified_name: str, entity_type: Optional[str], limit: int = 5
+    ) -> list[str]:
+        """Fetch neighbors (callers/callees) from the graph.
+
+        Args:
+            qualified_name: The entity's qualified name
+            entity_type: The entity type (Function, Class, File)
+            limit: Maximum number of neighbors to return
+
+        Returns:
+            List of neighbor qualified names
+        """
+        # Check cache first
+        if qualified_name in self._neighbor_cache:
+            return self._neighbor_cache[qualified_name]
+
+        # Files don't have caller/callee relationships in the same way
+        if entity_type == "File" or not self.connection:
+            self._neighbor_cache[qualified_name] = []
+            return []
+
+        neighbors = []
+        try:
+            conn = self.connection.connect()
+
+            # Query for callers and callees
+            if entity_type in ("Function", "Class"):
+                # Get callees (what this entity calls)
+                callees_query = f"""
+                    MATCH (n:{entity_type} {{qualified_name: $qname}})-[:CALLS]->(m)
+                    RETURN m.qualified_name
+                    LIMIT $limit
+                """
+                result = conn.execute(callees_query, {"qname": qualified_name, "limit": limit})
+                while result.has_next():
+                    row = result.get_next()
+                    if row[0]:
+                        neighbors.append(row[0])
+
+                # Get callers (what calls this entity)
+                remaining = limit - len(neighbors)
+                if remaining > 0:
+                    callers_query = f"""
+                        MATCH (n)-[:CALLS]->(m:{entity_type} {{qualified_name: $qname}})
+                        RETURN n.qualified_name
+                        LIMIT $limit
+                    """
+                    result = conn.execute(callers_query, {"qname": qualified_name, "limit": remaining})
+                    while result.has_next():
+                        row = result.get_next()
+                        if row[0] and row[0] not in neighbors:
+                            neighbors.append(row[0])
+
+        except Exception as e:
+            log.debug(f"Failed to fetch neighbors for {qualified_name}: {e}")
+
+        self._neighbor_cache[qualified_name] = neighbors
+        return neighbors
+
 
 # Auto-register provider
 ContextProviderRegistry.register("explored_areas", ExploredAreasProvider)

emdash_core/context/reranker.py

@@ -1,108 +1,42 @@
 """Re-ranker for filtering context items by query relevance.
 
-Uses a cross-encoder model to score context items against the current query,
-keeping only the most relevant items to save tokens in the LLM context.
+Uses a lightweight scoring system based on:
+1. Text matching (query terms vs entity names/paths/descriptions)
+2. Graph signals (pagerank, betweenness centrality)
+3. Session signals (recency, touch frequency)
+4. Longevity signals (items that keep appearing are important)
+5. File co-occurrence (files with multiple entities get boosted)
+
+This reranker requires zero external ML dependencies and runs in <10ms.
 """
 
 import os
 from typing import Optional
 
-# Disable tokenizers parallelism to avoid fork warnings when running in threads
-os.environ["TOKENIZERS_PARALLELISM"] = "false"
-
 from .models import ContextItem
+from .simple_reranker import simple_rerank_items, get_simple_rerank_scores
 from ..utils.logger import log
 
-# Model singleton to avoid reloading
-_reranker_model = None
-_model_load_attempted = False
-
-
-def get_reranker_model():
-    """Get or load the re-ranker model (singleton).
-
-    Returns:
-        CrossEncoder model or None if not available
-    """
-    global _reranker_model, _model_load_attempted
-
-    if _model_load_attempted:
-        return _reranker_model
-
-    _model_load_attempted = True
-
-    # Check if re-ranking is enabled
-    if os.getenv("CONTEXT_RERANK_ENABLED", "true").lower() != "true":
-        log.debug("Context re-ranking disabled via CONTEXT_RERANK_ENABLED")
-        return None
-
-    try:
-        from sentence_transformers import CrossEncoder
-
-        model_name = os.getenv(
-            "CONTEXT_RERANK_MODEL", "mixedbread-ai/mxbai-rerank-xsmall-v1"
-        )
-        log.info(f"Loading re-ranker model: {model_name}")
-        _reranker_model = CrossEncoder(model_name)
-        log.info("Re-ranker model loaded successfully")
-        return _reranker_model
-    except ImportError:
-        log.warning("sentence-transformers not installed, re-ranking disabled")
-        return None
-    except Exception as e:
-        log.warning(f"Failed to load re-ranker model: {e}")
-        return None
-
-
-def item_to_text(item: ContextItem) -> str:
-    """Convert a ContextItem to text for re-ranking.
-
-    Args:
-        item: Context item to convert
-
-    Returns:
-        Text representation for scoring
-    """
-    parts = [item.qualified_name]
-
-    if item.entity_type:
-        parts.append(f"({item.entity_type})")
-
-    if item.description:
-        parts.append(f": {item.description[:200]}")
-
-    if item.file_path:
-        # Just include the filename, not full path
-        filename = os.path.basename(item.file_path)
-        parts.append(f" [file: {filename}]")
-
-    return " ".join(parts)
-
 
 def rerank_context_items(
     items: list[ContextItem],
     query: str,
     top_k: Optional[int] = None,
     top_percent: Optional[float] = None,
+    connection=None,
 ) -> list[ContextItem]:
     """Re-rank context items by relevance to query.
 
-    Uses a cross-encoder model to score each item against the query,
-    then returns the top K or top N% most relevant items.
-
     Args:
         items: List of context items to re-rank
         query: The user's query/task description
         top_k: Keep top K items (default from env: CONTEXT_RERANK_TOP_K=20)
         top_percent: Keep top N% items (overrides top_k if set)
+        connection: Optional Kuzu connection for graph-based scoring
 
     Returns:
         Filtered and sorted list of context items (most relevant first)
     """
-    import time
-
-    original_count = len(items)
-
     if not items:
         return items
 
@@ -110,63 +44,31 @@ def rerank_context_items(
         log.debug("No query provided for re-ranking, returning original items")
         return items
 
-    model = get_reranker_model()
-    if model is None:
-        log.debug("Re-ranker model not available, returning original items")
+    # Check if re-ranking is enabled
+    if os.getenv("CONTEXT_RERANK_ENABLED", "true").lower() != "true":
+        log.debug("Context re-ranking disabled via CONTEXT_RERANK_ENABLED")
         return items
 
-    try:
-        start_time = time.time()
-
-        # Convert items to text for scoring
-        texts = [item_to_text(item) for item in items]
-
-        # Create query-document pairs
-        pairs = [(query, text) for text in texts]
-
-        # Score all pairs
-        scores = model.predict(pairs)
-
-        # Combine items with scores
-        scored_items = list(zip(items, scores))
-
-        # Sort by score descending
-        scored_items.sort(key=lambda x: x[1], reverse=True)
-
-        # Determine how many to keep
-        if top_percent is not None:
-            keep_count = max(1, int(len(items) * top_percent))
-        elif top_k is not None:
-            keep_count = min(top_k, len(items))
-        else:
-            # Default from environment
-            default_top_k = int(os.getenv("CONTEXT_RERANK_TOP_K", "20"))
-            keep_count = min(default_top_k, len(items))
-
-        duration_ms = (time.time() - start_time) * 1000
-
-        # Log statistics
-        if scored_items:
-            max_score = scored_items[0][1]
-            min_score = scored_items[-1][1]
-            filtered_count = original_count - keep_count
-            log.info(
-                f"Re-ranked context: {original_count} -> {keep_count} items "
-                f"(filtered {filtered_count}) in {duration_ms:.0f}ms | "
-                f"scores [{min_score:.3f}-{max_score:.3f}] | "
-                f"query: '{query[:40]}...'"
-            )
-
-        # Return top items (without scores)
-        return [item for item, score in scored_items[:keep_count]]
-
-    except Exception as e:
-        log.warning(f"Re-ranking failed: {e}, returning original items")
-        return items
+    # Determine effective top_k
+    if top_percent is not None:
+        effective_top_k = max(1, int(len(items) * top_percent))
+    elif top_k is not None:
+        effective_top_k = min(top_k, len(items))
+    else:
+        effective_top_k = int(os.getenv("CONTEXT_RERANK_TOP_K", "20"))
+
+    return simple_rerank_items(
+        items=items,
+        query=query,
+        connection=connection,
+        top_k=effective_top_k,
+    )
 
 
 def get_rerank_scores(
-    items: list[ContextItem], query: str
+    items: list[ContextItem],
+    query: str,
+    connection=None,
 ) -> list[tuple[ContextItem, float]]:
     """Get re-rank scores for context items without filtering.
 
@@ -175,6 +77,7 @@ def get_rerank_scores(
     Args:
         items: List of context items
         query: Query to score against
+        connection: Optional Kuzu connection for graph signals
 
     Returns:
         List of (item, score) tuples sorted by score descending
@@ -182,18 +85,6 @@ def get_rerank_scores(
     if not items or not query:
         return [(item, 0.0) for item in items]
 
-    model = get_reranker_model()
-    if model is None:
-        return [(item, 0.0) for item in items]
-
-    try:
-        texts = [item_to_text(item) for item in items]
-        pairs = [(query, text) for text in texts]
-        scores = model.predict(pairs)
-
-        scored = list(zip(items, scores))
-        scored.sort(key=lambda x: x[1], reverse=True)
-        return scored
-    except Exception as e:
-        log.warning(f"Failed to get rerank scores: {e}")
-        return [(item, 0.0) for item in items]
+    scored = get_simple_rerank_scores(items, query, connection)
+    # Return without component breakdown
+    return [(item, score) for item, score, _ in scored]