coderay 1.0.0__py3-none-any.whl

coderay/core/utils.py

@@ -0,0 +1,35 @@
+import hashlib
+from pathlib import Path
+
+
+def hash_content(content: str) -> str:
+    """Compute SHA-256 hex digest of content."""
+    return hashlib.sha256(content.encode("utf-8")).hexdigest()
+
+
+def read_from_path(path: Path) -> str:
+    """Read file as UTF-8 text with replacement for invalid bytes."""
+    return path.read_text(encoding="utf-8", errors="replace")
+
+
+def files_with_changed_content(
+    repo: Path,
+    paths: list[Path],
+    file_hashes: dict[str, str],
+) -> list[Path]:
+    """Return paths whose content hash differs from file_hashes."""
+    result: list[Path] = []
+    for p in paths:
+        try:
+            rel = str(p.relative_to(repo))
+        except ValueError:
+            result.append(p)
+            continue
+        try:
+            content = p.read_text(encoding="utf-8", errors="replace")
+            h = hash_content(content)
+            if file_hashes.get(rel) != h:
+                result.append(p)
+        except Exception:
+            result.append(p)
+    return result

coderay/embedding/base.py

@@ -0,0 +1,60 @@
+from __future__ import annotations
+
+import logging
+from abc import ABC, abstractmethod
+from typing import Any
+
+from coderay.core.config import get_embedding_dimensions
+
+logger = logging.getLogger(__name__)
+
+
+class Embedder(ABC):
+    """Abstract embedder: embed(texts) -> list of vectors."""
+
+    @property
+    @abstractmethod
+    def dimensions(self) -> int:
+        """Vector dimension (e.g. 384 for all-MiniLM-L6-v2)."""
+        ...
+
+    @abstractmethod
+    def embed(self, texts: list[str]) -> list[list[float]]:
+        """Embed a list of texts; returns one vector per text."""
+        ...
+
+
+def load_embedder_from_config(config: dict[str, Any]) -> Embedder:
+    """Build an Embedder from a config dict."""
+    emb = config.get("embedder") or {}
+    provider = (emb.get("provider") or "local").lower()
+
+    if provider == "local":
+        try:
+            from coderay.embedding.local import LocalEmbedder
+        except ImportError as exc:
+            raise ImportError(
+                "Local provider requires 'fastembed'. Install with: pip install coderay"
+            ) from exc
+        return LocalEmbedder(
+            model=emb.get("model") or "sentence-transformers/all-MiniLM-L6-v2",
+            dimensions=get_embedding_dimensions(config),
+        )
+
+    if provider == "openai":
+        try:
+            from coderay.embedding.openai import OpenAIEmbedder
+        except ImportError as exc:
+            raise ImportError(
+                "OpenAI provider requires the 'openai' package. "
+                "Install with: pip install coderay[openai]"
+            ) from exc
+        return OpenAIEmbedder(
+            model=emb.get("model") or "text-embedding-3-small",
+            dimensions=get_embedding_dimensions(config),
+            api_key=emb.get("api_key"),
+        )
+
+    raise ValueError(
+        f"Unknown embedder provider: {provider}. Supported: 'local', 'openai'."
+    )

coderay/embedding/local.py

@@ -0,0 +1,68 @@
+from __future__ import annotations
+
+import logging
+import os
+
+from coderay.embedding.base import Embedder
+
+logger = logging.getLogger(__name__)
+
+DEFAULT_MODEL = "sentence-transformers/all-MiniLM-L6-v2"
+DEFAULT_DIMENSIONS = 384
+
+# all-MiniLM-L6-v2 max_seq_length is 256 tokens (~1200 chars of code).
+# Truncating early avoids the tokenizer wasting time on text the model
+# will discard anyway.
+#
+# TODO: symbols exceeding this limit lose tail information. Future options:
+#   - Split long chunks into overlapping windows and average embeddings
+#   - Use a model with a larger context (e.g. nomic-embed-text, 8192 tokens)
+#   - Embed a signature+docstring summary instead of raw code for large symbols
+MAX_CHARS = 1500
+
+# Number of parallel ONNX workers (0 = auto based on CPU cores)
+_PARALLEL_WORKERS = int(os.environ.get("EMBED_WORKERS", 0)) or None
+
+
+class LocalEmbedder(Embedder):
+    """CPU-only embeddings via fastembed (ONNX Runtime)."""
+
+    def __init__(
+        self,
+        model: str = DEFAULT_MODEL,
+        dimensions: int = DEFAULT_DIMENSIONS,
+    ) -> None:
+        """Initialize with model name and vector dimensions."""
+        self._dimensions = dimensions
+        self._model_name = model
+        self._model = None
+
+    def _load_model(self):
+        """Lazily load the fastembed model on first use."""
+        from fastembed import TextEmbedding
+
+        logger.info("Loading local embedding model %s...", self._model_name)
+        self._model = TextEmbedding(model_name=self._model_name)
+
+    @property
+    def dimensions(self) -> int:
+        """Vector dimension (e.g. 384 for all-MiniLM-L6-v2)."""
+        return self._dimensions
+
+    def embed(self, texts: list[str]) -> list[list[float]]:
+        """Embed texts using fastembed; returns one vector per text."""
+        if not texts:
+            return []
+        if self._model is None:
+            self._load_model()
+
+        truncated = [t[:MAX_CHARS] if len(t) > MAX_CHARS else t for t in texts]
+        logger.info("Embedding %d chunks...", len(truncated))
+        embeddings = list(
+            self._model.embed(
+                truncated,
+                parallel=_PARALLEL_WORKERS,
+                batch_size=256,
+            )
+        )
+        return [e.tolist() for e in embeddings]

coderay/embedding/openai.py

@@ -0,0 +1,87 @@
+from __future__ import annotations
+
+import logging
+import time
+
+from coderay.embedding.base import Embedder
+
+logger = logging.getLogger(__name__)
+
+MAX_CHARS_PER_TEXT = 8000
+MAX_RETRIES = 3
+RETRY_BASE_DELAY = 1.0
+
+
+class OpenAIEmbedder(Embedder):
+    """OpenAI API embedder."""
+
+    def __init__(
+        self,
+        model: str = "text-embedding-3-small",
+        dimensions: int = 1536,
+        api_key: str | None = None,
+    ):
+        """Initialize with model, dimensions, and API key."""
+        import os
+
+        import openai
+
+        self._model = model
+        self._dimensions = dimensions
+        self._api_key = api_key or os.environ.get("OPENAI_API_KEY")
+        if not self._api_key:
+            raise ValueError(
+                "OpenAI API key required: set OPENAI_API_KEY or pass api_key"
+            )
+        self._client = openai.OpenAI(api_key=self._api_key)
+
+    @property
+    def dimensions(self) -> int:
+        """Vector dimension (e.g. 1536 for text-embedding-3-small)."""
+        return self._dimensions
+
+    def embed(self, texts: list[str]) -> list[list[float]]:
+        """Embed texts via the OpenAI API; returns one vector per text."""
+        if not texts:
+            return []
+
+        truncated = [
+            t[:MAX_CHARS_PER_TEXT] if len(t) > MAX_CHARS_PER_TEXT else t for t in texts
+        ]
+
+        batch_size = 100
+        all_vectors: list[list[float]] = []
+        for i in range(0, len(truncated), batch_size):
+            batch = truncated[i : i + batch_size]
+            vecs = self._embed_with_retry(batch)
+            all_vectors.extend(vecs)
+        return all_vectors
+
+    def _embed_with_retry(self, texts: list[str]) -> list[list[float]]:
+        """Call the API with exponential backoff on transient errors."""
+        import openai
+
+        for attempt in range(MAX_RETRIES):
+            try:
+                resp = self._client.embeddings.create(
+                    model=self._model,
+                    input=texts,
+                    dimensions=self._dimensions,
+                )
+                return [e.embedding for e in resp.data]
+            except (
+                openai.RateLimitError,
+                openai.APITimeoutError,
+                openai.InternalServerError,
+            ) as exc:
+                if attempt == MAX_RETRIES - 1:
+                    raise
+                delay = RETRY_BASE_DELAY * (2**attempt)
+                logger.warning(
+                    "OpenAI embed attempt %d failed (%s), retrying in %.1fs...",
+                    attempt + 1,
+                    exc,
+                    delay,
+                )
+                time.sleep(delay)
+        return []

coderay/graph/__init__.py

@@ -0,0 +1,19 @@
+from coderay.graph.builder import (
+    GRAPH_FILENAME,
+    build_and_save_graph,
+    build_graph,
+    load_graph,
+    save_graph,
+)
+from coderay.graph.code_graph import CodeGraph
+from coderay.graph.extractor import GraphExtractor
+
+__all__ = [
+    "GRAPH_FILENAME",
+    "CodeGraph",
+    "GraphExtractor",
+    "build_and_save_graph",
+    "build_graph",
+    "load_graph",
+    "save_graph",
+]

coderay/graph/builder.py

@@ -0,0 +1,128 @@
+from __future__ import annotations
+
+import json
+import logging
+from pathlib import Path
+from typing import Any
+
+from coderay.graph.code_graph import CodeGraph
+from coderay.graph.extractor import GraphExtractor
+
+logger = logging.getLogger(__name__)
+
+GRAPH_FILENAME = "graph.json"
+
+
+def build_graph(
+    repo_root: str | Path,
+    file_paths_and_contents: list[tuple[str, str]],
+    config: dict[str, Any] | None = None,
+) -> CodeGraph:
+    """Extract a CodeGraph from the given files.
+
+    Returns:
+        Built CodeGraph with resolved edges.
+    """
+    extractor = GraphExtractor(config=config)
+    graph = CodeGraph()
+    for file_path, content in file_paths_and_contents:
+        try:
+            nodes, edges = extractor.extract_from_file(file_path, content)
+            graph.add_nodes_and_edges(nodes, edges)
+        except Exception as exc:
+            logger.warning("Graph extraction failed for %s: %s", file_path, exc)
+    resolved = graph.resolve_edges()
+    logger.info(
+        "Graph built: %d nodes, %d edges (%d call edges resolved)",
+        graph.node_count,
+        graph.edge_count,
+        resolved,
+    )
+    return graph
+
+
+def save_graph(graph: CodeGraph, index_dir: str | Path) -> Path:
+    """Persist the graph to index_dir/graph.json."""
+    path = Path(index_dir) / GRAPH_FILENAME
+    path.parent.mkdir(parents=True, exist_ok=True)
+    path.write_text(json.dumps(graph.to_dict(), indent=2))
+    logger.info("Saved graph to %s", path)
+    return path
+
+
+def load_graph(index_dir: str | Path) -> CodeGraph | None:
+    """Load a previously-saved graph, or None if it doesn't exist."""
+    path = Path(index_dir) / GRAPH_FILENAME
+    if not path.is_file():
+        return None
+    try:
+        data = json.loads(path.read_text())
+        return CodeGraph.from_dict(data)
+    except Exception as exc:
+        logger.warning("Failed to load graph from %s: %s", path, exc)
+        return None
+
+
+def build_and_save_graph(
+    repo_root: str | Path,
+    index_dir: str | Path,
+    changed_paths: list[str] | None = None,
+) -> None:
+    """Build or incrementally update the graph, then save."""
+    from coderay.core.config import load_config
+    from coderay.state.machine import StateMachine
+
+    repo = Path(repo_root)
+    idx_dir = Path(index_dir)
+    config = load_config(idx_dir)
+
+    existing_graph = load_graph(idx_dir) if changed_paths else None
+    incremental = existing_graph is not None and changed_paths is not None
+
+    if changed_paths is not None:
+        paths_to_parse = changed_paths
+    else:
+        from coderay.chunking.registry import get_supported_extensions
+
+        supported = get_supported_extensions()
+        sm = StateMachine(idx_dir)
+        paths_to_parse = [
+            p for p in sm.file_hashes if any(p.endswith(ext) for ext in supported)
+        ]
+
+    files_with_content: list[tuple[str, str]] = []
+    for p in paths_to_parse:
+        full = repo / p
+        if full.is_file():
+            try:
+                content = full.read_text(encoding="utf-8", errors="replace")
+                files_with_content.append((p, content))
+            except Exception as e:
+                logger.warning("Could not read %s for graph: %s", p, e)
+
+    if incremental:
+        extractor = GraphExtractor(config=config)
+        for fp in paths_to_parse:
+            existing_graph.remove_file(fp)
+        for fp, content in files_with_content:
+            try:
+                nodes, edges = extractor.extract_from_file(fp, content)
+                existing_graph.add_nodes_and_edges(nodes, edges)
+            except Exception as exc:
+                logger.warning("Graph extraction failed for %s: %s", fp, exc)
+        existing_graph.resolve_edges()
+        graph = existing_graph
+        logger.info(
+            "Graph incremental update: re-parsed %d files",
+            len(files_with_content),
+        )
+    else:
+        graph = build_graph(repo_root, files_with_content, config=config)
+
+    save_graph(graph, idx_dir)
+    logger.info(
+        "Graph: saved %d nodes, %d edges from %d files",
+        graph.node_count,
+        graph.edge_count,
+        len(files_with_content) if not incremental else graph.node_count,
+    )