cicada-mcp 0.2.0__py3-none-any.whl → 0.3.0__py3-none-any.whl

cicada/extractors/keybert.py

@@ -0,0 +1,228 @@
+"""
+Keyword Extraction using KeyBERT
+Semantic keyword extraction using transformer-based embeddings
+"""
+
+import os
+import re
+import sys
+from collections import Counter
+from typing import Any
+
+# Disable tokenizers parallelism to avoid fork warnings
+# Must be set before importing transformers/keybert
+os.environ["TOKENIZERS_PARALLELISM"] = "false"
+
+from cicada.extractors.keyword import BaseKeywordExtractor
+from cicada.utils import extract_code_identifiers
+
+
+class KeyBERTExtractor(BaseKeywordExtractor):
+    """Extract keywords from text using KeyBERT semantic analysis."""
+
+    # Single model configuration
+    MODEL_NAME = "BAAI/bge-small-en-v1.5"  # 133MB, balanced performance
+
+    # Weighting strategy constants for keyword extraction
+    KEYBERT_CANDIDATE_MULTIPLIER = 3  # Extract 3x keywords for weighted reranking
+    CODE_IDENTIFIER_BOOST = 10  # 10x weight for exact code identifiers
+    CODE_SPLIT_WORD_BOOST = 3  # 3x weight for identifier components
+    BASE_SCORE_IDENTIFIER = 0.5  # Base score for identifiers not found by BERT
+    BASE_SCORE_SPLIT_WORD = 0.3  # Base score for split words not found by BERT
+
+    # Class variable to hold KeyBERT class (lazily loaded)
+    _KeyBERT: type | None = None
+
+    def __init__(self, verbose: bool = False):
+        super().__init__(verbose)
+        self.verbose = verbose
+
+        # Print message BEFORE the slow import
+        if self.verbose:
+            print(
+                f"Loading KeyBERT model ({self.MODEL_NAME})",
+                file=sys.stderr,
+            )
+            print("This can take up to a couple of minutes.", file=sys.stderr)
+
+        # Lazy import KeyBERT (only once per class)
+        if KeyBERTExtractor._KeyBERT is None:
+            try:
+                from keybert import KeyBERT
+
+                KeyBERTExtractor._KeyBERT = KeyBERT
+            except ImportError as e:
+                raise ImportError(
+                    "KeyBERT is not installed. Install it with:\n"
+                    "  uv add keybert\n"
+                    "or\n"
+                    "  pip install keybert"
+                ) from e
+
+        # Initialize KeyBERT with the model
+        try:
+            self.kw_model = KeyBERTExtractor._KeyBERT(model=self.MODEL_NAME)
+            if self.verbose:
+                print("✓ Model loaded successfully", file=sys.stderr)
+        except Exception as e:
+            raise RuntimeError(
+                f"Failed to load KeyBERT model '{self.MODEL_NAME}'. "
+                f"Ensure the model is downloaded and available. Error: {e}"
+            ) from e
+
+    def _calculate_term_frequencies(self, text: str) -> dict[str, int]:
+        """Calculate term frequencies for all words in the text.
+
+        Args:
+            text: Input text to analyze
+
+        Returns:
+            Dictionary mapping lowercase words to their raw frequency counts
+        """
+        tokens = re.findall(r"\b[a-zA-Z][a-zA-Z0-9_]*\b", text.lower())
+        term_freq = Counter(tokens)
+        return dict(term_freq)
+
+    def _apply_code_boosting(
+        self,
+        keyword_scores: dict[str, float],
+        code_identifiers: list[str],
+        code_split_words: list[str],
+        tf_scores: dict[str, int],
+    ) -> None:
+        """Apply boosting to code identifiers and split words in-place.
+
+        Args:
+            keyword_scores: Dictionary of keyword scores to modify in-place
+            code_identifiers: List of code identifiers to boost
+            code_split_words: List of split words from identifiers to boost
+            tf_scores: Term frequency scores for calculating base scores
+        """
+        # Apply code identifier boosting
+        code_identifiers_lower = [ident.lower() for ident in code_identifiers]
+        for identifier in code_identifiers_lower:
+            if identifier in keyword_scores:
+                keyword_scores[identifier] *= self.CODE_IDENTIFIER_BOOST
+            else:
+                # Add with base score × frequency if not found by KeyBERT
+                freq = tf_scores.get(identifier, 1)
+                keyword_scores[identifier] = (
+                    self.BASE_SCORE_IDENTIFIER * freq * self.CODE_IDENTIFIER_BOOST
+                )
+
+        # Apply split word boosting (but only if not already a code identifier)
+        code_split_words_lower = [word.lower() for word in code_split_words]
+        code_identifiers_set = set(code_identifiers_lower)  # For O(1) lookup
+        for word in code_split_words_lower:
+            # Skip words that are already code identifiers (avoid double-boosting)
+            if word in code_identifiers_set:
+                continue
+            if word in keyword_scores:
+                keyword_scores[word] *= self.CODE_SPLIT_WORD_BOOST
+            else:
+                freq = tf_scores.get(word, 1)
+                keyword_scores[word] = (
+                    self.BASE_SCORE_SPLIT_WORD * freq * self.CODE_SPLIT_WORD_BOOST
+                )
+
+    def _calculate_statistics(self, text: str) -> dict[str, int]:
+        """Calculate basic text statistics.
+
+        Args:
+            text: Input text to analyze
+
+        Returns:
+            Dictionary with basic statistics (total_tokens, total_words, unique_words, sentences)
+        """
+        words = text.split()
+        unique_words = {w.lower() for w in words if w.isalpha()}
+        sentences = text.count(".") + text.count("!") + text.count("?")
+
+        return {
+            "total_tokens": len(words),
+            "total_words": len([w for w in words if w.isalpha()]),
+            "unique_words": len(unique_words),
+            "sentences": max(1, sentences),
+        }
+
+    def extract_keywords(
+        self, text: str, top_n: int = 15, min_score: float = 0.0
+    ) -> dict[str, Any]:
+        """
+        Extract keywords using KeyBERT semantic analysis with code identifier emphasis and frequency weighting.
+
+        Weighting strategy:
+        - Semantic score × raw frequency (repetition increases score, document length doesn't matter)
+        - Full code identifiers (e.g., getUserData, snake_case): 10x weight
+        - Code split words (e.g., get, user, data): 3x weight
+
+        Args:
+            text: Input text to analyze
+            top_n: Number of top keywords to return
+            min_score: Minimum score threshold for keywords (filters out low-scoring terms)
+
+        Returns:
+            Dictionary with extracted keywords and analysis:
+            - top_keywords: List of (keyword, score) tuples, sorted by weighted score
+            - code_identifiers: Original identifiers (weighted 10x)
+            - code_split_words: Words extracted from identifiers (weighted 3x)
+            - stats: Basic text statistics
+        """
+        if not text or not text.strip():
+            return {
+                "top_keywords": [],
+                "code_identifiers": [],
+                "code_split_words": [],
+                "tf_scores": {},
+                "stats": {
+                    "total_tokens": 0,
+                    "total_words": 0,
+                    "unique_words": 0,
+                    "sentences": 0,
+                },
+            }
+
+        # 1. Extract code identifiers and their split words
+        code_identifiers, code_split_words = extract_code_identifiers(text)
+
+        # 2. Calculate term frequencies for all words (raw counts, not normalized)
+        tf_scores = self._calculate_term_frequencies(text)
+
+        # 3. Use KeyBERT to extract semantic keywords
+        try:
+            keybert_keywords: list[tuple[str, float]] = self.kw_model.extract_keywords(  # type: ignore[assignment]
+                text,
+                top_n=top_n * self.KEYBERT_CANDIDATE_MULTIPLIER,
+                keyphrase_ngram_range=(1, 1),  # Single words only
+            )
+        except Exception as e:
+            if self.verbose:
+                print(f"Warning: KeyBERT extraction failed: {e}", file=sys.stderr)
+            keybert_keywords = []
+
+        # 4. Build weighted keyword scores (semantic × frequency)
+        keyword_scores: dict[str, float] = {}
+
+        # Add KeyBERT keywords with semantic score × frequency
+        for keyword, semantic_score in keybert_keywords:
+            keyword_lower: str = keyword.lower()
+            freq = tf_scores.get(keyword_lower, 1)  # Default frequency of 1
+            keyword_scores[keyword_lower] = semantic_score * freq
+
+        # 5. Apply code identifier and split word boosting
+        self._apply_code_boosting(keyword_scores, code_identifiers, code_split_words, tf_scores)
+
+        # 6. Filter by minimum score threshold and sort by weighted score
+        filtered_scores = {k: v for k, v in keyword_scores.items() if v >= min_score}
+        top_keywords = sorted(filtered_scores.items(), key=lambda x: x[1], reverse=True)[:top_n]
+
+        # 7. Calculate basic statistics
+        stats = self._calculate_statistics(text)
+
+        return {
+            "top_keywords": top_keywords,
+            "code_identifiers": code_identifiers,
+            "code_split_words": code_split_words,
+            "tf_scores": tf_scores,
+            "stats": stats,
+        }

cicada/extractors/keyword.py

@@ -0,0 +1,191 @@
+import sys
+from collections import Counter
+from typing import Any
+
+from cicada.utils import extract_code_identifiers as util_extract_code_identifiers
+
+
+class BaseKeywordExtractor:
+    """Base class for keyword extraction."""
+
+    def __init__(self, verbose: bool = False):
+        self.verbose = verbose
+        self.STOPWORDS = {
+            "the",
+            "a",
+            "an",
+            "and",
+            "or",
+            "but",
+            "in",
+            "on",
+            "at",
+            "to",
+            "for",
+            "of",
+            "with",
+            "by",
+            "from",
+            "as",
+            "is",
+            "are",
+            "was",
+            "were",
+            "be",
+            "been",
+            "being",
+            "have",
+            "has",
+            "had",
+            "do",
+            "does",
+            "did",
+            "will",
+            "would",
+            "should",
+            "could",
+            "this",
+            "that",
+            "these",
+            "those",
+            "it",
+            "its",
+            "they",
+            "them",
+            "their",
+            "what",
+            "which",
+            "who",
+            "when",
+            "where",
+            "why",
+            "how",
+            "all",
+            "each",
+            "every",
+            "both",
+            "few",
+            "more",
+            "most",
+            "other",
+            "some",
+            "such",
+            "no",
+            "nor",
+            "not",
+            "only",
+            "own",
+            "same",
+            "so",
+            "than",
+            "too",
+            "very",
+            "can",
+            "just",
+            "up",
+            "out",
+        }
+
+    def _tokenize(self, text: str) -> list[str]:
+        """Tokenize text into words."""
+        import re
+
+        return re.findall(r"\b[a-zA-Z][a-zA-Z0-9_]*\b", text)
+
+    def extract_keywords_simple(self, text: str, top_n: int = 10) -> list[str]:
+        if not text or not text.strip():
+            return []
+        try:
+            results = self.extract_keywords(text, top_n=top_n)
+            return [keyword for keyword, _ in results["top_keywords"]]
+        except Exception as e:
+            if self.verbose:
+                print(f"Warning: Keyword extraction failed: {e}", file=sys.stderr)
+            return []
+
+    def _extract_keywords(
+        self,
+        words: list[str],
+        code_identifiers: list[str],
+        code_split_words: list[str],
+        top_n: int,
+        total_tokens: int,
+    ) -> tuple[list, dict, dict]:
+        """Extract keywords from a list of words.
+
+        Args:
+            words: Filtered words (after stopword removal)
+            code_identifiers: Extracted code identifiers
+            code_split_words: Words split from code identifiers
+            top_n: Number of top keywords to return
+            total_tokens: Total token count before filtering (for stats)
+        """
+        code_identifiers_lower = [ident.lower() for ident in code_identifiers]
+        all_keywords = words + (code_identifiers_lower * 10) + (code_split_words * 3)
+        keyword_freq = Counter(all_keywords)
+        top_keywords = keyword_freq.most_common(top_n)
+
+        total_words = len(all_keywords)
+        if total_words > 0:
+            tf_scores = {word: (freq / total_words) for word, freq in keyword_freq.items()}
+        else:
+            tf_scores = {}
+
+        stats = {
+            "total_tokens": total_tokens,
+            "total_words": len(words),
+            "unique_words": len(set(words)),
+        }
+
+        return top_keywords, tf_scores, stats
+
+    def extract_keywords(
+        self, text: str, top_n: int = 15, min_score: float = 0.0
+    ) -> dict[str, Any]:
+        raise NotImplementedError
+
+
+class RegularKeywordExtractor(BaseKeywordExtractor):
+    """Extract keywords using basic term frequency (TF) without lemmatization."""
+
+    def extract_keywords(
+        self, text: str, top_n: int = 15, min_score: float = 0.0
+    ) -> dict[str, Any]:
+        if not text or not text.strip():
+            return {
+                "top_keywords": [],
+                "regular_words": [],
+                "code_identifiers": [],
+                "code_split_words": [],
+                "tf_scores": {},
+                "stats": {
+                    "total_tokens": 0,
+                    "total_words": 0,
+                    "unique_words": 0,
+                },
+            }
+
+        code_identifiers, code_split_words = util_extract_code_identifiers(text)
+        tokens = self._tokenize(text)
+        total_tokens = len(tokens)
+        regular_words = []
+        for word in tokens:
+            word_lower = word.lower()
+            if len(word) > 2 and word_lower not in self.STOPWORDS:
+                regular_words.append(word_lower)
+
+        top_keywords, tf_scores, stats = self._extract_keywords(
+            regular_words, code_identifiers, code_split_words, top_n, total_tokens
+        )
+
+        # Filter by minimum score threshold (min_score is a frequency count for RegularKeywordExtractor)
+        filtered_keywords = [(word, score) for word, score in top_keywords if score >= min_score]
+
+        return {
+            "top_keywords": filtered_keywords,
+            "regular_words": list(set(regular_words))[:20],
+            "code_identifiers": code_identifiers,
+            "code_split_words": code_split_words,
+            "tf_scores": dict(sorted(tf_scores.items(), key=lambda x: x[1], reverse=True)[:10]),
+            "stats": stats,
+        }

cicada/extractors/module.py

@@ -2,6 +2,8 @@
 Module extraction logic.
 """
 
+from cicada.utils import extract_text_from_node
+
 from .base import extract_string_from_arguments
 
 
@@ -31,7 +33,7 @@ def _find_modules_recursive(node, source_code: bytes, modules: list):
 
         # Check if this is a defmodule call
         if target and arguments:
-            target_text = source_code[target.start_byte : target.end_byte].decode("utf-8")
+            target_text = extract_text_from_node(target, source_code)
 
             if target_text == "defmodule":
                 # Extract module name from arguments
@@ -39,9 +41,7 @@ def _find_modules_recursive(node, source_code: bytes, modules: list):
 
                 for arg_child in arguments.children:
                     if arg_child.type == "alias":
-                        module_name = source_code[arg_child.start_byte : arg_child.end_byte].decode(
-                            "utf-8"
-                        )
+                        module_name = extract_text_from_node(arg_child, source_code)
                         break
 
                 if module_name and do_block:
@@ -82,9 +82,7 @@ def _find_moduledoc_recursive(node, source_code: bytes) -> str | None:
             # Check if this is a moduledoc attribute
             for call_child in operand.children:
                 if call_child.type == "identifier":
-                    attr_name = source_code[call_child.start_byte : call_child.end_byte].decode(
-                        "utf-8"
-                    )
+                    attr_name = extract_text_from_node(call_child, source_code)
 
                     if attr_name == "moduledoc":
                         # Extract the documentation string from the arguments
@@ -102,9 +100,7 @@ def _find_moduledoc_recursive(node, source_code: bytes) -> str | None:
             is_defmodule = False
             for call_child in child.children:
                 if call_child.type == "identifier":
-                    target_text = source_code[call_child.start_byte : call_child.end_byte].decode(
-                        "utf-8"
-                    )
+                    target_text = extract_text_from_node(call_child, source_code)
                     if target_text == "defmodule":
                         is_defmodule = True
                         break

cicada/extractors/spec.py

@@ -2,62 +2,18 @@
 Type spec extraction logic.
 """
 
+from cicada.utils import extract_text_from_node
+
+from .common import _find_attribute_recursive
+
 
 def extract_specs(node, source_code: bytes) -> dict:
     """Extract all @spec attributes from a module body."""
     specs = {}
-    _find_specs_recursive(node, source_code, specs)
+    _find_attribute_recursive(node, source_code, specs, "spec", _parse_spec)
     return specs
 
 
-def _find_specs_recursive(node, source_code: bytes, specs: dict):
-    """Recursively find @spec declarations."""
-    # Look for unary_operator nodes (which represent @ attributes)
-    if node.type == "unary_operator":
-        operator = None
-        operand = None
-
-        for child in node.children:
-            if child.type == "@":
-                operator = child
-            elif child.type == "call":
-                operand = child
-
-        if operator and operand:
-            # Check if this is a spec attribute
-            for call_child in operand.children:
-                if call_child.type == "identifier":
-                    attr_name = source_code[call_child.start_byte : call_child.end_byte].decode(
-                        "utf-8"
-                    )
-
-                    if attr_name == "spec":
-                        # Extract the spec definition
-                        spec_info = _parse_spec(operand, source_code)
-                        if spec_info:
-                            key = f"{spec_info['name']}/{spec_info['arity']}"
-                            specs[key] = spec_info
-
-    # Recursively search children
-    for child in node.children:
-        # Don't recurse into nested defmodule or function definitions
-        if child.type == "call":
-            is_defmodule_or_def = False
-            for call_child in child.children:
-                if call_child.type == "identifier":
-                    target_text = source_code[call_child.start_byte : call_child.end_byte].decode(
-                        "utf-8"
-                    )
-                    if target_text in ["defmodule", "def", "defp"]:
-                        is_defmodule_or_def = True
-                        break
-
-            if is_defmodule_or_def:
-                continue
-
-        _find_specs_recursive(child, source_code, specs)
-
-
 def _parse_spec(spec_node, source_code: bytes) -> dict | None:
     """Parse a @spec attribute to extract function name, arity, parameter types, and return type."""
     # @spec is represented as: spec(function_signature)
@@ -81,9 +37,7 @@ def _parse_spec(spec_node, source_code: bytes) -> dict | None:
                             found_call = True
                         elif found_call and op_child.type not in ["::", "operator"]:
                             # This is the return type node (after :: operator)
-                            return_type = source_code[
-                                op_child.start_byte : op_child.end_byte
-                            ].decode("utf-8")
+                            return_type = extract_text_from_node(op_child, source_code)
 
                     if func_call:
                         func_name = None
@@ -91,9 +45,7 @@ def _parse_spec(spec_node, source_code: bytes) -> dict | None:
 
                         for fc_child in func_call.children:
                             if fc_child.type == "identifier":
-                                func_name = source_code[
-                                    fc_child.start_byte : fc_child.end_byte
-                                ].decode("utf-8")
+                                func_name = extract_text_from_node(fc_child, source_code)
                             elif fc_child.type == "arguments":
                                 param_types = _extract_param_types(fc_child, source_code)
 
@@ -117,7 +69,7 @@ def _extract_param_types(params_node, source_code: bytes) -> list[str]:
             continue
 
         # Get the type as a string
-        type_str = source_code[child.start_byte : child.end_byte].decode("utf-8")
+        type_str = extract_text_from_node(child, source_code)
         param_types.append(type_str)
 
     return param_types

cicada/format/__init__.py

@@ -0,0 +1,20 @@
+"""Format module for Cicada - handles formatting, colors, and ASCII art."""
+
+from .ascii_art import generate_gradient_ascii_art
+from .colors import BOLD, CYAN, GREEN, GREY, PRIMARY, RESET, SELECTED, YELLOW
+from .formatter import JSONFormatter, ModuleFormatter, main
+
+__all__ = [
+    "generate_gradient_ascii_art",
+    "BOLD",
+    "CYAN",
+    "GREEN",
+    "GREY",
+    "PRIMARY",
+    "RESET",
+    "SELECTED",
+    "YELLOW",
+    "ModuleFormatter",
+    "JSONFormatter",
+    "main",
+]

cicada/{ascii_art.py → format/ascii_art.py}

@@ -1,6 +1,6 @@
 """ASCII art and banner generation for Cicada CLI."""
 
-from cicada.colors import CYAN, RESET, YELLOW
+from .colors import CYAN, RESET, YELLOW
 
 
 def generate_gradient_ascii_art():