Update chain_app.py

chain_app.py

@@ -7,6 +7,7 @@ from chainlit.input_widget import Select, Slider
 import os
 import cohere 
 from huggingface_hub import InferenceClient
+from fireworks.client import Fireworks
 
 
 hf_token = os.environ.get("HF_TOKEN")
@@ -2197,34 +2198,17 @@ async def main(message: cl.Message):
     #     ).send()
         
     elif chat_profile == 'Llama-3.1-405B':        
-        completion = groq_client.chat.completions.create(
-            model="llama-3.1-405b-reasoning",
-            messages=[
-                {
-                    "role": "user",
-                    "content": message.content
-                }
-            ],
-            temperature=1,
-            max_tokens=1024,
-            top_p=1,
-            stream=True,
-            stop=None,
+        client = Fireworks(api_key="O9D2HQsat357QRQ2wEAC3buShjL8Ea4K7ndDqwkAsAdsDulz")
+        response = client.chat.completions.create(
+        model="accounts/fireworks/models/llama-v3p1-405b-instruct",
+        messages=[{
+           "role": "user",
+           "content": "Say this is a test",
+        }],
         )
-    
-        complete_content = ""
-    
-        # Iterate over each chunk
-        for chunk in completion:
-            # Retrieve the content from the current chunk
-            content = chunk.choices[0].delta.content
-            
-            # Check if the content is not None before concatenating it
-            if content is not None:
-                complete_content += content
-    
-        # Send the concatenated content as a message
-        await cl.Message(content=complete_content).send()
+        await cl.Message(
+            content=response.choices[0].message.content,
+        ).send()
 
     
     elif chat_profile == 'Llama-3.1-70B':