Minor fix

Author: christinaexyou

nemoguardrails/server/api.py

@@ -486,18 +486,24 @@ async def chat_completion(body: RequestBody, request: Request):
             messages = thread_messages + messages
 
         generation_options = body.options
+
+        # Initialize llm_params if not already set
+        if generation_options.llm_params is None:
+            generation_options.llm_params = {}
+
+        # Set OpenAI-compatible parameters in llm_params
         if body.max_tokens:
-            generation_options.max_tokens = body.max_tokens
+            generation_options.llm_params["max_tokens"] = body.max_tokens
         if body.temperature is not None:
-            generation_options.temperature = body.temperature
+            generation_options.llm_params["temperature"] = body.temperature
         if body.top_p is not None:
-            generation_options.top_p = body.top_p
+            generation_options.llm_params["top_p"] = body.top_p
         if body.stop:
-            generation_options.stop = body.stop
+            generation_options.llm_params["stop"] = body.stop
         if body.presence_penalty is not None:
-            generation_options.presence_penalty = body.presence_penalty
+            generation_options.llm_params["presence_penalty"] = body.presence_penalty
         if body.frequency_penalty is not None:
-            generation_options.frequency_penalty = body.frequency_penalty
+            generation_options.llm_params["frequency_penalty"] = body.frequency_penalty
 
         if (
             body.stream

tests/test_threads.py

@@ -140,4 +140,4 @@ def test_with_redis():
         },
     )
     res = response.json()
-    assert res["choices"]["message"][0]["content"] == "Hello again!"
+    assert res["choices"][0]["message"]["content"] == "Hello again!"