Add application files

app.py

@@ -1279,9 +1279,9 @@ After the tool completes, provide a summary of:
 # AI CHAT - With MCP Tool Support
 # ============================================================================
 async def chat_with_ai_async(message: str, history: list, hf_token: str):
-    """AI Chat with full MCP tool support - runs as async generator"""
+    """AI Chat powered by LLM with full MCP tool support"""
     if not knowledge_base["client"]["name"]:
-        yield history + [[message, "⚠️ Please complete Setup first."]], ""
+        yield history + [[message, "⚠️ Please complete Setup first. Enter your company name in the Setup tab."]], ""
         return
 
     if not message.strip():
@@ -1294,126 +1294,172 @@ async def chat_with_ai_async(message: str, history: list, hf_token: str):
         return
 
     client_name = knowledge_base["client"]["name"]
-    msg_lower = message.lower()
+    client_info = knowledge_base["client"].get("raw_research", "")
 
-    # Check if this needs AI agent (complex task) or can be answered locally
-    needs_agent = any(kw in msg_lower for kw in [
-        "search", "find", "research", "look up", "discover",
-        "draft", "write", "compose", "create email", "send",
-        "save", "add", "store",
-        "analyze", "suggest", "recommend", "strategy",
-        "new prospect", "find companies", "find contacts"
-    ])
-
-    if needs_agent:
-        # Use AI Agent with MCP tools
-        try:
-            agent = AutonomousMCPAgentHF(
-                mcp_registry=mcp_registry,
-                hf_token=token,
-                provider=HF_PROVIDER,
-                model=HF_MODEL
-            )
-
-            # Build context-aware prompt
-            prospects_info = "\n".join([f"- {p['name']} ({p.get('industry', 'Unknown')})" for p in knowledge_base["prospects"][:5]]) if knowledge_base["prospects"] else "None yet"
-            contacts_info = "\n".join([f"- {c['name']} at {c.get('company', 'Unknown')}" for c in knowledge_base["contacts"][:5]]) if knowledge_base["contacts"] else "None yet"
-
-            task = f"""You are a sales assistant for {client_name}.
-
-Current Knowledge Base:
-- Prospects: {len(knowledge_base['prospects'])}
-{prospects_info}
-- Contacts: {len(knowledge_base['contacts'])}
-{contacts_info}
-- Emails drafted: {len(knowledge_base['emails'])}
-
-User request: {message}
-
-Use the available tools to help. Available actions:
-- search_web: Search for company information
-- save_prospect: Save a new prospect
-- save_contact: Save a contact/decision maker
-- send_email: Draft an outreach email
-- save_fact: Save important facts about a company
-
-Complete the user's request and provide a helpful response."""
-
-            response_text = ""
-            current_history = history + [[message, "🤖 Working on it..."]]
-            yield current_history, ""
-
-            async for event in agent.run(task, max_iterations=10):
-                event_type = event.get("type")
-
-                if event_type == "tool_call":
-                    tool = event.get("tool", "")
-                    response_text += f"🔧 Using **{tool}**...\n"
-                    current_history = history + [[message, response_text]]
-                    yield current_history, ""
-
-                elif event_type == "tool_result":
-                    tool = event.get("tool", "")
-                    result = event.get("result", {})
-
-                    # Capture data from tool results (with deduplication)
-                    if tool == "save_prospect" and isinstance(result, dict):
-                        prospect_data = {
-                            "name": result.get("prospect_id", "Unknown"),
-                            "domain": result.get("domain", ""),
-                            "fit_score": 0,
-                            "research_complete": True,
-                            "discovered_at": datetime.now().strftime("%Y-%m-%d %H:%M")
-                        }
-                        # Deduplicate before adding
-                        merge_to_knowledge_base([prospect_data], [], [])
-
-                    elif tool == "save_contact" and isinstance(result, dict):
-                        contact_data = result
-                        merge_to_knowledge_base([], [contact_data], [])
-
-                    elif tool == "send_email" and isinstance(result, dict):
-                        merge_to_knowledge_base([], [], [result])
-
-                    response_text += f"✅ {tool} completed\n"
-                    current_history = history + [[message, response_text]]
-                    yield current_history, ""
-
-                elif event_type == "thought":
-                    thought = event.get("thought", "")
-                    if thought and not thought.startswith("[Processing"):
-                        response_text += f"\n{thought}\n"
-                        current_history = history + [[message, response_text]]
-                        yield current_history, ""
-
-                elif event_type == "agent_complete":
-                    final = event.get("final_answer", "")
-                    if final:
-                        response_text += f"\n---\n{final}"
-                    current_history = history + [[message, response_text]]
-                    yield current_history, ""
-                    return
-
-                elif event_type == "agent_error":
-                    response_text += f"\n⚠️ Error: {event.get('error', 'Unknown error')}"
-                    current_history = history + [[message, response_text]]
-                    yield current_history, ""
-                    return
-
-            # If we get here without returning
-            if not response_text:
-                response_text = "I processed your request but couldn't generate a response."
-            yield history + [[message, response_text]], ""
-            return
-
-        except Exception as e:
-            logger.error(f"Chat agent error: {e}")
-            yield history + [[message, f"⚠️ Error: {str(e)}"]], ""
-            return
+    # Always use LLM for all queries - this is a full AI assistant
+    try:
+        agent = AutonomousMCPAgentHF(
+            mcp_registry=mcp_registry,
+            hf_token=token,
+            provider=HF_PROVIDER,
+            model=HF_MODEL
+        )
 
-    # For simple queries, use local knowledge base lookup
-    response = get_local_response(message, client_name)
-    yield history + [[message, response]], ""
+        # Build comprehensive context with all knowledge base data
+        prospects_detail = ""
+        if knowledge_base["prospects"]:
+            for i, p in enumerate(knowledge_base["prospects"][:10], 1):
+                p_name = p.get('name', 'Unknown')
+                p_name_lower = p_name.lower()
+                # Get contacts for this prospect
+                p_contacts = [c for c in knowledge_base["contacts"]
+                             if p_name_lower in c.get("company", "").lower()
+                             or c.get("company", "").lower() in p_name_lower]
+                contacts_str = ", ".join([f"{c.get('name')} ({c.get('email')})" for c in p_contacts]) if p_contacts else "No contacts"
+                prospects_detail += f"{i}. {p_name} - {p.get('industry', 'Unknown industry')}, Fit: {p.get('fit_score', 'N/A')}\n"
+                prospects_detail += f"   Summary: {p.get('summary', 'No summary')[:100]}\n"
+                prospects_detail += f"   Contacts: {contacts_str}\n"
+        else:
+            prospects_detail = "No prospects discovered yet."
+
+        emails_detail = ""
+        if knowledge_base["emails"]:
+            for e in knowledge_base["emails"][:5]:
+                emails_detail += f"- To: {e.get('to')} | Subject: {e.get('subject', 'No subject')[:50]}\n"
+        else:
+            emails_detail = "No emails drafted yet."
+
+        task = f"""You are an AI sales assistant for {client_name}. You are a helpful, knowledgeable assistant that can answer any question about the sales pipeline, prospects, contacts, and help with various sales tasks.
+
+ABOUT {client_name}:
+{client_info[:500] if client_info else "No company research available yet."}
+
+CURRENT SALES PIPELINE:
+======================
+PROSPECTS ({len(knowledge_base['prospects'])}):
+{prospects_detail}
+
+CONTACTS ({len(knowledge_base['contacts'])}):
+{len(knowledge_base['contacts'])} decision makers found across prospects.
+
+DRAFTED EMAILS ({len(knowledge_base['emails'])}):
+{emails_detail}
+
+USER MESSAGE: {message}
+
+INSTRUCTIONS:
+- Answer the user's question helpfully and completely
+- If they ask about prospects, contacts, or emails, use the data above
+- If they ask you to search for something, use search_web tool
+- If they ask you to draft an email, create a professional, personalized email
+- If they ask for talking points, strategies, or recommendations, provide thoughtful, specific advice
+- If they ask to find similar companies or new prospects, use search_web to research
+- Be conversational and helpful - you're a knowledgeable sales assistant
+- Don't say "I don't have that capability" - try to help with whatever they ask
+- For follow-up questions, use context from the conversation
+
+Respond naturally and helpfully to the user's message."""
+
+        response_text = ""
+        current_history = history + [[message, "🤖 Thinking..."]]
+        yield current_history, ""
+
+        async for event in agent.run(task, max_iterations=12):
+            event_type = event.get("type")
+
+            if event_type == "tool_call":
+                tool = event.get("tool", "")
+                tool_input = event.get("input", {})
+                if tool == "search_web":
+                    query = tool_input.get("query", "") if isinstance(tool_input, dict) else ""
+                    response_text += f"🔍 Searching: {query[:50]}...\n"
+                elif tool == "send_email":
+                    response_text += f"✉️ Drafting email...\n"
+                else:
+                    response_text += f"🔧 Using {tool}...\n"
+                current_history = history + [[message, response_text]]
+                yield current_history, ""
+
+            elif event_type == "tool_result":
+                tool = event.get("tool", "")
+                result = event.get("result", {})
+
+                # Capture data from tool results (with deduplication)
+                if tool == "save_prospect" and isinstance(result, dict):
+                    prospect_data = {
+                        "name": result.get("company_name", result.get("prospect_id", "Unknown")),
+                        "domain": result.get("company_domain", result.get("domain", "")),
+                        "fit_score": result.get("fit_score", 75),
+                        "research_complete": True,
+                        "discovered_at": datetime.now().strftime("%Y-%m-%d %H:%M")
+                    }
+                    merge_to_knowledge_base([prospect_data], [], [])
+                    response_text += f"✅ Saved prospect: {prospect_data['name']}\n"
+
+                elif tool == "save_contact" and isinstance(result, dict):
+                    merge_to_knowledge_base([], [result], [])
+                    response_text += f"✅ Saved contact\n"
+
+                elif tool == "send_email" and isinstance(result, dict):
+                    merge_to_knowledge_base([], [], [result])
+                    response_text += f"✅ Email drafted\n"
+
+                elif tool == "search_web":
+                    count = result.get("count", 0) if isinstance(result, dict) else 0
+                    response_text += f"✅ Found {count} results\n"
+
+                current_history = history + [[message, response_text]]
+                yield current_history, ""
+
+            elif event_type == "thought":
+                thought = event.get("thought", "")
+                # Only show substantive thoughts, not processing messages
+                if thought and len(thought) > 50 and not thought.startswith("[Processing"):
+                    # This is likely the AI's actual response
+                    pass  # We'll get this in agent_complete
+
+            elif event_type == "agent_complete":
+                final = event.get("final_answer", "")
+                if final and "CX AI Agent" not in final and "Powered by AI" not in final:
+                    # Clean response - show just the final answer
+                    if response_text:
+                        response_text += "\n---\n\n"
+                    response_text += final
+                elif not response_text:
+                    response_text = "I've processed your request. Is there anything else you'd like to know?"
+                current_history = history + [[message, response_text]]
+                yield current_history, ""
+                return
+
+            elif event_type == "agent_error":
+                error = event.get("error", "Unknown error")
+                if "rate limit" in str(error).lower():
+                    response_text += "\n⚠️ Rate limit reached. Please wait a moment and try again."
+                else:
+                    response_text += f"\n⚠️ Error: {error}"
+                current_history = history + [[message, response_text]]
+                yield current_history, ""
+                return
+
+            elif event_type == "agent_max_iterations":
+                if not response_text:
+                    response_text = "I'm still processing your request. The task may be complex - please try a simpler question or try again."
+                current_history = history + [[message, response_text]]
+                yield current_history, ""
+                return
+
+        # If we get here without returning
+        if not response_text:
+            response_text = "I processed your request. Let me know if you need anything else!"
+        yield history + [[message, response_text]], ""
+
+    except Exception as e:
+        logger.error(f"Chat agent error: {e}")
+        error_msg = str(e)
+        if "rate limit" in error_msg.lower() or "429" in error_msg:
+            yield history + [[message, "⚠️ Rate limit reached. Please wait a moment and try again."]], ""
+        else:
+            yield history + [[message, f"⚠️ Error: {error_msg}"]], ""
 
 
 def chat_with_ai(message: str, history: list) -> tuple:
@@ -2192,54 +2238,67 @@ def create_app():
                     with gr.Column(scale=1):
                         gr.HTML("""
                         <div class="action-card">
-                            <h3>💡 Example Prompts</h3>
-                            <p style="font-size: 13px; color: var(--text-secondary);">Try these commands:</p>
+                            <h3>💡 Try These Prompts</h3>
+                            <p style="font-size: 12px; color: var(--text-secondary);">Your AI assistant can help with anything sales-related</p>
+                        </div>
+                        """)
+
+                        gr.HTML("""
+                        <div style="background: var(--bg-tertiary); padding: 12px; border-radius: 8px; margin-bottom: 8px;">
+                            <strong>🔍 Research & Discovery</strong>
+                            <ul style="margin: 8px 0 0 0; padding-left: 16px; font-size: 12px; line-height: 1.8;">
+                                <li>"Search for DTC fashion brands that recently raised Series A funding"</li>
+                                <li>"Find 5 subscription box companies in the wellness space"</li>
+                                <li>"Research what Allbirds is doing in sustainability this year"</li>
+                                <li>"What are the latest trends in e-commerce personalization?"</li>
+                            </ul>
                         </div>
                         """)
 
                         gr.HTML("""
                         <div style="background: var(--bg-tertiary); padding: 12px; border-radius: 8px; margin-bottom: 8px;">
-                            <strong>📋 Pipeline Overview</strong>
-                            <ul style="margin: 8px 0 0 0; padding-left: 16px; font-size: 12px; line-height: 1.6;">
-                                <li>"List all prospects"</li>
-                                <li>"Show me the contacts"</li>
-                                <li>"How many emails have been drafted?"</li>
-                                <li>"Give me a pipeline summary"</li>
+                            <strong>✉️ Email Drafting</strong>
+                            <ul style="margin: 8px 0 0 0; padding-left: 16px; font-size: 12px; line-height: 1.8;">
+                                <li>"Draft a cold outreach email to the CEO of Warby Parker about our analytics platform"</li>
+                                <li>"Write a follow-up email for my meeting request - it's been 5 days with no response"</li>
+                                <li>"Create a meeting invite email for Thursday 3pm, 30 mins, to discuss Q1 partnership"</li>
+                                <li>"Compose a breakup email for prospects who haven't responded after 3 touches"</li>
                             </ul>
                         </div>
                         """)
 
                         gr.HTML("""
                         <div style="background: var(--bg-tertiary); padding: 12px; border-radius: 8px; margin-bottom: 8px;">
-                            <strong>🔍 Research & Analysis</strong>
-                            <ul style="margin: 8px 0 0 0; padding-left: 16px; font-size: 12px; line-height: 1.6;">
-                                <li>"Research Acme Corp"</li>
-                                <li>"What do we know about Nike?"</li>
-                                <li>"Analyze Buffalo Jackson's fit score"</li>
-                                <li>"Find similar companies to Warby Parker"</li>
+                            <strong>📞 Call Preparation</strong>
+                            <ul style="margin: 8px 0 0 0; padding-left: 16px; font-size: 12px; line-height: 1.8;">
+                                <li>"I have a call with Glossier tomorrow - give me 10 talking points"</li>
+                                <li>"What objections might Patagonia raise and how should I handle them?"</li>
+                                <li>"Prepare a 5-minute pitch deck outline for my Everlane demo"</li>
+                                <li>"What questions should I ask to qualify Bombas as a prospect?"</li>
                             </ul>
                         </div>
                         """)
 
                         gr.HTML("""
                         <div style="background: var(--bg-tertiary); padding: 12px; border-radius: 8px; margin-bottom: 8px;">
-                            <strong>✉️ Custom Email Drafts</strong>
-                            <ul style="margin: 8px 0 0 0; padding-left: 16px; font-size: 12px; line-height: 1.6;">
-                                <li>"Draft a meeting request to Allbirds for Tuesday 2pm, 30 mins"</li>
-                                <li>"Write a follow-up email to Glossier about our Q4 partnership proposal"</li>
-                                <li>"Compose an intro email to Patagonia mentioning their sustainability initiatives"</li>
+                            <strong>📊 Pipeline & Analysis</strong>
+                            <ul style="margin: 8px 0 0 0; padding-left: 16px; font-size: 12px; line-height: 1.8;">
+                                <li>"Give me a summary of all prospects and their status"</li>
+                                <li>"Which prospect has the highest fit score and why?"</li>
+                                <li>"Compare the two prospects and tell me which to prioritize"</li>
+                                <li>"What's the total addressable market for our solution?"</li>
                             </ul>
                         </div>
                         """)
 
                         gr.HTML("""
                         <div style="background: var(--bg-tertiary); padding: 12px; border-radius: 8px;">
-                            <strong>💡 Sales Preparation</strong>
-                            <ul style="margin: 8px 0 0 0; padding-left: 16px; font-size: 12px; line-height: 1.6;">
-                                <li>"Give me talking points for the Everlane call"</li>
-                                <li>"Who is the decision maker at Bombas?"</li>
-                                <li>"What's Casper's industry and fit reason?"</li>
-                                <li>"Show the existing email draft for Away"</li>
+                            <strong>💡 Strategy & Insights</strong>
+                            <ul style="margin: 8px 0 0 0; padding-left: 16px; font-size: 12px; line-height: 1.8;">
+                                <li>"What's the best approach to sell to a company that just raised funding?"</li>
+                                <li>"How should I position our product against competitor X?"</li>
+                                <li>"Suggest a multi-touch outreach sequence for enterprise prospects"</li>
+                                <li>"What case studies should I reference when talking to retail brands?"</li>
                             </ul>
                         </div>
                         """)
@@ -2296,8 +2355,16 @@ def create_app():
         generate_packet_btn.click(fn=generate_handoff_packet, inputs=[prospect_dropdown], outputs=[handoff_output])
         refresh_dropdown_btn.click(fn=lambda: gr.Dropdown(choices=get_prospect_choices()), outputs=[prospect_dropdown])
 
-        send_btn.click(fn=chat_with_ai, inputs=[chat_input, chatbot], outputs=[chatbot, chat_input])
-        chat_input.submit(fn=chat_with_ai, inputs=[chat_input, chatbot], outputs=[chatbot, chat_input])
+        # Async chat wrapper that uses session token
+        async def chat_async_wrapper(message, history):
+            token = session_hf_token.get("token", "")
+            final_result = (history, "")
+            async for result in chat_with_ai_async(message, history, token):
+                final_result = result
+            return final_result
+
+        send_btn.click(fn=chat_async_wrapper, inputs=[chat_input, chatbot], outputs=[chatbot, chat_input])
+        chat_input.submit(fn=chat_async_wrapper, inputs=[chat_input, chatbot], outputs=[chatbot, chat_input])
 
     return demo