diffrag 0.2.3__py3-none-any.whl

diffrag/__init__.py

@@ -0,0 +1,101 @@
+"""diffrag — AI-powered code review with RAG for git diffs.
+
+Public API surface:
+
+    from diffrag import (
+        Settings,
+        GitRepository,
+        DiffResult,
+        FileDiff,
+        Hunk,
+        OpenAICompatClient,
+        OllamaClient,
+        RepoIndexer,
+        CodeReviewer,
+        CodeReview,
+        FileReview,
+    )
+"""
+
+from .ai.base import AIClient, ChatMessage, CompletionResponse, EmbeddingClient, EmbeddingResponse
+from .ai.ollama import OllamaClient
+from .ai.openai_compat import OpenAICompatClient
+from .config.settings import (
+    AISettings,
+    EmbeddingSettings,
+    IndexingSettings,
+    ReviewSettings,
+    Settings,
+    default_db_path,
+    user_config_path,
+)
+from .diff.git import GitRepository
+from .diff.models import DiffResult, FileDiff, Hunk
+from .exceptions import (
+    AIClientError,
+    AICodeReviewerError,
+    ConfigurationError,
+    GitError,
+    IndexingError,
+)
+from .indexing.indexer import ContextChunk, RepoIndexer
+from .indexing.splitter import (
+    EXTENSION_MAP,
+    LANGUAGE_SPLITTERS,
+    CodeSplitter,
+    LanguageSplitter,
+    PythonASTSplitter,
+    SplitChunk,
+)
+from .review.reviewer import CodeReview, CodeReviewer, FileReview
+
+try:
+    from importlib.metadata import PackageNotFoundError
+    from importlib.metadata import version as _version
+
+    __version__ = _version("diffrag")
+except PackageNotFoundError:
+    __version__ = "unknown"
+
+__all__ = [
+    # config
+    "AISettings",
+    "default_db_path",
+    "EmbeddingSettings",
+    "IndexingSettings",
+    "ReviewSettings",
+    "Settings",
+    "user_config_path",
+    # diff
+    "DiffResult",
+    "FileDiff",
+    "GitRepository",
+    "Hunk",
+    # ai
+    "AIClient",
+    "AIClientError",
+    "ChatMessage",
+    "CompletionResponse",
+    "EmbeddingClient",
+    "EmbeddingResponse",
+    "OllamaClient",
+    "OpenAICompatClient",
+    # indexing
+    "CodeSplitter",
+    "ContextChunk",
+    "EXTENSION_MAP",
+    "LANGUAGE_SPLITTERS",
+    "LanguageSplitter",
+    "PythonASTSplitter",
+    "RepoIndexer",
+    "SplitChunk",
+    # review
+    "CodeReview",
+    "CodeReviewer",
+    "FileReview",
+    # exceptions
+    "AICodeReviewerError",
+    "ConfigurationError",
+    "GitError",
+    "IndexingError",
+]

diffrag/__main__.py

@@ -0,0 +1,3 @@
+from diffrag.cli.main import cli
+
+cli()

diffrag/ai/__init__.py

@@ -0,0 +1,15 @@
+"""AI client protocols and concrete implementations."""
+
+from .base import AIClient, ChatMessage, CompletionResponse, EmbeddingClient, EmbeddingResponse
+from .ollama import OllamaClient
+from .openai_compat import OpenAICompatClient
+
+__all__ = [
+    "AIClient",
+    "ChatMessage",
+    "CompletionResponse",
+    "EmbeddingClient",
+    "EmbeddingResponse",
+    "OllamaClient",
+    "OpenAICompatClient",
+]

diffrag/ai/base.py

@@ -0,0 +1,106 @@
+"""Protocol definitions and shared data models for AI clients."""
+
+from dataclasses import dataclass
+from typing import Protocol, runtime_checkable
+
+
+@dataclass(frozen=True)
+class ChatMessage:
+    """A single message in a chat conversation.
+
+    Attributes:
+        role: Speaker role — one of ``"system"``, ``"user"``, or ``"assistant"``.
+        content: Text content of the message.
+    """
+
+    role: str
+    content: str
+
+
+@dataclass(frozen=True)
+class CompletionResponse:
+    """Response returned by a language model completion call.
+
+    Attributes:
+        content: Generated text from the model.
+        model: Model identifier reported by the server.
+        prompt_tokens: Number of tokens consumed by the prompt.
+        completion_tokens: Number of tokens in the generated response.
+    """
+
+    content: str
+    model: str
+    prompt_tokens: int
+    completion_tokens: int
+
+    @property
+    def total_tokens(self) -> int:
+        """Total tokens used (prompt + completion).
+
+        Returns:
+            Sum of prompt and completion token counts.
+        """
+        return self.prompt_tokens + self.completion_tokens
+
+
+@dataclass(frozen=True)
+class EmbeddingResponse:
+    """Response returned by an embedding model call.
+
+    Attributes:
+        embeddings: List of embedding vectors, one per input text.
+        model: Model identifier reported by the server.
+        total_tokens: Number of tokens processed.
+    """
+
+    embeddings: list[list[float]]
+    model: str
+    total_tokens: int
+
+
+@runtime_checkable
+class AIClient(Protocol):
+    """Protocol for language model completion clients.
+
+    Any object implementing this protocol can be used as the AI backend
+    for generating reviews and answers.
+    """
+
+    async def complete(
+        self,
+        messages: list[ChatMessage],
+        *,
+        temperature: float = 0.2,
+        max_tokens: int | None = None,
+    ) -> CompletionResponse:
+        """Send a chat completion request.
+
+        Args:
+            messages: Ordered list of messages forming the conversation.
+            temperature: Sampling temperature (0 = deterministic).
+            max_tokens: Maximum number of tokens to generate; ``None`` = model default.
+
+        Returns:
+            A :class:`CompletionResponse` with the generated text and usage info.
+        """
+        ...
+
+
+@runtime_checkable
+class EmbeddingClient(Protocol):
+    """Protocol for text embedding clients.
+
+    Any object implementing this protocol can be used to build and query
+    the vector index.
+    """
+
+    async def embed(self, texts: list[str]) -> EmbeddingResponse:
+        """Compute embeddings for a list of texts.
+
+        Args:
+            texts: Non-empty list of strings to embed.
+
+        Returns:
+            An :class:`EmbeddingResponse` with one vector per input text.
+        """
+        ...

diffrag/ai/ollama.py

@@ -0,0 +1,144 @@
+"""Native Ollama REST client using /api/chat and /api/embed."""
+
+import logging
+
+import httpx
+
+from ..exceptions import AIClientError
+from .base import ChatMessage, CompletionResponse, EmbeddingResponse
+
+log = logging.getLogger(__name__)
+
+
+class OllamaClient:
+    """Client for the native Ollama REST API (``/api/chat`` and ``/api/embed``).
+
+    Use this client when you need Ollama-specific parameters (e.g. ``num_ctx``).
+    For endpoints that expose an OpenAI-compatible ``/v1`` prefix, prefer
+    :class:`~diffrag.ai.openai_compat.OpenAICompatClient` instead.
+
+    Args:
+        base_url: URL of the Ollama server (default: ``http://localhost:11434``).
+        model: Model tag to use for requests.
+        timeout: Per-request timeout in seconds.
+        num_ctx: Context window size passed via Ollama's ``options`` field.
+    """
+
+    def __init__(
+        self,
+        base_url: str = "http://localhost:11434",
+        model: str = "llama3.2",
+        timeout: float = 600.0,
+        num_ctx: int | None = None,
+    ) -> None:
+        self._base_url = base_url.rstrip("/")
+        self._model = model
+        self._timeout = timeout
+        self._num_ctx = num_ctx
+
+    # ------------------------------------------------------------------
+    # AIClient protocol
+    # ------------------------------------------------------------------
+
+    async def complete(
+        self,
+        messages: list[ChatMessage],
+        *,
+        temperature: float = 0.2,
+        max_tokens: int | None = None,
+    ) -> CompletionResponse:
+        """Send a chat completion request to Ollama's ``/api/chat`` endpoint.
+
+        Args:
+            messages: Ordered list of messages forming the conversation.
+            temperature: Sampling temperature.
+            max_tokens: Maximum tokens to generate (mapped to ``num_predict``).
+
+        Returns:
+            A :class:`CompletionResponse` with the generated text and usage info.
+
+        Raises:
+            AIClientError: If the HTTP request fails or returns a non-2xx status.
+        """
+        options: dict = {"temperature": temperature}
+        if self._num_ctx is not None:
+            options["num_ctx"] = self._num_ctx
+        if max_tokens is not None:
+            options["num_predict"] = max_tokens
+
+        payload = {
+            "model": self._model,
+            "messages": [{"role": m.role, "content": m.content} for m in messages],
+            "stream": False,
+            "options": options,
+        }
+
+        log.debug("POST %s/api/chat model=%s", self._base_url, self._model)
+
+        data = await self._post("/api/chat", payload)
+
+        return CompletionResponse(
+            content=data["message"]["content"],
+            model=data.get("model", self._model),
+            prompt_tokens=data.get("prompt_eval_count", 0),
+            completion_tokens=data.get("eval_count", 0),
+        )
+
+    # ------------------------------------------------------------------
+    # EmbeddingClient protocol
+    # ------------------------------------------------------------------
+
+    async def embed(self, texts: list[str]) -> EmbeddingResponse:
+        """Compute embeddings via Ollama's ``/api/embed`` endpoint.
+
+        Args:
+            texts: Non-empty list of strings to embed.
+
+        Returns:
+            An :class:`EmbeddingResponse` with one vector per input text.
+
+        Raises:
+            AIClientError: If the HTTP request fails or returns a non-2xx status.
+        """
+        payload = {"model": self._model, "input": texts}
+
+        log.debug("POST %s/api/embed model=%s n=%d", self._base_url, self._model, len(texts))
+
+        data = await self._post("/api/embed", payload)
+
+        return EmbeddingResponse(
+            embeddings=data["embeddings"],
+            model=data.get("model", self._model),
+            total_tokens=0,
+        )
+
+    # ------------------------------------------------------------------
+    # Internal helpers
+    # ------------------------------------------------------------------
+
+    async def _post(self, path: str, payload: dict) -> dict:
+        """Execute a POST request and return the parsed JSON body.
+
+        Args:
+            path: URL path relative to ``base_url`` (must start with ``/``).
+            payload: JSON-serialisable request body.
+
+        Returns:
+            Parsed JSON response as a dict.
+
+        Raises:
+            AIClientError: On network errors or non-2xx HTTP responses.
+        """
+        url = f"{self._base_url}{path}"
+        try:
+            async with httpx.AsyncClient(timeout=self._timeout) as client:
+                response = await client.post(url, json=payload)
+                response.raise_for_status()
+                return response.json()
+        except httpx.HTTPStatusError as exc:
+            raise AIClientError(
+                f"Ollama request to {url} failed with status {exc.response.status_code}: "
+                f"{exc.response.text[:200]}"
+            ) from exc
+        except httpx.RequestError as exc:
+            raise AIClientError(f"Network error reaching {url}: {exc}") from exc

diffrag/ai/openai_compat.py

@@ -0,0 +1,193 @@
+"""OpenAI-compatible REST client (works with OpenWebUI, Ollama /v1, OpenAI, etc.)."""
+
+import json
+import logging
+
+import httpx
+
+from ..exceptions import AIClientError
+from .base import ChatMessage, CompletionResponse, EmbeddingResponse
+
+log = logging.getLogger(__name__)
+
+
+class OpenAICompatClient:
+    """HTTP client for any OpenAI-compatible ``/chat/completions`` + ``/embeddings`` API.
+
+    A single instance can be used for both completion and embedding requests.
+    Point it at OpenWebUI, Ollama's ``/v1`` prefix, or the real OpenAI endpoint.
+
+    Args:
+        base_url: Root URL of the API (e.g. ``http://localhost:3000``).
+            The ``/v1`` prefix should **not** be included; it is appended automatically.
+        model: Default model name sent with every request.
+        api_key: Bearer token for the ``Authorization`` header. Empty string disables auth.
+        timeout: Per-request timeout in seconds.
+    """
+
+    def __init__(
+        self,
+        base_url: str,
+        model: str,
+        api_key: str = "",
+        timeout: float = 120.0,
+    ) -> None:
+        self._base_url = base_url.rstrip("/")
+        self._model = model
+        self._api_key = api_key
+        self._timeout = timeout
+
+    # ------------------------------------------------------------------
+    # AIClient protocol
+    # ------------------------------------------------------------------
+
+    async def complete(
+        self,
+        messages: list[ChatMessage],
+        *,
+        temperature: float = 0.2,
+        max_tokens: int | None = None,
+    ) -> CompletionResponse:
+        """Send a chat completion request using streaming to avoid gateway timeouts.
+
+        The request uses ``stream: true`` so the gateway receives bytes as soon as
+        the model starts generating, preventing idle-timeout 504 errors.  The full
+        response is accumulated internally; callers receive a single
+        :class:`CompletionResponse` as before.
+
+        Usage statistics are requested via ``stream_options: {include_usage: true}``.
+        Servers that do not support this option silently omit usage; token counts
+        default to 0 in that case.
+
+        Args:
+            messages: Ordered list of messages forming the conversation.
+            temperature: Sampling temperature.
+            max_tokens: Maximum tokens to generate; ``None`` uses the model default.
+
+        Returns:
+            A :class:`CompletionResponse` with the generated text and usage info.
+
+        Raises:
+            AIClientError: If the HTTP request fails or returns a non-2xx status.
+        """
+        url = f"{self._base_url}/v1/chat/completions"
+        payload: dict = {
+            "model": self._model,
+            "messages": [{"role": m.role, "content": m.content} for m in messages],
+            "temperature": temperature,
+            "stream": True,
+            "stream_options": {"include_usage": True},
+        }
+        if max_tokens is not None:
+            payload["max_tokens"] = max_tokens
+
+        log.debug("POST %s/v1/chat/completions model=%s (streaming)", self._base_url, self._model)
+
+        content = ""
+        model = self._model
+        prompt_tokens = completion_tokens = 0
+        try:
+            async with (
+                httpx.AsyncClient(timeout=self._timeout) as client,
+                client.stream("POST", url, json=payload, headers=self._headers()) as response,
+            ):
+                if not response.is_success:
+                    await response.aread()
+                response.raise_for_status()
+                async for line in response.aiter_lines():
+                    if not line.startswith("data: "):
+                        continue
+                    chunk_str = line[6:]
+                    if chunk_str == "[DONE]":
+                        break
+                    try:
+                        chunk = json.loads(chunk_str)
+                    except json.JSONDecodeError:
+                        continue
+                    model = chunk.get("model", model)
+                    for choice in chunk.get("choices", []):
+                        content += choice.get("delta", {}).get("content", "")
+                    if usage := chunk.get("usage"):
+                        prompt_tokens = usage.get("prompt_tokens", 0)
+                        completion_tokens = usage.get("completion_tokens", 0)
+        except httpx.HTTPStatusError as exc:
+            raise AIClientError(
+                f"API request to {url} failed with status {exc.response.status_code}: "
+                f"{exc.response.text[:200]}"
+            ) from exc
+        except httpx.RequestError as exc:
+            raise AIClientError(f"Network error reaching {url}: {exc}") from exc
+
+        return CompletionResponse(
+            content=content,
+            model=model,
+            prompt_tokens=prompt_tokens,
+            completion_tokens=completion_tokens,
+        )
+
+    # ------------------------------------------------------------------
+    # EmbeddingClient protocol
+    # ------------------------------------------------------------------
+
+    async def embed(self, texts: list[str]) -> EmbeddingResponse:
+        """Compute embeddings for a list of texts.
+
+        Args:
+            texts: Non-empty list of strings to embed.
+
+        Returns:
+            An :class:`EmbeddingResponse` with one vector per input text.
+
+        Raises:
+            AIClientError: If the HTTP request fails or returns a non-2xx status.
+        """
+        payload = {"model": self._model, "input": texts}
+
+        log.debug("POST %s/v1/embeddings model=%s n=%d", self._base_url, self._model, len(texts))
+
+        data = await self._post("/v1/embeddings", payload)
+        embeddings = [item["embedding"] for item in data["data"]]
+        usage = data.get("usage", {})
+
+        return EmbeddingResponse(
+            embeddings=embeddings,
+            model=data.get("model", self._model),
+            total_tokens=usage.get("total_tokens", 0),
+        )
+
+    # ------------------------------------------------------------------
+    # Internal helpers
+    # ------------------------------------------------------------------
+
+    def _headers(self) -> dict[str, str]:
+        headers: dict[str, str] = {"Content-Type": "application/json"}
+        if self._api_key:
+            headers["Authorization"] = f"Bearer {self._api_key}"
+        return headers
+
+    async def _post(self, path: str, payload: dict) -> dict:
+        """Execute a POST request and return the parsed JSON body.
+
+        Args:
+            path: URL path relative to ``base_url`` (must start with ``/``).
+            payload: JSON-serialisable request body.
+
+        Returns:
+            Parsed JSON response as a dict.
+
+        Raises:
+            AIClientError: On network errors or non-2xx HTTP responses.
+        """
+        url = f"{self._base_url}{path}"
+        try:
+            async with httpx.AsyncClient(timeout=self._timeout) as client:
+                response = await client.post(url, json=payload, headers=self._headers())
+                response.raise_for_status()
+                return response.json()
+        except httpx.HTTPStatusError as exc:
+            raise AIClientError(
+                f"API request to {url} failed with status {exc.response.status_code}: "
+                f"{exc.response.text[:200]}"
+            ) from exc
+        except httpx.RequestError as exc:
+            raise AIClientError(f"Network error reaching {url}: {exc}") from exc

diffrag/cli/__init__.py

@@ -0,0 +1,5 @@
+"""Click-based command-line interface."""
+
+from .main import cli
+
+__all__ = ["cli"]