Update app.py

app.py

@@ -2,6 +2,7 @@ import os
 from dotenv import load_dotenv
 import streamlit as st
 import openai
+from tavily import TavilySearchResults  # Import TavilySearchResults
 
 # Load environment variables
 load_dotenv()
@@ -20,30 +21,60 @@ if "messages" not in st.session_state:
 for msg in st.session_state["messages"]:
     st.chat_message(msg["role"]).write(msg["content"])
 
+# Initialize TavilySearchResults tool
+tavily_tool = TavilySearchResults(
+    max_results=7,
+    search_depth="advanced",
+    topic="news",
+    days=1,
+    include_answer=True,
+    include_raw_content=True,
+    exclude_domains=['example.com'],
+)
+
 # Input from user
 if prompt := st.chat_input():
     # Append user message to session state
     st.session_state["messages"].append({"role": "user", "content": prompt})
     st.chat_message("user").write(prompt)
-    
+
     try:
-        # OpenAI ChatCompletion API call
-        response = openai.chat.completions.create(
-            model="gpt-4o",
-            messages=[
-                {"role": "system", "content": "You are an helpful assistant "},
-                *st.session_state["messages"]
-            ],
-            max_tokens=150,  # Adjust as needed
-            temperature=0.7  # Adjust as needed
-        )
-        
-        # Extract response message content
-        msg = response.choices[0].message.content.strip()
-        
-        # Append assistant message to session state
-        st.session_state["messages"].append({"role": "assistant", "content": msg})
-        st.chat_message("assistant").write(msg)
+        # Check if user query requires Tavily
+        if "search for" in prompt.lower() or "find" in prompt.lower():
+            # Use TavilySearchResults tool
+            tavily_response = tavily_tool.run(prompt)
+
+            # Format Tavily results
+            if tavily_response:
+                formatted_results = "\n\n".join(
+                    [f"**{i+1}.** {result['title']}\n{result['content']}" for i, result in enumerate(tavily_response)]
+                )
+                tavily_message = f"Here are the top results I found:\n\n{formatted_results}"
+            else:
+                tavily_message = "I couldn't find any relevant results."
+
+            # Append Tavily response to session state
+            st.session_state["messages"].append({"role": "assistant", "content": tavily_message})
+            st.chat_message("assistant").write(tavily_message)
+
+        else:
+            # GPT-4 response for general queries
+            response = openai.chat.completions.create(
+                model="gpt-4o",
+                messages=[
+                    {"role": "system", "content": "You are a helpful assistant."},
+                    *st.session_state["messages"]
+                ],
+                max_tokens=150,
+                temperature=0.7
+            )
+
+            # Extract GPT response message
+            gpt_message = response.choices[0].message.content.strip()
+
+            # Append GPT response to session state
+            st.session_state["messages"].append({"role": "assistant", "content": gpt_message})
+            st.chat_message("assistant").write(gpt_message)
 
     except Exception as e:
         st.error(f"An error occurred: {e}")