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

emdash_core/context/longevity.py

@@ -0,0 +1,197 @@
+"""Longevity tracking for context items.
+
+Tracks which entities appear repeatedly across reranking calls.
+Items that keep appearing are likely important and get boosted.
+
+This uses an in-memory cache that resets on process restart.
+For persistence, the cache could be stored in the graph database.
+"""
+
+import math
+import time
+from dataclasses import dataclass, field
+from typing import Optional
+
+
+@dataclass
+class LongevityRecord:
+    """Track an entity's appearance history."""
+
+    qualified_name: str
+    appearance_count: int = 0
+    first_seen: float = field(default_factory=time.time)
+    last_seen: float = field(default_factory=time.time)
+
+    def record_appearance(self) -> None:
+        """Record a new appearance of this entity."""
+        self.appearance_count += 1
+        self.last_seen = time.time()
+
+    def get_longevity_score(self, now: Optional[float] = None) -> float:
+        """Calculate longevity score based on appearance count.
+
+        Longevity = items that have appeared in context frame more than once.
+        No time-based decay - if it keeps appearing, it's important.
+
+        Score formula (log scale for diminishing returns):
+        - 1 appearance = 0.0 (first time, no longevity yet)
+        - 2 appearances = 0.37
+        - 3 appearances = 0.50
+        - 5 appearances = 0.62
+        - 10 appearances = 0.77
+        - 20 appearances = 0.90
+
+        Args:
+            now: Current timestamp (unused, kept for API compatibility)
+
+        Returns:
+            Score between 0.0 and 1.0
+        """
+        if self.appearance_count <= 1:
+            return 0.0  # First appearance = no longevity
+
+        # Log scale for diminishing returns
+        # Subtract 1 so first repeat (count=2) starts contributing
+        return min(1.0, math.log(self.appearance_count) / 3)
+
+
+class LongevityTracker:
+    """Tracks entity appearances across reranking calls."""
+
+    def __init__(self, max_entries: int = 1000):
+        """Initialize the tracker.
+
+        Args:
+            max_entries: Maximum number of entities to track (LRU eviction)
+        """
+        self._records: dict[str, LongevityRecord] = {}
+        self._max_entries = max_entries
+
+    def record_appearance(self, qualified_name: str) -> None:
+        """Record that an entity appeared in reranking.
+
+        Args:
+            qualified_name: The entity's qualified name
+        """
+        if qualified_name in self._records:
+            self._records[qualified_name].record_appearance()
+        else:
+            # Evict oldest entries if at capacity
+            if len(self._records) >= self._max_entries:
+                self._evict_oldest()
+
+            self._records[qualified_name] = LongevityRecord(
+                qualified_name=qualified_name,
+                appearance_count=1,
+            )
+
+    def record_batch(self, qualified_names: list[str]) -> None:
+        """Record appearances for multiple entities.
+
+        Args:
+            qualified_names: List of entity qualified names
+        """
+        for qname in qualified_names:
+            self.record_appearance(qname)
+
+    def get_longevity_score(self, qualified_name: str) -> float:
+        """Get the longevity score for an entity.
+
+        Args:
+            qualified_name: The entity's qualified name
+
+        Returns:
+            Score between 0.0 and 1.0 (0.0 if never seen)
+        """
+        record = self._records.get(qualified_name)
+        if record is None:
+            return 0.0
+        return record.get_longevity_score()
+
+    def get_appearance_count(self, qualified_name: str) -> int:
+        """Get how many times an entity has appeared.
+
+        Args:
+            qualified_name: The entity's qualified name
+
+        Returns:
+            Number of appearances (0 if never seen)
+        """
+        record = self._records.get(qualified_name)
+        return record.appearance_count if record else 0
+
+    def _evict_oldest(self) -> None:
+        """Evict the oldest (least recently seen) entries."""
+        if not self._records:
+            return
+
+        # Sort by last_seen and remove bottom 10%
+        sorted_records = sorted(
+            self._records.items(),
+            key=lambda x: x[1].last_seen,
+        )
+
+        evict_count = max(1, len(sorted_records) // 10)
+        for qname, _ in sorted_records[:evict_count]:
+            del self._records[qname]
+
+    def clear(self) -> None:
+        """Clear all longevity records."""
+        self._records.clear()
+
+    def get_stats(self) -> dict:
+        """Get statistics about the tracker.
+
+        Returns:
+            Dictionary with tracker statistics
+        """
+        if not self._records:
+            return {
+                "total_entities": 0,
+                "total_appearances": 0,
+                "avg_appearances": 0,
+                "max_appearances": 0,
+            }
+
+        appearances = [r.appearance_count for r in self._records.values()]
+        return {
+            "total_entities": len(self._records),
+            "total_appearances": sum(appearances),
+            "avg_appearances": sum(appearances) / len(appearances),
+            "max_appearances": max(appearances),
+        }
+
+
+# Global tracker instance (shared across reranking calls)
+_global_tracker: Optional[LongevityTracker] = None
+
+
+def get_longevity_tracker() -> LongevityTracker:
+    """Get the global longevity tracker (creates if needed)."""
+    global _global_tracker
+    if _global_tracker is None:
+        _global_tracker = LongevityTracker()
+    return _global_tracker
+
+
+def record_reranked_items(qualified_names: list[str]) -> None:
+    """Record that items appeared in a reranking result.
+
+    Call this after reranking to update longevity scores.
+
+    Args:
+        qualified_names: List of qualified names that were reranked
+    """
+    get_longevity_tracker().record_batch(qualified_names)
+
+
+def get_longevity_score(qualified_name: str) -> float:
+    """Get the longevity score for an entity.
+
+    Args:
+        qualified_name: The entity's qualified name
+
+    Returns:
+        Score between 0.0 and 1.0
+    """
+    return get_longevity_tracker().get_longevity_score(qualified_name)

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]