Spaces:
Runtime error
Runtime error
| """LLM-powered agent using LiteLLM for model-agnostic inference. | |
| Works with any LiteLLM-supported provider: | |
| - ``"anthropic/claude-sonnet-4-20250514"`` | |
| - ``"openai/gpt-4o"`` | |
| - ``"ollama/llama3.1:70b"`` | |
| - ``"hosted_vllm/Qwen/Qwen3-32B"`` (pass ``api_base=`` kwarg) | |
| """ | |
| from __future__ import annotations | |
| import copy | |
| from typing import Any, Literal | |
| from open_range.agents.observation import format_observation | |
| from open_range.agents.parsing import extract_command | |
| from open_range.agents.prompts import BLUE_SYSTEM_PROMPT, RED_SYSTEM_PROMPT | |
| class LLMRangeAgent: | |
| """Generic agent powered by any LiteLLM model. | |
| Satisfies the :class:`RangeAgent` protocol. | |
| Args: | |
| model: LiteLLM model string (default ``"anthropic/claude-sonnet-4-20250514"``). | |
| temperature: Sampling temperature. | |
| max_tokens: Maximum tokens per completion. | |
| **litellm_kwargs: Extra kwargs forwarded to ``litellm.completion`` | |
| (e.g. ``api_base``, ``api_key``). | |
| """ | |
| def __init__( | |
| self, | |
| model: str = "anthropic/claude-sonnet-4-20250514", | |
| temperature: float | None = 0.3, | |
| max_tokens: int = 512, | |
| bootstrap_messages: list[dict[str, Any]] | None = None, | |
| system_suffix: str = "", | |
| **litellm_kwargs: Any, | |
| ) -> None: | |
| self.model = model | |
| self.temperature = temperature | |
| self.max_tokens = max_tokens | |
| self.bootstrap_messages = copy.deepcopy(bootstrap_messages or []) | |
| self.system_suffix = system_suffix.strip() | |
| self.litellm_kwargs = litellm_kwargs | |
| self.messages: list[dict[str, Any]] = [] | |
| self.role: str = "red" | |
| self.last_response_text: str = "" | |
| self.last_command: str = "" | |
| def reset(self, briefing: str, role: Literal["red", "blue"]) -> None: | |
| """Initialize conversation history with role-specific system prompt.""" | |
| self.role = role | |
| system = RED_SYSTEM_PROMPT if role == "red" else BLUE_SYSTEM_PROMPT | |
| if self.system_suffix: | |
| system = f"{system}\n\n{self.system_suffix}" | |
| self.messages = [ | |
| {"role": "system", "content": system}, | |
| ] | |
| self.messages.extend(copy.deepcopy(self.bootstrap_messages)) | |
| self.messages.append({"role": "user", "content": briefing}) | |
| self.last_response_text = "" | |
| self.last_command = "" | |
| def act(self, observation: Any) -> str: | |
| """Call the LLM with the conversation history and return a command. | |
| Appends the observation as a user message (unless it was already | |
| appended via ``reset``), calls litellm, extracts the shell command | |
| from the response, and returns it. | |
| """ | |
| import litellm | |
| observation_text = format_observation(observation) | |
| # Append observation only if it wasn't already the last user message | |
| if self.messages and self.messages[-1]["role"] != "user": | |
| self.messages.append({"role": "user", "content": observation_text}) | |
| kwargs: dict[str, Any] = { | |
| "model": self.model, | |
| "messages": self.messages, | |
| "max_tokens": self.max_tokens, | |
| "drop_params": True, | |
| **self.litellm_kwargs, | |
| } | |
| # Codex deployments commonly reject temperature; omit it when unsupported. | |
| if self.temperature is not None and "codex" not in self.model.lower(): | |
| kwargs["temperature"] = self.temperature | |
| response = litellm.completion(**kwargs) | |
| text = response.choices[0].message.content.strip() | |
| self.messages.append({"role": "assistant", "content": text}) | |
| self.last_response_text = text | |
| self.last_command = extract_command(text) | |
| return self.last_command | |