| """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), |
| "llm_mode": kwargs.get("llm_mode", "expert"), |
| "attachments": [], |
| } |
| try: |
| resp = requests.post( |
| f"{self.base_url}/v1/completions", |
| json=payload, |
| timeout=self.timeout, |
| verify=False, |
| ) |
|
|
| |
| 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", {}) |
| |
| self._last_prompt_tokens = usage.get("prompt_tokens", 0) |
| self._last_completion_tokens = usage.get("completion_tokens", 0) |
| self._last_model_name = data.get("model", self.model) |
| 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 |
|
|
|
|
| |
| _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 |
|
|
