|
| 1 | +from __future__ import annotations |
| 2 | + |
| 3 | +import asyncio |
| 4 | +import json |
| 5 | +import os |
| 6 | +from dataclasses import dataclass |
| 7 | +from typing import Any, Dict, Optional |
| 8 | + |
| 9 | +import requests |
| 10 | + |
| 11 | + |
| 12 | +def _run_async(coro): |
| 13 | + """ |
| 14 | + Helper to run async code from sync context. |
| 15 | + Mirrors lab_facade._run_async behavior. |
| 16 | + """ |
| 17 | + try: |
| 18 | + asyncio.get_running_loop() |
| 19 | + running = True |
| 20 | + except RuntimeError: |
| 21 | + running = False |
| 22 | + |
| 23 | + if not running: |
| 24 | + return asyncio.run(coro) |
| 25 | + |
| 26 | + loop = asyncio.get_event_loop() |
| 27 | + if loop.is_closed(): |
| 28 | + return asyncio.run(coro) |
| 29 | + return loop.run_until_complete(coro) |
| 30 | + |
| 31 | + |
| 32 | +@dataclass |
| 33 | +class GenerationModel: |
| 34 | + """ |
| 35 | + Simple, library-agnostic text generation interface. |
| 36 | + Implementations should provide: |
| 37 | + - generate(prompt, system_prompt=None) -> str |
| 38 | + - a_generate(prompt, system_prompt=None) -> str (async) |
| 39 | + """ |
| 40 | + |
| 41 | + def generate(self, prompt: str, system_prompt: Optional[str] = None) -> str: # pragma: no cover - interface |
| 42 | + raise NotImplementedError |
| 43 | + |
| 44 | + async def a_generate(self, prompt: str, system_prompt: Optional[str] = None) -> str: |
| 45 | + # Default async implementation delegates to sync in a thread |
| 46 | + loop = asyncio.get_running_loop() |
| 47 | + return await loop.run_in_executor(None, self.generate, prompt, system_prompt) |
| 48 | + |
| 49 | + |
| 50 | +class LocalHTTPGenerationModel(GenerationModel): |
| 51 | + """ |
| 52 | + Generation model that talks to a local HTTP server with an OpenAI-compatible |
| 53 | + /v1/chat/completions endpoint (e.g., vLLM, sglang, or TransformerLab inference server). |
| 54 | + """ |
| 55 | + |
| 56 | + def __init__(self, base_url: str, model: str, api_key: str = "dummy") -> None: |
| 57 | + self.base_url = base_url.rstrip("/") |
| 58 | + self.model = model |
| 59 | + self.api_key = api_key or "dummy" |
| 60 | + |
| 61 | + def generate(self, prompt: str, system_prompt: Optional[str] = None) -> str: |
| 62 | + url = f"{self.base_url}/chat/completions" |
| 63 | + messages: list[Dict[str, Any]] = [] |
| 64 | + if system_prompt: |
| 65 | + messages.append({"role": "system", "content": system_prompt}) |
| 66 | + messages.append({"role": "user", "content": prompt}) |
| 67 | + |
| 68 | + headers = { |
| 69 | + "Content-Type": "application/json", |
| 70 | + "Authorization": f"Bearer {self.api_key}", |
| 71 | + } |
| 72 | + payload: Dict[str, Any] = { |
| 73 | + "model": self.model, |
| 74 | + "messages": messages, |
| 75 | + } |
| 76 | + |
| 77 | + resp = requests.post(url, headers=headers, data=json.dumps(payload), timeout=60) |
| 78 | + resp.raise_for_status() |
| 79 | + data = resp.json() |
| 80 | + try: |
| 81 | + return data["choices"][0]["message"]["content"] |
| 82 | + except Exception as exc: # noqa: BLE001 |
| 83 | + raise RuntimeError(f"Unexpected response from local generation server: {data}") from exc |
| 84 | + |
| 85 | + |
| 86 | +def load_generation_model(config: Optional[Dict[str, Any] | str] = None) -> GenerationModel: |
| 87 | + """ |
| 88 | + Load a simple generation model wrapper based on a configuration object. |
| 89 | + This function is intentionally library-agnostic (no deepeval, no LangChain). |
| 90 | + Current support: |
| 91 | + - provider=\"local\": talks to a local HTTP server exposing an OpenAI-style |
| 92 | + /v1/chat/completions endpoint (e.g., TransformerLab inference server). |
| 93 | + Args: |
| 94 | + config: Either: |
| 95 | + - dict with \"provider\" and optionally \"base_url\", \"model\", \"api_key\" (e.g. {\"provider\": \"local\", \"model\": \"MyModel\", \"base_url\": \"http://localhost:8338/v1\"}) |
| 96 | + - JSON string representing such a dict |
| 97 | + - simple string provider name (e.g. \"local\") |
| 98 | + Returns: |
| 99 | + GenerationModel implementation. |
| 100 | + """ |
| 101 | + # Normalize config to dict |
| 102 | + if config is None: |
| 103 | + config_dict: Dict[str, Any] = {"provider": "local"} |
| 104 | + elif isinstance(config, str): |
| 105 | + # Try to parse as JSON first; if that fails, treat as provider name |
| 106 | + try: |
| 107 | + config_dict = json.loads(config) |
| 108 | + if not isinstance(config_dict, dict): |
| 109 | + config_dict = {"provider": str(config)} |
| 110 | + except json.JSONDecodeError: |
| 111 | + config_dict = {"provider": config} |
| 112 | + elif isinstance(config, dict): |
| 113 | + config_dict = dict(config) |
| 114 | + else: |
| 115 | + raise TypeError("config must be a dict, JSON string, provider string, or None") |
| 116 | + |
| 117 | + provider = str(config_dict.get("provider", "local")).lower() |
| 118 | + |
| 119 | + if provider == "local": |
| 120 | + # Prefer config; fall back to env vars, then defaults |
| 121 | + base_url = config_dict.get("base_url") or os.environ.get("TFL_LOCAL_MODEL_BASE_URL", "http://localhost:8338/v1") |
| 122 | + model_name = config_dict.get("model") or os.environ.get("TFL_LOCAL_MODEL_NAME", "default") |
| 123 | + api_key = config_dict.get("api_key") or os.environ.get("TFL_LOCAL_MODEL_API_KEY", "dummy") |
| 124 | + return LocalHTTPGenerationModel(base_url=base_url, model=model_name, api_key=api_key) |
| 125 | + |
| 126 | + raise NotImplementedError( |
| 127 | + f"Provider '{provider}' is not yet supported by load_generation_model. " |
| 128 | + "For now, only provider='local' is implemented." |
| 129 | + ) |
0 commit comments