feat(clients): add GenericOpenAIAgent for any OpenAI Chat Completions API endpoint

.env.example

@@ -40,6 +40,22 @@ GROQ_API_KEY=your_groq_api_key_here
 # OpenRouter Model Selection (optional)
 OPENROUTER_MODEL=openai/gpt-4o-mini
 
+# =============================================================================
+# GENERIC OPENAI CHAT COMPLETIONS API CLIENT
+# =============================================================================
+
+# API key for any OpenAI Chat Completions API compatible endpoint (e.g. Poe)
+# Get a Poe API key from: https://creator.poe.com/
+OPENAI_COMPATIBLE_API_KEY=your_api_key_here
+
+# Base URL of the target endpoint
+# Example for Poe: https://api.poe.com/llm/v1
+OPENAI_COMPATIBLE_BASE_URL=https://api.poe.com/llm/v1
+
+# Model name to use via the generic client (default: gpt-4o)
+# Example for Poe: MiniMax-Text-01, Claude-3-7-Sonnet, etc.
+OPENAI_COMPATIBLE_MODEL=gpt-4o
+
 # Set it to a positive integer(window size in lines) to trim log blocks that only differ in timestamps. This saves input tokens and reduces context length.
 # For example, if LOG_TRIM=2, the code will compare every 1 line, and then every 2 lines for duplication.
 LOG_TRIM=0

clients/README.md

@@ -13,6 +13,7 @@ These clients are some baselines that we have implemented and evaluated to help
 - [ReAct](/clients/react.py): A naive LLM agent that uses the ReAct framework.
 - [FLASH](/clients/flash.py): A naive LLM agent that uses status supervision and hindsight integration components to ensure the high reliability of workflow execution.
 - [OpenRouter](/clients/openrouter.py): A naive OpenRouter LLM agent with only shell access.
+- [Generic OpenAI](/clients/generic_openai.py): A generic agent that works with any provider exposing the [OpenAI Chat Completions API](https://platform.openai.com/docs/api-reference/chat) (`/v1/chat/completions`), such as [Poe](https://creator.poe.com/docs/external-applications/openai-compatible-api), vLLM, LocalAI, standard OpenAI deployments, or other compatible services. The `base_url` and model are fully configurable via environment variables.
 
 ### Using the vLLM Client
 
@@ -84,6 +85,19 @@ cp .env.example .env
 - `OPENROUTER_MODEL`: OpenRouter model to use (default: `openai/gpt-4o-mini`)
 - `USE_WANDB`: Enable Weights & Biases logging (default: `false`)
 
+### Generic OpenAI-Compatible Client
+The [Generic OpenAI client](/clients/generic_openai.py) works with any provider that implements the OpenAI Chat Completions API (`/v1/chat/completions`), such as [Poe](https://creator.poe.com/docs/external-applications/openai-compatible-api), vLLM, LocalAI, or standard OpenAI.
+
+Set the following environment variables:
+- `OPENAI_COMPATIBLE_API_KEY`: API key for your target endpoint (required)
+- `OPENAI_COMPATIBLE_BASE_URL`: Base URL of your target endpoint, e.g. `https://api.poe.com/llm/v1` (required)
+- `OPENAI_COMPATIBLE_MODEL`: Model name to use, e.g. `MiniMax-Text-01` (default: `gpt-4o`)
+
+Then run:
+```bash
+python clients/generic_openai.py
+```
+
 ### Keyless Authentication
 
 The script [`gpt_azure_identity.py`](/clients/gpt_azure_identity.py) supports keyless authentication for **securely** accessing Azure OpenAI endpoints. It supports two authentication methods:

clients/generic_openai.py

@@ -0,0 +1,100 @@
+"""Generic OpenAI-compatible chat client (with shell access) for AIOpsLab.
+
+This agent works with any provider that implements the OpenAI Chat Completions
+API endpoint (/v1/chat/completions), such as Poe
+(https://creator.poe.com/docs/external-applications/openai-compatible-api),
+standard OpenAI deployments, vLLM, LocalAI, or other compatible services.
+
+Configure the endpoint and model via environment variables or constructor arguments:
+ OPENAI_COMPATIBLE_API_KEY — API key for the target endpoint
+ OPENAI_COMPATIBLE_BASE_URL — Base URL of the target endpoint (e.g. https://api.poe.com/llm/v1)
+ OPENAI_COMPATIBLE_MODEL — Model name to use (e.g. MiniMax-Text-01)
+"""
+
+import os
+import asyncio
+import wandb
+from aiopslab.orchestrator import Orchestrator
+from aiopslab.orchestrator.problems.registry import ProblemRegistry
+from clients.utils.llm import GenericOpenAIClient
+from clients.utils.templates import DOCS_SHELL_ONLY
+from dotenv import load_dotenv
+
+# Load environment variables from the .env file
+load_dotenv()
+
+
+class GenericOpenAIAgent:
+ def __init__(
+ self,
+ base_url: str | None = None,
+ model: str | None = None,
+ api_key: str | None = None,
+ ):
+ self.history = []
+ self.llm = GenericOpenAIClient(
+ base_url=base_url,
+ model=model,
+ api_key=api_key,
+ )
+
+ def init_context(self, problem_desc: str, instructions: str, apis: 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 (GenericOpenAI - {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 = GenericOpenAIAgent()
+
+ orchestrator = Orchestrator()
+ orchestrator.register_agent(agent, name="generic-openai")
+
+ 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

@@ -5,6 +5,7 @@
 from clients.deepseek import DeepSeekAgent
 from clients.vllm import vLLMAgent
 from clients.openrouter import OpenRouterAgent
+from clients.generic_openai import GenericOpenAIAgent
 
 class AgentRegistry:
  """Registry for agent implementations."""
@@ -16,6 +17,7 @@ def __init__(self):
  "deepseek": DeepSeekAgent,
  "vllm": vLLMAgent,
  "openrouter": OpenRouterAgent,
+ "generic": GenericOpenAIAgent,
  }
 
  def register(self, name, agent_cls):

clients/utils/llm.py

@@ -330,6 +330,84 @@ def run(self, payload: list[dict[str, str]]) -> list[str]:
  return response
 
 
+class GenericOpenAIClient:
+ """Generic client for any OpenAI Chat Completions compatible endpoint.
+
+ Uses the standard Chat Completions API (client.chat.completions.create),
+ making it compatible with any provider that implements the OpenAI Chat
+ Completions spec — including Poe, OpenRouter, vLLM, LocalAI, DeepSeek,
+ and standard OpenAI deployments, as well as Azure- or other cloud-hosted
+ gateways that expose an OpenAI-compatible `/v1/chat/completions` endpoint
+ via `base_url`.
+
+ Note: Native Azure OpenAI endpoints typically require the AzureOpenAI client
+ with an `azure_endpoint` and `api_version`, and are not used via `base_url`
+ in this class unless they are fronted by such a compatibility gateway.
+ Environment variables:
+ OPENAI_COMPATIBLE_API_KEY: API key for the target endpoint.
+ OPENAI_COMPATIBLE_BASE_URL: Base URL of the target endpoint.
+ OPENAI_COMPATIBLE_MODEL: Model name to use (default: gpt-4o).
+
+ All three can be overridden by passing explicit arguments to the constructor.
+ """
+
+ def __init__(
+ self,
+ base_url: Optional[str] = None,
+ model: Optional[str] = None,
+ api_key: Optional[str] = None,
+ max_tokens: int = 16000,
+ ):
+ self.cache = Cache()
+ self.model = model or os.getenv("OPENAI_COMPATIBLE_MODEL", "gpt-4o")
+ self.max_tokens = max_tokens
+ resolved_base_url = base_url or os.getenv("OPENAI_COMPATIBLE_BASE_URL")
+ if not resolved_base_url:
+ raise ValueError(
+ "base_url must be provided either as a constructor argument or via "
+ "the OPENAI_COMPATIBLE_BASE_URL environment variable."
+ )
+ resolved_api_key = api_key or os.getenv("OPENAI_COMPATIBLE_API_KEY")
+ if not resolved_api_key:
+ raise ValueError(
+ "api_key must be provided either as a constructor argument or via "
+ "the OPENAI_COMPATIBLE_API_KEY environment variable."
+ )
+ self.client = OpenAI(api_key=resolved_api_key, base_url=resolved_base_url)
+
+ 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:
+ response = self.client.chat.completions.create(
+ messages=payload, # type: ignore
+ model=self.model,
+ max_tokens=self.max_tokens,
+ temperature=0.5,
+ top_p=0.95,
+ frequency_penalty=0.0,
+ presence_penalty=0.0,
+ n=1,
+ timeout=60,
+ stop=[],
+ )
+ except Exception as e:
+ print(f"Exception: {repr(e)}")
+ raise e
+
+ 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
+
+
 class LLaMAClient:
  """Abstraction for Meta's LLaMA-3 model."""