codexa 0.4.0__py3-none-any.whl

semantic_code_intelligence/llm/ollama_provider.py

@@ -0,0 +1,122 @@
+"""Ollama LLM provider — integration with the Ollama local model server."""
+
+from __future__ import annotations
+
+import json
+from typing import Any
+from urllib.error import URLError
+from urllib.request import Request, urlopen
+
+from semantic_code_intelligence.llm.provider import (
+    LLMMessage,
+    LLMProvider,
+    LLMResponse,
+    MessageRole,
+)
+from semantic_code_intelligence.utils.logging import get_logger
+
+logger = get_logger("llm.ollama")
+
+
+class OllamaProvider(LLMProvider):
+    """LLM provider for the Ollama local model server.
+
+    Communicates via HTTP with the Ollama REST API.
+    No external packages required beyond the standard library.
+    """
+
+    def __init__(
+        self,
+        model: str = "llama3",
+        base_url: str = "http://localhost:11434",
+        temperature: float = 0.2,
+        max_tokens: int = 2048,
+    ) -> None:
+        self._model = model
+        self._base_url = base_url.rstrip("/")
+        self._temperature = temperature
+        self._max_tokens = max_tokens
+
+    @property
+    def name(self) -> str:
+        return "ollama"
+
+    def _api_call(self, endpoint: str, payload: dict[str, Any]) -> dict[str, Any]:
+        """Make a POST request to the Ollama API."""
+        url = f"{self._base_url}{endpoint}"
+        data = json.dumps(payload).encode("utf-8")
+        req = Request(url, data=data, headers={"Content-Type": "application/json"})
+
+        with urlopen(req, timeout=120) as resp:  # noqa: S310 — localhost only
+            result: dict[str, Any] = json.loads(resp.read().decode("utf-8"))
+            return result
+
+    def complete(self, prompt: str, **kwargs: Any) -> LLMResponse:
+        temperature = kwargs.get("temperature", self._temperature)
+        payload: dict[str, Any] = {
+            "model": self._model,
+            "prompt": prompt,
+            "stream": False,
+            "options": {
+                "temperature": temperature,
+                "num_predict": kwargs.get("max_tokens", self._max_tokens),
+            },
+        }
+
+        logger.debug("Ollama generate request: model=%s", self._model)
+        result = self._api_call("/api/generate", payload)
+
+        return LLMResponse(
+            content=result.get("response", ""),
+            model=self._model,
+            provider=self.name,
+            usage={
+                "prompt_tokens": result.get("prompt_eval_count", 0),
+                "completion_tokens": result.get("eval_count", 0),
+                "total_tokens": (
+                    result.get("prompt_eval_count", 0)
+                    + result.get("eval_count", 0)
+                ),
+            },
+        )
+
+    def chat(self, messages: list[LLMMessage], **kwargs: Any) -> LLMResponse:
+        temperature = kwargs.get("temperature", self._temperature)
+        payload: dict[str, Any] = {
+            "model": self._model,
+            "messages": [m.to_dict() for m in messages],
+            "stream": False,
+            "options": {
+                "temperature": temperature,
+                "num_predict": kwargs.get("max_tokens", self._max_tokens),
+            },
+        }
+
+        logger.debug("Ollama chat request: model=%s, messages=%d", self._model, len(messages))
+        result = self._api_call("/api/chat", payload)
+
+        msg = result.get("message", {})
+        return LLMResponse(
+            content=msg.get("content", ""),
+            model=self._model,
+            provider=self.name,
+            usage={
+                "prompt_tokens": result.get("prompt_eval_count", 0),
+                "completion_tokens": result.get("eval_count", 0),
+                "total_tokens": (
+                    result.get("prompt_eval_count", 0)
+                    + result.get("eval_count", 0)
+                ),
+            },
+        )
+
+    def is_available(self) -> bool:
+        """Check whether the Ollama server is reachable."""
+        try:
+            url = f"{self._base_url}/api/tags"
+            req = Request(url)
+            with urlopen(req, timeout=5) as resp:  # noqa: S310 — localhost only
+                ok: bool = resp.status == 200
+                return ok
+        except (URLError, OSError):
+            return False

semantic_code_intelligence/llm/openai_provider.py

@@ -0,0 +1,100 @@
+"""OpenAI LLM provider — integration with the OpenAI Chat Completions API."""
+
+from __future__ import annotations
+
+from typing import Any
+
+from semantic_code_intelligence.llm.provider import (
+    LLMMessage,
+    LLMProvider,
+    LLMResponse,
+    MessageRole,
+)
+from semantic_code_intelligence.utils.logging import get_logger
+
+logger = get_logger("llm.openai")
+
+
+class OpenAIProvider(LLMProvider):
+    """LLM provider for the OpenAI API (GPT-3.5, GPT-4, etc.).
+
+    Requires the ``openai`` package and a valid API key.
+    """
+
+    def __init__(
+        self,
+        api_key: str,
+        model: str = "gpt-3.5-turbo",
+        base_url: str | None = None,
+        temperature: float = 0.2,
+        max_tokens: int = 2048,
+    ) -> None:
+        self._api_key = api_key
+        self._model = model
+        self._base_url = base_url
+        self._temperature = temperature
+        self._max_tokens = max_tokens
+        self._client: Any = None
+
+    @property
+    def name(self) -> str:
+        return "openai"
+
+    def _get_client(self) -> Any:
+        """Lazily initialise the OpenAI client."""
+        if self._client is None:
+            try:
+                import openai
+            except ImportError as exc:
+                raise ImportError(
+                    "The 'openai' package is required for OpenAIProvider. "
+                    "Install it with: pip install openai"
+                ) from exc
+
+            kwargs: dict[str, Any] = {"api_key": self._api_key}
+            if self._base_url:
+                kwargs["base_url"] = self._base_url
+            self._client = openai.OpenAI(**kwargs)
+        return self._client
+
+    def complete(self, prompt: str, **kwargs: Any) -> LLMResponse:
+        messages = [LLMMessage(role=MessageRole.USER, content=prompt)]
+        return self.chat(messages, **kwargs)
+
+    def chat(self, messages: list[LLMMessage], **kwargs: Any) -> LLMResponse:
+        client = self._get_client()
+        temperature = kwargs.get("temperature", self._temperature)
+        max_tokens = kwargs.get("max_tokens", self._max_tokens)
+
+        api_messages = [m.to_dict() for m in messages]
+
+        logger.debug(
+            "OpenAI chat request: model=%s, messages=%d", self._model, len(api_messages)
+        )
+
+        response = client.chat.completions.create(
+            model=self._model,
+            messages=api_messages,
+            temperature=temperature,
+            max_tokens=max_tokens,
+        )
+
+        choice = response.choices[0]
+        usage = {}
+        if response.usage:
+            usage = {
+                "prompt_tokens": response.usage.prompt_tokens,
+                "completion_tokens": response.usage.completion_tokens,
+                "total_tokens": response.usage.total_tokens,
+            }
+
+        return LLMResponse(
+            content=choice.message.content or "",
+            model=response.model,
+            provider=self.name,
+            usage=usage,
+            raw=response.model_dump() if hasattr(response, "model_dump") else {},
+        )
+
+    def is_available(self) -> bool:
+        return bool(self._api_key)

semantic_code_intelligence/llm/provider.py

@@ -0,0 +1,92 @@
+"""LLM provider abstraction — base class and data types for LLM integration."""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass, field
+from enum import Enum
+from typing import Any
+
+
+class MessageRole(str, Enum):
+    """Role of a message in a conversation."""
+
+    SYSTEM = "system"
+    USER = "user"
+    ASSISTANT = "assistant"
+
+
+@dataclass
+class LLMMessage:
+    """A single message in an LLM conversation."""
+
+    role: MessageRole
+    content: str
+
+    def to_dict(self) -> dict[str, str]:
+        return {"role": self.role.value, "content": self.content}
+
+
+@dataclass
+class LLMResponse:
+    """Response from an LLM provider."""
+
+    content: str
+    model: str = ""
+    provider: str = ""
+    usage: dict[str, int] = field(default_factory=dict)
+    raw: dict[str, Any] = field(default_factory=dict)
+
+    def to_dict(self) -> dict[str, Any]:
+        return {
+            "content": self.content,
+            "model": self.model,
+            "provider": self.provider,
+            "usage": self.usage,
+        }
+
+
+class LLMProvider(ABC):
+    """Abstract base class for LLM providers.
+
+    Subclasses must implement `complete()` and `chat()`.
+    """
+
+    @property
+    @abstractmethod
+    def name(self) -> str:
+        """Provider name (e.g. 'openai', 'ollama')."""
+        ...
+
+    @abstractmethod
+    def complete(self, prompt: str, **kwargs: Any) -> LLMResponse:
+        """Generate a completion for a single prompt.
+
+        Args:
+            prompt: The text prompt to complete.
+            **kwargs: Provider-specific options (temperature, max_tokens, etc.).
+
+        Returns:
+            An LLMResponse containing the generated text.
+        """
+        ...
+
+    @abstractmethod
+    def chat(self, messages: list[LLMMessage], **kwargs: Any) -> LLMResponse:
+        """Generate a response for a multi-turn conversation.
+
+        Args:
+            messages: Conversation history as a list of LLMMessage.
+            **kwargs: Provider-specific options.
+
+        Returns:
+            An LLMResponse containing the assistant's reply.
+        """
+        ...
+
+    def is_available(self) -> bool:
+        """Check whether the provider is configured and reachable.
+
+        Default: True. Subclasses may override for connectivity checks.
+        """
+        return True

semantic_code_intelligence/llm/rate_limiter.py

@@ -0,0 +1,164 @@
+"""LLM rate limiter — sliding-window rate limiting for API calls.
+
+Enforces requests-per-minute (RPM) and tokens-per-minute (TPM) limits
+using a sliding window of recent events.  Callers can either block
+until capacity is available or receive a ``RateLimitExceeded`` error.
+"""
+
+from __future__ import annotations
+
+import threading
+import time
+from dataclasses import dataclass, field
+from typing import Any
+
+
+class RateLimitExceeded(Exception):
+    """Raised when a rate limit has been exceeded."""
+
+    def __init__(self, message: str = "Rate limit exceeded", retry_after: float = 0.0) -> None:
+        super().__init__(message)
+        self.retry_after = retry_after
+
+
+@dataclass
+class RateLimiterStats:
+    """Rate limiter statistics."""
+
+    total_requests: int = 0
+    total_tokens: int = 0
+    rejected_requests: int = 0
+    current_rpm: int = 0
+    current_tpm: int = 0
+
+    def to_dict(self) -> dict[str, Any]:
+        return {
+            "total_requests": self.total_requests,
+            "total_tokens": self.total_tokens,
+            "rejected_requests": self.rejected_requests,
+            "current_rpm": self.current_rpm,
+            "current_tpm": self.current_tpm,
+        }
+
+
+@dataclass
+class _Event:
+    """Internal record of a single API call."""
+
+    timestamp: float
+    tokens: int = 0
+
+
+class RateLimiter:
+    """Sliding-window rate limiter for LLM API calls.
+
+    Parameters
+    ----------
+    rpm : int
+        Maximum requests per minute.  0 = unlimited.
+    tpm : int
+        Maximum tokens per minute.  0 = unlimited.
+    blocking : bool
+        If ``True``, :meth:`acquire` will sleep until capacity is
+        available instead of raising :class:`RateLimitExceeded`.
+    """
+
+    def __init__(
+        self,
+        rpm: int = 0,
+        tpm: int = 0,
+        blocking: bool = True,
+    ) -> None:
+        self._rpm = rpm
+        self._tpm = tpm
+        self._blocking = blocking
+        self._events: list[_Event] = []
+        self._lock = threading.Lock()
+        self._stats = RateLimiterStats()
+
+    # ------------------------------------------------------------------
+    # Public API
+    # ------------------------------------------------------------------
+
+    def acquire(self, estimated_tokens: int = 0) -> None:
+        """Acquire permission to make an API call.
+
+        If the rate limit would be exceeded and ``blocking`` is True,
+        this method sleeps until capacity is available.  Otherwise it
+        raises :class:`RateLimitExceeded`.
+        """
+        while True:
+            with self._lock:
+                self._prune()
+                rpm_ok = self._check_rpm()
+                tpm_ok = self._check_tpm(estimated_tokens)
+                if rpm_ok and tpm_ok:
+                    # Record the request event (token count updated later via record_usage)
+                    self._events.append(_Event(timestamp=time.monotonic(), tokens=estimated_tokens))
+                    self._stats.total_requests += 1
+                    return
+
+                # Calculate wait time
+                wait = self._wait_time()
+
+            if not self._blocking:
+                self._stats.rejected_requests += 1
+                raise RateLimitExceeded(
+                    f"Rate limit exceeded (RPM={self._rpm}, TPM={self._tpm})",
+                    retry_after=wait,
+                )
+
+            # Sleep outside the lock
+            time.sleep(min(wait, 1.0))
+
+    def record_usage(self, tokens: int) -> None:
+        """Record the actual token usage after a response is received."""
+        with self._lock:
+            self._stats.total_tokens += tokens
+            # Update the last event's token count with actual usage
+            if self._events:
+                self._events[-1].tokens = tokens
+
+    @property
+    def stats(self) -> RateLimiterStats:
+        """Return current rate limiter statistics."""
+        with self._lock:
+            self._prune()
+            self._stats.current_rpm = len(self._events)
+            self._stats.current_tpm = sum(e.tokens for e in self._events)
+        return self._stats
+
+    @property
+    def is_enabled(self) -> bool:
+        """Return True if any rate limit is configured."""
+        return self._rpm > 0 or self._tpm > 0
+
+    # ------------------------------------------------------------------
+    # Internal helpers
+    # ------------------------------------------------------------------
+
+    def _prune(self) -> None:
+        """Remove events older than 60 seconds."""
+        cutoff = time.monotonic() - 60.0
+        self._events = [e for e in self._events if e.timestamp > cutoff]
+
+    def _check_rpm(self) -> bool:
+        """Check whether adding one more request is within the RPM limit."""
+        if self._rpm <= 0:
+            return True
+        return len(self._events) < self._rpm
+
+    def _check_tpm(self, estimated_tokens: int) -> bool:
+        """Check whether adding tokens is within the TPM limit."""
+        if self._tpm <= 0:
+            return True
+        current = sum(e.tokens for e in self._events)
+        return (current + estimated_tokens) <= self._tpm
+
+    def _wait_time(self) -> float:
+        """Estimate how long to wait before capacity is available."""
+        if not self._events:
+            return 0.1
+        oldest = self._events[0].timestamp
+        elapsed = time.monotonic() - oldest
+        return max(60.0 - elapsed, 0.1)