several fixes, removed logging

app/ui_components.py

@@ -41,7 +41,7 @@ def display_info_icons():
         if time.time() - st.session_state.info_icons_time > 10 or ("messages" in st.session_state and len(st.session_state.messages) > 0):
             st.session_state.info_icons_displayed = False
 
-def extract_data_from_markdown(text: Union[str, bytes, io.BytesIO]) -> Union[str, bytes, None]:
+def extract_data_from_markdown(text: Union[str, bytes, io.BytesIO]) -> Union[str, bytes, io.BytesIO, None]:
     if isinstance(text, io.BytesIO):
         return text
     if isinstance(text, bytes):
@@ -49,16 +49,24 @@ def extract_data_from_markdown(text: Union[str, bytes, io.BytesIO]) -> Union[str
     pattern = r'```(csv|excel)\n(.*?)\n```'
     match = re.search(pattern, text, re.DOTALL)
     if match:
-        return match.group(2).strip()
+        data_type = match.group(1)
+        data = match.group(2).strip()
+        if data_type == 'excel':
+            return io.BytesIO(data.encode())
+        return data
     return None
 
 def format_data(data: Union[str, bytes, io.BytesIO], format_type: str):
     try:
         if isinstance(data, io.BytesIO):
+            if format_type == 'excel':
+                return pd.read_excel(data, engine='openpyxl')
             data.seek(0)
-            return pd.read_excel(data, engine='openpyxl')
+            return pd.read_csv(data)
         elif isinstance(data, bytes):
-            return pd.read_excel(io.BytesIO(data), engine='openpyxl')
+            if format_type == 'excel':
+                return pd.read_excel(io.BytesIO(data), engine='openpyxl')
+            return pd.read_csv(io.BytesIO(data))
         else:
             if format_type == 'csv':
                 csv_data = []

main.py

@@ -1,7 +1,6 @@
 import streamlit as st
 import json
 import asyncio
-import logging
 from app.streamlit_web_scraper_chat import StreamlitWebScraperChat
 from app.ui_components import display_info_icons, display_message, extract_data_from_markdown, format_data
 from app.utils import loading_animation, get_loading_message
@@ -11,6 +10,8 @@ import pandas as pd
 import base64
 from google_auth_oauthlib.flow import Flow
 import io
+from io import BytesIO
+import re
 from src.utils.google_sheets_utils import SCOPES, get_redirect_uri, display_google_sheets_button, initiate_google_auth
 
 def handle_oauth_callback():
@@ -28,18 +29,88 @@ def handle_oauth_callback():
         except Exception as e:
             st.error(f"Error during OAuth callback: {str(e)}")
 
+def serialize_bytesio(obj):
+    if isinstance(obj, BytesIO):
+        return {
+            "_type": "BytesIO",
+            "data": base64.b64encode(obj.getvalue()).decode('utf-8')
+        }
+    raise TypeError(f"Object of type {obj.__class__.__name__} is not JSON serializable")
+
+def deserialize_bytesio(obj):
+    if isinstance(obj, dict) and "_type" in obj and obj["_type"] == "BytesIO":
+        return BytesIO(base64.b64decode(obj["data"]))
+    return obj
+
+def save_chat_history(chat_history):
+    with open("chat_history.json", "w") as f:
+        json.dump(chat_history, f, default=serialize_bytesio)
+
+def load_chat_history():
+    try:
+        with open("chat_history.json", "r") as f:
+            return json.load(f, object_hook=deserialize_bytesio)
+    except FileNotFoundError:
+        return {}
+
 def safe_process_message(web_scraper_chat, message):
     if message is None or message.strip() == "":
         return "I'm sorry, but I didn't receive any input. Could you please try again?"
     try:
         response = web_scraper_chat.process_message(message)
-        if isinstance(response, tuple) and len(response) == 2 and isinstance(response[1], pd.DataFrame):
-            csv_string, df = response
-            st.text("CSV Data:")
-            st.code(csv_string, language="csv")
-            st.text("Interactive Table:")
-            st.dataframe(df)
-            return csv_string
+        st.write("Debug: Response type:", type(response))
+        
+        if isinstance(response, tuple):
+            st.write("Debug: Response is a tuple")
+            if len(response) == 2 and isinstance(response[1], pd.DataFrame):
+                st.write("Debug: CSV data detected")
+                csv_string, df = response
+                st.text("CSV Data:")
+                st.code(csv_string, language="csv")
+                st.text("Interactive Table:")
+                st.dataframe(df)
+                
+                csv_buffer = BytesIO()
+                df.to_csv(csv_buffer, index=False)
+                csv_buffer.seek(0)
+                st.download_button(
+                    label="Download CSV",
+                    data=csv_buffer,
+                    file_name="data.csv",
+                    mime="text/csv"
+                )
+                
+                return csv_string
+            elif len(response) == 2 and isinstance(response[0], BytesIO):
+                st.write("Debug: Excel data detected")
+                excel_buffer, df = response
+                st.text("Excel Data:")
+                st.dataframe(df)
+                
+                excel_buffer.seek(0)
+                st.download_button(
+                    label="Download Original Excel file",
+                    data=excel_buffer,
+                    file_name="data_original.xlsx",
+                    mime="application/vnd.openxmlformats-officedocument.spreadsheetml.sheet"
+                )
+                
+                excel_data = BytesIO()
+                with pd.ExcelWriter(excel_data, engine='xlsxwriter') as writer:
+                    df.to_excel(writer, index=False, sheet_name='Sheet1')
+                excel_data.seek(0)
+                
+                st.download_button(
+                    label="Download Excel (from DataFrame)",
+                    data=excel_data,
+                    file_name="data_from_df.xlsx",
+                    mime="application/vnd.openxmlformats-officedocument.spreadsheetml.sheet"
+                )
+                
+                return ("Excel data displayed and available for download.", excel_buffer)
+        else:
+            st.write("Debug: Response is not a tuple")
+        
         return response
     except AttributeError as e:
         if "'NoneType' object has no attribute 'lower'" in str(e):
@@ -47,19 +118,9 @@ def safe_process_message(web_scraper_chat, message):
         else:
             raise e
     except Exception as e:
+        st.write("Debug: Exception occurred:", str(e))
         return f"An unexpected error occurred: {str(e)}. Please try again or contact support if the issue persists."
 
-def load_chat_history():
-    try:
-        with open("chat_history.json", "r") as f:
-            return json.load(f)
-    except FileNotFoundError:
-        return {}
-
-def save_chat_history(chat_history):
-    with open("chat_history.json", "w") as f:
-        json.dump(chat_history, f)
-
 def get_date_group(date_str):
     date = datetime.strptime(date_str, "%Y-%m-%d")
     today = datetime.now().date()
@@ -95,13 +156,6 @@ async def list_ollama_models():
         st.error(f"Error fetching Ollama models: {str(e)}")
         return []
 
-def setup_logging(enable_logging):
-    if enable_logging:
-        logging.basicConfig(level=logging.DEBUG, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
-        return logging.getLogger(__name__)
-    else:
-        return logging.getLogger(__name__)
-
 def load_css():
     with open("app/styles.css", "r") as f:
         st.markdown(f"<style>{f.read()}</style>", unsafe_allow_html=True)
@@ -124,19 +178,49 @@ def render_message(role, content, avatar_path):
 
 def display_message_with_sheets_upload(message, message_index):
     content = message["content"]
-    if isinstance(content, (str, bytes, io.BytesIO)):
+    if isinstance(content, (str, bytes, BytesIO)):
         data = extract_data_from_markdown(content)
         if data is not None:
-            if isinstance(data, io.BytesIO) or (isinstance(content, str) and 'excel' in content.lower()):
-                df = format_data(data, 'excel')
-            else:
-                df = format_data(data, 'csv')
-            
-            if df is not None:
-                st.dataframe(df)
-                display_google_sheets_button(df)
-            else:
-                st.warning("Failed to display data as a table. Showing raw content:")
+            try:
+                is_excel = isinstance(data, BytesIO) or (isinstance(content, str) and 'excel' in content.lower())
+                if is_excel:
+                    df = format_data(data, 'excel')
+                else:
+                    df = format_data(data, 'csv')
+                
+                if df is not None:
+                    st.dataframe(df)
+                    
+                    if not is_excel:
+                        csv_buffer = BytesIO()
+                        df.to_csv(csv_buffer, index=False)
+                        csv_buffer.seek(0)
+                        st.download_button(
+                            label="📥 Download as CSV",
+                            data=csv_buffer,
+                            file_name="data.csv",
+                            mime="text/csv",
+                            key=f"csv_download_{message_index}"
+                        )
+                    else:
+                        excel_buffer = BytesIO()
+                        with pd.ExcelWriter(excel_buffer, engine='xlsxwriter') as writer:
+                            df.to_excel(writer, index=False, sheet_name='Sheet1')
+                        excel_buffer.seek(0)
+                        st.download_button(
+                            label="📥 Download as Excel",
+                            data=excel_buffer,
+                            file_name="data.xlsx",
+                            mime="application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
+                            key=f"excel_download_{message_index}"
+                        )
+                    
+                    display_google_sheets_button(df, f"sheets_upload_{message_index}")
+                else:
+                    st.warning("Failed to display data as a table. Showing raw content:")
+                    st.code(content)
+            except Exception as e:
+                st.error(f"Error processing data: {str(e)}")
                 st.code(content)
         else:
             st.markdown(content)
@@ -159,12 +243,6 @@ def main():
     user_avatar_path = "app/icons/man.png"
     ai_avatar_path = "app/icons/skull.png"
 
-    if 'enable_logging' not in st.session_state:
-        st.session_state.enable_logging = False
-
-    logger = setup_logging(st.session_state.enable_logging)
-    logger.debug("Starting CyberScraper 2077")
-
     if 'chat_history' not in st.session_state:
         st.session_state.chat_history = load_chat_history()
     if 'current_chat_id' not in st.session_state or st.session_state.current_chat_id not in st.session_state.chat_history:
@@ -186,12 +264,6 @@ def main():
     with st.sidebar:
         st.title("Conversation History")
 
-        st.session_state.enable_logging = st.toggle("Enable Logging", st.session_state.enable_logging)
-        if st.session_state.enable_logging:
-            st.info("Logging is enabled. Check your console for log messages.")
-        else:
-            st.info("Logging is disabled.")
-
         # Model selection
         st.subheader("Select Model")
         default_models = ["gpt-4o-mini", "gpt-3.5-turbo"]
@@ -296,14 +368,9 @@ def main():
     prompt = st.chat_input("Enter the URL to scrape or ask a question regarding the data", key="user_input")
 
     if prompt:
-        if st.session_state.enable_logging:
-            logger.debug(f"Received prompt: {prompt}")
         st.session_state.chat_history[st.session_state.current_chat_id]["messages"].append({"role": "user", "content": prompt})
-        save_chat_history(st.session_state.chat_history)
         
         if not st.session_state.web_scraper_chat:
-            if st.session_state.enable_logging:
-                logger.debug("Initializing web_scraper_chat")
             st.session_state.web_scraper_chat = initialize_web_scraper_chat()
 
         with st.chat_message("assistant"):
@@ -313,12 +380,14 @@ def main():
                     st.session_state.web_scraper_chat,
                     prompt
                 )
+                st.write("Debug: Full response type:", type(full_response))
                 if full_response is not None:
-                    st.session_state.chat_history[st.session_state.current_chat_id]["messages"].append({"role": "assistant", "content": full_response})
+                    if isinstance(full_response, tuple) and len(full_response) == 2 and isinstance(full_response[1], BytesIO):
+                        st.session_state.chat_history[st.session_state.current_chat_id]["messages"].append({"role": "assistant", "content": full_response[0]})
+                    else:
+                        st.session_state.chat_history[st.session_state.current_chat_id]["messages"].append({"role": "assistant", "content": full_response})
                     save_chat_history(st.session_state.chat_history)
             except Exception as e:
-                if st.session_state.enable_logging:
-                    logger.error(f"An unexpected error occurred: {str(e)}")
                 st.error(f"An unexpected error occurred: {str(e)}")
             
             st.rerun()

src/ollama_models.py

@@ -1,19 +1,14 @@
 import requests
 from typing import List, Dict, Any
-import logging
 import os
 import json
 
 class OllamaModel:
     def __init__(self, model_name: str):
         self.model_name = model_name
-        self.logger = logging.getLogger(__name__)
-        self.logger.setLevel(logging.DEBUG)
         self.base_url = os.getenv('OLLAMA_BASE_URL', 'http://localhost:11434')
 
     async def generate(self, prompt: str, system_prompt: str = "") -> str:
-        self.logger.debug(f"Generating with Ollama model: {self.model_name}")
-        self.logger.debug(f"Prompt (first 500 chars): {prompt[:500]}...")
         try:
             response = requests.post(
                 f"{self.base_url}/api/generate",
@@ -35,26 +30,22 @@ class OllamaModel:
                         if 'response' in data:
                             full_response += data['response']
                     except json.JSONDecodeError:
-                        self.logger.warning(f"Failed to parse JSON: {line}")
+                        print(f"Error decoding JSON: {line}")
             
-            self.logger.debug(f"Ollama response (first 500 chars): {full_response[:500]}...")
             return full_response
         except Exception as e:
-            self.logger.error(f"Error generating with Ollama: {str(e)}")
+            print(f"An error occurred: {str(e)}")
             raise
 
     @staticmethod
     async def list_models() -> List[str]:
-        logger = logging.getLogger(__name__)
         base_url = os.getenv('OLLAMA_BASE_URL', 'http://localhost:11434')
         try:
             response = requests.get(f"{base_url}/api/tags")
             response.raise_for_status()
             models = response.json()
-            logger.debug(f"Available Ollama models: {models['models']}")
             return [model['name'] for model in models['models']]
         except Exception as e:
-            logger.error(f"Error listing Ollama models: {str(e)}")
             return []
 
 class OllamaModelManager:

src/utils/google_sheets_utils.py

@@ -8,11 +8,9 @@ import pandas as pd
 from datetime import datetime
 import os
 import json
-import logging
 import hashlib
-
-logging.basicConfig(level=logging.DEBUG)
-logger = logging.getLogger(__name__)
+import re
+from io import BytesIO
 
 SCOPES = ['https://www.googleapis.com/auth/spreadsheets', 'https://www.googleapis.com/auth/drive.file']
 TOKEN_FILE = 'token.json'
@@ -38,18 +36,16 @@ def get_google_sheets_credentials():
     if os.path.exists(TOKEN_FILE):
         try:
             creds = Credentials.from_authorized_user_file(TOKEN_FILE, SCOPES)
-            logger.debug("Loaded credentials from token file")
         except Exception as e:
-            logger.error(f"Error loading credentials from token file: {str(e)}")
+            print(f"Error loading credentials from file: {str(e)}")
     
     if not creds or not creds.valid:
         if creds and creds.expired and creds.refresh_token:
             try:
                 creds.refresh(Request())
-                logger.debug("Refreshed expired credentials")
                 save_credentials(creds)
             except Exception as e:
-                logger.error(f"Error refreshing credentials: {str(e)}")
+                print(f"Error refreshing credentials: {str(e)}")
                 creds = None
         else:
             creds = None
@@ -58,31 +54,41 @@ def get_google_sheets_credentials():
         if 'google_auth_token' in st.session_state:
             try:
                 creds = Credentials.from_authorized_user_info(json.loads(st.session_state['google_auth_token']), SCOPES)
-                logger.debug("Loaded credentials from session state")
                 save_credentials(creds)
             except Exception as e:
-                logger.error(f"Error loading credentials from session state: {str(e)}")
-    
+                print(f"Error creating credentials from session state: {str(e)}")
     return creds
 
 def save_credentials(creds):
     try:
         with open(TOKEN_FILE, 'w') as token:
             token.write(creds.to_json())
-        logger.debug("Saved credentials to token file")
     except Exception as e:
-        logger.error(f"Error saving credentials to token file: {str(e)}")
+        print(f"Error saving credentials: {str(e)}")
+
+def clean_data_for_sheets(df):
+    def clean_value(val):
+        if pd.isna(val):
+            return ""
+        if isinstance(val, (int, float)):
+            return str(val)
+        return str(val).replace('\n', ' ').replace('\r', '')
+
+    for col in df.columns:
+        df[col] = df[col].map(clean_value)
+
+    if 'comments' in df.columns:
+        df['comments'] = df['comments'].astype(str)
 
-def upload_to_google_sheets(df):
+    return df
+
+def upload_to_google_sheets(data):
     creds = get_google_sheets_credentials()
     if not creds:
-        logger.error("Failed to obtain valid credentials.")
         return None
 
     try:
-        service = build('sheets', 'v4', credentials=creds)
-        logger.debug("Built Sheets service")
-        
+        service = build('sheets', 'v4', credentials=creds)        
         spreadsheet = {
             'properties': {
                 'title': f"CyberScraper Data {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}"
@@ -90,38 +96,38 @@ def upload_to_google_sheets(df):
         }
         spreadsheet = service.spreadsheets().create(body=spreadsheet, fields='spreadsheetId').execute()
         spreadsheet_id = spreadsheet.get('spreadsheetId')
-        logger.debug(f"Created new spreadsheet with ID: {spreadsheet_id}")
+
+        if isinstance(data, pd.DataFrame):
+            df = clean_data_for_sheets(data)
+        else:
+            return None
 
         values = [df.columns.tolist()] + df.values.tolist()
         body = {'values': values}
         result = service.spreadsheets().values().update(
             spreadsheetId=spreadsheet_id, range='Sheet1',
             valueInputOption='RAW', body=body).execute()
-        logger.debug(f"Updated spreadsheet. Cells updated: {result.get('updatedCells')}")
-
         return spreadsheet_id
     except HttpError as error:
-        logger.error(f"An HTTP error occurred: {error}")
+        print(f"An HTTP error occurred: {error}")
         return None
     except Exception as e:
-        logger.error(f"An unexpected error occurred: {str(e)}")
+        print(f"An error occurred: {str(e)}")
         return None
 
-def display_google_sheets_button(df):
-    df_hash = hash(str(df))
-
+def display_google_sheets_button(data, unique_key):
     creds = get_google_sheets_credentials()
     if not creds:
         auth_button = '🔑 Authorize Google Sheets'
-        if st.button(auth_button, key=f"auth_sheets_{df_hash}", help="Authorize access to Google Sheets"):
+        if st.button(auth_button, key=f"auth_sheets_{unique_key}", help="Authorize access to Google Sheets"):
             initiate_google_auth()
     else:
         upload_button = '✅ Upload to Google Sheets'
-        if st.button(upload_button, key=f"upload_{df_hash}", help="Upload data to Google Sheets"):
+        if st.button(upload_button, key=f"upload_{unique_key}", help="Upload data to Google Sheets"):
             with st.spinner("Uploading to Google Sheets..."):
-                spreadsheet_id = upload_to_google_sheets(df)
+                spreadsheet_id = upload_to_google_sheets(data)
                 if spreadsheet_id:
                     st.success(f"Data uploaded successfully. Spreadsheet ID: {spreadsheet_id}")
                     st.markdown(f"[Open Spreadsheet](https://docs.google.com/spreadsheets/d/{spreadsheet_id})")
                 else:
-                    st.error("Failed to upload data to Google Sheets.")
+                    st.error("Failed to upload data to Google Sheets. Check the console for error details.")

src/web_extractor.py

@@ -2,7 +2,8 @@ import asyncio
 from typing import Dict, Any, Optional, List, Tuple
 import json
 import pandas as pd
-from io import StringIO
+from io import StringIO, BytesIO
+import base64
 import re
 from functools import lru_cache
 import hashlib
@@ -17,7 +18,6 @@ from langchain.prompts import PromptTemplate
 from langchain.schema.runnable import RunnableSequence
 from langchain.text_splitter import RecursiveCharacterTextSplitter
 import tiktoken
-import logging
 import csv
 from bs4 import BeautifulSoup, Comment
 
@@ -46,8 +46,6 @@ class WebExtractor:
             length_function=self.num_tokens_from_string,
         )
         self.max_tokens = 128000 if model_name == "gpt-4o-mini" else 16385
-        self.logger = logging.getLogger(__name__)
-        self.logger.setLevel(logging.DEBUG)
         self.query_cache = {}
         self.content_hash = None
 
@@ -156,9 +154,7 @@ class WebExtractor:
 
         return text
 
-    async def _extract_info(self, query: str) -> str:
-        self.logger.debug(f"Extracting info with model: {self.model}")
-        
+    async def _extract_info(self, query: str) -> str:        
         if not self.preprocessed_content:
             return "Please provide a URL first before asking for information."
 
@@ -179,15 +175,12 @@ class WebExtractor:
             extracted_data = await self._cached_api_call(content_hash, query)
         else:
             chunks = self.optimized_text_splitter(self.preprocessed_content)
-            self.logger.debug(f"Content split into {len(chunks)} chunks")
             all_extracted_data = []
             for i, chunk in enumerate(chunks):
                 chunk_data = await self._cached_api_call(self._hash_content(chunk), query)
                 all_extracted_data.append(chunk_data)
             extracted_data = self._merge_json_chunks(all_extracted_data)
 
-        self.logger.debug(f"Extracted data (first 500 chars): {extracted_data[:500]}...")
-
         formatted_result = self._format_result(extracted_data, query)
         self.query_cache[cache_key] = formatted_result
         return formatted_result
@@ -199,7 +192,7 @@ class WebExtractor:
             csv_string, df = self._format_as_csv(extracted_data)
             return f"```csv\n{csv_string}\n```", df
         elif 'excel' in query.lower():
-            return self._format_as_excel_and_save(extracted_data)
+            return self._format_as_excel(extracted_data)
         elif 'sql' in query.lower():
             return self._format_as_sql(extracted_data)
         elif 'html' in query.lower():
@@ -220,7 +213,7 @@ class WebExtractor:
                 else:
                     merged_data.append(data)
             except json.JSONDecodeError:
-                self.logger.error(f"Failed to parse JSON chunk: {chunk[:100]}...")
+                print(f"Error decoding JSON chunk: {chunk[:100]}...")
         return json.dumps(merged_data)
 
     def _format_as_json(self, data: str) -> str:
@@ -260,15 +253,13 @@ class WebExtractor:
             
             return csv_string, df
         except json.JSONDecodeError as e:
-            self.logger.error(f"JSON Decode Error: {str(e)}")
             error_msg = f"Error: Invalid JSON data. Raw data: {data[:500]}..."
             return error_msg, pd.DataFrame()
         except Exception as e:
-            self.logger.error(f"Unexpected error in _format_as_csv: {str(e)}")
             error_msg = f"Error: Failed to convert data to CSV. {str(e)}"
             return error_msg, pd.DataFrame()
 
-    def _format_as_excel_and_save(self, data: str) -> str:
+    def _format_as_excel(self, data: str) -> Tuple[BytesIO, pd.DataFrame]:
         json_pattern = r'```json\s*([\s\S]*?)\s*```'
         match = re.search(json_pattern, data)
         if match:
@@ -276,17 +267,21 @@ class WebExtractor:
         try:
             parsed_data = json.loads(data)
             if not parsed_data:
-                return "No data to convert to Excel."
+                return BytesIO(b"No data to convert to Excel."), pd.DataFrame()
             
             df = pd.DataFrame(parsed_data)
-            output_filename = "output.xlsx"
-            with pd.ExcelWriter(output_filename, engine='xlsxwriter') as writer:
-                df.to_excel(writer, index=False)
-            return f"Excel data saved to {output_filename}"
+            excel_buffer = BytesIO()
+            with pd.ExcelWriter(excel_buffer, engine='xlsxwriter') as writer:
+                df.to_excel(writer, index=False, sheet_name='Sheet1')
+            excel_buffer.seek(0)
+            
+            return excel_buffer, df
         except json.JSONDecodeError:
-            return f"Error: Invalid JSON data. Raw data: {data[:500]}..."
+            error_msg = f"Error: Invalid JSON data. Raw data: {data[:500]}..."
+            return BytesIO(error_msg.encode()), pd.DataFrame()
         except Exception as e:
-            return f"Error: Failed to convert data to Excel. {str(e)}"
+            error_msg = f"Error: Failed to convert data to Excel. {str(e)}"
+            return BytesIO(error_msg.encode()), pd.DataFrame()
 
     def _format_as_sql(self, data: str) -> str:
         json_pattern = r'```json\s*([\s\S]*?)\s*```'