feat: add LiteLLM client and agent

clients/litellm_agent.py

@@ -0,0 +1,116 @@
+"""LiteLLM-backed agent (with shell access) for AIOpsLab.
+
+LiteLLM (https://docs.litellm.ai/) gives one interface to 100+ model
+providers, so a problem can be evaluated against models the registry has no
+dedicated client for, including Anthropic and Gemini.
+
+Two ways to use it, from the same model name:
+
+1. Direct, with no extra infrastructure. LiteLLM routes each model to its own
+   provider using that provider's key:
+
+       LITELLM_MODEL=anthropic/claude-opus-4-7
+       ANTHROPIC_API_KEY=...
+
+2. Through a self-hosted LiteLLM gateway, which keeps every provider key
+   server-side and adds centralized cost tracking, budgets, rate limiting and
+   fallbacks across a long benchmark run:
+
+       LITELLM_MODEL=claude-opus-4-7
+       LITELLM_API_BASE=http://localhost:4000
+       LITELLM_API_KEY=...
+
+Configure via environment variables or constructor arguments:
+    LITELLM_MODEL    — model to route (default: gpt-4o)
+    LITELLM_API_BASE — optional gateway URL; unset means call the provider directly
+    LITELLM_API_KEY  — gateway key when LITELLM_API_BASE is set, otherwise unset
+                       so LiteLLM reads the provider's own variable
+"""
+
+import os
+import asyncio
+import wandb
+from aiopslab.orchestrator import Orchestrator
+from aiopslab.orchestrator.problems.registry import ProblemRegistry
+from clients.utils.llm import LiteLLMClient
+from clients.utils.templates import DOCS_SHELL_ONLY
+from dotenv import load_dotenv
+
+# Load environment variables from the .env file
+load_dotenv()
+
+
+class LiteLLMAgent:
+    def __init__(
+        self,
+        model: str | None = None,
+        api_base: str | None = None,
+        api_key: str | None = None,
+    ):
+        self.history = []
+        self.llm = LiteLLMClient(
+            model=model,
+            api_base=api_base,
+            api_key=api_key,
+        )
+
+    def init_context(self, problem_desc: str, instructions: str, apis: dict[str, str]):
+        """Initialize the context for the agent."""
+
+        self.shell_api = self._filter_dict(apis, lambda k, _: "exec_shell" in k)
+        self.submit_api = self._filter_dict(apis, lambda k, _: "submit" in k)
+        stringify_apis = lambda apis: "\n\n".join(
+            [f"{k}\n{v}" for k, v in apis.items()]
+        )
+
+        self.system_message = DOCS_SHELL_ONLY.format(
+            prob_desc=problem_desc,
+            shell_api=stringify_apis(self.shell_api),
+            submit_api=stringify_apis(self.submit_api),
+        )
+
+        self.task_message = instructions
+
+        self.history.append({"role": "system", "content": self.system_message})
+        self.history.append({"role": "user", "content": self.task_message})
+
+    async def get_action(self, input) -> str:
+        """Wrapper to interface the agent with AIOpsLab.
+
+        Args:
+            input (str): The input from the orchestrator/environment.
+
+        Returns:
+            str: The response from the agent.
+        """
+        self.history.append({"role": "user", "content": input})
+        response = self.llm.run(self.history)
+        model_name = self.llm.model
+        print(f"===== Agent (LiteLLM - {model_name}) ====\n{response[0]}")
+        self.history.append({"role": "assistant", "content": response[0]})
+        return response[0]
+
+    def _filter_dict(self, dictionary, filter_func):
+        return {k: v for k, v in dictionary.items() if filter_func(k, v)}
+
+
+if __name__ == "__main__":
+    # Load use_wandb from environment variable with a default of False
+    use_wandb = os.getenv("USE_WANDB", "false").lower() == "true"
+
+    if use_wandb:
+        wandb.init(project="AIOpsLab", entity="AIOpsLab")
+
+    problems = ProblemRegistry().PROBLEM_REGISTRY
+    for pid in problems:
+        agent = LiteLLMAgent()
+
+        orchestrator = Orchestrator()
+        orchestrator.register_agent(agent, name="litellm")
+
+        problem_desc, instructs, apis = orchestrator.init_problem(pid)
+        agent.init_context(problem_desc, instructs, apis)
+        asyncio.run(orchestrator.start_problem(max_steps=30))
+
+    if use_wandb:
+        wandb.finish()

clients/registry.py

@@ -6,6 +6,7 @@
 from clients.vllm import vLLMAgent
 from clients.openrouter import OpenRouterAgent
 from clients.generic_openai import GenericOpenAIAgent
+from clients.litellm_agent import LiteLLMAgent
 
 class AgentRegistry:
     """Registry for agent implementations."""
@@ -18,6 +19,7 @@ def __init__(self):
             "vllm": vLLMAgent,
             "openrouter": OpenRouterAgent,
             "generic": GenericOpenAIAgent,
+            "litellm": LiteLLMAgent,
         }
 
     def register(self, name, agent_cls):

clients/utils/llm.py

@@ -446,3 +446,114 @@ def run(self, payload: list[dict[str, str]]) -> list[str]:
             self.cache.add_to_cache(payload, response)
             self.cache.save_cache()
         return response
+
+
+class LiteLLMClient:
+    """Abstraction for any model reachable through LiteLLM.
+
+    LiteLLM (https://docs.litellm.ai/) provides one interface to 100+ model
+    providers, which is useful here because the registry has no Anthropic or
+    Gemini client: reaching those models today means either routing through
+    OpenRouter or standing up an OpenAI-compatible endpoint for
+    GenericOpenAIClient, which requires a base_url and refuses to construct
+    without one.
+
+    This client needs no endpoint. Used directly it calls each provider with
+    that provider's own key (ANTHROPIC_API_KEY and friends), so evaluating a
+    problem across vendors needs no extra infrastructure:
+
+        LITELLM_MODEL=anthropic/claude-opus-4-7
+
+    Pointing it at a self-hosted LiteLLM gateway instead keeps every provider
+    key server-side and adds centralized cost tracking, budgets, rate limiting
+    and fallbacks, which matter over a long benchmark run:
+
+        LITELLM_MODEL=claude-opus-4-7
+        LITELLM_API_BASE=http://localhost:4000
+        LITELLM_API_KEY=sk-...
+
+    Environment variables:
+        LITELLM_MODEL:    model to route (default: gpt-4o).
+        LITELLM_API_BASE: optional gateway URL. Unset means call the provider
+                          directly.
+        LITELLM_API_KEY:  gateway key when LITELLM_API_BASE is set; otherwise
+                          leave unset so LiteLLM reads the provider's own
+                          variable.
+
+    All three can be overridden by passing explicit constructor arguments.
+    """
+
+    # LiteLLM's own prefix meaning "forward this to my gateway rather than
+    # resolving the provider yourself". Without it LiteLLM infers the provider
+    # from the model name and calls it directly, silently bypassing the
+    # gateway. See https://docs.litellm.ai/docs/providers/litellm_proxy
+    PROXY_PREFIX = "litellm_proxy/"
+
+    def __init__(
+        self,
+        model: Optional[str] = None,
+        api_base: Optional[str] = None,
+        api_key: Optional[str] = None,
+        max_tokens: int = 16000,
+        use_cache: bool = True,
+    ):
+        self.cache = Cache() if use_cache else None
+        self.model = model or os.getenv("LITELLM_MODEL", GPT_MODEL)
+        self.api_base = (
+            (api_base or os.getenv("LITELLM_API_BASE") or "").strip().rstrip("/")
+        )
+        self.api_key = (api_key or os.getenv("LITELLM_API_KEY") or "").strip()
+        self.max_tokens = max_tokens
+
+    def route(self) -> str:
+        """Model name to send, prefixed when a gateway should forward it."""
+        if not self.api_base or self.model.startswith(self.PROXY_PREFIX):
+            return self.model
+        return f"{self.PROXY_PREFIX}{self.model}"
+
+    def inference(self, payload: list[dict[str, str]]) -> list[str]:
+        if self.cache is not None:
+            cache_result = self.cache.get_from_cache(payload)
+            if cache_result is not None:
+                return cache_result
+
+        try:
+            import litellm
+        except ImportError as e:
+            raise ImportError(
+                "The 'litellm' package is required for LiteLLMClient. "
+                "Install it with: poetry install --with clients"
+            ) from e
+
+        request: Dict = {
+            "model": self.route(),
+            "messages": payload,
+            "max_tokens": self.max_tokens,
+            "temperature": 0.5,
+            "top_p": 0.95,
+            "n": 1,
+            "timeout": 60,
+            # Providers reject each other's parameters, so let LiteLLM drop
+            # what a given model does not support instead of failing the call.
+            # Without this the shared settings above cannot serve every model.
+            "drop_params": True,
+        }
+        if self.api_base:
+            request["api_base"] = self.api_base
+        if self.api_key:
+            request["api_key"] = self.api_key
+
+        try:
+            response = litellm.completion(**request)
+        except Exception as e:
+            print(f"Exception: {repr(e)}")
+            raise
+
+        return [c.message.content for c in response.choices]  # type: ignore
+
+    def run(self, payload: list[dict[str, str]]) -> list[str]:
+        response = self.inference(payload)
+        if self.cache is not None:
+            self.cache.add_to_cache(payload, response)
+            self.cache.save_cache()
+        return response

pyproject.toml

@@ -41,6 +41,7 @@ vllm = "^0.7.3"
 transformers = ">=4.49,<6.0"
 groq = "^0.28.0"
 flwr = "^1.19.0"
+litellm = ">=1.92.0,<1.101.0"
 
 # Developer tooling (testing, etc.)
 [tool.poetry.group.dev.dependencies]

tests/clients/__init__.py

@@ -0,0 +1,2 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.