Update chain_app.py

chain_app.py

@@ -5,16 +5,22 @@ from groq import Groq
 import requests
 from chainlit.input_widget import Select, Slider
 import os
+import cohere 
 
 hf_token = os.environ.get("HF_TOKEN")
 openai_api_key = os.environ.get('OPENAI_API_KEY')
 groq_api_key = os.environ.get('GROQ_API_KEY')
+cohere_api_key = os.enviorn.get('COHERE_API_KEY')
 
 hf_text_client = Client("Artin2009/text-generation", hf_token=hf_token)
 # hf_image_client = Client('Artin2009/image-generation')
 openai_client = OpenAI(api_key=openai_api_key)
 groq_client = Groq(api_key=groq_api_key)
 
+co = cohere.Client(
+  api_key=cohere_api_key, # This is your trial API key
+) 
+
 # API_URL = "https://api-inference.huggingface.co/models/PartAI/TookaBERT-Large"
 # headers = {"Authorization": f"Bearer {hf_token}"}
 
@@ -70,6 +76,10 @@ async def chat_profile():
             name="Llama-3-70B",
             markdown_description="Meta Open Source model Llama-2 with 70B parameters",
         ),
+        cl.ChatProfile(
+            name='Aya-23B',
+            markdown_description='Cohere open sourced AI model with 23B parameters'
+        ),
         cl.ChatProfile(
             name="Llama-3-8B",
             markdown_description="Meta Open Source model Llama-2 with 7B parameters",
@@ -677,6 +687,19 @@ async def main(message: cl.Message):
     #         content=output[0]['sequence']
     #     ).send()
 
+    elif chat_profile == 'Aya-23B':
+        stream = co.chat_stream( 
+        model='c4ai-aya-23',
+        message=message.content,
+        temperature=0.3,
+        # chat_history=[{"role": "User", "message": "Hello"}, {"role": "Chatbot", "message": "Hello! How can I help you today?"}, {"role": "User", "message": "Hi"}, {"role": "User", "message": "hello"}],
+        prompt_truncation='OFF',
+        connectors=[],
+        ) 
+
+        for event in stream:
+            if event.event_type == "text-generation":
+               await cl.Message(content=event.text).send()
 @cl.on_settings_update
 async def setup_agent(settings):
     print("on_settings_update", settings)