Initial commit: Original Python LLMTester project

.core/providers/lmstudio.py

@@ -0,0 +1,100 @@
+from __future__ import annotations
+
+import json
+import time
+from typing import Dict, List, Optional
+
+import requests
+
+from config import get_env_setting
+from .base import ModelInfo, Provider, ProviderError
+
+DEFAULT_BASE_URL = "http://localhost:1234"
+
+
+class LMStudioProvider(Provider):
+    name = "LM Studio"
+
+    def __init__(self, base_url: Optional[str] = None) -> None:
+        self.base_url = base_url or get_env_setting("LM_STUDIO_BASE_URL", DEFAULT_BASE_URL)
+        self._models_url = f"{self.base_url}/v1/models"
+        self._chat_url = f"{self.base_url}/v1/chat/completions"
+
+    def list_models(self) -> List[ModelInfo]:
+        try:
+            response = requests.get(self._models_url, timeout=10)
+            response.raise_for_status()
+        except requests.exceptions.RequestException as exc:
+            raise ProviderError(f"Failed to fetch models from LM Studio: {exc}") from exc
+
+        try:
+            data = response.json()
+        except json.JSONDecodeError as exc:
+            raise ProviderError(f"Failed to decode model list: {exc}") from exc
+
+        models = data.get("data", [])
+        result: List[ModelInfo] = []
+        if isinstance(models, list):
+            for item in models:
+                if isinstance(item, dict):
+                    model_id = str(item.get("id") or item.get("name") or "")
+                else:
+                    model_id = str(item)
+                if not model_id:
+                    continue
+                result.append(ModelInfo(id=model_id, display_name=model_id))
+
+        default_model = data.get("default_model") or data.get("model") or data.get("id")
+        if default_model and all(m.id != default_model for m in result):
+            result.insert(0, ModelInfo(id=str(default_model), display_name=str(default_model)))
+
+        if not result:
+            raise ProviderError("LM Studio returned no models.")
+
+        return result
+
+    def run_prompt(self, model_id: str, prompt: str, *, temperature: float) -> Dict[str, object]:
+        headers = {"Content-Type": "application/json"}
+        payload = {
+            "model": model_id,
+            "messages": [{"role": "user", "content": prompt}],
+            "temperature": temperature,
+        }
+
+        start_time = time.time()
+        try:
+            response = requests.post(
+                self._chat_url,
+                headers=headers,
+                data=json.dumps(payload),
+                timeout=180,
+            )
+            response.raise_for_status()
+        except requests.exceptions.RequestException as exc:
+            return {"error": f"API request failed: {exc}"}
+
+        end_time = time.time()
+
+        try:
+            data = response.json()
+        except json.JSONDecodeError as exc:
+            return {"error": f"Failed to parse response JSON: {exc}\nRaw: {response.text}"}
+
+        response_time = end_time - start_time
+        content = data.get("choices", [{}])[0].get("message", {}).get("content", "Error: No content found.")
+        usage = data.get("usage", {})
+        total_tokens = usage.get("completion_tokens", 0)
+        time_to_first_token = usage.get("prompt_eval_duration", 0) / 1_000_000_000
+        stop_reason = data.get("choices", [{}])[0].get("finish_reason", "N/A")
+
+        tokens_per_second = total_tokens / response_time if response_time > 0 else 0
+
+        return {
+            "response": content,
+            "metrics": {
+                "tokens_per_second": round(tokens_per_second, 2),
+                "total_tokens": total_tokens,
+                "time_to_first_token": round(time_to_first_token, 2),
+                "stop_reason": stop_reason,
+            },
+        }