Merge branch 'main' into tlsdc/log_prob

Author: TLSDC

src/agentlab/agents/generic_agent/__init__.py

@@ -16,7 +16,9 @@
     AGENT_4o,
     AGENT_4o_MINI,
     AGENT_CLAUDE_SONNET_35,
+    AGENT_CLAUDE_SONNET_35_VISION,
     AGENT_4o_VISION,
+    AGENT_4o_MINI_VISION,
     AGENT_o3_MINI,
     AGENT_o1_MINI,
 )

src/agentlab/llm/chat_api.py

@@ -147,6 +147,13 @@ def make_model(self):
                 n_retry_server=self.n_retry_server,
                 log_probs=self.log_probs,
             )
+        elif self.backend == "vllm":
+            return VLLMChatModel(
+                model_name=self.model_name,
+                temperature=self.temperature,
+                max_tokens=self.max_new_tokens,
+                n_retry_server=self.n_retry_server,
+            )
         else:
             raise ValueError(f"Backend {self.backend} is not supported")
 
@@ -440,3 +447,27 @@ def __init__(
 
         client = InferenceClient(model=model_url, token=token)
         self.llm = partial(client.text_generation, max_new_tokens=max_new_tokens, details=log_probs)
+
+
+class VLLMChatModel(ChatModel):
+    def __init__(
+        self,
+        model_name,
+        api_key=None,
+        temperature=0.5,
+        max_tokens=100,
+        n_retry_server=4,
+        min_retry_wait_time=60,
+    ):
+        super().__init__(
+            model_name=model_name,
+            api_key=api_key,
+            temperature=temperature,
+            max_tokens=max_tokens,
+            max_retry=n_retry_server,
+            min_retry_wait_time=min_retry_wait_time,
+            api_key_env_var="VLLM_API_KEY",
+            client_class=OpenAI,
+            client_args={"base_url": "http://0.0.0.0:8000/v1"},
+            pricing_func=None,
+        )