Merge pull request #455 from ishitachauhan12/main

Add auth-key support and validation to vLLM

Author: zainhoda

src/vanna/vllm/vllm.py

@@ -17,6 +17,11 @@ def __init__(self, config=None):
         else:
             self.model = config["model"]
 
+        if "auth-key" in config:
+            self.auth_key = config["auth-key"]
+        else:
+            self.auth_key = None
+
     def system_message(self, message: str) -> any:
         return {"role": "system", "content": message}
 
@@ -67,7 +72,17 @@ def submit_prompt(self, prompt, **kwargs) -> str:
             "messages": prompt,
         }
 
-        response = requests.post(url, json=data)
+        if self.auth_key is not None:
+            headers = { 
+            'Content-Type': 'application/json', 
+            'Authorization': f'Bearer {self.auth_key}' 
+            }
+
+            response = requests.post(url, headers=headers,json=data)
+
+
+        else:
+            response = requests.post(url, json=data)
 
         response_dict = response.json()