Update app.py

app.py

@@ -1,6 +1,7 @@
 # app.py
 
 import os
+import re
 import requests
 import uuid
 import datetime
@@ -10,7 +11,8 @@ import shutil
 import secrets
 import time
 import json
-from typing import List, Tuple, Any, Dict
+from typing import List, Tuple, Any, Dict, Set, Generator
+from urllib.parse import urljoin, urlparse
 
 # Third-party libraries
 import gradio as gr
@@ -24,34 +26,29 @@ from agent_prompt import PromptLibrary, SystemAuditor
 
 # --- CONFIGURATION ---
 class Config:
-    """Centralized configuration for the Maestro application."""
     HF_MODEL = os.getenv("HF_MODEL", "mistralai/Mixtral-8x7B-Instruct-v0.1")
     HF_TOKEN = os.getenv("HF_TOKEN")
     VERBOSE = os.getenv("VERBOSE", "True").lower() == "true"
-    MAX_NEW_TOKENS_REPORT = 4096
-    MAX_NEW_TOKENS_CHAT = 1024
+    MAX_NEW_TOKENS_REPORT = 8192
+    MAX_NEW_TOKENS_CHAT = 2048
     REQUESTS_TIMEOUT = 15
     USER_AGENT = 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/125.0.0.0 Safari/537.36'
-
+    
 # --- UTILITIES ---
 def log(message: str) -> None:
     if Config.VERBOSE:
-        print(f"[LOG] {datetime.datetime.now(datetime.timezone.utc).isoformat()} - {message}")
+        print(f"[{datetime.datetime.now(datetime.timezone.utc).isoformat()}] {message}")
 
 class SessionManager:
-    """A context manager for creating and cleaning up session-specific temporary directories."""
+    # ... (No changes from previous version)
     def __init__(self, session_id: str):
         self.session_id = session_id
         self.temp_dir = os.path.join(tempfile.gettempdir(), f"session_{session_id}")
-
     def __enter__(self) -> str:
         os.makedirs(self.temp_dir, exist_ok=True)
-        log(f"Session '{self.session_id}' started. Temp dir: {self.temp_dir}")
         return self.temp_dir
-
     def __exit__(self, exc_type, exc_val, exc_tb):
         shutil.rmtree(self.temp_dir, ignore_errors=True)
-        log(f"Session '{self.session_id}' ended. Temp dir cleaned up.")
 
 # --- CORE APPLICATION ENGINE ---
 class MaestroEngine:
@@ -65,104 +62,174 @@ class MaestroEngine:
             nltk.download('punkt', quiet=True)
         log("MaestroEngine initialized.")
 
+    # --- File Processors ---
     def _read_pdf(self, file_path: str) -> str:
         try:
             reader = PdfReader(file_path)
             return "\n".join(page.extract_text() or "" for page in reader.pages)
-        except Exception as e:
-            log(f"Error reading PDF {os.path.basename(file_path)}: {e}")
-            return f"Error reading PDF: {e}"
+        except Exception as e: return f"Error reading PDF {os.path.basename(file_path)}: {e}"
 
     def _process_zip(self, zip_path: str, temp_dir: str) -> str:
+        # ... (Identical to previous correct version)
         extracted_texts = []
         extract_path = os.path.join(temp_dir, "zip_extract")
         os.makedirs(extract_path, exist_ok=True)
         try:
             with zipfile.ZipFile(zip_path, 'r') as zf:
                 for member in zf.infolist():
-                    if member.filename.endswith('.pdf'):
-                        zf.extract(member, extract_path)
-                        extracted_texts.append(self._read_pdf(os.path.join(extract_path, member.filename)))
+                    if member.is_dir() or member.filename.startswith('__MACOSX'): continue
+                    member_path = zf.extract(member, path=extract_path)
+                    if member.filename.endswith('.pdf'): extracted_texts.append(self._read_pdf(member_path))
                     elif member.filename.endswith('.txt'):
-                        extracted_texts.append(zf.read(member).decode('utf-8', errors='ignore'))
+                         with open(member_path, 'r', encoding='utf-8', errors='ignore') as f: extracted_texts.append(f.read())
             return "\n\n".join(extracted_texts)
+        except Exception as e: return f"Error processing ZIP {os.path.basename(zip_path)}: {e}"
+        
+    # --- Google Workspace Integration ---
+    def _get_google_drive_id(self, url: str) -> str | None:
+        match = re.search(r"/file/d/([a-zA-Z0-9_-]+)", url)
+        return match.group(1) if match else None
+
+    def _download_google_drive_file(self, file_id: str, temp_dir: str) -> Tuple[str | None, str | None]:
+        URL = "https://docs.google.com/uc?export=download&id="
+        try:
+            response = requests.get(URL + file_id, stream=True, timeout=30)
+            response.raise_for_status()
+            disposition = response.headers.get('content-disposition', '')
+            fname = re.findall('filename="(.+)"', disposition)
+            filename = fname[0] if fname else f"{file_id}_download"
+            file_path = os.path.join(temp_dir, filename)
+            with open(file_path, "wb") as f:
+                for chunk in response.iter_content(chunk_size=8192):
+                    f.write(chunk)
+            return file_path, None
         except Exception as e:
-            log(f"Error processing ZIP {os.path.basename(zip_path)}: {e}")
-            return f"Error processing ZIP: {e}"
+            return None, f"Failed to download Google Drive file ({file_id}): {e}"
 
-    def process_data_sources(self, session_id: str, url: str, text: str, files: List[Any]) -> Tuple[str, List[str]]:
-        """Orchestrates data ingestion from all provided sources."""
+    def _fetch_google_doc(self, url: str) -> Tuple[str | None, str | None]:
+        try:
+            doc_id = re.search(r"/document/d/([a-zA-Z0-9_-]+)", url).group(1)
+            export_url = f"https://docs.google.com/document/d/{doc_id}/export?format=txt"
+            response = requests.get(export_url, timeout=Config.REQUESTS_TIMEOUT)
+            response.raise_for_status()
+            return response.text, None
+        except Exception as e: return None, f"Failed to fetch Google Doc: {e}"
+
+    # --- Web Crawler ---
+    def _crawl_website(self, start_url: str, max_depth: int) -> Tuple[str, List[str]]:
+        if not start_url: return "", []
+        visited: Set[str] = set()
+        to_visit: List[Tuple[str, int]] = [(start_url, 0)]
+        all_text, errors = [], []
+        base_netloc = urlparse(start_url).netloc
+
+        while to_visit and len(visited) < 50: # Safety break
+            current_url, depth = to_visit.pop(0)
+            if current_url in visited or depth > max_depth: continue
+            
+            log(f"Crawling (depth {depth}): {current_url}")
+            visited.add(current_url)
+            
+            try:
+                response = requests.get(current_url, headers={'User-Agent': Config.USER_AGENT}, timeout=Config.REQUESTS_TIMEOUT)
+                response.raise_for_status()
+                soup = BeautifulSoup(response.content, 'html.parser')
+                all_text.append(soup.get_text(separator="\n", strip=True))
+                
+                if depth < max_depth:
+                    for link in soup.find_all('a', href=True):
+                        abs_url = urljoin(current_url, link['href'])
+                        if urlparse(abs_url).netloc == base_netloc and abs_url not in visited:
+                            to_visit.append((abs_url, depth + 1))
+            except Exception as e:
+                errors.append(f"Crawl Error on {current_url}: {e}")
+        return "\n\n".join(all_text), errors
+        
+    # --- Main Orchestrator ---
+    def process_data_sources(self, session_id: str, url: str, crawl_depth: int, text: str, file_paths: List[str]) -> Tuple[str, List[str]]:
         all_content, errors = [], []
         with SessionManager(session_id) as temp_dir:
-            if url:
-                try:
-                    response = requests.get(url, headers={'User-Agent': Config.USER_AGENT}, timeout=Config.REQUESTS_TIMEOUT)
-                    response.raise_for_status()
-                    all_content.append(BeautifulSoup(response.content, 'html.parser').get_text(separator="\n", strip=True))
-                except Exception as e:
-                    errors.append(f"URL Fetch Error: {e}")
-            if text:
-                all_content.append(text)
-            if files:
-                for file_obj in files:
-                    file_path = os.path.join(temp_dir, os.path.basename(file_obj.name))
-                    with open(file_path, "wb") as f:
-                        shutil.copyfileobj(file_obj, f)
-                    ext = os.path.splitext(file_obj.name)[1].lower()
-                    if ext == '.pdf':
-                        all_content.append(self._read_pdf(file_path))
+            # 1. Direct Text
+            if text: all_content.append(text)
+            
+            # 2. Uploaded Files
+            if file_paths:
+                for path in file_paths:
+                    if not path: continue
+                    filename = os.path.basename(path)
+                    ext = os.path.splitext(filename)[1].lower()
+                    if ext == '.pdf': all_content.append(self._read_pdf(path))
                     elif ext == '.txt':
-                        all_content.append(open(file_path, 'r', encoding='utf-8').read())
-                    elif ext == '.zip':
-                        all_content.append(self._process_zip(file_path, temp_dir))
-                    else:
-                        errors.append(f"Unsupported file type: {file_obj.name}")
+                        with open(path, 'r', encoding='utf-8', errors='ignore') as f: all_content.append(f.read())
+                    elif ext == '.zip': all_content.append(self._process_zip(path, temp_dir))
+                    else: errors.append(f"Unsupported file type: {filename}")
+            
+            # 3. URL Processing (Crawler, Google Drive, Google Docs)
+            if url:
+                if "drive.google.com/file/d/" in url:
+                    file_id = self._get_google_drive_id(url)
+                    if file_id:
+                        path, err = self._download_google_drive_file(file_id, temp_dir)
+                        if err: errors.append(err)
+                        elif path: all_content.append(self._read_pdf(path)) # Assuming PDF for simplicity
+                    else: errors.append(f"Invalid Google Drive URL: {url}")
+                elif "docs.google.com/document/d/" in url:
+                    content, err = self._fetch_google_doc(url)
+                    if err: errors.append(err)
+                    else: all_content.append(content)
+                else: # Standard web crawling
+                    content, crawl_errors = self._crawl_website(url, crawl_depth)
+                    all_content.append(content)
+                    errors.extend(crawl_errors)
+                    
         return "\n\n---\n\n".join(all_content), errors
 
-    def _query_llm(self, prompt: str, max_tokens: int) -> str:
+    # --- LLM Interaction ---
+    def _query_llm_stream(self, system_prompt: str, user_prompt: str, max_tokens: int) -> Generator[str, None, None]:
+        full_prompt = f"<s>[INST] {system_prompt}\n\n{user_prompt} [/INST]"
         try:
-            response = self.client.text_generation(prompt, max_new_tokens=max_tokens, temperature=0.7, top_p=0.95)
-            return response.strip()
+            for token in self.client.text_generation(full_prompt, max_new_tokens=max_tokens, stream=True, temperature=0.7, top_p=0.95):
+                yield token
         except Exception as e:
-            log(f"LLM query failed: {e}")
-            return f"Error communicating with the model: {e}"
+            log(f"LLM stream query failed: {e}")
+            yield f"Error communicating with the model: {e}"
 
-    def run_rag_query(self, query: str, context: str) -> str:
-        prompt = f"Context:\n---\n{context}\n---\nBased only on the context provided, answer the following question:\nQuestion: {query}"
-        return self._query_llm(prompt, Config.MAX_NEW_TOKENS_CHAT)
+    def run_rag_query(self, query: str, context: str, purpose: str) -> Generator[str, None, None]:
+        system_prompt = PromptLibrary.AGENT_PREFIX.format(
+            dynamic_timestamp_utc=datetime.datetime.now(datetime.timezone.utc).isoformat(),
+            user_purpose=purpose or "Direct Question & Answer"
+        )
+        user_prompt = f"Based *only* on the context provided below, answer the user's question.\n\nCONTEXT:\n---\n{context}\n---\n\nQUESTION: {query}"
+        yield from self._query_llm_stream(system_prompt, user_prompt, Config.MAX_NEW_TOKENS_CHAT)
 
     def generate_report(self, report_type: str, context: str, objective: str) -> str:
+        system_prompt = PromptLibrary.AGENT_PREFIX.format(
+            dynamic_timestamp_utc=datetime.datetime.now(datetime.timezone.utc).isoformat(),
+            user_purpose=objective
+        )
         if report_type == "Narrative Prose Report":
-            prompt = PromptLibrary.NARRATIVE_PROSE_REPORT.format(
-                task_objective=objective,
-                knowledge_base=context
-            )
-            return self._query_llm(prompt, Config.MAX_NEW_TOKENS_REPORT)
-        elif report_type == "Technical JSON Report":
-            prompt = PromptLibrary.TECHNICAL_JSON_REPORT.format(
-                task_objective=objective,
-                baseline_knowledge="Previously established facts",
-                new_information=context
-            )
-            raw_response = self._query_llm(prompt, Config.MAX_NEW_TOKENS_REPORT)
-            clean_json_str = raw_response.replace("```json", "").replace("```", "").strip()
-            try:
-                return json.dumps(json.loads(clean_json_str), indent=2)
-            except json.JSONDecodeError:
-                log(f"Failed to parse LLM response as JSON. Raw response: {raw_response}")
-                return '{"error": "The model did not return valid JSON.", "raw_response": ' + json.dumps(raw_response) + '}'
-        return "Invalid report type selected."
+            user_prompt = PromptLibrary.NARRATIVE_PROSE_REPORT.format(task_objective=objective, knowledge_base=context)
+        else: # Technical JSON Report
+            user_prompt = PromptLibrary.TECHNICAL_JSON_REPORT.format(task_objective=objective, knowledge_base=context)
+        
+        # Reports are not streamed
+        full_prompt = f"<s>[INST] {system_prompt}\n\n{user_prompt} [/INST]"
+        response = self.client.text_generation(full_prompt, max_new_tokens=Config.MAX_NEW_TOKENS_REPORT)
+        
+        if report_type == "Technical JSON Report":
+            clean_json_str = re.sub(r'```json\s*|\s*```', '', response).strip()
+            try: return json.dumps(json.loads(clean_json_str), indent=2)
+            except json.JSONDecodeError: return json.dumps({"error": "Model returned invalid JSON", "raw_response": response})
+        return response
 
 # --- GRADIO APPLICATION ---
 class GradioApp:
-    """Manages the Gradio UI and application workflow."""
     def __init__(self, engine: MaestroEngine):
         self.engine = engine
         self.app = self._build_ui()
 
     def _build_ui(self) -> gr.Blocks:
-        with gr.Blocks(theme=gr.themes.Monochrome(), title="Maestro AI Engine") as app:
-            # State management
+        with gr.Blocks(theme=gr.themes.Soft(primary_hue="emerald", secondary_hue="green"), title="Maestro AI Engine") as app:
             session_id = gr.State(lambda: secrets.token_hex(16))
             processed_data = gr.State("")
 
@@ -170,88 +237,81 @@ class GradioApp:
             
             with gr.Tabs():
                 with gr.TabItem("① Data Ingestion"):
+                    user_purpose = gr.Textbox(label="High-Level Goal / Purpose", placeholder="e.g., 'Research AI impact on agriculture for a market report'")
                     with gr.Row():
-                        with gr.Column():
-                            url_input = gr.Textbox(label="Scrape from URL")
-                            text_input = gr.Textbox(label="Paste Text", lines=10)
-                        file_upload = gr.Files(label="Upload Files (.pdf, .txt, .zip)", type="file")
+                        with gr.Column(scale=3):
+                            url_input = gr.Textbox(label="Ingest from URL", placeholder="Enter a standard URL, Google Drive, or Google Docs link")
+                            text_input = gr.Textbox(label="Paste Text", lines=8)
+                        with gr.Column(scale=2):
+                            crawl_depth = gr.Slider(label="Web Crawl Depth", minimum=0, maximum=5, value=1, step=1, info="For standard URLs only. 0=current page, 1=page + links, etc.")
+                            file_upload = gr.File(label="Upload Files (.pdf, .txt, .zip)", file_count="multiple", type="filepath")
                     process_button = gr.Button("🚀 Process All Sources", variant="primary")
-                    ingestion_summary = gr.Textbox(label="Ingestion Summary", interactive=False)
-                    error_log = gr.Textbox(label="Errors", interactive=False)
+                    ingestion_summary = gr.Textbox(label="Ingestion Summary", interactive=False, lines=2)
+                    error_log = gr.Textbox(label="Errors & Warnings", interactive=False, lines=3)
                 
                 with gr.TabItem("② Reporting & Synthesis"):
-                    report_objective = gr.Textbox(label="Report Objective", placeholder="e.g., 'Synthesize findings on AI in agriculture'")
-                    report_type = gr.Dropdown(label="Select Report Type", choices=["Narrative Prose Report", "Technical JSON Report"])
-                    generate_button = gr.Button("Generate Report", variant="primary")
+                    with gr.Row():
+                        report_type = gr.Dropdown(label="Select Report Type", choices=["Narrative Prose Report", "Technical JSON Report"], value="Narrative Prose Report")
+                        generate_button = gr.Button("Generate Report", variant="primary")
                     with gr.Tabs():
-                        with gr.TabItem("Narrative Output"):
-                            report_output_md = gr.Markdown()
-                        with gr.TabItem("JSON Output"):
-                            report_output_json = gr.JSON()
+                        with gr.TabItem("Narrative Output"): report_output_md = gr.Markdown()
+                        with gr.TabItem("JSON Output"): report_output_json = gr.JSON()
 
                 with gr.TabItem("③ Direct Chat Q&A"):
-                    chatbot = gr.Chatbot(label="Chat Interface", height=550)
-                    msg_input = gr.Textbox(label="Your Question", placeholder="Ask a question about the processed data...")
-                    msg_input.submit(self._chat_workflow, [msg_input, chatbot, processed_data], [msg_input, chatbot])
+                    chatbot = gr.Chatbot(label="Chat Interface", height=550, bubble_full_width=False)
+                    msg_input = gr.Textbox(label="Your Question", placeholder="Ask a question about the processed data...", scale=4)
+                    msg_input.submit(self._chat_workflow, [msg_input, chatbot, processed_data, user_purpose], [msg_input, chatbot])
 
-            # --- Workflow Connections ---
-            process_button.click(self._ingest_workflow, [session_id, url_input, text_input, file_upload], [processed_data, ingestion_summary, error_log])
-            
-            # **FIXED LINE**
-            # Pass session_id to the reporting workflow instead of the non-existent auditor state
-            generate_button.click(self._reporting_workflow, [session_id, report_type, processed_data, report_objective], [report_output_md, report_output_json])
+            process_button.click(self._ingest_workflow, [session_id, url_input, crawl_depth, text_input, file_upload, process_button], [processed_data, ingestion_summary, error_log, process_button])
+            generate_button.click(self._reporting_workflow, [session_id, report_type, processed_data, user_purpose, generate_button], [report_output_md, report_output_json, generate_button])
 
         return app
 
-    def _ingest_workflow(self, s_id, url, text, files):
+    def _ingest_workflow(self, s_id, url, depth, text, files, btn):
+        yield gr.update(value="⚙️ Processing...", interactive=False)
         log(f"Starting ingestion for session {s_id}...")
-        data, errors = self.engine.process_data_sources(s_id, url, text, files)
-        summary = f"Processing complete. {len(data)} characters ingested. {len(errors)} errors encountered."
-        return data, summary, "\n".join(errors)
+        data, errors = self.engine.process_data_sources(s_id, url, depth, text, files)
+        summary = f"Processing complete. {len(data):,} characters ingested from all sources. {len(errors)} errors encountered."
+        yield data, summary, "\n".join(errors), gr.update(value="🚀 Process All Sources", interactive=True)
 
-    def _chat_workflow(self, message, history, context):
+    def _chat_workflow(self, message, history, context, purpose):
         if not context:
             history.append((message, "Error: No data has been ingested. Please process data in Tab 1 first."))
             return "", history
-        response = self.engine.run_rag_query(message, context)
-        history.append((message, response))
-        return "", history
-
-    # **FIXED METHOD**
-    # The method now accepts session_id (s_id) and creates the auditor instance internally.
-    def _reporting_workflow(self, s_id: str, r_type: str, context: str, objective: str):
-        auditor = SystemAuditor(session_id=s_id) # Create auditor instance on-the-fly
+        
+        history.append((message, ""))
+        log(f"Starting RAG query. Purpose: {purpose}")
+        full_response = ""
+        for token in self.engine.run_rag_query(message, context, purpose):
+            full_response += token
+            history[-1] = (message, full_response)
+            yield "", history
 
+    def _reporting_workflow(self, s_id, r_type, context, objective, btn):
+        yield gr.update(value="Generating...", interactive=False), None, None
+        auditor = SystemAuditor(session_id=s_id)
         if not context:
-            md_out = "### Error: No data has been ingested. Please process data in Tab 1 first."
-            log(auditor.format_prompt_log(f"REPORTING FAILED: No context data. Objective: {objective}"))
-            return md_out, None
-        
+            md_out = "### Error: No data ingested. Please process sources in Tab 1."
+            yield gr.update(value="Generate Report", interactive=True), md_out, None
+            return
+
         start_time = time.time()
-        log(auditor.format_prompt_log(f"Generating report of type '{r_type}' with objective: '{objective}'"))
+        log(auditor.format_prompt_log(f"Generating report: '{r_type}'"))
         response = self.engine.generate_report(r_type, context, objective)
         latency = (time.time() - start_time) * 1000
-
-        log(auditor.format_response_log(response, latency, 1, 0.95)) # Log the event
+        log(auditor.format_response_log(response, latency, 1, 0.95))
 
         if r_type == "Narrative Prose Report":
-            return response, None
-        elif r_type == "Technical JSON Report":
-            try:
-                # The engine already returns a JSON string or an error object
-                json_data = json.loads(response)
-                return None, json_data
-            except json.JSONDecodeError:
-                return None, {"error": "Failed to decode the final JSON response from the engine."}
+            yield gr.update(interactive=True), response, None
+        else: # Technical JSON Report
+            try: yield gr.update(interactive=True), None, json.loads(response)
+            except json.JSONDecodeError: yield gr.update(interactive=True), None, {"error": "Could not parse JSON from model."}
 
-    def launch(self):
-        self.app.launch(debug=Config.VERBOSE)
+    def launch(self): self.app.launch(debug=Config.VERBOSE, share=False)
 
-# --- MAIN EXECUTION BLOCK ---
 if __name__ == "__main__":
     if not Config.HF_TOKEN:
         print("FATAL: Hugging Face token (HF_TOKEN) not found in environment variables.")
-        print("Please set your token, e.g., `export HF_TOKEN='hf_...'`")
     else:
         log("Instantiating Maestro Engine and launching Gradio App...")
         maestro_engine = MaestroEngine()