From a0c8007a1336e98930ecec86448a374468b735cd Mon Sep 17 00:00:00 2001
From: Amit Singh <142410046+AmitSinghShorthillsAI@users.noreply.github.com>
Date: Tue, 17 Sep 2024 01:11:18 +0530
Subject: [PATCH] Add temperature support to vLLM

Added temperature support for hosted LLM using vLLM

Changes made:
* Introduced default temperature of 0.7 in __init__ method
* Updated JSON payload in submit_prompt method to include temperature


This change allows users to control the randomness of the model's output.
If not specified, it defaults to 0.7, providing a balance between
creativity and coherence in the generated text.
---
 src/vanna/vllm/vllm.py | 7 +++++++
 1 file changed, 7 insertions(+)

diff --git a/src/vanna/vllm/vllm.py b/src/vanna/vllm/vllm.py
index 53990821..3227d83a 100644
--- a/src/vanna/vllm/vllm.py
+++ b/src/vanna/vllm/vllm.py
@@ -22,6 +22,12 @@ def __init__(self, config=None):
         else:
             self.auth_key = None
 
+        if "temperature" in config:
+            self.temperature = config["temperature"]
+        else:
+            # default temperature - can be overrided using config
+            self.temperature = 0.7
+
     def system_message(self, message: str) -> any:
         return {"role": "system", "content": message}
 
@@ -68,6 +74,7 @@ def submit_prompt(self, prompt, **kwargs) -> str:
         url = f"{self.host}/v1/chat/completions"
         data = {
             "model": self.model,
+            "temperature": self.temperature,
             "stream": False,
             "messages": prompt,
         }