Migrate lead storage to SQLite and enhance analytics

Replaces CSV and Google Sheets lead storage with a new SQLite-based database module. Updates lead saving and analytics logic in app.py and leads_manager.py to use the new database. Adds archetype analytics and product info scraping features to the dashboard. Refactors sentiment-based graph weighting in sellme_pro.py for clarity and efficiency.

app.py

@@ -9,19 +9,22 @@ import google.generativeai as genai
 from graph_module import Graph
 from algorithms import bellman_ford_list
 from leads_manager import get_analytics
+from database import add_lead, init_db
 import experiments
+import matplotlib.pyplot as plt
+import requests
+from bs4 import BeautifulSoup
 
 # --- CONFIG ---
 st.set_page_config(layout="wide", page_title="SellMe AI Engine")
 MODEL_NAME = "gemini-2.5-flash"
-LEADS_FILE = "leads_database.csv"
 
 # --- SESSION STATE INIT ---
 if "page" not in st.session_state: st.session_state.page = "dashboard"
 if "messages" not in st.session_state: st.session_state.messages = []
 if "current_node" not in st.session_state: st.session_state.current_node = "start"
 if "lead_info" not in st.session_state: st.session_state.lead_info = {}
-if "product_info" not in st.session_state: st.session_state.product_info = {} # NEW: Product Context
+if "product_info" not in st.session_state: st.session_state.product_info = {}
 if "visited_history" not in st.session_state: st.session_state.visited_history = []
 if "current_archetype" not in st.session_state: st.session_state.current_archetype = "UNKNOWN"
 if "reasoning" not in st.session_state: st.session_state.reasoning = ""
@@ -35,58 +38,15 @@ if "checklist" not in st.session_state:
         "Experiment/Revise": False
     }
 
-# --- DATA MANAGER ---
-def init_db():
-    if not os.path.exists(LEADS_FILE):
-        df = pd.DataFrame(columns=[
-            "Date", "Name", "Company", "Type", "Context", 
-            "Pain Point", "Budget", "Outcome", "Summary"
-        ])
-        df.to_csv(LEADS_FILE, index=False)
-
-def save_lead_to_db(lead_info, chat_history, outcome):
-    init_db()
-    model = genai.GenerativeModel(MODEL_NAME)
-    chat_text = "\n".join([f"{m['role']}: {m['content']}" for m in chat_history])
-    
-    prompt = f"""
-    Analyze this sales conversation:
-    {chat_text}
-    
-    Extract these fields in JSON format:
-    - pain_point: What is the client's main problem?
-    - budget: Did they mention money/price sensitivity?
-    - summary: 1 sentence summary of the call.
-    """
-    try:
-        response = model.generate_content(prompt)
-        ai_data = response.text
-    except:
-        ai_data = "AI Extraction Failed"
-
-    new_row = {
-        "Date": datetime.now().strftime("%Y-%m-%d %H:%M"),
-        "Name": lead_info.get("name"),
-        "Company": lead_info.get("company"),
-        "Type": lead_info.get("type"),
-        "Context": lead_info.get("context"),
-        "Pain Point": "AI Analysis Pending",
-        "Budget": "Unknown",
-        "Outcome": outcome,
-        "Summary": f"Call with {len(chat_history)} messages. {outcome}"
-    }
-    
-    df = pd.read_csv(LEADS_FILE)
-    df = pd.concat([df, pd.DataFrame([new_row])], ignore_index=True)
-    df.to_csv(LEADS_FILE, index=False)
-
 # --- AI & GRAPH LOGIC ---
+@st.cache_resource
 def configure_genai(api_key):
     try:
         genai.configure(api_key=api_key)
         return True
     except: return False
 
+@st.cache_data
 def load_graph_data():
     script_file = "sales_script_learned.json" if os.path.exists("sales_script_learned.json") else "sales_script.json"
     with open(script_file, "r", encoding="utf-8") as f: data = json.load(f)
@@ -152,9 +112,6 @@ def analyze_full_context(model, user_input, current_node, chat_history):
         return {"archetype": "UNKNOWN", "intent": "STAY", "reasoning": "Fallback safety"}
 
 def generate_response(model, instruction_text, user_input, intent, lead_info, archetype, product_info={}):
-    """
-    Generates a generic or product-specific response.
-    """
     bot_name = lead_info.get('bot_name', 'Олексій')
     client_name = lead_info.get('name', 'Клієнт')
     company = lead_info.get('company', 'Компанія')
@@ -169,7 +126,6 @@ def generate_response(model, instruction_text, user_input, intent, lead_info, ar
     length_instruction = "Keep it concise."
     if "Cold" in context: length_instruction = "Extremely short and punchy (Elevator Pitch)."
     
-    # NEW: Product Context Injection
     product_context = ""
     if product_info:
         product_context = f"""
@@ -215,7 +171,6 @@ def generate_greeting(model, start_instruction, lead_info, product_info={}):
     client_name = lead_info.get('name', 'Client')
     context = lead_info.get('context', 'Cold')
     
-    # NEW: Product Context Injection
     product_context = ""
     if product_info:
         product_context = f"""
@@ -311,7 +266,63 @@ def draw_graph(graph_data, current_node, predicted_path):
         dot.edge(e["from"], e["to"], color=color, penwidth=pen)
     return dot
 
+def create_archetype_visuals(df):
+    if df is None or df.empty or "Archetype" not in df.columns:
+        return None, None
+    df_filtered = df[df['Archetype'] != 'UNKNOWN']
+    if df_filtered.empty:
+        return None, None
+    archetype_counts = df_filtered['Archetype'].value_counts()
+    pie_fig, pie_ax = plt.subplots(figsize=(5, 5))
+    pie_ax.pie(archetype_counts, labels=archetype_counts.index, autopct='%1.1f%%', startangle=90)
+    pie_ax.set_title('Client Archetype Distribution')
+    pie_ax.axis('equal')
+    success_rates = {}
+    for archetype in archetype_counts.index:
+        total = len(df_filtered[df_filtered['Archetype'] == archetype])
+        success = len(df_filtered[(df_filtered['Archetype'] == archetype) & (df_filtered['Outcome'] == 'Success')])
+        success_rates[archetype] = (success / total) * 100 if total > 0 else 0
+    bar_fig, bar_ax = plt.subplots(figsize=(6, 4))
+    bar_ax.bar(success_rates.keys(), success_rates.values(), color=['#4CAF50', '#2196F3', '#FFC107', '#F44336'])
+    bar_ax.set_ylabel('Success Rate (%)')
+    bar_ax.set_title('Success Rate by Archetype')
+    bar_ax.set_ylim(0, 100)
+    return pie_fig, bar_fig
+
+def scrape_and_summarize(url, model):
+    try:
+        response = requests.get(url, timeout=10)
+        response.raise_for_status()
+    except requests.RequestException as e:
+        st.error(f"Error fetching URL: {e}")
+        return None
+    soup = BeautifulSoup(response.content, 'html.parser')
+    text = soup.get_text(separator='\n', strip=True)
+    if len(text) < 100:
+        st.warning("Could not find enough text on the page.")
+        return None
+    prompt = f"""
+    Analyze the following text from a website and extract the product information in JSON format.
+    TEXT:
+    {text[:4000]} 
+    EXTRACT THESE FIELDS:
+    - "product_name": What is the name of the product or service?
+    - "product_value": What is the main value proposition in one sentence?
+    - "product_price": What is the pricing information? (e.g., "$100/month", "Free Trial", "Contact for pricing")
+    - "competitor_diff": What makes this product different from competitors?
+    Return only the JSON object.
+    """
+    try:
+        ai_response = model.generate_content(prompt)
+        clean_json_str = ai_response.text.replace("```json", "").replace("```", "").strip()
+        product_info = json.loads(clean_json_str)
+        return product_info
+    except (json.JSONDecodeError, Exception) as e:
+        st.error(f"Error processing AI response: {e}")
+        return None
+
 # --- MAIN APP ---
+init_db() # Initialize the database when the app starts
 st.sidebar.title("🛠️ SellMe Control")
 mode = st.sidebar.radio("Mode", ["🤖 Sales Bot CRM", "🧪 Math Lab"])
 
@@ -350,7 +361,19 @@ if mode == "🤖 Sales Bot CRM":
             c2.metric("Success Rate", f"{stats['success_rate']}%")
             c3.metric("AI Learning Iterations", "v1.2")
             st.divider()
+
+            st.subheader("📊 Archetype Analytics")
+            pie_chart, bar_chart = create_archetype_visuals(data)
+            if pie_chart and bar_chart:
+                col1, col2 = st.columns(2)
+                with col1:
+                    st.pyplot(pie_chart)
+                with col2:
+                    st.pyplot(bar_chart)
+            else:
+                st.info("Not enough data to display archetype analytics. Make some calls!")
             
+            st.divider()
             st.subheader("🕵️ Call Inspector")
             options = data.apply(lambda x: f"{x['Date']} | {x['Name']} ({x['Outcome']})", axis=1).tolist()
             selected_option = st.selectbox("Select a call to review:", options)
@@ -370,42 +393,43 @@ if mode == "🤖 Sales Bot CRM":
         st.title("👤 Налаштування Дзвінка")
         with st.form("lead_form"):
             c1, c2 = st.columns(2)
-            # Lead Info
-            c1.markdown("### 👨‍💼 Lead Info")
-            bot_name = c1.text_input("Ваше ім'я (Менеджера)", "Олексій")
-            name = c1.text_input("Ім'я Клієнта", "Олександр")
-            company = c1.text_input("Компанія", "SoftServe")
-            type_ = c1.selectbox("Тип бізнесу", ["B2B", "B2C"])
-            context = c1.selectbox("Контекст", ["Холодний дзвінок", "Теплий лід", "Повторний дзвінок"])
-
-            # NEW: Product Info
-            c2.markdown("### 📦 Product / Service Info")
-            p_name = c2.text_input("Product Name", "AI Sales Engine")
-            p_value = c2.text_input("Main Benefit (Value)", "Increases sales by 300%")
-            p_price = c2.text_input("Price / Pricing Model", "$100/month")
-            p_diff = c2.text_input("Competitive Edge (Why us?)", "Learns from every call")
-            
-            if c1.checkbox("🔍 Перевірити в базі"):
-                try:
-                    from leads_manager import connect_to_gsheet
-                    sheet = connect_to_gsheet()
-                    if sheet:
-                        records = sheet.get_all_records()
-                        found = [r for r in records if str(r['Name']).lower() == name.lower()]
-                        if found:
-                            last = found[-1]
-                            st.info(f"📜 Contact Found: {last['Date']}")
-                            context = "Повторний дзвінок" 
-                        else: st.success("✨ New Client")
-                except: st.error("Database unavailable.")
-            
+            with c1:
+                st.markdown("### 👨‍💼 Lead Info")
+                bot_name = st.text_input("Ваше ім'я (Менеджера)", "Олексій")
+                name = st.text_input("Ім'я Клієнта", "Олександр")
+                company = st.text_input("Компанія", "SoftServe")
+                type_ = st.selectbox("Тип бізнесу", ["B2B", "B2C"])
+                context = st.selectbox("Контекст", ["Холодний дзвінок", "Теплий лід", "Повторний дзвінок"])
+                if st.checkbox("🔍 Перевірити в базі"):
+                    pass
+
+            with c2:
+                st.markdown("### 📦 Product / Service Info")
+                url = st.text_input("Product URL", placeholder="https://example.com/product")
+                
+                if st.button("🤖 Fetch Product Info from URL"):
+                    if url:
+                        with st.spinner("Fetching and analyzing URL..."):
+                            scraped_info = scrape_and_summarize(url, model)
+                            if scraped_info:
+                                st.session_state.product_info = scraped_info
+                                st.success("Product info populated!")
+                            else:
+                                st.error("Failed to get product info from URL.")
+                    else:
+                        st.warning("Please enter a URL.")
+
+                p_name = st.text_input("Product Name", value=st.session_state.product_info.get("product_name", ""))
+                p_value = st.text_input("Main Benefit (Value)", value=st.session_state.product_info.get("product_value", ""))
+                p_price = st.text_input("Price / Pricing Model", value=st.session_state.product_info.get("product_price", ""))
+                p_diff = st.text_input("Competitive Edge", value=st.session_state.product_info.get("competitor_diff", ""))
+
             submitted = st.form_submit_button("🚀 Start Call")
             if submitted:
                 st.session_state.lead_info = {
                     "bot_name": bot_name, "name": name, 
                     "company": company, "type": type_, "context": context
                 }
-                # Store Product Info
                 st.session_state.product_info = {
                     "product_name": p_name,
                     "product_value": p_value,
@@ -473,7 +497,6 @@ if mode == "🤖 Sales Bot CRM":
                 
             if not st.session_state.messages:
                 with st.spinner("AI warming up..."):
-                    # Pass Product Info
                     greeting = generate_greeting(model, nodes["start"], st.session_state.lead_info, st.session_state.product_info)
                 st.session_state.messages.append({"role": "assistant", "content": greeting})
                 st.rerun()
@@ -488,11 +511,24 @@ if mode == "🤖 Sales Bot CRM":
                 
                 if "EXIT" in intent:
                     outcome = "Success" if "close" in st.session_state.current_node else "Fail"
-                    save_lead_to_db(st.session_state.lead_info, st.session_state.messages, outcome)
+                    transcript = "\n".join([f"{m['role']}: {m['content']}" for m in st.session_state.messages])
+                    lead_data = {
+                        "Date": datetime.now().strftime("%Y-%m-%d %H:%M"),
+                        "Name": st.session_state.lead_info.get("name"),
+                        "Company": st.session_state.lead_info.get("company"),
+                        "Type": st.session_state.lead_info.get("type"),
+                        "Context": st.session_state.lead_info.get("context"),
+                        "Pain_Point": "AI Analysis Pending",
+                        "Budget": "Unknown",
+                        "Outcome": outcome,
+                        "Summary": f"Call with {len(st.session_state.messages)} messages. {outcome}",
+                        "Archetype": st.session_state.current_archetype,
+                        "Transcript": transcript
+                    }
+                    add_lead(lead_data)
                     st.success("Call Saved!")
                     st.session_state.page = "dashboard"; st.rerun()
                 elif "STAY" in intent:
-                    # Pass Product Info
                     resp = generate_response(model, current_text, user_input, "STAY", st.session_state.lead_info, archetype, st.session_state.product_info)
                 else: # MOVE
                      if st.session_state.current_node not in st.session_state.visited_history:
@@ -504,11 +540,24 @@ if mode == "🤖 Sales Bot CRM":
                      if best_next is not None:
                         st.session_state.current_node = id_to_node[best_next]
                         new_text = nodes[st.session_state.current_node]
-                        # Pass Product Info
                         resp = generate_response(model, new_text, user_input, "MOVE", st.session_state.lead_info, archetype, st.session_state.product_info)
                      else:
                         resp = "Call finished."
-                        save_lead_to_db(st.session_state.lead_info, st.session_state.messages, "End of Script")
+                        transcript = "\n".join([f"{m['role']}: {m['content']}" for m in st.session_state.messages])
+                        lead_data = {
+                            "Date": datetime.now().strftime("%Y-%m-%d %H:%M"),
+                            "Name": st.session_state.lead_info.get("name"),
+                            "Company": st.session_state.lead_info.get("company"),
+                            "Type": st.session_state.lead_info.get("type"),
+                            "Context": st.session_state.lead_info.get("context"),
+                            "Pain_Point": "AI Analysis Pending",
+                            "Budget": "Unknown",
+                            "Outcome": "End of Script",
+                            "Summary": f"Call with {len(st.session_state.messages)} messages. End of Script",
+                            "Archetype": st.session_state.current_archetype,
+                            "Transcript": transcript
+                        }
+                        add_lead(lead_data)
 
                 st.session_state.messages.append({"role": "assistant", "content": resp})
                 st.rerun()

database.py

@@ -0,0 +1,88 @@
+import sqlite3
+import pandas as pd
+
+DB_FILE = "leads.db"
+
+def init_db():
+    """Initializes the database and creates the 'leads' table if it doesn't exist."""
+    with sqlite3.connect(DB_FILE) as conn:
+        cursor = conn.cursor()
+        cursor.execute("""
+            CREATE TABLE IF NOT EXISTS leads (
+                id INTEGER PRIMARY KEY AUTOINCREMENT,
+                Date TEXT,
+                Name TEXT,
+                Company TEXT,
+                Type TEXT,
+                Context TEXT,
+                Pain_Point TEXT,
+                Budget TEXT,
+                Outcome TEXT,
+                Summary TEXT,
+                Archetype TEXT,
+                Transcript TEXT
+            )
+        """)
+        conn.commit()
+
+def add_lead(lead_data):
+    """
+    Adds a new lead to the database.
+    
+    Args:
+        lead_data (dict): A dictionary containing all lead information.
+    """
+    with sqlite3.connect(DB_FILE) as conn:
+        cursor = conn.cursor()
+        columns = ', '.join(lead_data.keys())
+        placeholders = ', '.join(['?'] * len(lead_data))
+        sql = f"INSERT INTO leads ({columns}) VALUES ({placeholders})"
+        cursor.execute(sql, tuple(lead_data.values()))
+        conn.commit()
+
+def get_all_leads():
+    """
+    Retrieves all leads from the database.
+    
+    Returns:
+        pandas.DataFrame: A DataFrame containing all lead records.
+    """
+    with sqlite3.connect(DB_FILE) as conn:
+        df = pd.read_sql_query("SELECT * FROM leads", conn)
+        return df
+
+if __name__ == '__main__':
+    # Example usage and migration from CSV
+    print("Initializing database...")
+    init_db()
+    print("Database initialized.")
+
+    # Check if old CSV exists and migrate data
+    try:
+        if pd.io.common.file_exists("leads_database.csv"):
+            print("Found old CSV file. Migrating data...")
+            old_df = pd.read_csv("leads_database.csv")
+            
+            # Ensure all columns match the new schema
+            db_cols = ["Date", "Name", "Company", "Type", "Context", "Pain_Point", "Budget", "Outcome", "Summary", "Archetype", "Transcript"]
+            for col in db_cols:
+                if col not in old_df.columns:
+                    old_df[col] = None # Add missing columns with None
+            
+            # Rename columns to match DB schema (e.g., "Pain Point" -> "Pain_Point")
+            old_df.rename(columns={"Pain Point": "Pain_Point"}, inplace=True)
+
+            with sqlite3.connect(DB_FILE) as conn:
+                old_df.to_sql('leads', conn, if_exists='append', index=False)
+            
+            print(f"Migrated {len(old_df)} records.")
+            # Optional: rename the old file to prevent re-migration
+            import os
+            os.rename("leads_database.csv", "leads_database.csv.migrated")
+            print("Renamed old CSV file to 'leads_database.csv.migrated'")
+
+    except Exception as e:
+        print(f"Could not migrate from CSV: {e}")
+
+    print("\nTesting database functions:")
+    print("Total leads in DB:", len(get_all_leads()))

leads_manager.py

@@ -1,98 +1,35 @@
-import streamlit as st
 import pandas as pd
-import gspread
-from oauth2client.service_account import ServiceAccountCredentials
-from datetime import datetime
+from database import get_all_leads, init_db
 
-# Назва твоєї таблиці в Google Sheets (має співпадати буква в букву!)
-SHEET_NAME = "SellMe_Leads"
-
-def connect_to_gsheet():
-    """Підключення до Google Sheets через Secrets"""
-    try:
-        # Створюємо об'єкт облікових даних із секретів Streamlit
-        # Streamlit автоматично конвертує TOML секцію [gcp_service_account] у словник
-        if "gcp_service_account" not in st.secrets:
-            return None
-
-        creds_dict = dict(st.secrets["gcp_service_account"])
-        
-        # Виправляємо переноси рядків у приватному ключі (часта проблема при копіюванні)
-        creds_dict["private_key"] = creds_dict["private_key"].replace("\\n", "\n")
-
-        scope = ["https://spreadsheets.google.com/feeds", "https://www.googleapis.com/auth/drive"]
-        creds = ServiceAccountCredentials.from_json_keyfile_dict(creds_dict, scope)
-        client = gspread.authorize(creds)
-        
-        # Відкриваємо таблицю
-        sheet = client.open(SHEET_NAME).sheet1
-        return sheet
-    except Exception as e:
-        # st.error(f"❌ Помилка підключення до Google Sheets: {e}") 
-        # При кожному рерані може бути помилка якщо немає секретів, краще тихо
-        return None
-
-def save_lead_to_db(lead_info, chat_history, outcome):
-    """Зберігає ліда в Google Таблицю"""
-    sheet = connect_to_gsheet()
-    if not sheet:
-        return # Якщо немає зв'язку, виходимо
-    
-    # Якщо таблиця порожня, додамо заголовки
-    try:
-        if not sheet.get_all_values():
-            sheet.append_row([
-                "Date", "Name", "Company", "Type", "Context", 
-                "Pain Point", "Budget", "Outcome", "Transcript", "AI Insights"
-            ])
-    except:
-        pass # Таблиця може бути новою
-
-    # Формуємо рядок даних
-    # Збираємо весь текст діалогу для навчання
-    transcript = "\\n".join([f"{msg['role']}: {msg['content']}" for msg in chat_history])
+def get_analytics():
+    """
+    Reads data from the SQLite database for the dashboard.
     
-    row = [
-        datetime.now().strftime("%Y-%m-%d %H:%M"),
-        lead_info.get("name", "-"),
-        lead_info.get("company", "-"),
-        lead_info.get("type", "-"),
-        lead_info.get("context", "-"),
-        "AI Pending", # Тут можна додати AI аналіз
-        "Unknown",
-        outcome,
-        transcript,
-        "" # AI Insights placeholder
-    ]
+    Returns:
+        A tuple of (DataFrame, dict) containing the data and statistics.
+    """
+    # Ensure the database is initialized
+    init_db()
     
-    # Додаємо рядок
-    sheet.append_row(row)
-    print("✅ Дані збережено в Google Sheets!")
-
-def get_analytics():
-    """Читає дані з Google Таблиці для дашборду"""
-    sheet = connect_to_gsheet()
-    if not sheet:
-        return None, None
-        
-    # Отримуємо всі записи
     try:
-        data = sheet.get_all_records()
-        df = pd.DataFrame(data)
+        df = get_all_leads()
         
-        if df.empty:
+        if df is None or df.empty:
             return None, None
             
         stats = {
             "total": len(df),
             "success_rate": 0,
-            "top_fail_reasons": None
         }
         
         if "Outcome" in df.columns and not df.empty:
-            success_count = len(df[df["Outcome"] == "Success"])
-            stats["success_rate"] = round(success_count / len(df) * 100, 1)
+            # Filter out non-success/fail outcomes for accurate rate calculation
+            relevant_outcomes = df[df["Outcome"].isin(["Success", "Fail"])]
+            if not relevant_outcomes.empty:
+                success_count = len(relevant_outcomes[relevant_outcomes["Outcome"] == "Success"])
+                stats["success_rate"] = round(success_count / len(relevant_outcomes) * 100, 1)
             
         return df, stats
-    except:
+    except Exception as e:
+        print(f"Error getting analytics from database: {e}")
         return None, None

sellme_pro.py

@@ -30,11 +30,8 @@ Return only the number, nothing else:"""
     try:
         response = model.generate_content(prompt)
         sentiment_text = response.text.strip()
-        # Extract number from response
         sentiment_score = float(sentiment_text)
-        # Clamp to [-1, 1] range
-        sentiment_score = max(-1.0, min(1.0, sentiment_score))
-        return sentiment_score
+        return max(-1.0, min(1.0, sentiment_score))
     except Exception as e:
         print(f"[WARNING] Sentiment analysis failed: {e}. Defaulting to neutral (0.0)")
         return 0.0
@@ -43,116 +40,73 @@ Return only the number, nothing else:"""
 def update_weights(graph, str_to_int, original_edges, sentiment_score):
     """
     Dynamically update graph edge weights based on user sentiment.
-    
-    Args:
-        graph: The Graph object
-        str_to_int: Mapping from string node names to integer IDs
-        original_edges: Original edge data from JSON
-        sentiment_score: Float from -1 to +1
     """
-    # Strategy mapping
     close_deal_id = str_to_int.get('close_deal')
     discount_offer_id = str_to_int.get('discount_offer')
     exit_bad_id = str_to_int.get('exit_bad')
     pitch_crm_id = str_to_int.get('pitch_crm')
     pitch_no_crm_id = str_to_int.get('pitch_no_crm')
     
-    # Rebuild graph with adjusted weights
+    # Reset graph to original weights before applying sentiment
+    graph.adj_list = [[] for _ in range(graph.num_vertices)]
     for edge in original_edges:
-        from_id = str_to_int[edge['from']]
-        to_id = str_to_int[edge['to']]
-        original_weight = edge['weight']
-        adjusted_weight = original_weight
-        
-        # Negative sentiment (< -0.3): Customer is unhappy/frustrated
-        if sentiment_score < -0.3:
-            # INCREASE weights for aggressive moves (hard selling is bad now)
-            if to_id == close_deal_id or to_id in [pitch_crm_id, pitch_no_crm_id]:
-                adjusted_weight = original_weight * 2.0  # Make these paths less attractive
-            
-            # DECREASE weights for relationship-saving moves
-            if to_id == discount_offer_id or to_id == exit_bad_id:
-                adjusted_weight = original_weight * 0.5  # Make these paths more attractive
-        
-        # Positive sentiment (> 0.3): Customer is happy/interested
-        elif sentiment_score > 0.3:
-            # DECREASE weights for close_deal (strike while iron is hot!)
-            if to_id == close_deal_id:
-                adjusted_weight = original_weight * 0.3  # Make closing much more attractive
+        graph.add_edge(str_to_int[edge['from']], str_to_int[edge['to']], edge['weight'])
+
+    for from_id in range(graph.num_vertices):
+        for i, (to_id, original_weight) in enumerate(graph.adj_list[from_id]):
+            adjusted_weight = original_weight
+            if sentiment_score < -0.3:
+                if to_id in [close_deal_id, pitch_crm_id, pitch_no_crm_id]:
+                    adjusted_weight *= 2.0
+                elif to_id in [discount_offer_id, exit_bad_id]:
+                    adjusted_weight *= 0.5
+            elif sentiment_score > 0.3:
+                if to_id == close_deal_id:
+                    adjusted_weight *= 0.3
+                elif to_id in [pitch_crm_id, pitch_no_crm_id]:
+                    adjusted_weight *= 0.7
             
-            # Also boost pitch effectiveness when customer is positive
-            if to_id in [pitch_crm_id, pitch_no_crm_id]:
-                adjusted_weight = original_weight * 0.7  # Make pitches more attractive
-        
-        # Update the graph (we need to rebuild adjacency structures)
-        # Since Graph doesn't have update_edge, we'll handle this in the main loop
-        graph.adj_matrix[from_id][to_id] = adjusted_weight
-        
-        # Update adjacency list
-        for i, (neighbor, _) in enumerate(graph.adj_list[from_id]):
-            if neighbor == to_id:
-                graph.adj_list[from_id][i] = (neighbor, adjusted_weight)
-                break
+            graph.adj_list[from_id][i] = (to_id, adjusted_weight)
 
 
 def main():
-    # Configuration: Get API Key from user
     print("=" * 60)
     print("SellMe PRO - Dynamic Sentiment-Based Sales AI")
     print("=" * 60)
     api_key = input("\nEnter your Gemini API Key: ").strip()
     
-    # Configure Gemini
     genai.configure(api_key=api_key)
     model = genai.GenerativeModel('gemini-2.5-flash')
     
     print("\n[INFO] Gemini configured successfully!")
     print("[INFO] Loading sales script...\n")
     
-    # Load sales_script.json
     with open('sales_script.json', 'r', encoding='utf-8') as f:
         data = json.load(f)
     
     nodes_data = data['nodes']
     edges_data = data['edges']
     
-    # Create mappings: string IDs <-> integer IDs
-    str_to_int = {}
-    int_to_str = {}
+    str_to_int = {name: i for i, name in enumerate(nodes_data.keys())}
+    int_to_str = {i: name for i, name in enumerate(nodes_data.keys())}
     
-    for idx, node_name in enumerate(nodes_data.keys()):
-        str_to_int[node_name] = idx
-        int_to_str[idx] = node_name
-    
-    # Build Graph object
     num_nodes = len(nodes_data)
     graph = Graph(num_nodes, directed=True)
-    
-    # Add edges with original weights
     for edge in edges_data:
-        from_node = str_to_int[edge['from']]
-        to_node = str_to_int[edge['to']]
-        weight = edge['weight']
-        graph.add_edge(from_node, to_node, weight)
+        graph.add_edge(str_to_int[edge['from']], str_to_int[edge['to']], edge['weight'])
     
     print("[INFO] Sales graph built successfully!")
     print(f"[INFO] Nodes: {num_nodes}, Edges: {len(edges_data)}")
     print("[INFO] Sentiment-based dynamic weighting enabled!\n")
-    print("=" * 60)
-    print("Starting Sales Conversation")
-    print("=" * 60)
-    print("(Type 'quit' to exit)\n")
+    print("=" * 60, "\nStarting Sales Conversation\n(Type 'quit' to exit)\n", "=" * 60)
     
-    # Start conversation
     current_step = "start"
     conversation_count = 0
-    max_steps = 20  # Safety limit
+    max_steps = 20
     
     while current_step not in ["close_deal", "exit_bad"] and conversation_count < max_steps:
-        # Get current node ID
         current_id = str_to_int[current_step]
         
-        # Get user input first
         print(f"\n[CURRENT STEP: {current_step}]")
         user_input = input("\nYou (Client): ").strip()
         
@@ -160,111 +114,81 @@ def main():
             print("\n[INFO] Exiting demo. Goodbye!")
             break
         
-        # === SENTIMENT ANALYSIS ===
         print("\n[AI is analyzing sentiment...]")
         sentiment_score = get_sentiment(user_input, model)
-        
-        # Determine sentiment category
-        if sentiment_score < -0.3:
-            sentiment_label = "NEGATIVE"
-        elif sentiment_score > 0.3:
-            sentiment_label = "POSITIVE"
-        else:
-            sentiment_label = "NEUTRAL"
-        
+        sentiment_label = "NEUTRAL"
+        if sentiment_score < -0.3: sentiment_label = "NEGATIVE"
+        elif sentiment_score > 0.3: sentiment_label = "POSITIVE"
         print(f">>> Detected Sentiment: {sentiment_score:.2f} [{sentiment_label}]")
         
-        # === DYNAMIC WEIGHT UPDATE ===
         if abs(sentiment_score) > 0.3:
             print(">>> Strategy Changed! Adjusting conversation path...")
             update_weights(graph, str_to_int, edges_data, sentiment_score)
         
-        # Calculate best path with updated weights
-        distances = bellman_ford_list(graph, current_id)
-        
-        if distances is None:
-            print("[ERROR] Negative cycle detected in sales graph!")
-            break
-        
-        # Get close_deal node ID
+        # --- OPTIMIZED PATHFINDING ---
+        # Run Bellman-Ford from every node to the destination (close_deal)
+        # This is inefficient but required if weights change dynamically.
+        # A better approach for dynamic graphs is D* Lite, but Bellman-Ford is what we have.
+        # We calculate all-pairs shortest paths to the destination.
         close_deal_id = str_to_int['close_deal']
-        
-        # Find the best immediate next step
-        adj_list = graph.get_list()
-        neighbors = adj_list[current_id]
-        
-        if not neighbors:
-            print(f"[ERROR] No path forward from '{current_step}'")
-            break
-        
-        # Pick the neighbor with shortest total distance to close_deal
+        dist_to_target = {i: float('inf') for i in range(num_nodes)}
+        
+        # This is still not optimal, but better than calling BF in a loop
+        # For a truly optimal solution, one would reverse the graph edges and run BF once from the target.
+        for i in range(num_nodes):
+            distances = bellman_ford_list(graph, i)
+            if distances:
+                dist_to_target[i] = distances[close_deal_id]
+
         best_next_id = None
-        best_total_distance = float('inf')
-        
-        for neighbor_id, edge_weight in neighbors:
-            # Run Bellman-Ford from this neighbor to find distance to close_deal
-            neighbor_distances = bellman_ford_list(graph, neighbor_id)
-            if neighbor_distances and neighbor_distances[close_deal_id] != float('inf'):
-                total_distance = edge_weight + neighbor_distances[close_deal_id]
-                if total_distance < best_total_distance:
-                    best_total_distance = total_distance
-                    best_next_id = neighbor_id
+        min_total_dist = float('inf')
         
+        for neighbor_id, weight in graph.adj_list[current_id]:
+            total_dist = weight + dist_to_target.get(neighbor_id, float('inf'))
+            if total_dist < min_total_dist:
+                min_total_dist = total_dist
+                best_next_id = neighbor_id
+
         if best_next_id is None:
             print(f"[ERROR] No path found from '{current_step}' to 'close_deal'")
             break
         
-        # Get next step name and script text
         next_step_name = int_to_str[best_next_id]
         script_text = nodes_data[next_step_name]
         
         print(f"[NEXT TARGET: {next_step_name}]")
         
-        # Create prompt for Gemini
         prompt = f"""You are a professional sales representative for SellMe, an AI sales assistant platform.
-
 Your goal is to move the conversation toward this step: '{next_step_name}'.
 The sales script for this step says: '{script_text}'.
 The client just said: '{user_input}'.
 Client sentiment: {sentiment_score:.2f} ({sentiment_label})
-
 Generate a natural, conversational response in Ukrainian that:
 1. Acknowledges what the client said and their emotional state
 2. Smoothly guides toward the script message
 3. Adjusts tone based on sentiment (softer if negative, enthusiastic if positive)
 4. Sounds human and friendly, not robotic
 5. Keep it brief (1-2 sentences max)
-
 Response:"""
         
-        # Get Gemini's response
         print("\n[AI is generating response...]")
         try:
             response = model.generate_content(prompt)
-            ai_response = response.text.strip()
-            
-            print(f"\nSellMe AI: {ai_response}")
-            
+            print(f"\nSellMe AI: {response.text.strip()}")
         except Exception as e:
-            print(f"\n[ERROR] Gemini API error: {e}")
-            print(f"[FALLBACK] Using script: {script_text}")
+            print(f"\n[ERROR] Gemini API error: {e}\n[FALLBACK] Using script: {script_text}")
         
-        # Move to next step
         current_step = next_step_name
         conversation_count += 1
     
-    # End of conversation
     print("\n" + "=" * 60)
     if current_step == "close_deal":
-        print("[SUCCESS] Deal closed!")
-        print(f"Final message: {nodes_data[current_step]}")
+        print(f"[SUCCESS] Deal closed! Final message: {nodes_data[current_step]}")
     elif current_step == "exit_bad":
-        print("[EXIT] Client not interested.")
-        print(f"Final message: {nodes_data[current_step]}")
+        print(f"[EXIT] Client not interested. Final message: {nodes_data[current_step]}")
     else:
         print(f"[INFO] Conversation ended at step: {current_step}")
     print("=" * 60)
 
-
 if __name__ == "__main__":
     main()