From b2032d268a586d5cee090b11bbeb89f79d56d1d0 Mon Sep 17 00:00:00 2001 From: Prodman Devokadev Date: Mon, 21 Sep 2026 14:43:53 +0530 Subject: [PATCH] feat: add LiteLLM client and agent --- clients/litellm_agent.py | 116 ++++++++++++++ clients/registry.py | 2 + clients/utils/llm.py | 111 +++++++++++++ pyproject.toml | 1 + tests/clients/__init__.py | 2 + tests/clients/test_litellm_client.py | 227 +++++++++++++++++++++++++++ 6 files changed, 459 insertions(+) create mode 100644 clients/litellm_agent.py create mode 100644 tests/clients/__init__.py create mode 100644 tests/clients/test_litellm_client.py diff --git a/clients/litellm_agent.py b/clients/litellm_agent.py new file mode 100644 index 00000000..299db949 --- /dev/null +++ b/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() diff --git a/clients/registry.py b/clients/registry.py index 82cfd358..e623283d 100644 --- a/clients/registry.py +++ b/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): diff --git a/clients/utils/llm.py b/clients/utils/llm.py index 48a86cb0..61c796da 100644 --- a/clients/utils/llm.py +++ b/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 diff --git a/pyproject.toml b/pyproject.toml index 8bdd51e1..5ef66a96 100644 --- a/pyproject.toml +++ b/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] diff --git a/tests/clients/__init__.py b/tests/clients/__init__.py new file mode 100644 index 00000000..59e481eb --- /dev/null +++ b/tests/clients/__init__.py @@ -0,0 +1,2 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. diff --git a/tests/clients/test_litellm_client.py b/tests/clients/test_litellm_client.py new file mode 100644 index 00000000..5c1bedf4 --- /dev/null +++ b/tests/clients/test_litellm_client.py @@ -0,0 +1,227 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +import sys +import types +import unittest +from unittest.mock import MagicMock, patch + +from clients.utils.llm import LiteLLMClient + + +def _fake_litellm(content="hello", raises=None): + """Return a stub litellm module and the dict capturing its request.""" + captured = {} + + def completion(**kwargs): + captured.update(kwargs) + if raises is not None: + raise raises + message = MagicMock() + message.content = content + choice = MagicMock() + choice.message = message + response = MagicMock() + response.choices = [choice] + return response + + module = types.ModuleType("litellm") + module.completion = completion + return module, captured + + +class TestLiteLLMClientConfig(unittest.TestCase): + """Model, endpoint and key resolution.""" + + def test_defaults_to_gpt_4o_with_no_endpoint(self): + with patch.dict("os.environ", {}, clear=True): + client = LiteLLMClient(use_cache=False) + self.assertEqual(client.model, "gpt-4o") + self.assertEqual(client.api_base, "") + self.assertEqual(client.api_key, "") + + def test_reads_the_environment(self): + env = { + "LITELLM_MODEL": "anthropic/claude-opus-4-7", + "LITELLM_API_BASE": "http://localhost:4000", + "LITELLM_API_KEY": "sk-gateway", + } + with patch.dict("os.environ", env, clear=True): + client = LiteLLMClient(use_cache=False) + self.assertEqual(client.model, "anthropic/claude-opus-4-7") + self.assertEqual(client.api_base, "http://localhost:4000") + self.assertEqual(client.api_key, "sk-gateway") + + def test_constructor_arguments_win(self): + env = {"LITELLM_MODEL": "from-env", "LITELLM_API_BASE": "http://env:4000"} + with patch.dict("os.environ", env, clear=True): + client = LiteLLMClient( + model="explicit", api_base="http://explicit:4000", use_cache=False + ) + self.assertEqual(client.model, "explicit") + self.assertEqual(client.api_base, "http://explicit:4000") + + def test_trailing_slash_is_trimmed(self): + with patch.dict("os.environ", {}, clear=True): + client = LiteLLMClient(api_base="http://localhost:4000/", use_cache=False) + self.assertEqual(client.api_base, "http://localhost:4000") + + def test_no_endpoint_is_required(self): + """Unlike GenericOpenAIClient, constructing without a base_url is valid. + + That is the point of this client: reaching a provider needs no + OpenAI-compatible endpoint to exist first. + """ + with patch.dict("os.environ", {}, clear=True): + LiteLLMClient(use_cache=False) # must not raise + + +class TestLiteLLMClientRouting(unittest.TestCase): + """Whether a configured gateway is actually used.""" + + def test_direct_mode_sends_the_bare_model_name(self): + with patch.dict("os.environ", {}, clear=True): + client = LiteLLMClient(model="anthropic/claude-opus-4-7", use_cache=False) + self.assertEqual(client.route(), "anthropic/claude-opus-4-7") + + def test_gateway_mode_prefixes_the_model(self): + """Regression guard for a silently bypassed gateway. + + Without the litellm_proxy/ prefix LiteLLM infers the provider from the + model name and calls it directly, so api_base is ignored and the call + fails on missing provider credentials. + """ + with patch.dict("os.environ", {}, clear=True): + client = LiteLLMClient( + model="gemini-2.5-flash", + api_base="http://localhost:4000", + use_cache=False, + ) + self.assertEqual(client.route(), "litellm_proxy/gemini-2.5-flash") + + def test_an_explicit_prefix_is_not_doubled(self): + with patch.dict("os.environ", {}, clear=True): + client = LiteLLMClient( + model="litellm_proxy/gpt-4o", + api_base="http://localhost:4000", + use_cache=False, + ) + self.assertEqual(client.route(), "litellm_proxy/gpt-4o") + + +class TestLiteLLMClientInference(unittest.TestCase): + """The request LiteLLM actually receives.""" + + def _run(self, client, payload=None): + module, captured = _fake_litellm() + with patch.dict(sys.modules, {"litellm": module}): + result = client.inference(payload or [{"role": "user", "content": "hi"}]) + return result, captured + + def test_returns_the_message_content(self): + with patch.dict("os.environ", {}, clear=True): + client = LiteLLMClient(use_cache=False) + result, _ = self._run(client) + self.assertEqual(result, ["hello"]) + + def test_forwards_the_payload_unchanged(self): + with patch.dict("os.environ", {}, clear=True): + client = LiteLLMClient(use_cache=False) + payload = [ + {"role": "system", "content": "you are an SRE"}, + {"role": "user", "content": "the service is down"}, + ] + _, captured = self._run(client, payload) + self.assertEqual(captured["messages"], payload) + + def test_drop_params_is_on(self): + """Providers reject each other's parameters; without this the shared + temperature/top_p/n settings cannot serve every model.""" + with patch.dict("os.environ", {}, clear=True): + client = LiteLLMClient(use_cache=False) + _, captured = self._run(client) + self.assertTrue(captured["drop_params"]) + + def test_direct_mode_sends_no_endpoint_or_key(self): + with patch.dict("os.environ", {}, clear=True): + client = LiteLLMClient(model="gpt-4o", use_cache=False) + _, captured = self._run(client) + self.assertNotIn("api_base", captured) + self.assertNotIn("api_key", captured) + + def test_gateway_mode_forwards_endpoint_and_key(self): + with patch.dict("os.environ", {}, clear=True): + client = LiteLLMClient( + model="gpt-4o", + api_base="http://localhost:4000", + api_key="sk-gateway", + use_cache=False, + ) + _, captured = self._run(client) + self.assertEqual(captured["api_base"], "http://localhost:4000") + self.assertEqual(captured["api_key"], "sk-gateway") + self.assertEqual(captured["model"], "litellm_proxy/gpt-4o") + + def test_max_tokens_is_configurable(self): + with patch.dict("os.environ", {}, clear=True): + client = LiteLLMClient(max_tokens=256, use_cache=False) + _, captured = self._run(client) + self.assertEqual(captured["max_tokens"], 256) + + def test_provider_errors_propagate(self): + with patch.dict("os.environ", {}, clear=True): + client = LiteLLMClient(use_cache=False) + module, _ = _fake_litellm(raises=RuntimeError("provider exploded")) + with patch.dict(sys.modules, {"litellm": module}): + with self.assertRaises(RuntimeError): + client.inference([{"role": "user", "content": "hi"}]) + + def test_missing_package_explains_how_to_install(self): + with patch.dict("os.environ", {}, clear=True): + client = LiteLLMClient(use_cache=False) + with patch.dict(sys.modules, {"litellm": None}): + with self.assertRaises(ImportError) as ctx: + client.inference([{"role": "user", "content": "hi"}]) + self.assertIn("poetry install", str(ctx.exception)) + + +class TestLiteLLMClientCaching(unittest.TestCase): + """Caching matches the other clients' behaviour.""" + + def test_run_stores_the_response(self): + with patch.dict("os.environ", {}, clear=True): + client = LiteLLMClient(use_cache=False) + client.cache = MagicMock() + client.cache.get_from_cache.return_value = None + + module, _ = _fake_litellm() + with patch.dict(sys.modules, {"litellm": module}): + result = client.run([{"role": "user", "content": "hi"}]) + + self.assertEqual(result, ["hello"]) + client.cache.add_to_cache.assert_called_once() + client.cache.save_cache.assert_called_once() + + def test_a_cache_hit_skips_the_provider(self): + with patch.dict("os.environ", {}, clear=True): + client = LiteLLMClient(use_cache=False) + client.cache = MagicMock() + client.cache.get_from_cache.return_value = ["cached"] + + # No litellm module installed: a cache hit must not need one. + with patch.dict(sys.modules, {"litellm": None}): + result = client.inference([{"role": "user", "content": "hi"}]) + + self.assertEqual(result, ["cached"]) + + +class TestLiteLLMAgentRegistration(unittest.TestCase): + def test_agent_is_registered(self): + from clients.registry import AgentRegistry + from clients.litellm_agent import LiteLLMAgent + + self.assertIs(AgentRegistry().AGENT_REGISTRY["litellm"], LiteLLMAgent) + + +if __name__ == "__main__": + unittest.main()