codebase-cortex 0.1.0__py3-none-any.whl

codebase_cortex/config.py

@@ -0,0 +1,150 @@
+"""Configuration and LLM factory for Codebase Cortex."""
+
+from __future__ import annotations
+
+import os
+from dataclasses import dataclass, field
+from pathlib import Path
+
+from dotenv import load_dotenv
+from langchain_core.language_models import BaseChatModel
+
+CORTEX_DIR_NAME = ".cortex"
+
+
+def find_cortex_dir(start: Path | None = None) -> Path:
+    """Find the .cortex directory, starting from the given path or cwd.
+
+    Returns the .cortex path (may not exist yet — for `cortex init`).
+    """
+    start = start or Path.cwd()
+    return start.resolve() / CORTEX_DIR_NAME
+
+
+DEFAULT_MODELS: dict[str, str] = {
+    "google": "gemini-2.5-flash-lite",
+    "anthropic": "claude-sonnet-4-20250514",
+    "openrouter": "",  # no sensible default — user must choose
+}
+
+RECOMMENDED_MODELS: dict[str, list[str]] = {
+    "google": [
+        "gemini-2.5-flash-lite",
+        "gemini-3-flash-preview",
+        "gemini-2.5-pro",
+    ],
+    "anthropic": [
+        "claude-sonnet-4-20250514",
+        "claude-haiku-4-5-20251001",
+    ],
+    "openrouter": [
+        "anthropic/claude-sonnet-4",
+        "google/gemini-2.5-flash-lite",
+        "google/gemini-3-flash-preview",
+    ],
+}
+
+
+@dataclass
+class Settings:
+    """Application settings loaded from .cortex/.env in the target repo."""
+
+    llm_provider: str = "google"
+    llm_model: str = ""
+    google_api_key: str = ""
+    anthropic_api_key: str = ""
+    openrouter_api_key: str = ""
+    github_token: str = ""
+    repo_path: Path = field(default_factory=lambda: Path.cwd())
+    cortex_dir: Path = field(default_factory=lambda: find_cortex_dir())
+    oauth_callback_port: int = 9876
+
+    @property
+    def data_dir(self) -> Path:
+        return self.cortex_dir
+
+    @property
+    def notion_token_path(self) -> Path:
+        return self.cortex_dir / "notion_tokens.json"
+
+    @property
+    def faiss_index_dir(self) -> Path:
+        return self.cortex_dir / "faiss_index"
+
+    @property
+    def page_cache_path(self) -> Path:
+        return self.cortex_dir / "page_cache.json"
+
+    @property
+    def env_path(self) -> Path:
+        return self.cortex_dir / ".env"
+
+    @classmethod
+    def from_env(cls, repo_path: Path | None = None) -> Settings:
+        """Load settings from .cortex/.env in the given or current directory."""
+        repo = (repo_path or Path.cwd()).resolve()
+        cortex_dir = repo / CORTEX_DIR_NAME
+        env_file = cortex_dir / ".env"
+
+        if env_file.exists():
+            load_dotenv(env_file)
+
+        provider = os.getenv("LLM_PROVIDER", "google")
+        return cls(
+            llm_provider=provider,
+            llm_model=os.getenv("LLM_MODEL", DEFAULT_MODELS.get(provider, "")),
+            google_api_key=os.getenv("GOOGLE_API_KEY", ""),
+            anthropic_api_key=os.getenv("ANTHROPIC_API_KEY", ""),
+            openrouter_api_key=os.getenv("OPENROUTER_API_KEY", ""),
+            github_token=os.getenv("GITHUB_TOKEN", ""),
+            repo_path=repo,
+            cortex_dir=cortex_dir,
+        )
+
+    @property
+    def is_initialized(self) -> bool:
+        return self.env_path.exists()
+
+
+def get_llm(settings: Settings | None = None, model: str | None = None) -> BaseChatModel:
+    """Create an LLM instance based on settings.
+
+    Args:
+        settings: Application settings. Loads from env if not provided.
+        model: Override model name. Uses LLM_MODEL from env, then provider default.
+    """
+    if settings is None:
+        settings = Settings.from_env()
+
+    resolved_model = model or settings.llm_model
+
+    if settings.llm_provider == "anthropic":
+        from langchain_anthropic import ChatAnthropic
+
+        return ChatAnthropic(
+            model=resolved_model or DEFAULT_MODELS["anthropic"],
+            api_key=settings.anthropic_api_key,
+        )
+
+    if settings.llm_provider == "openrouter":
+        from langchain_openai import ChatOpenAI
+
+        if not resolved_model:
+            raise ValueError(
+                "LLM_MODEL is required for OpenRouter. "
+                "Set it in .cortex/.env (e.g. LLM_MODEL=anthropic/claude-sonnet-4)"
+            )
+
+        return ChatOpenAI(
+            model=resolved_model,
+            api_key=settings.openrouter_api_key,
+            base_url="https://openrouter.ai/api/v1",
+        )
+
+    # Default: Google Gemini
+    from langchain_google_genai import ChatGoogleGenerativeAI
+
+    return ChatGoogleGenerativeAI(
+        model=resolved_model or DEFAULT_MODELS["google"],
+        google_api_key=settings.google_api_key,
+    )

codebase_cortex/embeddings/clustering.py

@@ -0,0 +1,140 @@
+"""HDBSCAN topic clustering for Knowledge Map generation."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from collections import defaultdict
+
+import hdbscan
+import numpy as np
+
+from codebase_cortex.embeddings.indexer import CodeChunk
+
+
+@dataclass
+class TopicCluster:
+    """A cluster of related code chunks representing a topic."""
+
+    cluster_id: int
+    label: str
+    chunks: list[CodeChunk] = field(default_factory=list)
+    centroid: np.ndarray | None = field(default=None, repr=False)
+
+    @property
+    def file_paths(self) -> list[str]:
+        """Unique file paths in this cluster."""
+        return sorted(set(c.file_path for c in self.chunks))
+
+    @property
+    def size(self) -> int:
+        return len(self.chunks)
+
+
+@dataclass
+class TopicClusterer:
+    """Clusters code embeddings into topics using HDBSCAN.
+
+    HDBSCAN is density-based — it automatically determines the number
+    of clusters and marks sparse points as noise (cluster_id = -1).
+    """
+
+    min_cluster_size: int = 3
+    min_samples: int = 2
+
+    def cluster(
+        self,
+        embeddings: np.ndarray,
+        chunks: list[CodeChunk],
+    ) -> list[TopicCluster]:
+        """Cluster embeddings and return topic groups.
+
+        Args:
+            embeddings: Array of shape (n, dimension).
+            chunks: Corresponding code chunks.
+
+        Returns:
+            List of TopicCluster, excluding noise cluster.
+        """
+        if len(embeddings) < self.min_cluster_size:
+            # Too few chunks to cluster meaningfully
+            return [TopicCluster(
+                cluster_id=0,
+                label=self._generate_label(chunks),
+                chunks=list(chunks),
+            )] if chunks else []
+
+        clusterer = hdbscan.HDBSCAN(
+            min_cluster_size=self.min_cluster_size,
+            min_samples=self.min_samples,
+            metric="euclidean",
+        )
+        labels = clusterer.fit_predict(embeddings.astype(np.float64))
+
+        # Group chunks by cluster
+        cluster_map: dict[int, list[tuple[CodeChunk, np.ndarray]]] = defaultdict(list)
+        for label, chunk, emb in zip(labels, chunks, embeddings):
+            cluster_map[int(label)].append((chunk, emb))
+
+        topics = []
+        for cluster_id, items in sorted(cluster_map.items()):
+            if cluster_id == -1:
+                continue  # Skip noise
+
+            cluster_chunks = [item[0] for item in items]
+            cluster_embeddings = np.array([item[1] for item in items])
+            centroid = cluster_embeddings.mean(axis=0)
+
+            topics.append(TopicCluster(
+                cluster_id=cluster_id,
+                label=self._generate_label(cluster_chunks),
+                chunks=cluster_chunks,
+                centroid=centroid,
+            ))
+
+        return topics
+
+    @staticmethod
+    def _generate_label(chunks: list[CodeChunk]) -> str:
+        """Generate a descriptive label from chunk metadata.
+
+        Uses the most common directory and chunk names to create
+        a human-readable topic label.
+        """
+        if not chunks:
+            return "Unknown"
+
+        # Find most common directory
+        dirs = defaultdict(int)
+        names = defaultdict(int)
+        for c in chunks:
+            parts = c.file_path.split("/")
+            if len(parts) > 1:
+                dirs[parts[0]] += 1
+            names[c.name] += 1
+
+        top_dir = max(dirs, key=dirs.get) if dirs else ""
+        # Pick top 2 most common names
+        top_names = sorted(names, key=names.get, reverse=True)[:2]
+
+        if top_dir:
+            return f"{top_dir}: {', '.join(top_names)}"
+        return ", ".join(top_names)
+
+    def to_markdown(self, topics: list[TopicCluster]) -> str:
+        """Render topic clusters as a Markdown Knowledge Map."""
+        if not topics:
+            return "No topics identified yet.\n"
+
+        lines = ["# Knowledge Map\n"]
+        lines.append(f"*{sum(t.size for t in topics)} code chunks across {len(topics)} topics*\n")
+
+        for topic in sorted(topics, key=lambda t: t.size, reverse=True):
+            lines.append(f"## {topic.label}")
+            lines.append(f"*{topic.size} chunks across {len(topic.file_paths)} files*\n")
+            for fp in topic.file_paths[:10]:
+                lines.append(f"- `{fp}`")
+            if len(topic.file_paths) > 10:
+                lines.append(f"- ... and {len(topic.file_paths) - 10} more files")
+            lines.append("")
+
+        return "\n".join(lines)

codebase_cortex/embeddings/indexer.py

@@ -0,0 +1,208 @@
+"""Sentence-transformers embedding pipeline for code chunks."""
+
+from __future__ import annotations
+
+import re
+from pathlib import Path
+from dataclasses import dataclass, field
+
+import numpy as np
+
+# Lazy-loaded to avoid slow import at startup
+_model = None
+MODEL_NAME = "all-MiniLM-L6-v2"
+
+# File extensions to index
+CODE_EXTENSIONS = {
+    ".py", ".js", ".ts", ".tsx", ".jsx", ".java", ".go", ".rs",
+    ".rb", ".php", ".c", ".cpp", ".h", ".hpp", ".cs", ".swift",
+    ".kt", ".scala", ".sh", ".bash", ".yml", ".yaml", ".toml",
+    ".json", ".md", ".rst", ".txt",
+}
+
+# Directories to skip
+SKIP_DIRS = {
+    ".git", ".venv", "venv", "node_modules", "__pycache__",
+    ".pytest_cache", "dist", "build", ".eggs", ".tox",
+    ".mypy_cache", ".ruff_cache",
+}
+
+# Max file size to index (100KB)
+MAX_FILE_SIZE = 100_000
+
+
+def _get_model():
+    """Lazy-load the sentence-transformers model."""
+    global _model
+    if _model is None:
+        from sentence_transformers import SentenceTransformer
+        _model = SentenceTransformer(MODEL_NAME)
+    return _model
+
+
+@dataclass
+class CodeChunk:
+    """A chunk of code with metadata for embedding."""
+
+    file_path: str
+    chunk_type: str  # "function" | "class" | "module" | "section"
+    name: str
+    content: str
+    start_line: int
+    end_line: int
+
+
+@dataclass
+class EmbeddingIndexer:
+    """Indexes code chunks using sentence-transformers.
+
+    Walks a repository, extracts meaningful code chunks,
+    and generates embeddings for similarity search.
+    """
+
+    repo_path: Path
+    chunks: list[CodeChunk] = field(default_factory=list)
+
+    def collect_chunks(self) -> list[CodeChunk]:
+        """Walk the repo and extract code chunks from all indexable files."""
+        self.chunks = []
+        for file_path in self._iter_files():
+            try:
+                content = file_path.read_text(encoding="utf-8", errors="ignore")
+                if not content.strip():
+                    continue
+                rel_path = str(file_path.relative_to(self.repo_path))
+                chunks = self._chunk_file(rel_path, content)
+                self.chunks.extend(chunks)
+            except (OSError, UnicodeDecodeError):
+                continue
+        return self.chunks
+
+    def embed_chunks(self, chunks: list[CodeChunk] | None = None) -> np.ndarray:
+        """Generate embeddings for code chunks.
+
+        Args:
+            chunks: Chunks to embed. Uses self.chunks if not provided.
+
+        Returns:
+            numpy array of shape (n_chunks, embedding_dim).
+        """
+        chunks = chunks or self.chunks
+        if not chunks:
+            return np.array([])
+
+        model = _get_model()
+        texts = [self._chunk_to_text(c) for c in chunks]
+        embeddings = model.encode(texts, show_progress_bar=False, convert_to_numpy=True)
+        return embeddings
+
+    def embed_texts(self, texts: list[str]) -> np.ndarray:
+        """Embed arbitrary text strings (for query embedding)."""
+        if not texts:
+            return np.array([])
+        model = _get_model()
+        return model.encode(texts, show_progress_bar=False, convert_to_numpy=True)
+
+    def _iter_files(self):
+        """Yield indexable files from the repo."""
+        for path in self.repo_path.rglob("*"):
+            if any(skip in path.parts for skip in SKIP_DIRS):
+                continue
+            if not path.is_file():
+                continue
+            if path.suffix not in CODE_EXTENSIONS:
+                continue
+            if path.stat().st_size > MAX_FILE_SIZE:
+                continue
+            yield path
+
+    def _chunk_file(self, rel_path: str, content: str) -> list[CodeChunk]:
+        """Split a file into meaningful chunks."""
+        if rel_path.endswith(".py"):
+            return self._chunk_python(rel_path, content)
+        # For non-Python files, chunk by sections or as whole module
+        return self._chunk_by_sections(rel_path, content)
+
+    def _chunk_python(self, rel_path: str, content: str) -> list[CodeChunk]:
+        """Extract Python functions and classes as chunks."""
+        chunks = []
+        lines = content.split("\n")
+
+        # Regex patterns for top-level definitions
+        func_pattern = re.compile(r"^(async\s+)?def\s+(\w+)")
+        class_pattern = re.compile(r"^class\s+(\w+)")
+
+        current_def = None
+        current_start = 0
+        current_name = ""
+        current_type = ""
+
+        for i, line in enumerate(lines):
+            func_match = func_pattern.match(line)
+            class_match = class_pattern.match(line)
+
+            if func_match or class_match:
+                # Save previous definition
+                if current_def is not None:
+                    chunk_content = "\n".join(lines[current_start:i])
+                    if chunk_content.strip():
+                        chunks.append(CodeChunk(
+                            file_path=rel_path,
+                            chunk_type=current_type,
+                            name=current_name,
+                            content=chunk_content,
+                            start_line=current_start + 1,
+                            end_line=i,
+                        ))
+
+                if func_match:
+                    current_type = "function"
+                    current_name = func_match.group(2)
+                else:
+                    current_type = "class"
+                    current_name = class_match.group(1)
+                current_def = True
+                current_start = i
+
+        # Save last definition
+        if current_def is not None:
+            chunk_content = "\n".join(lines[current_start:])
+            if chunk_content.strip():
+                chunks.append(CodeChunk(
+                    file_path=rel_path,
+                    chunk_type=current_type,
+                    name=current_name,
+                    content=chunk_content,
+                    start_line=current_start + 1,
+                    end_line=len(lines),
+                ))
+
+        # If no definitions found, treat whole file as module chunk
+        if not chunks and content.strip():
+            chunks.append(CodeChunk(
+                file_path=rel_path,
+                chunk_type="module",
+                name=rel_path,
+                content=content[:3000],
+                start_line=1,
+                end_line=len(lines),
+            ))
+
+        return chunks
+
+    def _chunk_by_sections(self, rel_path: str, content: str) -> list[CodeChunk]:
+        """Chunk non-Python files as whole modules (truncated if large)."""
+        lines = content.split("\n")
+        return [CodeChunk(
+            file_path=rel_path,
+            chunk_type="module",
+            name=rel_path,
+            content=content[:3000],
+            start_line=1,
+            end_line=len(lines),
+        )]
+
+    @staticmethod
+    def _chunk_to_text(chunk: CodeChunk) -> str:
+        """Convert a chunk to a text string suitable for embedding."""
+        return f"{chunk.file_path} ({chunk.chunk_type}: {chunk.name})\n{chunk.content[:2000]}"

codebase_cortex/embeddings/store.py

@@ -0,0 +1,126 @@
+"""FAISS index management for code embeddings."""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+from dataclasses import dataclass, field
+
+import faiss
+import numpy as np
+
+from codebase_cortex.embeddings.indexer import CodeChunk
+
+
+@dataclass
+class SearchResult:
+    """A single search result from the FAISS index."""
+
+    chunk: CodeChunk
+    distance: float
+    score: float  # 1 / (1 + distance), higher = more similar
+
+
+@dataclass
+class FAISSStore:
+    """Manages a FAISS vector index for code embeddings.
+
+    Stores embeddings in a flat L2 index with chunk metadata
+    persisted alongside in a JSON sidecar file.
+    """
+
+    index_dir: Path
+    index: faiss.IndexFlatL2 | None = field(default=None, repr=False)
+    chunks: list[CodeChunk] = field(default_factory=list)
+    _dimension: int = 384  # all-MiniLM-L6-v2 output dimension
+
+    def build(self, embeddings: np.ndarray, chunks: list[CodeChunk]) -> None:
+        """Build a new index from embeddings and chunks.
+
+        Args:
+            embeddings: Array of shape (n, dimension).
+            chunks: Corresponding code chunks (must match embeddings length).
+        """
+        if len(embeddings) == 0:
+            self.index = faiss.IndexFlatL2(self._dimension)
+            self.chunks = []
+            return
+
+        self._dimension = embeddings.shape[1]
+        self.index = faiss.IndexFlatL2(self._dimension)
+        self.index.add(embeddings.astype(np.float32))
+        self.chunks = list(chunks)
+
+    def search(self, query_embedding: np.ndarray, k: int = 5) -> list[SearchResult]:
+        """Search the index for the k most similar chunks.
+
+        Args:
+            query_embedding: Array of shape (1, dimension) or (dimension,).
+            k: Number of results to return.
+
+        Returns:
+            List of SearchResult sorted by similarity (highest first).
+        """
+        if self.index is None or self.index.ntotal == 0:
+            return []
+
+        query = query_embedding.reshape(1, -1).astype(np.float32)
+        k = min(k, self.index.ntotal)
+        distances, indices = self.index.search(query, k)
+
+        results = []
+        for dist, idx in zip(distances[0], indices[0]):
+            if idx < 0 or idx >= len(self.chunks):
+                continue
+            results.append(SearchResult(
+                chunk=self.chunks[idx],
+                distance=float(dist),
+                score=1.0 / (1.0 + float(dist)),
+            ))
+        return results
+
+    def save(self) -> None:
+        """Persist the index and chunk metadata to disk."""
+        self.index_dir.mkdir(parents=True, exist_ok=True)
+
+        if self.index is not None:
+            faiss.write_index(self.index, str(self.index_dir / "index.faiss"))
+
+        metadata = [
+            {
+                "file_path": c.file_path,
+                "chunk_type": c.chunk_type,
+                "name": c.name,
+                "content": c.content[:500],  # Truncate for storage
+                "start_line": c.start_line,
+                "end_line": c.end_line,
+            }
+            for c in self.chunks
+        ]
+        (self.index_dir / "chunks.json").write_text(json.dumps(metadata, indent=2))
+
+    def load(self) -> bool:
+        """Load an existing index from disk.
+
+        Returns:
+            True if loaded successfully, False if no index exists.
+        """
+        index_path = self.index_dir / "index.faiss"
+        chunks_path = self.index_dir / "chunks.json"
+
+        if not index_path.exists() or not chunks_path.exists():
+            return False
+
+        self.index = faiss.read_index(str(index_path))
+        self._dimension = self.index.d
+
+        metadata = json.loads(chunks_path.read_text())
+        self.chunks = [CodeChunk(**m) for m in metadata]
+        return True
+
+    @property
+    def size(self) -> int:
+        """Number of vectors in the index."""
+        if self.index is None:
+            return 0
+        return self.index.ntotal