Update llm_node.py

This commit is contained in:
Sam Khoze
2024-06-15 23:32:24 +05:30
committed by GitHub
parent 7cdc1ef471
commit edcff6a89a
+19 -14
View File
@@ -19,21 +19,26 @@ class LLM_node:
} }
CATEGORY = "DeepFuze" CATEGORY = "DeepFuze"
RETURN_TYPES = ("STRING",) RETURN_TYPES = ("NEW_STRING",)
RETURN_NAMES = ("text",) RETURN_NAMES = ("LLM_RESPONSE",)
FUNCTION = "run_llm" FUNCTION = "run_llm"
def run_llm(self, system_prompt, user_query, model_name,temperature,api_key,max_tokens,timeout): def run_llm(self, system_prompt, user_query, model_name,temperature,api_key,max_tokens,timeout):
client = OpenAI(api_key=api_key) client = OpenAI(api_key=api_key)
response = client.chat.completions.create( try:
model=model_name, response = client.chat.completions.create(
messages=[ model=model_name,
{"role": "system", "content": system_prompt}, messages=[
{"role": "user", "content": user_query} {"role": "system", "content": system_prompt},
], {"role": "user", "content": user_query}
# stop = stop.split(","), ],
temperature=temperature, # stop = stop.split(","),
max_tokens=max_tokens, temperature=temperature,
timeout=timeout max_tokens=max_tokens,
) timeout=timeout
return (text,) ).to_dict()
print(response)
result = response["choices"][0]["message"]["content"]
except Exception as e:
raise ValueError(e)
return (result,)