Update app.py

app.py

@@ -1,26 +1,41 @@
+import logging
 import os
 import streamlit as st
+from dotenv import load_dotenv
 import openai
-from langchain.embeddings.openai import OpenAIEmbeddings
-from langchain.vectorstores.faiss import FAISS
-from langchain.document_loaders import TextLoader
-from langchain.chains import LLMChain
-from langchain.prompts.chat import ChatPromptTemplate, MessagesPlaceholder
-from langchain.chat_models import ChatOpenAI
-from langchain.schema import HumanMessage, AIMessage
-from serpapi import GoogleSearch
-import logging
+from langchain_openai import ChatOpenAI
+from langchain_community.vectorstores import FAISS
+from langchain_openai import OpenAIEmbeddings
+from langchain.prompts import ChatPromptTemplate, MessagesPlaceholder
+from langchain.agents import tool, AgentExecutor
+from langchain.agents.output_parsers.openai_tools import OpenAIToolsAgentOutputParser
+from langchain.agents.format_scratchpad.openai_tools import format_to_openai_tool_messages
+from langchain_core.messages import AIMessage, HumanMessage
+from langchain_community.document_loaders import TextLoader
+from langchain_text_splitters import CharacterTextSplitter
+import serpapi
+import requests
 
-# Configure logging
-logging.basicConfig(level=logging.ERROR)
+# Initialize logging
+logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
 # Load environment variables
+load_dotenv()
+
+# Define and validate API keys
 openai_api_key = os.getenv("OPENAI_API_KEY")
 serper_api_key = os.getenv("SERPER_API_KEY")
+
+if not openai_api_key or not serper_api_key:
+    logger.error("API keys are not set properly.")
+    raise ValueError("API keys for OpenAI and SERPER must be set in the .env file.")
+
+# Initialize OpenAI client
 openai.api_key = openai_api_key
 
 # Load knowledge base
+@st.cache_resource
 def load_knowledge_base():
     try:
         loader = TextLoader("./data_source/time_to_rethink_trust_book.md")
@@ -35,12 +50,8 @@ def load_knowledge_base():
 knowledge_base = load_knowledge_base()
 
 # Initialize embeddings and FAISS index
-try:
-    embeddings = OpenAIEmbeddings()
-    db = FAISS.from_documents(knowledge_base, embeddings)
-except Exception as e:
-    logger.error(f"Error initializing FAISS index: {e}")
-    raise e
+embeddings = OpenAIEmbeddings()
+db = FAISS.from_documents(knowledge_base, embeddings)
 
 # Define search function for knowledge base
 def search_knowledge_base(query):
@@ -54,8 +65,11 @@ def search_knowledge_base(query):
 # SERPER API Google Search function
 def google_search(query):
     try:
-        search_client = GoogleSearch({"q": query, "api_key": serper_api_key})
-        results = search_client.get_dict()
+        search_client = serpapi.Client(api_key=serper_api_key)
+        results = search_client.search({
+            "engine": "google",
+            "q": query,
+        })
         snippets = [result["snippet"] for result in results.get("organic_results", [])]
         return snippets
     except requests.exceptions.HTTPError as http_err:
@@ -70,42 +84,104 @@ def rag_response(query):
     try:
         retrieved_docs = search_knowledge_base(query)
         context = "\n".join(doc.page_content for doc in retrieved_docs)
-        google_results = google_search(query)
-        combined_context = context + "\n" + "\n".join(google_results)
-
-        prompt = f"Context:\n{combined_context}\n\nQuestion: {query}\nAnswer:"
-        llm = ChatOpenAI(model="gpt-4", temperature=0.5, api_key=openai_api_key)
+        prompt = f"Context:\n{context}\n\nQuestion: {query}\nAnswer:"
+        llm = ChatOpenAI(model="gpt-4o", temperature=0.5, api_key=openai_api_key)
         response = llm.invoke(prompt)
         return response.content
     except Exception as e:
         logger.error(f"Error generating RAG response: {e}")
         return "Error occurred during RAG response generation"
 
+# Define tools
+@tool
+def knowledge_base_tool(query: str):
+    """Query the knowledge base and retrieve a response."""
+    return rag_response(query)
+
+@tool
+def google_search_tool(query: str):
+    """Perform a Google search using the SERPER API."""
+    return google_search(query)
+
+tools = [knowledge_base_tool, google_search_tool]
+
+# Create the prompt template
+prompt_message = """
+Act as an expert copywriter who specializes in creating compelling marketing copy using AI technologies.
+Engage in a friendly and informative conversation based on the knowledge base.
+Only proceed to create sales materials when the user explicitly requests it.
+Work together with the user to update the outcome of the sales material.
+"""
+
+prompt_template = ChatPromptTemplate.from_messages([
+    ("system", prompt_message),
+    MessagesPlaceholder(variable_name="chat_history"),
+    ("user", "{input}"),
+    MessagesPlaceholder(variable_name="agent_scratchpad"),
+])
+
+# Create Langchain Agent
+llm = ChatOpenAI(model="gpt-4o", temperature=0.5)
+llm_with_tools = llm.bind_tools(tools)
+
+# Define the agent pipeline
+agent = (
+    {
+        "input": lambda x: x["input"],
+        "agent_scratchpad": lambda x: format_to_openai_tool_messages(x["intermediate_steps"]),
+        "chat_history": lambda x: x["chat_history"],
+    }
+    | prompt_template
+    | llm_with_tools
+    | OpenAIToolsAgentOutputParser()
+)
+
+# Instantiate an AgentExecutor
+agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=True)
+
+# Streamlit app
+st.title("AI Copywriting Assistant")
+
 # Initialize chat history
 if 'chat_history' not in st.session_state:
     st.session_state.chat_history = []
 
-# Function to handle chat responses
-def chatbot_response(message):
-    try:
-        response = rag_response(message)
-        st.session_state.chat_history.append(HumanMessage(content=message))
-        st.session_state.chat_history.append(AIMessage(content=response))
-        return response
-    except Exception as e:
-        logger.error(f"Error generating chatbot response: {e}")
-        return "Error occurred during response generation"
+# Display chat history
+for message in st.session_state.chat_history:
+    with st.chat_message(message["role"]):
+        st.markdown(message["content"])
 
-# Streamlit UI setup
-st.title("Instant Insight-2-Action")
+# Chat input
+if prompt := st.chat_input("Type your message here..."):
+    # Add user message to chat history
+    st.session_state.chat_history.append({"role": "user", "content": prompt})
+    
+    # Display user message
+    with st.chat_message("user"):
+        st.markdown(prompt)
 
-prompt = st.chat_input("Type your prompt here...", key="unique_chat_input_key")
-if prompt:
-    response = chatbot_response(prompt)
+    # Generate AI response
+    with st.chat_message("assistant"):
+        message_placeholder = st.empty()
+        full_response = ""
+        
+        try:
+            # Generate response using the agent executor
+            output = agent_executor.invoke({
+                "input": prompt,
+                "chat_history": st.session_state.chat_history
+            })
+            full_response = output["output"]
+            
+            # Display the response word by word
+            for chunk in full_response.split():
+                full_response += chunk + " "
+                message_placeholder.markdown(full_response + "▌")
+            message_placeholder.markdown(full_response)
+        except Exception as e:
+            logger.error(f"Error generating response: {e}")
+            full_response = "I apologize, but an error occurred while generating the response. Please try again."
+            message_placeholder.markdown(full_response)
 
-# Display chat history
-for msg in st.session_state.chat_history:
-    if isinstance(msg, HumanMessage):
-        st.write(f"You: {msg.content}")
-    elif isinstance(msg, AIMessage):
-        st.write(f"Bot: {msg.content}")
+    # Add AI response to chat history
+    st.session_state.chat_history.append({"role": "assistant", "content": full_response})