Update mistral_api.py

mistral_api.py

@@ -26,28 +26,38 @@ client = InferenceClient("mistralai/Mistral-Small-24B-Instruct-2501")#, token=os
 metaprompt_filepath = "meta_prompt.txt"
 META_PROMPT = load_file(metaprompt_filepath)
 
-def get_response(message, system_message=META_PROMPT, max_tokens=800, temperature=0) -> str:
-    """Chat completion using HuggingFace's Inference API"""
-    
-    messages = [{"role":"system", "content":system_message}]
-    print("Got to get_response")
-    # uncomment the lines below and add a parameter history: list[tuple[str, str]] to support history
-    #if history:
-    #    for val in history:
-    #        if val[0]:
-    #            messages.append({"role":"user", "content":val[0]})
-    #       if val[1]:
-    #            messages.append({"role": "assistant", "content": val[1]})
-
-    messages.append({"role": "user", "content": message})
-    response = ""
-    for message in client.chat_completion(
-        messages,
+def get_response(messages: list[dict], temperature=0, max_tokens=800) -> str:
+    """Chat completion using HuggingFace's Inference API."""
+    response = client.chat_completion(
+        messages=messages,
         max_tokens=max_tokens,
-        stream=True,
-        temperature=temperature
-    ):
-        token = message.choices[0].delta.content
-        response += str(token)
+        temperature=temperature,
+        stream=True  # Enables streaming
+    )
+    return response.["choices"][0]["delta"]["content"]
+
+# def get_response(message, system_message=META_PROMPT, max_tokens=800, temperature=0) -> str:
+#     """Chat completion using HuggingFace's Inference API"""
+    
+#     messages = [{"role":"system", "content":system_message}]
+#     print("Got to get_response")
+#     # uncomment the lines below and add a parameter history: list[tuple[str, str]] to support history
+#     #if history:
+#     #    for val in history:
+#     #        if val[0]:
+#     #            messages.append({"role":"user", "content":val[0]})
+#     #       if val[1]:
+#     #            messages.append({"role": "assistant", "content": val[1]})
+
+#     messages.append({"role": "user", "content": message})
+#     response = ""
+#     for message in client.chat_completion(
+#         messages,
+#         max_tokens=max_tokens,
+#         stream=True,
+#         temperature=temperature
+#     ):
+#         token = message.choices[0].delta.content
+#         response += str(token)
     
-    yield response
+#     yield response