Update app.py

app.py

@@ -92,61 +92,49 @@ def search_web(query):
     except Exception as e:
         return f"An error occurred: {e}"
 
-def run_conversation(user_prompt):
-    # Step 1: send the conversation and available functions to the model
-    messages = [
-        {
-            "role": "system",
-            "content": "You are a function calling LLM that uses the data extracted from the get_game_score function to answer questions around NBA game scores. Include the team and their opponent in your response."
-        },
-        {
-            "role": "user",
-            "content": user_prompt,
-        }
-    ]
-    tools = [
-        {
-            "type": "internet",
-            "internet": {
-                "allow": ["search_web"],
-                "description": "Search the web for information.",
-                "parameters": {
-                    "type": "object",
-                    "properties": {
-                        "query": {
-                            "type": "string",
-                            "description": "The search query.",
-                        }
-                    },
-                    "required": ["query"],
-                },
-            },
-        }
-    ]
-    response = client.chat.completions.create(
-        model=model_option,
-        messages=messages,
-        tools=tools,
-        tool_choice="auto",  
-        max_tokens=max_tokens
-    )
-
-    for chunk in response:
-        if chunk.choices[0].delta.content:
-            yield chunk.choices[0].delta.content
+full_response = None  # Initialize full_response to None
 
-if prompt := st.text_input("Enter your prompt here..."):
+if prompt := st.chat_input("Enter your prompt here..."):
     st.session_state.messages.append({"role": "user", "content": prompt})
 
     with st.chat_message("user", avatar="🕺"):  
         st.markdown(prompt)
 
     try:
-        chat_completion = run_conversation(prompt)
-
-        with st.chat_message("assistant", avatar="🤖"):
-            chat_responses_generator = generate_chat_responses(chat_completion)
-            for response in chat_responses_generator:
-                st.markdown(response)
+        if "search for" in prompt.lower():
+            query = prompt.lower().replace("search for", "").strip()
+            search_results = search_web(query)
+            formatted_results = "\n\n".join([f"Title: {result['title']}\nURL: {result['url']}\nSnippet: {result['snippet']}" for result in search_results])
+            st.session_state.messages.append({"role": "assistant", "content": formatted_results})
+        else:
+            chat_completion = client.chat.completions.create(
+                model=model_option,
+                messages=[
+                    {
+                        "role": "system",
+                        "content": "You are a helpful assistant that can access the web.",
+                    },
+                    {"role": m["role"], "content": m["content"]}
+                    for m in st.session_state.messages
+                ],
+                max_tokens=max_tokens,
+                stream=True,
+            )
+
+            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="🚨")
+
+    # Check if full_response is defined before using it
+    if full_response is not None:
+        if isinstance(full_response, str):
+            st.session_state.messages.append(
+                {"role": "assistant", "content": full_response}
+            )
+        else:
+            combined_response = "\n".join(str(item) for item in full_response)
+            st.session_state.messages.append(
+                {"role": "assistant", "content": combined_response}
+            )