Update app.py

app.py

@@ -9,6 +9,7 @@ import json
 _ = load_dotenv(find_dotenv())
 st.set_page_config(page_icon="💬", layout="wide", page_title="Groq Chat Bot...")
 
+
 def icon(emoji: str):
     """Shows an emoji as a Notion-style page icon."""
     st.write(
@@ -16,6 +17,7 @@ def icon(emoji: str):
         unsafe_allow_html=True,
     )
 
+
 icon("📣")
 
 st.subheader("Groq Chat Streamlit App", divider="rainbow", anchor=False)
@@ -24,6 +26,73 @@ client = Groq(
     api_key=os.environ['GROQ_API_KEY'],
 )
 
+def get_tool_owner_info():
+    owner_info = {
+        "date_time": datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
+    }
+    return json.dumps(owner_info)
+
+def run_conversation(user_prompt, messages):
+    tools = [
+        {
+            "type": "function",
+            "function": {
+                "name": "time_date",
+                "description": "The tool will return information about the time and date to the AI.",
+                "parameters": {},
+            },
+        }
+    ]
+    response = client.chat.completions.create(
+        model=model_option,
+        messages=messages,
+        tools=tools,
+        tool_choice="auto",  
+        max_tokens=max_tokens
+    )
+
+    response_message = response.choices[0].message
+    tool_calls = response_message.tool_calls
+
+    if tool_calls:
+        available_functions = {
+            "time_date": get_tool_owner_info
+        }
+
+        messages.append(response_message)
+
+        for tool_call in tool_calls:
+            function_name = tool_call.function.name
+            function_to_call = available_functions[function_name]
+            function_args = json.loads(tool_call.function.arguments)
+            function_response = function_to_call(**function_args)
+            messages.append(
+                {
+                    "tool_call_id": tool_call.id,
+                    "role": "tool",
+                    "name": function_name,
+                    "content": function_response,
+                }
+            )
+
+        second_response = client.chat.completions.create(
+            model=model_option,
+            messages=messages
+        )
+
+        return second_response.choices[0].message.content
+    else:
+        return response_message
+
+
+# Initialize chat history and selected model
+if "messages" not in st.session_state:
+    st.session_state.messages = []
+
+if "selected_model" not in st.session_state:
+    st.session_state.selected_model = None
+
+# Define model details
 models = {
     "mixtral-8x7b-32768": {
         "name": "Mixtral-8x7b-Instruct-v0.1",
@@ -34,6 +103,7 @@ models = {
     "gemma-7b-it": {"name": "Gemma-7b-it", "tokens": 8192, "developer": "Google"},
 }
 
+# Layout for model selection and max_tokens slider
 col1, col2 = st.columns(2)
 
 with col1:
@@ -41,15 +111,10 @@ with col1:
         "Choose a model:",
         options=list(models.keys()),
         format_func=lambda x: models[x]["name"],
-        index=0,
+        index=0,  # Default to the first model in the list
     )
 
-if "messages" not in st.session_state:
-    st.session_state.messages = []
-
-if "selected_model" not in st.session_state:
-    st.session_state.selected_model = None
-
+# Detect model change and clear chat history if model has changed
 if st.session_state.selected_model != model_option:
     st.session_state.messages = []
     st.session_state.selected_model = model_option
@@ -57,110 +122,55 @@ if st.session_state.selected_model != model_option:
 max_tokens_range = models[model_option]["tokens"]
 
 with col2:
+    # Adjust max_tokens slider dynamically based on the selected model
     max_tokens = st.slider(
         "Max Tokens:",
-        min_value=512,
+        min_value=512,  # Minimum value to allow some flexibility
         max_value=max_tokens_range,
+        # Default value or max allowed if less
         value=min(32768, max_tokens_range),
         step=512,
         help=f"Adjust the maximum number of tokens (words) for the model's response. Max for selected model: {max_tokens_range}",
     )
 
+# Display chat messages from history on app rerun
 for message in st.session_state.messages:
     avatar = "🤖" if message["role"] == "assistant" else "🕺"
     with st.chat_message(message["role"], avatar=avatar):
         st.markdown(message["content"])
 
-def generate_chat_responses(user_prompt):
-    """Fetches response from the Groq API using the run_conversation function."""
-    response = run_conversation(user_prompt)
-    yield response  # Yield the response content
-
-def run_conversation(user_prompt):
-    messages=[
-        {
-            "role": "system",
-            "content": "You are a helpful assistant named ChattyBot."
-        },
-        {
-            "role": "user",
-            "content": user_prompt,
-        }
-    ]
-    tools = [
-        {
-            "type": "function",
-            "function": {
-                "name": "time_date",
-                "description": "The tool will return information about the time and date to the AI.",
-                "parameters": {},
-            },
-        }
-    ]
-    try:
-        response = client.chat.completions.create(
-            model=model_option,
-            messages=messages,
-            tools=tools,
-            tool_choice="auto",
-            max_tokens=4096
-        )
-
-        response_message = response.choices[0].delta
-        tool_calls = response_message.tool_calls
-
-        if tool_calls:
-            available_functions = {
-                "time_date": get_tool_owner_info
-            }
-
-            messages.append(response_message)
-
-            for tool_call in tool_calls:
-                function_name = tool_call.function.name
-                function_to_call = available_functions[function_name]
-                function_args = json.loads(tool_call.function.arguments)
-                function_response = function_to_call(**function_args)
-                messages.append(
-                    {
-                        "tool_call_id": tool_call.id,
-                        "role": "tool",
-                        "name": function_name,
-                        "content": function_response,
-                    }
-                )
-
-            second_response = client.chat.completions.create(
-                model=model_option,
-                messages=messages
-            )
 
-            return second_response.choices[0].delta.content
-        else:
-            return response_message.content
-    except Exception as e:
-        st.error(e, icon="🚨")
-        return None
+def generate_chat_responses(chat_completion) -> Generator[str, None, None]:
+    """Yield chat response content from the Groq API response."""
+    for chunk in chat_completion:
+        if chunk.choices[0].delta.content:
+            yield chunk.choices[0].delta.content
 
-def get_tool_owner_info():
-    owner_info = {
-        "date_time": datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
-    }
-    return json.dumps(owner_info)
 
 if prompt := st.chat_input("Enter your prompt here..."):
     st.session_state.messages.append({"role": "user", "content": prompt})
 
-    with st.chat_message("user", avatar=""):
+    with st.chat_message("user", avatar="🕺"):
         st.markdown(prompt)
 
+    # Fetch response from Groq API
     try:
-        # Use generate_chat_responses with user prompt
-        with st.chat_message("assistant", avatar=""):
-            chat_responses_generator = generate_chat_responses(prompt)
+        chat_completion = client.chat.completions.create(
+            model=model_option,
+            messages=[
+                {"role": m["role"], "content": m["content"]}
+                for m in st.session_state.messages
+            ],
+            max_tokens=max_tokens,
+            stream=True,
+        )
+
+        # Use the generator function with st.write_stream
+        with st.chat_message("assistant", avatar="🤖"):
+            chat_responses_generator = generate_chat_responses(chat_completion)
             full_response = st.write_stream(chat_responses_generator)
     except Exception as e:
-        st.error(e, icon="")
+        st.error(e, icon="🚨")
 
     # Append the full response to session_state.messages
     if isinstance(full_response, str):
@@ -172,4 +182,5 @@ if prompt := st.chat_input("Enter your prompt here..."):
         combined_response = "\n".join(str(item) for item in full_response)
         st.session_state.messages.append(
             {"role": "assistant", "content": combined_response}
-        )
+        )
+