rex/llm_client.py
Rama master · réplica de lectura del repositorio original.
"""Language model client with pluggable backends (Transformers, OpenAI, Echo)."""
from __future__ import annotations
import contextvars
import json
import sys
import threading
import types as _types
from collections.abc import Iterator, Sequence
from dataclasses import dataclass
from importlib import import_module
from importlib.util import find_spec
from typing import Any, Protocol, cast
from rex.assistant_errors import ConfigurationError
from rex.config import OPENROUTER_DEFAULT_BASE_URL, AppConfig, load_config
from rex.logging_utils import get_logger
from rex.retry import RetryConfig, with_retry
logger = get_logger(__name__)
def _module_available(name: str) -> bool:
if name in sys.modules:
return True
return find_spec(name) is not None
TORCH_AVAILABLE = _module_available("torch")
TRANSFORMERS_AVAILABLE = _module_available("transformers")
OPENAI_AVAILABLE = _module_available("openai")
OLLAMA_AVAILABLE = _module_available("ollama")
ANTHROPIC_AVAILABLE = _module_available("anthropic")
@dataclass
class GenerationConfig:
max_new_tokens: int
temperature: float
top_p: float
top_k: int
seed: int
class LLMStrategy(Protocol):
name: str
def generate(
self,
prompt: str,
config: GenerationConfig,
*,
messages: list[dict[str, Any]] | None = None,
**kwargs: Any,
) -> str: ...
def stream(
self,
prompt: str,
config: GenerationConfig,
*,
messages: list[dict[str, Any]] | None = None,
**kwargs: Any,
) -> Iterator[str]:
raise NotImplementedError
class EchoStrategy:
name = "echo"
def __init__(self, model_name: str) -> None:
self.model_name = model_name
def generate(
self,
prompt: str,
_config: GenerationConfig,
*,
messages: list[dict[str, Any]] | None = None,
**_kwargs: Any,
) -> str:
text = prompt.strip()
if not text and messages:
text = "\n".join(m.get("content", "").strip() for m in messages).strip()
return f"[{self.model_name}] {text or '(silence)'}"
def stream(
self,
prompt: str,
_config: GenerationConfig,
*,
messages: list[dict[str, Any]] | None = None,
**_kwargs: Any,
) -> Iterator[str]:
text = prompt.strip()
if not text and messages:
text = "\n".join(m.get("content", "").strip() for m in messages).strip()
words = (text or "(silence)").split()
for index, word in enumerate(words):
suffix = " " if index < len(words) - 1 else ""
yield f"{word}{suffix}"
class OfflineTransformersStrategy:
name = "offline-transformers"
def __init__(self, model_name: str) -> None:
self.model_name = model_name
def generate(
self,
prompt: str,
_config: GenerationConfig,
*,
messages: list[dict[str, Any]] | None = None,
**_kwargs: Any,
) -> str:
text = prompt.strip()
if not text and messages:
text = "\n".join(m.get("content", "").strip() for m in messages).strip()
return f"(offline) {text or ''}".strip() or "(offline)"
def stream(
self,
prompt: str,
config: GenerationConfig,
*,
messages: list[dict[str, Any]] | None = None,
**_kwargs: Any,
) -> Iterator[str]:
yield self.generate(prompt, config, messages=messages)
@dataclass(frozen=True)
class _TransformersRuntime:
torch: Any
pipeline: Any
tokenizer: Any
class TransformersStrategy:
name = "transformers"
def __init__(self, model_name: str) -> None:
if not (TORCH_AVAILABLE and TRANSFORMERS_AVAILABLE):
raise ConfigurationError("Transformers backend requires `torch` and `transformers`.")
try:
torch = import_module("torch")
transformers = import_module("transformers")
except ImportError as exc:
raise ConfigurationError(
"Transformers backend requires `torch` and `transformers`."
) from exc
from rex.runtime.warm import (
WarmComponentSpec,
default_idle_timeout,
get_global_warm_runtime,
warm_component_key,
)
device = 0 if torch.cuda.is_available() else -1
manager = get_global_warm_runtime()
component_name = warm_component_key("llm", model_name, device)
def _load() -> _TransformersRuntime:
tokenizer = transformers.AutoTokenizer.from_pretrained(model_name)
model = transformers.AutoModelForCausalLM.from_pretrained(model_name)
pipeline = transformers.pipeline(
"text-generation",
model=model,
tokenizer=tokenizer,
device=device,
)
if tokenizer.pad_token_id is None:
tokenizer.pad_token_id = tokenizer.eos_token_id
return _TransformersRuntime(torch=torch, pipeline=pipeline, tokenizer=tokenizer)
manager.register_if_absent(
WarmComponentSpec(
name=component_name,
loader=_load,
estimated_cost_mb=2048.0,
idle_timeout_s=default_idle_timeout(),
)
)
manager.warm(component_name)
self._warm_manager = manager
self._warm_component_name = component_name
def _runtime(self) -> _TransformersRuntime:
return cast(_TransformersRuntime, self._warm_manager.get(self._warm_component_name))
@property
def _torch(self) -> Any:
return self._runtime().torch
@property
def pipeline(self) -> Any:
return self._runtime().pipeline
@property
def tokenizer(self) -> Any:
return self._runtime().tokenizer
def generate(
self,
prompt: str,
config: GenerationConfig,
*,
messages: list[dict[str, Any]] | None = None,
**kwargs: Any,
) -> str:
with self._warm_manager.acquire(self._warm_component_name) as value:
runtime = cast(_TransformersRuntime, value)
torch = runtime.torch
torch.manual_seed(config.seed)
if torch.cuda.is_available():
torch.cuda.manual_seed_all(config.seed)
do_sample = config.temperature > 0 or config.top_p < 1.0
generate_kwargs = {
"max_new_tokens": config.max_new_tokens,
"do_sample": do_sample,
"pad_token_id": runtime.tokenizer.pad_token_id,
}
if do_sample:
generate_kwargs.update(
{
"temperature": max(config.temperature, 1e-4),
"top_p": config.top_p,
"top_k": config.top_k,
}
)
outputs = runtime.pipeline(prompt, **generate_kwargs)
text = outputs[0]["generated_text"]
return text[len(prompt) :].strip() or "(silence)"
class OpenAIStrategy:
name = "openai"
def __init__(
self,
model_name: str,
client_factory,
retry_config: RetryConfig | None = None,
) -> None:
self.model_name = model_name
self._client_factory = client_factory
self._cached_client = None
self.tools: list[Any] = [] # Will be set by LanguageModel
self._retry_config = retry_config or RetryConfig()
def _get_client(self) -> Any:
if self._cached_client is None:
self._cached_client = self._client_factory()
return self._cached_client
def generate(
self,
prompt: str,
config: GenerationConfig,
*,
messages=None,
tools=None,
user: str | None = None,
) -> Any:
payload = messages or [{"role": "user", "content": prompt}]
# Prepare API call parameters
api_params: dict[str, Any] = {
"model": self.model_name,
"messages": payload,
"temperature": config.temperature,
"max_tokens": config.max_new_tokens,
"top_p": config.top_p,
}
# Pass a stable user key so OpenClaw can maintain session state.
if user is not None:
api_params["user"] = user
# Add tools if available
if tools or self.tools:
api_params["tools"] = tools or self.tools
# Don't force tool use, let model decide
# api_params["tool_choice"] = "auto"
def _call() -> str:
response = self._get_client().chat.completions.create(**api_params)
message = response.choices[0].message
# When tool_calls are present, serialize to JSON so generate() always returns str
if hasattr(message, "tool_calls") and message.tool_calls:
tool_calls_data = [
{
"id": tc.id,
"type": "function",
"function": {
"name": tc.function.name,
"arguments": tc.function.arguments,
},
}
for tc in message.tool_calls
]
return json.dumps({"__tool_calls__": tool_calls_data})
content = getattr(message, "content", "") or ""
return content.strip() or "(silence)"
return with_retry(_call, config=self._retry_config)
def stream(
self,
prompt: str,
config: GenerationConfig,
*,
messages=None,
tools=None,
user: str | None = None,
) -> Iterator[str]:
payload = messages or [{"role": "user", "content": prompt}]
api_params: dict[str, Any] = {
"model": self.model_name,
"messages": payload,
"temperature": config.temperature,
"max_tokens": config.max_new_tokens,
"top_p": config.top_p,
"stream": True,
}
if user is not None:
api_params["user"] = user
if tools or self.tools:
api_params["tools"] = tools or self.tools
stream = self._get_client().chat.completions.create(**api_params)
for chunk in stream:
choices = getattr(chunk, "choices", None)
if not choices:
continue
delta = getattr(choices[0], "delta", None)
if delta is None:
continue
content = getattr(delta, "content", None)
if isinstance(content, str):
if content:
yield content
continue
if isinstance(content, list):
for item in content:
text = getattr(item, "text", None)
if text is None and isinstance(item, dict):
text = item.get("text")
if text:
yield str(text)
class OllamaStrategy:
name = "ollama"
def __init__(
self,
model_name: str,
api_key: str | None = None,
base_url: str = "http://localhost:11434",
use_cloud: bool = False,
retry_config: RetryConfig | None = None,
) -> None:
if not OLLAMA_AVAILABLE:
raise ConfigurationError("Ollama backend requires the `ollama` package.")
try:
ollama = import_module("ollama")
client_cls = ollama.Client
except Exception as exc:
raise ConfigurationError("Ollama backend requires the `ollama` package.") from exc
self.model_name = model_name
self.base_url = base_url
self.use_cloud = use_cloud
self.api_key = api_key
self._ollama = ollama
self._client_cls = client_cls
self._retry_config = retry_config or RetryConfig()
if use_cloud and not api_key:
raise ConfigurationError("Ollama cloud requires an API key.")
try:
if use_cloud:
self._client = None # cloud handled via ollama.chat
else:
self._client = self._client_cls(host=self.base_url)
logger.info("Ollama local client initialized at %s", self.base_url)
except Exception as exc:
raise ConfigurationError(f"Failed to initialize Ollama client: {exc}") from exc
def generate(
self,
prompt: str,
config: GenerationConfig,
*,
messages: list[dict[str, Any]] | None = None,
**kwargs: Any,
) -> str:
payload = messages or [{"role": "user", "content": prompt}]
options = {
"temperature": config.temperature,
"top_p": config.top_p,
"top_k": config.top_k,
"seed": config.seed,
}
def _call() -> str:
if self.use_cloud:
response = self._ollama.chat(
model=self.model_name,
messages=payload,
options=options,
)
else:
response = self._client.chat( # type: ignore[union-attr]
model=self.model_name,
messages=payload,
options=options,
)
# Log usage when token counts are available in the response
try:
from rex.llm_usage import record_usage
if isinstance(response, dict):
eval_count = int(response.get("eval_count") or 0)
prompt_eval_count = int(response.get("prompt_eval_count") or 0)
else:
eval_count = int(getattr(response, "eval_count", None) or 0)
prompt_eval_count = int(getattr(response, "prompt_eval_count", None) or 0)
record_usage(
model=self.model_name,
prompt_tokens=prompt_eval_count,
completion_tokens=eval_count,
)
except Exception as _log_exc: # noqa: BLE001
logger.debug("Usage logging failed: %s", _log_exc)
if isinstance(response, dict):
content = response.get("message", {}).get("content", "")
else:
message = getattr(response, "message", None)
content = str(getattr(message, "content", "") or "") if message else ""
return content.strip() or "(silence)"
try:
return with_retry(_call, config=self._retry_config)
except Exception as exc:
# Connection errors: server not running or unreachable
if isinstance(exc, ConnectionRefusedError):
logger.error("Ollama connection failed: %s", exc)
return "[Ollama: connection failed — is Ollama running?]"
# httpx.ConnectError — ollama uses httpx internally; avoid top-level import
exc_cls = type(exc)
if exc_cls.__name__ == "ConnectError" and exc_cls.__module__ == "httpx":
logger.error("Ollama connection failed: %s", exc)
return "[Ollama: connection failed — is Ollama running?]"
# 404 model-not-found (ollama.ResponseError carries status_code)
if getattr(exc, "status_code", None) == 404:
logger.error("Ollama model '%s' not found: %s", self.model_name, exc)
return (
f"[Ollama: model '{self.model_name}' not found"
f" — run: ollama pull {self.model_name}]"
)
# All other errors
logger.error("Ollama unexpected error (model=%s): %s", self.model_name, exc)
return f"[Ollama: unexpected error: {exc}]"
def stream(
self,
prompt: str,
config: GenerationConfig,
*,
messages=None,
**_kwargs: Any,
) -> Iterator[str]:
payload = messages or [{"role": "user", "content": prompt}]
options = {
"temperature": config.temperature,
"top_p": config.top_p,
"top_k": config.top_k,
"seed": config.seed,
}
try:
if self.use_cloud:
stream = self._ollama.chat(
model=self.model_name,
messages=payload,
options=options,
stream=True,
)
else:
stream = self._client.chat( # type: ignore[union-attr]
model=self.model_name,
messages=payload,
options=options,
stream=True,
)
for chunk in stream:
content = ""
if isinstance(chunk, dict):
content = str(chunk.get("message", {}).get("content", "") or "")
else:
message = getattr(chunk, "message", None)
content = str(getattr(message, "content", "") or "")
if content:
yield content
except Exception as exc:
logger.warning("Ollama stream interrupted (model=%s): %s", self.model_name, exc)
return
class AnthropicStrategy:
name = "anthropic"
def __init__(
self,
model_name: str,
client_factory,
retry_config: RetryConfig | None = None,
) -> None:
self.model_name = model_name
self._client_factory = client_factory
self._cached_client = None
self._retry_config = retry_config or RetryConfig()
def _get_client(self) -> Any:
if self._cached_client is None:
self._cached_client = self._client_factory()
return self._cached_client
def generate(
self,
prompt: str,
config: GenerationConfig,
*,
messages: list[dict[str, Any]] | None = None,
**kwargs: Any,
) -> str:
payload = messages or [{"role": "user", "content": prompt}]
# Anthropic requires alternating user/assistant turns; filter system messages
system_parts = [m["content"] for m in payload if m.get("role") == "system"]
user_messages = [m for m in payload if m.get("role") != "system"]
system_text = "\n".join(system_parts) if system_parts else None
api_params: dict[str, Any] = {
"model": self.model_name,
"max_tokens": config.max_new_tokens,
"messages": user_messages,
}
if system_text:
api_params["system"] = system_text
def _call() -> str:
response = self._get_client().messages.create(**api_params)
content = response.content[0].text if response.content else ""
return content.strip() or "(silence)"
return with_retry(_call, config=self._retry_config)
def stream(
self,
prompt: str,
config: GenerationConfig,
*,
messages=None,
**_kwargs: Any,
) -> Iterator[str]:
payload = messages or [{"role": "user", "content": prompt}]
system_parts = [m["content"] for m in payload if m.get("role") == "system"]
user_messages = [m for m in payload if m.get("role") != "system"]
system_text = "\n".join(system_parts) if system_parts else None
api_params: dict[str, Any] = {
"model": self.model_name,
"max_tokens": config.max_new_tokens,
"messages": user_messages,
}
if system_text:
api_params["system"] = system_text
try:
stream_manager = self._get_client().messages.stream(**api_params)
with stream_manager as stream:
for text in stream.text_stream:
if text:
yield str(text)
except Exception as exc:
logger.warning("Anthropic stream interrupted (model=%s): %s", self.model_name, exc)
return
_STRATEGIES: dict[str, type[Any]] = {
EchoStrategy.name: EchoStrategy,
TransformersStrategy.name: TransformersStrategy,
}
def register_strategy(name: str, strategy: type[Any]) -> None:
_STRATEGIES[name] = strategy
class LanguageModel:
def __init__(self, config: AppConfig | None = None, **overrides) -> None:
self.config = config or load_config()
llm_config = self.config.llm
if llm_config is None:
raise ConfigurationError("LLM configuration is unavailable.")
self.provider = str(overrides.get("provider") or llm_config.llm_provider).lower()
if self.provider == "openai":
self.model_name = (
overrides.get("model") or self.config.openai_model or self.config.llm_model
)
elif self.provider == "openrouter":
self.model_name = (
overrides.get("model") or self.config.openrouter_model or self.config.llm_model
)
elif self.provider == "anthropic":
self.model_name = (
overrides.get("model") or self.config.anthropic_model or self.config.llm_model
)
else:
self.model_name = overrides.get("model") or self.config.llm_model
self.api_key: str | None
self.api_base_url: str | None
if self.provider == "openrouter":
self.api_key = overrides.get("openrouter_api_key") or self.config.openrouter_api_key
self.api_base_url = OPENROUTER_DEFAULT_BASE_URL
else:
self.api_key = overrides.get("openai_api_key") or self.config.openai_api_key
self.api_base_url = overrides.get("base_url") or self.config.openai_base_url
self.anthropic_api_key = overrides.get("anthropic_api_key") or self.config.anthropic_api_key
self.generation = GenerationConfig(
max_new_tokens=overrides.get("max_new_tokens", self.config.llm_max_tokens),
temperature=overrides.get("temperature", self.config.llm_temperature),
top_p=overrides.get("top_p", self.config.llm_top_p),
top_k=overrides.get("top_k", self.config.llm_top_k),
seed=overrides.get("seed", self.config.llm_seed),
)
self._openai_client = None
self._anthropic_client = None
self._tools: list[Any] = [] # OpenAI tool definitions
self._tool_functions: dict[str, Any] = {} # Map of tool name -> callable
from rex.runtime.warm import get_global_warm_runtime
# Establish process-wide warm-cache policy before constructing model strategies.
get_global_warm_runtime(self.config)
self._request_model: contextvars.ContextVar[str | None] = contextvars.ContextVar(
f"askrex_llm_model_{id(self)}", default=None
)
self._strategy_lock = threading.Lock()
self.strategy = self._init_strategy(self.model_name)
self._strategy_cache: dict[str, LLMStrategy] = {self.model_name: self.strategy}
def _ensure_anthropic_client(self) -> Any:
if self._anthropic_client is not None:
return self._anthropic_client
if not self.anthropic_api_key:
raise ConfigurationError("Missing Anthropic API key.")
if not ANTHROPIC_AVAILABLE:
raise ConfigurationError("Anthropic backend requires the `anthropic` package.")
try:
Anthropic = import_module("anthropic").Anthropic
except Exception as exc:
raise ConfigurationError("Anthropic backend requires the `anthropic` package.") from exc
self._anthropic_client = Anthropic(api_key=self.anthropic_api_key)
return self._anthropic_client
def _init_strategy(self, model_name: str | None = None) -> LLMStrategy:
model_name = model_name or self.model_name
if self.provider in {"openai", "openrouter"}:
return cast(LLMStrategy, OpenAIStrategy(model_name, self._ensure_openai_client))
if self.provider == "anthropic":
return cast(LLMStrategy, AnthropicStrategy(model_name, self._ensure_anthropic_client))
if self.provider == "ollama":
return cast(
LLMStrategy,
OllamaStrategy(
model_name,
api_key=self.config.ollama_api_key,
base_url=self.config.ollama_base_url,
use_cloud=self.config.ollama_use_cloud,
),
)
strategy_cls = _STRATEGIES.get(self.provider)
if strategy_cls is None:
logger.warning("Unknown LLM provider '%s'. Falling back to echo mode.", self.provider)
return cast(LLMStrategy, EchoStrategy(model_name))
if strategy_cls is TransformersStrategy and not (
TORCH_AVAILABLE and TRANSFORMERS_AVAILABLE
):
logger.warning(
"Transformers missing; using offline fallback for model '%s'.", model_name
)
return cast(LLMStrategy, OfflineTransformersStrategy(model_name))
try:
return cast(LLMStrategy, strategy_cls(model_name))
except Exception as exc:
logger.warning(
"LLM backend init failed (%s). Falling back to echo. (%s)", self.provider, exc
)
return cast(LLMStrategy, EchoStrategy(model_name))
def set_request_model(self, model_name: str) -> contextvars.Token[str | None]:
"""Set a model override for the current request context only."""
return self._request_model.set(model_name)
def reset_request_model(self, token: contextvars.Token[str | None]) -> None:
"""Restore the previous request-local model override."""
self._request_model.reset(token)
def active_model_name(self) -> str:
"""Return the request-local model when set, otherwise the configured base model."""
request_model = getattr(self, "_request_model", None)
if request_model is None:
return self.model_name
return request_model.get() or self.model_name
def _strategy_for_request(self) -> LLMStrategy:
model_name = self.active_model_name()
strategy_cache = getattr(self, "_strategy_cache", None)
strategy_lock = getattr(self, "_strategy_lock", None)
if strategy_cache is None or strategy_lock is None:
return self.strategy
cached = strategy_cache.get(model_name)
if cached is not None:
return cast(LLMStrategy, cached)
with strategy_lock:
cached = strategy_cache.get(model_name)
if cached is None:
cached = self._init_strategy(model_name)
if hasattr(cached, "tools"):
cached.tools = self._tools
strategy_cache[model_name] = cached
return cast(LLMStrategy, cached)
def _ensure_openai_client(self):
if self._openai_client is not None:
return self._openai_client
base_url = self.api_base_url
api_key = self.api_key
if self.provider == "openrouter":
if not api_key:
raise ConfigurationError("Missing OpenRouter API key.")
elif not api_key and base_url:
api_key = "local"
if not api_key:
raise ConfigurationError("Missing OpenAI API key.")
if not OPENAI_AVAILABLE:
raise ConfigurationError("OpenAI backend requires the `openai` package.")
try:
OpenAI = import_module("openai").OpenAI
except Exception as exc:
raise ConfigurationError("OpenAI backend requires the `openai` package.") from exc
client_kwargs: dict[str, Any] = {"api_key": api_key}
if base_url:
client_kwargs["base_url"] = base_url
if self.provider == "openrouter":
client_kwargs["default_headers"] = {"X-OpenRouter-Title": "AskRex Assistant"}
self._openai_client = OpenAI(**client_kwargs)
return self._openai_client
def register_tool(
self, name: str, description: str, parameters: dict[str, Any], function: Any
) -> None:
"""Register a tool/function that the LLM can call.
Args:
name: Function name (e.g., "web_search")
description: What the function does
parameters: JSON schema for parameters
function: Callable to execute when tool is invoked
"""
tool_def = {
"type": "function",
"function": {
"name": name,
"description": description,
"parameters": parameters,
},
}
self._tools.append(tool_def)
self._tool_functions[name] = function
# Update strategy if it supports tools
if hasattr(self.strategy, "tools"):
self.strategy.tools = self._tools
logger.info(f"Registered tool: {name}")
def _execute_tool_call(self, tool_call: Any) -> str:
"""Execute a tool call and return the result."""
import json
function_name = tool_call.function.name
function_args = json.loads(tool_call.function.arguments)
logger.info(f"Executing tool: {function_name} with args: {function_args}")
if function_name not in self._tool_functions:
return f"Error: Unknown function {function_name}"
from rex.mobile_api.action_context import ( # noqa: PLC0415
MobileActionDeniedError,
authorized_mobile_tool,
)
try:
with authorized_mobile_tool(function_name, arguments=function_args):
result = self._tool_functions[function_name](**function_args)
return str(result) if result is not None else "No results found"
except MobileActionDeniedError:
raise
except Exception as e:
logger.error("Tool execution failed: %s", type(e).__name__)
return f"Error executing {function_name}"
def _format_messages(self, messages: Sequence[dict[str, str]]) -> str:
return "\n".join(
f"{m.get('role', 'user').capitalize()}: {m.get('content', '').strip()}"
for m in messages
if isinstance(m, dict)
).strip()
def generate(
self,
prompt: str | None = None,
*,
messages: Sequence[dict[str, Any]] | None = None,
config: GenerationConfig | None = None,
max_tool_rounds: int = 3,
) -> str:
if messages is not None:
prompt_text = self._format_messages(messages)
normalized_messages: list[dict[str, Any]] = [
{"role": str(m.get("role", "")), "content": str(m.get("content", ""))}
for m in messages
if isinstance(m, dict)
]
elif isinstance(prompt, str):
prompt_text = prompt
normalized_messages = [{"role": "user", "content": prompt}]
else:
raise ValueError("Prompt or messages must be provided.")
if not prompt_text.strip():
raise ValueError("Prompt must not be empty.")
strategy = self._strategy_for_request()
# Derive a stable user key for OpenAI/OpenClaw session persistence.
_user_key: str | None = None
if self.provider == "openai":
uid = getattr(self.config, "user_id", "default")
profile = getattr(self.config, "active_profile", "default")
_user_key = uid if uid and uid != "default" else profile
# Tool calling loop (only for OpenAI strategy with tools)
if self.provider == "openai" and self._tools:
current_messages = list(normalized_messages)
for round_num in range(max_tool_rounds):
try:
_tool_extra: dict[str, Any] = {}
if _user_key is not None:
_tool_extra["user"] = _user_key
result = strategy.generate(
prompt_text,
config or self.generation,
messages=current_messages,
tools=self._tools,
**_tool_extra,
)
except TypeError:
# Fallback for strategies that don't support tools parameter
result = strategy.generate(
prompt_text, config or self.generation, messages=current_messages
)
# Check if result is a JSON-serialized tool call response
_tc_parsed: dict[str, Any] | None = None
if isinstance(result, str):
try:
_maybe = json.loads(result)
if isinstance(_maybe, dict) and "__tool_calls__" in _maybe:
_tc_parsed = _maybe
except (json.JSONDecodeError, ValueError):
pass
if _tc_parsed is not None:
tc_list: list[dict[str, Any]] = _tc_parsed["__tool_calls__"]
logger.info(
f"Tool call round {round_num + 1}: Model requested {len(tc_list)} tool(s)"
)
# Add assistant message to conversation
current_messages.append(
{
"role": "assistant",
"content": "",
"tool_calls": tc_list,
}
)
# Execute each tool call and add results
for tc_data in tc_list:
tool_call = _types.SimpleNamespace(
id=tc_data["id"],
function=_types.SimpleNamespace(
name=tc_data["function"]["name"],
arguments=tc_data["function"]["arguments"],
),
)
tool_result = self._execute_tool_call(tool_call)
current_messages.append(
{
"role": "tool",
"tool_call_id": tc_data["id"],
"content": tool_result,
}
)
# Continue loop to get final response with tool results
continue
# No tool calls - return final response
return result.strip() if isinstance(result, str) else str(result).strip()
# Max rounds reached
logger.warning(f"Max tool rounds ({max_tool_rounds}) reached")
return "I apologize, but I encountered too many tool calls. Please try rephrasing your question."
# No tools or not OpenAI - use original simple generation
try:
try:
extra: dict[str, Any] = {}
if _user_key is not None:
extra["user"] = _user_key
result = strategy.generate(
prompt_text,
config or self.generation,
messages=normalized_messages,
**extra,
)
return result.strip() if isinstance(result, str) else str(result).strip()
except TypeError:
result = strategy.generate(prompt_text, config or self.generation)
return result.strip() if isinstance(result, str) else str(result).strip()
except Exception as exc:
from rex.dep_errors import LMStudioConnectionError, is_connection_error
if (
is_connection_error(exc)
and self.provider == "openai"
and self.config.openai_base_url
):
url = self.config.openai_base_url
raise LMStudioConnectionError(
f"Rex can't reach LM Studio at {url}. Is LM Studio running?"
) from exc
raise
def stream(
self,
prompt: str | None = None,
*,
messages: Sequence[dict[str, Any]] | None = None,
config: GenerationConfig | None = None,
) -> Iterator[str]:
generation_config = config or self.generation
if messages is not None:
prompt_text = self._format_messages(messages)
normalized_messages = [
{"role": str(m.get("role", "")), "content": str(m.get("content", ""))}
for m in messages
if isinstance(m, dict)
]
elif isinstance(prompt, str):
prompt_text = prompt
normalized_messages = [{"role": "user", "content": prompt}]
else:
raise ValueError("Prompt or messages must be provided.")
if not prompt_text.strip():
raise ValueError("Prompt must not be empty.")
strategy = self._strategy_for_request()
def _generate_once() -> Iterator[str]:
if messages is not None:
completion = self.generate(messages=normalized_messages, config=generation_config)
else:
completion = self.generate(prompt_text, config=generation_config)
if completion:
yield completion
extra: dict[str, Any] = {}
if self.provider == "openai":
uid = getattr(self.config, "user_id", "default")
profile = getattr(self.config, "active_profile", "default")
user_key = uid if uid and uid != "default" else profile
if user_key:
extra["user"] = user_key
if self._tools:
extra["tools"] = self._tools
stream_fn = getattr(strategy, "stream", None)
if not callable(stream_fn):
return _generate_once()
try:
token_iterator = stream_fn(
prompt_text,
generation_config,
messages=normalized_messages,
**extra,
)
except TypeError:
try:
token_iterator = stream_fn(prompt_text, generation_config)
except NotImplementedError:
return _generate_once()
except NotImplementedError:
return _generate_once()
def _guarded_stream() -> Iterator[str]:
try:
yield from token_iterator
except NotImplementedError:
yield from _generate_once()
return _guarded_stream()
__all__ = [
"LanguageModel",
"GenerationConfig",
"register_strategy",
"TORCH_AVAILABLE",
"TRANSFORMERS_AVAILABLE",
"RetryConfig",
]