Update app.py

app.py

@@ -80,11 +80,11 @@ def extract_main_content(soup):
     """
     if not soup:
         return ""
-        
+
     # Remove script and style elements
     for element in soup(['script', 'style', 'header', 'footer', 'nav', 'ads', 'sidebar']):
         element.decompose()
-    
+
     # First try to find content in main content areas
     main_content_tags = soup.find_all(['article', 'main', 'div.content', 'div.post', 'div.entry-content'])
     if main_content_tags:
@@ -97,14 +97,14 @@ def extract_main_content(soup):
         else:
             # Fallback to body content
             content = soup.body.get_text(strip=True, separator=' ') if soup.body else soup.get_text(strip=True, separator=' ')
-    
+
     # Clean up the text
     content = ' '.join(content.split())
     content = re.sub(r'\s+', ' ', content)  # Remove multiple spaces
     content = re.sub(r'[\n\r\t]', ' ', content)  # Remove newlines and tabs
-    
-    # Limit content length to avoid token limits (adjust as needed)
-    return content[:5000]
+
+    # Return the content
+    return content
 
 def get_page_metadata(soup):
     """
@@ -115,15 +115,15 @@ def get_page_metadata(soup):
         'description': '',
         'keywords': ''
     }
-    
+
     if not soup:
         return metadata
-        
+
     # Get title
     title_tag = soup.find('title')
     if title_tag and title_tag.string:
         metadata['title'] = title_tag.string.strip()
-    
+
     # Get meta description (try multiple variants)
     meta_desc = (
         soup.find('meta', attrs={'name': 'description'}) or
@@ -132,99 +132,108 @@ def get_page_metadata(soup):
     )
     if meta_desc:
         metadata['description'] = meta_desc.get('content', '').strip()
-    
+
     # Get meta keywords
     meta_keywords = soup.find('meta', attrs={'name': 'keywords'})
     if meta_keywords:
         metadata['keywords'] = meta_keywords.get('content', '').strip()
-    
+
     # Get OG title if main title is empty
     if not metadata['title']:
         og_title = soup.find('meta', attrs={'property': 'og:title'})
         if og_title:
             metadata['title'] = og_title.get('content', '').strip()
-    
+
     return metadata
 
 def generate_summary(bookmark):
     """
-    Generate a comprehensive summary for a bookmark using available content and LLM.
+    Generate a comprehensive summary for a bookmark using available content and LLM via the Groq Cloud API.
     """
     logger.info(f"Generating summary for bookmark: {bookmark.get('url')}")
-    
+
     try:
-        # Get the HTML soup object from the bookmark if it exists
+        # Get the HTML soup object from the bookmark
         soup = BeautifulSoup(bookmark.get('html_content', ''), 'html.parser')
-        
-        # Step 1: Extract all available information
+
+        # Extract metadata and main content
         metadata = get_page_metadata(soup)
         main_content = extract_main_content(soup)
-        
-        # Step 2: Generate summary using LLM with all available content
-        try:
-            # Prepare comprehensive context for LLM
-            available_content = []
-            if metadata['title']:
-                available_content.append(f"Title: {metadata['title']}")
-            if metadata['description']:
-                available_content.append(f"Description: {metadata['description']}")
-            if metadata['keywords']:
-                available_content.append(f"Keywords: {metadata['keywords']}")
-            if main_content:
-                available_content.append(f"Main Content: {main_content}")
-            
-            if not available_content:
-                logger.warning("No content available for summary generation")
-                bookmark['summary'] = bookmark.get('title', 'No summary available.')
-                return bookmark
-                
-            prompt = f"""
-            Analyze and summarize this webpage based on the following information:
-            
-            {' | '.join(available_content)}
-            
-            Please provide a concise summary (2-3 sentences) focusing on:
-            1. The main purpose or topic of the page
-            2. Key information or features
-            3. Target audience or use case (if apparent)
-            
-            Be factual and objective.
-            """
-            
-            response = openai.ChatCompletion.create(
-                model='llama3-8b-8192',
-                messages=[
-                    {"role": "system", "content": "You are a helpful assistant that creates concise webpage summaries."},
-                    {"role": "user", "content": prompt}
-                ],
-                max_tokens=150,
-                temperature=0.5,
-            )
-            
-            summary = response['choices'][0]['message']['content'].strip()
-            logger.info("Successfully generated LLM summary")
-            bookmark['summary'] = summary
-            return bookmark
-            
-        except Exception as e:
-            logger.error(f"Error generating LLM summary: {e}")
-            # Fallback mechanisms in order of preference
-            if metadata['description']:
-                logger.info("Falling back to meta description")
-                bookmark['summary'] = metadata['description']
-            elif main_content:
-                logger.info("Falling back to truncated main content")
-                bookmark['summary'] = ' '.join(main_content.split()[:50]) + '...'
-            elif metadata['title']:
-                logger.info("Falling back to title")
-                bookmark['summary'] = metadata['title']
-            else:
-                bookmark['summary'] = bookmark.get('title', 'No summary available.')
+
+        # Prepare content for the prompt
+        available_content = []
+        if metadata['title']:
+            available_content.append(f"Title: {metadata['title']}")
+        if metadata['description']:
+            available_content.append(f"Description: {metadata['description']}")
+        if metadata['keywords']:
+            available_content.append(f"Keywords: {metadata['keywords']}")
+        if main_content:
+            available_content.append(f"Main Content: {main_content}")
+
+        if not available_content:
+            logger.warning("No content available for summary generation")
+            bookmark['summary'] = bookmark.get('title', 'No summary available.')
             return bookmark
-            
+
+        # Estimate token count and trim content if necessary
+        max_total_tokens = 8000  # Adjust based on model's maximum context length
+        prompt_tokens_estimate = len(' '.join(available_content).split()) + 200  # 200 tokens reserved for response
+        if prompt_tokens_estimate > max_total_tokens:
+            # Trim main content
+            allowable_content_tokens = max_total_tokens - 200  # Reserve 200 tokens for response
+            main_content_tokens = len(main_content.split())
+            if main_content_tokens > allowable_content_tokens:
+                main_content = ' '.join(main_content.split()[:allowable_content_tokens])
+                logger.info("Trimmed main content to fit within token limits.")
+
+            # Update available content
+            available_content[-1] = f"Main Content: {main_content}"
+
+        # Construct the prompt
+        prompt = f"""
+Analyze and summarize the following webpage content:
+
+{' '.join(available_content)}
+
+Provide a concise summary (2-3 sentences) focusing on:
+- The main purpose or topic of the page.
+- Key information or features.
+- Target audience or use case (if apparent).
+
+Be factual and objective.
+"""
+
+        # Call the LLM via Groq Cloud API
+        response = openai.ChatCompletion.create(
+            model='llama3-8b-8192',  # Use the model as per your Groq Cloud API configuration
+            messages=[
+                {"role": "system", "content": "You are a helpful assistant that creates concise webpage summaries."},
+                {"role": "user", "content": prompt}
+            ],
+            max_tokens=200,  # Adjust as necessary to accommodate longer summaries
+            temperature=0.5,
+        )
+
+        summary = response['choices'][0]['message']['content'].strip()
+        logger.info("Successfully generated LLM summary")
+        bookmark['summary'] = summary
+        return bookmark
+
     except Exception as e:
-        logger.error(f"Error in generate_summary: {e}")
-        bookmark['summary'] = bookmark.get('title', 'No summary available.')
+        logger.error(f"Error generating summary: {e}")
+        # Fallback mechanisms
+        if metadata['description']:
+            logger.info("Falling back to meta description")
+            bookmark['summary'] = metadata['description']
+        elif main_content:
+            logger.info("Falling back to main content")
+            bookmark['summary'] = ' '.join(main_content.split()[:50]) + '...'
+        elif metadata['title']:
+            logger.info("Falling back to title")
+            bookmark['summary'] = metadata['title']
+        else:
+            bookmark['summary'] = 'No summary available.'
         return bookmark
 
 def parse_bookmarks(file_content):
@@ -313,67 +322,75 @@ async def process_bookmarks_async(bookmarks_list):
 
 def assign_category(bookmark):
     """
-    Assign a category to a bookmark based on its content.
+    Assign a category to a bookmark using the LLM based on its summary via the Groq Cloud API.
     """
     if bookmark.get('dead_link'):
         bookmark['category'] = 'Dead Link'
         logger.info(f"Assigned category 'Dead Link' to bookmark: {bookmark.get('url')}")
         return bookmark
 
-    summary = bookmark.get('summary', '').lower()
-    assigned_category = 'Uncategorized'
-
-    # Keywords associated with each category
-    category_keywords = {
-        "Social Media": ["social media", "networking", "friends", "connect", "posts", "profile"],
-        "News and Media": ["news", "journalism", "media", "headlines", "breaking news"],
-        "Education and Learning": ["education", "learning", "courses", "tutorial", "university", "academy", "study"],
-        "Entertainment": ["entertainment", "movies", "tv shows", "games", "comics", "fun"],
-        "Shopping and E-commerce": ["shopping", "e-commerce", "buy", "sell", "marketplace", "deals", "store"],
-        "Finance and Banking": ["finance", "banking", "investment", "money", "economy", "stock", "trading"],
-        "Technology": ["technology", "tech", "gadgets", "software", "computers", "innovation"],
-        "Health and Fitness": ["health", "fitness", "medical", "wellness", "exercise", "diet"],
-        "Travel and Tourism": ["travel", "tourism", "destinations", "hotels", "flights", "vacation"],
-        "Food and Recipes": ["food", "recipes", "cooking", "cuisine", "restaurant", "dining"],
-        "Sports": ["sports", "scores", "teams", "athletics", "matches", "leagues"],
-        "Arts and Culture": ["arts", "culture", "museum", "gallery", "exhibition", "artistic"],
-        "Government and Politics": ["government", "politics", "policy", "election", "public service"],
-        "Business and Economy": ["business", "corporate", "industry", "economy", "markets"],
-        "Science and Research": ["science", "research", "experiment", "laboratory", "study", "scientific"],
-        "Personal Blogs and Journals": ["blog", "journal", "personal", "diary", "thoughts", "opinions"],
-        "Job Search and Careers": ["jobs", "careers", "recruitment", "resume", "employment", "hiring"],
-        "Music and Audio": ["music", "audio", "songs", "albums", "artists", "bands"],
-        "Videos and Movies": ["video", "movies", "film", "clips", "trailers", "cinema"],
-        "Reference and Knowledge Bases": ["reference", "encyclopedia", "dictionary", "wiki", "knowledge", "information"],
-    }
+    summary = bookmark.get('summary', '')
+    if not summary:
+        bookmark['category'] = 'Uncategorized'
+        return bookmark
 
-    for category, keywords in category_keywords.items():
-        for keyword in keywords:
-            if re.search(r'\b' + re.escape(keyword) + r'\b', summary):
-                assigned_category = category
-                logger.info(f"Assigned category '{assigned_category}' to bookmark: {bookmark.get('url')}")
-                break
-        if assigned_category != 'Uncategorized':
-            break
-
-    bookmark['category'] = assigned_category
-    if assigned_category == 'Uncategorized':
-        logger.info(f"No matching category found for bookmark: {bookmark.get('url')}")
-    return bookmark
+    # Prepare the prompt
+    categories_str = ', '.join([f'"{cat}"' for cat in CATEGORIES if cat != 'Dead Link'])
+    prompt = f"""
+Based on the following summary, assign the most appropriate category from the list below.
+
+Summary:
+{summary}
+
+Categories:
+{categories_str}
+
+Respond with only the category name.
+"""
+
+    try:
+        response = openai.ChatCompletion.create(
+            model='llama3-8b-8192',  # Use the model as per your Groq Cloud API configuration
+            messages=[
+                {"role": "system", "content": "You categorize webpages based on their content."},
+                {"role": "user", "content": prompt}
+            ],
+            max_tokens=10,
+            temperature=0,
+        )
+
+        category = response['choices'][0]['message']['content'].strip().strip('"')
+
+        # Validate the category
+        if category in CATEGORIES:
+            bookmark['category'] = category
+            logger.info(f"Assigned category '{category}' to bookmark: {bookmark.get('url')}")
+        else:
+            bookmark['category'] = 'Uncategorized'
+            logger.warning(f"Invalid category '{category}' returned by LLM for bookmark: {bookmark.get('url')}")
+
+        return bookmark
+
+    except Exception as e:
+        logger.error(f"Error assigning category: {e}")
+        bookmark['category'] = 'Uncategorized'
+        return bookmark
 
 def vectorize_and_index(bookmarks_list):
     """
-    Create vector embeddings for bookmarks and build FAISS index.
+    Create vector embeddings for bookmarks and build FAISS index with ID mapping.
     """
     logger.info("Vectorizing summaries and building FAISS index")
     try:
         summaries = [bookmark['summary'] for bookmark in bookmarks_list]
         embeddings = embedding_model.encode(summaries)
         dimension = embeddings.shape[1]
-        faiss_idx = faiss.IndexFlatL2(dimension)
-        faiss_idx.add(np.array(embeddings))
-        logger.info("FAISS index built successfully")
-        return faiss_idx, embeddings
+        index = faiss.IndexIDMap(faiss.IndexFlatL2(dimension))
+        # Assign unique IDs to each bookmark
+        ids = np.array([bookmark['id'] for bookmark in bookmarks_list], dtype=np.int64)
+        index.add_with_ids(np.array(embeddings).astype('float32'), ids)
+        logger.info("FAISS index built successfully with IDs")
+        return index
     except Exception as e:
         logger.error(f"Error in vectorizing and indexing: {e}")
         raise
@@ -400,6 +417,13 @@ def display_bookmarks():
             card_style = "border: 2px solid var(--success-color);"
             text_style = "color: var(--text-color);"
 
+        # Escape HTML content to prevent XSS attacks
+        from html import escape
+        title = escape(title)
+        url = escape(url)
+        summary = escape(summary)
+        category = escape(category)
+
         card_html = f'''
         <div class="card" style="{card_style}; padding: 10px; margin: 10px; border-radius: 5px;">
             <div class="card-content">
@@ -421,7 +445,7 @@ def process_uploaded_file(file):
     """
     global bookmarks, faiss_index
     logger.info("Processing uploaded file")
-    
+
     if file is None:
         logger.warning("No file uploaded")
         return "Please upload a bookmarks HTML file.", '', gr.update(choices=[]), display_bookmarks()
@@ -442,6 +466,10 @@ def process_uploaded_file(file):
         logger.warning("No bookmarks found in the uploaded file")
         return "No bookmarks found in the uploaded file.", '', gr.update(choices=[]), display_bookmarks()
 
+    # Assign unique IDs to bookmarks
+    for idx, bookmark in enumerate(bookmarks):
+        bookmark['id'] = idx
+
     # Asynchronously fetch bookmark info
     try:
         asyncio.run(process_bookmarks_async(bookmarks))
@@ -455,48 +483,52 @@ def process_uploaded_file(file):
         assign_category(bookmark)
 
     try:
-        faiss_index, embeddings = vectorize_and_index(bookmarks)
+        faiss_index = vectorize_and_index(bookmarks)
     except Exception as e:
         logger.error(f"Error building FAISS index: {e}")
         return "Error building search index.", '', gr.update(choices=[]), display_bookmarks()
 
     message = f"✅ Successfully processed {len(bookmarks)} bookmarks."
     logger.info(message)
-    
+
     # Generate displays and updates
     bookmark_html = display_bookmarks()
-    choices = [f"{i+1}. {bookmark['title']} (Category: {bookmark['category']})" 
+    choices = [f"{i+1}. {bookmark['title']} (Category: {bookmark['category']})"
               for i, bookmark in enumerate(bookmarks)]
 
     return message, bookmark_html, gr.update(choices=choices), bookmark_html
 
 def delete_selected_bookmarks(selected_indices):
     """
-    Delete selected bookmarks.
+    Delete selected bookmarks and remove their vectors from the FAISS index.
     """
     global bookmarks, faiss_index
     if not selected_indices:
         return "⚠️ No bookmarks selected.", gr.update(choices=[]), display_bookmarks()
 
-    indices = [int(s.split('.')[0])-1 for s in selected_indices]
-    indices = sorted(indices, reverse=True)
-    for idx in indices:
+    ids_to_delete = []
+    indices_to_delete = []
+    for s in selected_indices:
+        idx = int(s.split('.')[0]) - 1
         if 0 <= idx < len(bookmarks):
+            bookmark_id = bookmarks[idx]['id']
+            ids_to_delete.append(bookmark_id)
+            indices_to_delete.append(idx)
             logger.info(f"Deleting bookmark at index {idx + 1}")
-            bookmarks.pop(idx)
 
-    if bookmarks:
-        faiss_index, embeddings = vectorize_and_index(bookmarks)
-    else:
-        faiss_index = None
+    # Remove vectors from FAISS index
+    if faiss_index is not None and ids_to_delete:
+        faiss_index.remove_ids(np.array(ids_to_delete, dtype=np.int64))
+
+    # Remove bookmarks from the list (reverse order to avoid index shifting)
+    for idx in sorted(indices_to_delete, reverse=True):
+        bookmarks.pop(idx)
 
     message = "🗑️ Selected bookmarks deleted successfully."
     logger.info(message)
-    
-    # Update choices and display
-    choices = [f"{i+1}. {bookmark['title']} (Category: {bookmark['category']})" 
-              for i, bookmark in enumerate(bookmarks)]
-    
+    choices = [f"{i+1}. {bookmark['title']} (Category: {bookmark['category']})"
+               for i, bookmark in enumerate(bookmarks)]
+
     return message, gr.update(choices=choices), display_bookmarks()
 
 def edit_selected_bookmarks_category(selected_indices, new_category):
@@ -516,11 +548,11 @@ def edit_selected_bookmarks_category(selected_indices, new_category):
 
     message = "✏️ Category updated for selected bookmarks."
     logger.info(message)
-    
+
     # Update choices and display
-    choices = [f"{i+1}. {bookmark['title']} (Category: {bookmark['category']})" 
+    choices = [f"{i+1}. {bookmark['title']} (Category: {bookmark['category']})"
               for i, bookmark in enumerate(bookmarks)]
-    
+
     return message, gr.update(choices=choices), display_bookmarks()
 
 def export_bookmarks():
@@ -553,40 +585,52 @@ def export_bookmarks():
 
 def chatbot_response(user_query):
     """
-    Generate chatbot response using Groq Cloud API.
+    Generate chatbot response using the FAISS index and embeddings.
     """
-    if not GROQ_API_KEY:
-        logger.warning("GROQ_API_KEY not set.")
-        return "⚠️ API key not set. Please set the GROQ_API_KEY environment variable."
-
-    if not bookmarks:
+    if not bookmarks or faiss_index is None:
         logger.warning("No bookmarks available for chatbot")
         return "⚠️ No bookmarks available. Please upload and process your bookmarks first."
 
     logger.info(f"Chatbot received query: {user_query}")
 
     try:
-        max_bookmarks = 50
-        bookmark_data = ""
-        for idx, bookmark in enumerate(bookmarks[:max_bookmarks]):
-            bookmark_data += f"{idx+1}. Title: {bookmark['title']}\nURL: {bookmark['url']}\nSummary: {bookmark['summary']}\n\n"
+        # Encode the user query
+        query_vector = embedding_model.encode([user_query]).astype('float32')
+
+        # Search the FAISS index
+        k = 5  # Number of results to return
+        distances, ids = faiss_index.search(query_vector, k)
+        ids = ids.flatten()
 
+        # Retrieve the bookmarks
+        id_to_bookmark = {bookmark['id']: bookmark for bookmark in bookmarks}
+        matching_bookmarks = [id_to_bookmark.get(id) for id in ids if id in id_to_bookmark]
+
+        if not matching_bookmarks:
+            return "No relevant bookmarks found for your query."
+
+        # Format the response
+        bookmarks_info = "\n".join([
+            f"Title: {bookmark['title']}\nURL: {bookmark['url']}\nSummary: {bookmark['summary']}"
+            for bookmark in matching_bookmarks
+        ])
+
+        # Use the LLM via Groq Cloud API to generate a response
         prompt = f"""
-        You are an assistant that helps users find relevant bookmarks from their collection based on their queries.
+A user asked: "{user_query}"
 
-        User Query:
-        {user_query}
+Based on the bookmarks below, provide a helpful answer to the user's query, referencing the relevant bookmarks.
 
-        Bookmarks:
-        {bookmark_data}
+Bookmarks:
+{bookmarks_info}
 
-        Please identify the most relevant bookmarks that match the user's query. Provide a concise list including the index, title, URL, and a brief summary.
-        """
+Provide a concise and helpful response.
+"""
 
         response = openai.ChatCompletion.create(
-            model='llama3-8b-8192',
+            model='llama3-8b-8192',  # Use the model as per your Groq Cloud API configuration
             messages=[
-                {"role": "system", "content": "You help users find relevant bookmarks based on their queries."},
+                {"role": "system", "content": "You assist users by finding relevant information from their bookmarks."},
                 {"role": "user", "content": prompt}
             ],
             max_tokens=500,
@@ -680,7 +724,7 @@ def build_app():
                     delete_button = gr.Button("🗑️ Delete Selected")
                     edit_category_button = gr.Button("✏️ Edit Category")
                     export_button = gr.Button("💾 Export")
-                
+
                 download_link = gr.HTML(label="📥 Download")
 
             # Set up event handlers
@@ -720,4 +764,4 @@ def build_app():
         print(f"Error building the app: {e}")
 
 if __name__ == "__main__":
-    build_app()
+    build_app()