Update app.py

app.py

@@ -1,17 +1,17 @@
-import os
-import torch
 import streamlit as st
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
 from langchain_core.prompts import PromptTemplate
-from langchain.chains import LLMChain
-from langchain_community.llms import HuggingFacePipeline
-from langchain.agents import AgentExecutor, Tool, ZeroShotAgent
+from langchain_core.runnables import RunnableSequence
+from langchain_huggingface import HuggingFacePipeline
+from langchain.agents import create_react_agent, AgentExecutor, Tool
 from langchain.memory import ConversationBufferMemory
-import pandas as pd
-from sqlalchemy import create_engine
-import smtplib
-from email.mime.text import MIMEText
-from email.mime.multipart import MIMEMultipart
+
+# Mock lead data
+LEADS = [
+    {"name": "John Doe", "email": "john@example.com", "company": "TechCorp"},
+    {"name": "Jane Smith", "email": "jane@example.com", "company": "InnoSoft"},
+    {"name": "Bob Johnson", "email": "[email protected]", "company": "DataTech"},
+]
 
 # Set up the open-source LLM
 @st.cache_resource
@@ -29,14 +29,9 @@ def load_model():
 
 local_llm = load_model()
 
-# Set up the database connection
-db_connection_string = "sqlite:///leads.db"  # Replace with your actual database connection string
-engine = create_engine(db_connection_string)
-
 # Define the tools for the agent
 def search_leads(query):
-    df = pd.read_sql(f"SELECT * FROM leads WHERE name LIKE '%{query}%'", engine)
-    return df.to_dict(orient='records')
+    return [lead for lead in LEADS if query.lower() in lead['name'].lower()]
 
 def send_email(to_email, subject, body):
     # For demo purposes, we'll just print the email details
@@ -49,7 +44,7 @@ tools = [
     Tool(
         name="Search Leads",
         func=search_leads,
-        description="Useful for searching leads in the database"
+        description="Useful for searching leads by name"
     ),
     Tool(
         name="Send Email",
@@ -59,7 +54,7 @@ tools = [
 ]
 
 # Set up the agent
-prefix = """You are an AI CyberSecurity Program Advisor. Your goal is to engage with leads and get them to book a video call for an in-person sales meeting. You have access to a database of leads and can send emails.
+prefix = """You are an AI CyberSecurity Program Advisor. Your goal is to engage with leads and get them to book a video call for an in-person sales meeting. You have access to a list of leads and can send emails.
 
 You have access to the following tools:"""
 
@@ -69,17 +64,15 @@ suffix = """Begin!
 Human: {human_input}
 AI: Let's approach this step-by-step:"""
 
-prompt = ZeroShotAgent.create_prompt(
-    tools,
-    prefix=prefix,
-    suffix=suffix,
-    input_variables=["human_input", "chat_history"]
+prompt = PromptTemplate(
+    template=prefix + "{agent_scratchpad}" + suffix,
+    input_variables=["human_input", "chat_history", "agent_scratchpad"]
 )
 
-llm_chain = LLMChain(llm=local_llm, prompt=prompt)
+llm_chain = RunnableSequence(prompt, local_llm)
 memory = ConversationBufferMemory(memory_key="chat_history")
 
-agent = ZeroShotAgent(llm_chain=llm_chain, tools=tools, verbose=True)
+agent = create_react_agent(local_llm, tools, prompt)
 agent_executor = AgentExecutor.from_agent_and_tools(
     agent=agent, tools=tools, verbose=True, memory=memory
 )
@@ -97,9 +90,9 @@ if lead_name:
         st.write(f"No lead found with the name {lead_name}")
     else:
         lead = lead_info[0]
-        st.write(f"Lead found: {lead['name']} (Email: {lead['email']})")
+        st.write(f"Lead found: {lead['name']} (Email: {lead['email']}, Company: {lead['company']})")
         
-        initial_message = f"Hello {lead['name']}, I'd like to discuss our cybersecurity program with you. Are you available for a quick video call?"
+        initial_message = f"Hello {lead['name']}, I'd like to discuss our cybersecurity program with {lead['company']}. Are you available for a quick video call?"
         
         if st.button("Engage with Lead"):
             with st.spinner("AI is generating a response..."):