File size: 4,018 Bytes
af8ac78 e1c98c1 af8ac78 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 | """Local LLM wrapper — wraps the local completions endpoint as a LangChain LLM."""
from __future__ import annotations
from typing import Any, Optional
import requests, sys, os
sys.path.insert(0, os.path.dirname(os.path.dirname(__file__)))
import config
from crewai.llms.base_llm import BaseLLM
from pydantic import Field
class LocalLLM(BaseLLM):
"""CrewAI-compatible wrapper for the local OpenAI-compatible completions endpoint."""
def __init__(self, **kwargs):
kwargs.setdefault("model", config.LLM_MODEL_ID)
kwargs.setdefault("base_url", config.LLM_BASE_URL)
super().__init__(**kwargs)
self.max_tokens = kwargs.get("max_tokens", config.LLM_MAX_TOKENS)
self.temperature = kwargs.get("temperature", config.LLM_TEMPERATURE)
self.top_p = kwargs.get("top_p", config.LLM_TOP_P)
self.timeout = kwargs.get("timeout", config.LLM_TIMEOUT)
self.use_kv_cache = kwargs.get("use_kv_cache", False)
def call(self, messages: list[dict], callbacks: list[Any] | None = None, **kwargs: Any) -> str:
payload = {
"model_id": self.model,
"prompt": messages,
"max_tokens": kwargs.get("max_tokens") or self.max_tokens,
"temperature": self.temperature,
"top_p": self.top_p,
"use_kv_cache": self.use_kv_cache,
"use_gpu": kwargs.get("use_gpu", False),
"cpu_threads": kwargs.get("cpu_threads", 2),
"attachments": [],
}
try:
resp = requests.post(
f"{self.base_url}/v1/completions",
json=payload,
timeout=self.timeout,
verify=False,
)
# ── Handle 503 "Server busy" explicitly ───────────────────────────
if resp.status_code == 503:
err_body = resp.json() if resp.content else {}
retry_hint = err_body.get("retry_after", 10)
reason = err_body.get("error", "The inference server is busy.")
self._last_prompt_tokens = 0
self._last_completion_tokens = 0
return (
f"[LLM BUSY] {reason} "
f"The server is processing another request. "
f"Please try again in {retry_hint} seconds."
)
resp.raise_for_status()
data = resp.json()
usage = data.get("usage", {})
# Store token counts so callers can read them without CrewAI usage_metrics
self._last_prompt_tokens = usage.get("prompt_tokens", 0)
self._last_completion_tokens = usage.get("completion_tokens", 0)
if usage:
self._track_token_usage_internal(usage)
return data["choices"][0]["text"].strip()
except requests.exceptions.ConnectionError:
self._last_prompt_tokens = 0
self._last_completion_tokens = 0
return "[LLM OFFLINE] Cannot connect to the inference server. Is nvidia_llm.py running?"
except requests.exceptions.Timeout:
self._last_prompt_tokens = 0
self._last_completion_tokens = 0
return (
f"[LLM TIMEOUT] The inference server did not respond within {self.timeout}s. "
"The server may be busy. Please try again shortly."
)
except Exception as e:
self._last_prompt_tokens = 0
self._last_completion_tokens = 0
return f"[LLM ERROR] {e}"
def supports_function_calling(self) -> bool:
return False
def supports_stop_words(self) -> bool:
return False
# Singleton instance
_llm_instance: LocalLLM | None = None
def get_llm() -> LocalLLM:
global _llm_instance
if _llm_instance is None:
_llm_instance = LocalLLM(model=config.LLM_MODEL_ID)
return _llm_instance
|