pr-context-engine 0.1.0__py3-none-any.whl

src/llm/gemini_provider.py

@@ -0,0 +1,33 @@
+"""Google Gemini-backed LLMProvider implementation (model: gemini-2.5-flash)."""
+import logging
+
+from google import genai
+
+from src.llm.base import LLMProvider
+
+logger = logging.getLogger(__name__)
+
+_MODEL = "gemini-2.5-flash"
+
+
+class GeminiProvider(LLMProvider):
+    """Calls the Gemini API via google-genai to satisfy the LLMProvider contract."""
+
+    def __init__(self, api_key: str) -> None:
+        """Build a Gemini client from an explicit API key."""
+        self._client = genai.Client(api_key=api_key)
+
+    def generate(self, prompt: str) -> str:
+        """Send `prompt` to Gemini and return the completion text."""
+        logger.info("Requesting Gemini completion (model=%s)", _MODEL)
+        response = self._client.models.generate_content(
+            model=_MODEL,
+            contents=prompt,
+        )
+        try:
+            text = response.text
+        except ValueError as exc:
+            raise RuntimeError("Gemini blocked the response (safety/recitation filter)") from exc
+        if not text:
+            raise RuntimeError("Gemini returned an empty completion")
+        return text

src/llm/groq_provider.py

@@ -0,0 +1,30 @@
+"""Groq-backed LLMProvider implementation (model: llama-3.3-70b-versatile)."""
+import logging
+
+from groq import Groq
+
+from src.llm.base import LLMProvider
+
+logger = logging.getLogger(__name__)
+
+_MODEL = "llama-3.3-70b-versatile"
+
+
+class GroqProvider(LLMProvider):
+    """Calls the Groq chat-completions API to satisfy the LLMProvider contract."""
+
+    def __init__(self, api_key: str) -> None:
+        """Build a Groq client from an explicit API key."""
+        self._client = Groq(api_key=api_key)
+
+    def generate(self, prompt: str) -> str:
+        """Send `prompt` as a single user message and return the completion text."""
+        logger.info("Requesting Groq completion (model=%s)", _MODEL)
+        response = self._client.chat.completions.create(
+            model=_MODEL,
+            messages=[{"role": "user", "content": prompt}],
+        )
+        content = response.choices[0].message.content
+        if content is None:
+            raise RuntimeError("Groq returned an empty completion")
+        return content

src/llm/ollama_provider.py

@@ -0,0 +1,41 @@
+"""Ollama-backed LLMProvider for local development (default model: qwen2.5-coder:7b)."""
+import logging
+
+import requests
+
+from src.llm.base import LLMProvider
+
+logger = logging.getLogger(__name__)
+
+_DEFAULT_BASE_URL = "http://localhost:11434"
+_DEFAULT_MODEL = "qwen2.5-coder:7b"
+
+
+class OllamaProvider(LLMProvider):
+    """Calls a local Ollama instance via REST to satisfy the LLMProvider contract.
+
+    No API key required — runs entirely offline. Suitable for dev/iteration.
+    """
+
+    def __init__(
+        self,
+        base_url: str = _DEFAULT_BASE_URL,
+        model: str = _DEFAULT_MODEL,
+    ) -> None:
+        """Configure the Ollama endpoint and model."""
+        self._base_url = base_url.rstrip("/")
+        self._model = model
+
+    def generate(self, prompt: str) -> str:
+        """POST the prompt to Ollama's /api/generate and return the response text."""
+        logger.info("Requesting Ollama completion (model=%s)", self._model)
+        response = requests.post(
+            f"{self._base_url}/api/generate",
+            json={"model": self._model, "prompt": prompt, "stream": False},
+            timeout=120,
+        )
+        response.raise_for_status()
+        text = response.json().get("response", "")
+        if not text:
+            raise RuntimeError("Ollama returned an empty response")
+        return text