From 747169156c88d3f9137b6d2d784052b934b50156 Mon Sep 17 00:00:00 2001 From: ishita Date: Thu, 23 May 2024 22:31:42 +0530 Subject: [PATCH] Add auth-key support and validation to Vllm --- src/vanna/vllm/vllm.py | 17 ++++++++++++++++- 1 file changed, 16 insertions(+), 1 deletion(-) diff --git a/src/vanna/vllm/vllm.py b/src/vanna/vllm/vllm.py index 0dd67e4f..53990821 100644 --- a/src/vanna/vllm/vllm.py +++ b/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()