""" MozeAI Document Studio - Complete Application Intelligent Document Workspace with AI Memory, Streaming, and Multi-Turn Context """ import streamlit as st from groq import Groq import requests import re import numpy as np from bs4 import BeautifulSoup import PyPDF2 import docx from io import StringIO, BytesIO import csv import json from datetime import datetime import pytz import time import hashlib from collections import defaultdict import os from difflib import unified_diff import uuid from typing import Dict, List, Optional, Callable, Any, Tuple from dataclasses import dataclass, field, asdict # ============================================================================ # DOCUMENT GENERATION LIBRARIES # ============================================================================ from pptx import Presentation from pptx.util import Inches, Pt from pptx.enum.text import PP_ALIGN from pptx.dml.color import RGBColor from docx import Document as WordDocument from docx.shared import Inches as DocInches, Pt as DocPt from docx.enum.text import WD_ALIGN_PARAGRAPH # ============================================================================ # DATA CLASSES FOR INTELLIGENCE CORE # ============================================================================ @dataclass class ParsedInstruction: """Structured representation of user instruction""" intent: str # "improve", "analyze", "transform", "generate", "create", "edit" target_audience: Optional[str] = None tone: str = "neutral" scope: str = "full" domain: str = "general" constraints: List[str] = field(default_factory=list) reasoning: str = "" confidence: float = 0.0 needs_clarification: bool = False clarification_questions: List[str] = field(default_factory=list) extracted_entities: Dict[str, Any] = field(default_factory=dict) @dataclass class EditPlan: """Execution plan for document editing""" strategy: str steps: List[str] = field(default_factory=list) constraints: List[str] = field(default_factory=list) target_metrics: Dict[str, Any] = field(default_factory=dict) rationale: str = "" estimated_tokens: int = 0 @dataclass class DocumentProfile: """Comprehensive document analysis result""" structure: Dict[str, Any] = field(default_factory=dict) content: Dict[str, Any] = field(default_factory=dict) quality: Dict[str, Any] = field(default_factory=dict) suggestions: List[str] = field(default_factory=list) strengths: List[str] = field(default_factory=list) metadata: Dict[str, Any] = field(default_factory=dict) @dataclass class EditResult: """Result of document edit operation""" edited_document: str = "" changes_made: Dict[str, Any] = field(default_factory=dict) reasoning: str = "" successful: bool = False streaming_complete: bool = True execution_time_ms: int = 0 @dataclass class ConversationTurn: """Single conversation turn with context""" user_query: str = "" document_snapshot: Dict[str, Any] = field(default_factory=dict) assistant_response: str = "" edits_made: Dict[str, Any] = field(default_factory=dict) timestamp: datetime = field(default_factory=datetime.now) turn_id: str = field(default_factory=lambda: hashlib.md5(str(time.time()).encode()).hexdigest()[:8]) # ============================================================================ # DOCUMENT WORKSPACE CLASS # ============================================================================ class DocumentWorkspace: """Manages the active document with version control and change tracking""" def __init__(self): self.current_document = { "id": str(uuid.uuid4()), "title": "Untitled Document", "content": "", "type": "text", "created_at": datetime.now().isoformat(), "modified_at": datetime.now().isoformat(), "versions": [], "changes": [], "metadata": { "word_count": 0, "char_count": 0, "reading_time": 0, "style": "general", "language": "english" } } self.version_history = [] self.pending_changes = [] self.suggestion_mode = False self.track_changes = True def update_document(self, new_content, change_description=""): """Update document with change tracking""" old_content = self.current_document["content"] if old_content == new_content: return False if self.track_changes: self.save_version(f"Before: {change_description}") changes = self._calculate_changes(old_content, new_content) self.current_document["content"] = new_content self.current_document["modified_at"] = datetime.now().isoformat() self._update_metadata() change_record = { "id": len(self.current_document["changes"]), "timestamp": datetime.now().isoformat(), "description": change_description, "changes": changes, "type": "edit" } self.current_document["changes"].append(change_record) if self.track_changes: self.save_version(f"After: {change_description}") return True def _calculate_changes(self, old_text, new_text): """Calculate specific changes between versions""" changes = [] old_lines = old_text.split('\n') new_lines = new_text.split('\n') diff = list(unified_diff(old_lines, new_lines, lineterm='')) for line in diff: if line.startswith('+') and not line.startswith('+++'): changes.append({"type": "addition", "text": line[1:]}) elif line.startswith('-') and not line.startswith('---'): changes.append({"type": "deletion", "text": line[1:]}) return changes def _update_metadata(self): """Update document metadata""" content = self.current_document["content"] words = len(content.split()) chars = len(content) self.current_document["metadata"]["word_count"] = words self.current_document["metadata"]["char_count"] = chars self.current_document["metadata"]["reading_time"] = max(1, words // 200) def save_version(self, description=""): """Save current state as version""" version = { "id": len(self.version_history), "timestamp": datetime.now().isoformat(), "content": self.current_document["content"], "description": description, "metadata": self.current_document["metadata"].copy() } self.version_history.append(version) if len(self.version_history) > 50: self.version_history = self.version_history[-50:] return version def restore_version(self, version_id): """Restore a previous version""" if version_id < len(self.version_history): version = self.version_history[version_id] self.update_document(version["content"], f"Restored version {version_id}") return True return False def analyze_document(self): """Perform comprehensive document analysis""" content = self.current_document["content"] analysis = { "structure": self._analyze_structure(), "readability": self._analyze_readability(), "grammar_issues": self._check_grammar(), "style_analysis": self._analyze_style(), "suggestions": self._generate_suggestions() } return analysis def _analyze_structure(self): """Analyze document structure""" content = self.current_document["content"] lines = content.split('\n') headings = [] paragraphs = 0 lists = 0 for line in lines: if line.strip().startswith('#'): headings.append(line.strip()) elif len(line.strip()) > 20: paragraphs += 1 elif line.strip().startswith(('-', '*', '•')): lists += 1 return { "headings": headings, "paragraph_count": paragraphs, "list_items": lists, "total_lines": len(lines) } def _analyze_readability(self): """Calculate readability scores""" content = self.current_document["content"] sentences = re.split(r'[.!?]+', content) words = content.split() if len(sentences) == 0 or len(words) == 0: return {"score": 0, "level": "Unknown"} avg_words_per_sentence = len(words) / len(sentences) if avg_words_per_sentence < 10: score = 90 level = "Very Easy" elif avg_words_per_sentence < 15: score = 70 level = "Easy" elif avg_words_per_sentence < 20: score = 50 level = "Medium" elif avg_words_per_sentence < 25: score = 30 level = "Difficult" else: score = 10 level = "Very Difficult" return {"score": score, "level": level} def _check_grammar(self): """Basic grammar checking""" content = self.current_document["content"].lower() issues = [] common_errors = [ (r'\b(i)\s+(am|is|are|was|were)\s+(\w+ed)\b', "Passive voice detected"), (r'\b(very|really|quite|extremely)\s+(\w+)\b', "Consider removing intensifier"), (r'\b(there is|there are)\s+(\w+)\s+that\b', "Wordy construction"), ] for pattern, message in common_errors: if re.search(pattern, content): issues.append(message) return issues[:5] def _analyze_style(self): """Analyze writing style""" content = self.current_document["content"] style = "general" if re.search(r'\b(according to|citation|reference|study|research)\b', content, re.I): style = "academic" elif re.search(r'\b(proposal|budget|timeline|deliverable|stakeholder)\b', content, re.I): style = "business" elif re.search(r'\b(algorithm|function|class|import|def|return)\b', content): style = "technical" elif re.search(r'\b(chapter|scene|character|dialogue)\b', content, re.I): style = "creative" return {"detected_style": style, "confidence": 0.8} def _generate_suggestions(self): """Generate improvement suggestions""" content = self.current_document["content"] suggestions = [] if len(content.split()) < 100: suggestions.append("Consider expanding the document with more details") structure = self._analyze_structure() if len(structure["headings"]) == 0 and len(content) > 500: suggestions.append("Add headings to improve document structure") readability = self._analyze_readability() if readability["score"] < 30: suggestions.append("Simplify sentences to improve readability") return suggestions # ============================================================================ # 1. CONVERSATION MANAGER # ============================================================================ class ConversationManager: """Maintains persistent, multi-turn conversation memory with document state tracking""" def __init__(self, max_history: int = 10): self.max_history = max_history self.turns: List[ConversationTurn] = [] self.intent_summary: Optional[str] = None self.cumulative_edits: Dict[str, Any] = { "total_edits": 0, "sections_affected": defaultdict(int), "first_interaction": None, "last_interaction": None } def add_user_message(self, query: str, document_state: dict) -> None: snapshot = { "word_count": document_state.get("word_count", 0), "char_count": document_state.get("char_count", 0), "title": document_state.get("title", "Untitled"), "content_preview": document_state.get("content", "")[:200], "has_content": bool(document_state.get("content", "")) } turn = ConversationTurn( user_query=query, document_snapshot=snapshot, assistant_response="", edits_made={} ) self.turns.append(turn) if self.cumulative_edits["first_interaction"] is None: self.cumulative_edits["first_interaction"] = datetime.now() self.cumulative_edits["last_interaction"] = datetime.now() if len(self.turns) > self.max_history: self.turns = self.turns[-self.max_history:] def add_assistant_message(self, response: str, edits_made: dict) -> None: if self.turns: self.turns[-1].assistant_response = response self.turns[-1].edits_made = edits_made self.cumulative_edits["total_edits"] += 1 for section in edits_made.get("sections_modified", []): self.cumulative_edits["sections_affected"][section] += 1 def get_conversation_context(self) -> str: if not self.turns: return "No previous conversation." context_parts = ["## Conversation History\n"] for i, turn in enumerate(self.turns[-self.max_history:], 1): context_parts.append(f"**Turn {i}:**") context_parts.append(f"User: \"{turn.user_query[:200]}\"") context_parts.append(f"Document: {turn.document_snapshot.get('title', 'Untitled')} " f"({turn.document_snapshot.get('word_count', 0)} words)") if turn.edits_made: changes_desc = ", ".join(turn.edits_made.get("key_changes", [])[:3]) if changes_desc: context_parts.append(f"Result: {changes_desc}") context_parts.append("") if self.cumulative_edits["total_edits"] > 1: context_parts.append(f"**Cumulative:** {self.cumulative_edits['total_edits']} edits across " f"{len(self.cumulative_edits['sections_affected'])} sections") return "\n".join(context_parts) def get_document_evolution(self) -> List[Dict]: evolution = [] for turn in self.turns: evolution.append({ "turn_id": turn.turn_id, "timestamp": turn.timestamp.isoformat(), "query": turn.user_query[:100], "document_state": turn.document_snapshot, "changes": turn.edits_made }) return evolution def summarize_intent(self, llm_client=None) -> str: if not self.turns: return "No conversation to summarize" queries = [turn.user_query for turn in self.turns[-5:]] if llm_client and len(queries) > 1: try: prompt = f"""Based on these user queries about document editing, what is the user's OVERARCHING intent? Queries: {chr(10).join(f'- {q}' for q in queries)} Summarize in one sentence what the user is trying to achieve:""" messages = [{"role": "user", "content": prompt}] response = llm_client.chat.completions.create( model="llama-3.3-70b-versatile", messages=messages, max_tokens=100, temperature=0.3 ) self.intent_summary = response.choices[0].message.content.strip() return self.intent_summary except Exception: pass keywords = [] for q in queries: words = q.lower().split()[:5] keywords.extend(words) unique_keywords = list(set(keywords))[:5] self.intent_summary = f"User is focused on: {', '.join(unique_keywords)}" return self.intent_summary def clear(self) -> None: self.turns = [] self.intent_summary = None self.cumulative_edits = { "total_edits": 0, "sections_affected": defaultdict(int), "first_interaction": None, "last_interaction": None } # ============================================================================ # 2. INSTRUCTION PARSER # ============================================================================ class InstructionParser: """Extract semantic meaning from user instructions using AI""" def __init__(self, llm_client=None): self.llm_client = llm_client self.confidence_threshold = 0.7 def parse(self, instruction: str, document: dict) -> ParsedInstruction: if self.llm_client: try: return self._parse_with_ai(instruction, document) except Exception as e: print(f"AI parsing failed: {e}") return self._parse_with_regex(instruction, document) def _parse_with_ai(self, instruction: str, document: dict) -> ParsedInstruction: system_prompt = """You are an instruction parser for a document editing AI. Given a user instruction, extract semantic intent. Output ONLY valid JSON with this structure: { "intent": "improve|analyze|transform|generate|create|edit", "target_audience": "string or null", "tone": "formal|casual|academic|persuasive|neutral", "scope": "full|introduction|section|conclusion|paragraph", "domain": "business|academic|technical|creative|general", "constraints": ["list of specific requirements"], "confidence": 0.95, "reasoning": "Why this interpretation?", "needs_clarification": false, "clarification_questions": ["Question if needed?"], "extracted_entities": {"key": "value"} }""" user_prompt = f"""Instruction: "{instruction}" Document context: {document.get('title', 'Untitled')} ({document.get('word_count', 0)} words) Parse this instruction and include clarification_questions if the instruction is ambiguous (confidence < 0.7):""" try: messages = [ {"role": "system", "content": system_prompt}, {"role": "user", "content": user_prompt} ] response = self.llm_client.chat.completions.create( model="llama-3.3-70b-versatile", messages=messages, max_tokens=500, temperature=0.2 ) result_text = response.choices[0].message.content.strip() json_match = re.search(r'\{.*\}', result_text, re.DOTALL) if json_match: data = json.loads(json_match.group()) return ParsedInstruction( intent=data.get("intent", "edit"), target_audience=data.get("target_audience"), tone=data.get("tone", "neutral"), scope=data.get("scope", "full"), domain=data.get("domain", "general"), constraints=data.get("constraints", []), reasoning=data.get("reasoning", ""), confidence=data.get("confidence", 0.5), needs_clarification=data.get("needs_clarification", False), clarification_questions=data.get("clarification_questions", []), extracted_entities=data.get("extracted_entities", {}) ) except Exception as e: print(f"AI parsing error: {e}") return self._parse_with_regex(instruction, document) def _parse_with_regex(self, instruction: str, document: dict) -> ParsedInstruction: inst_lower = instruction.lower() intent = "edit" if any(word in inst_lower for word in ["improve", "enhance", "better"]): intent = "improve" elif any(word in inst_lower for word in ["analyze", "review", "check"]): intent = "analyze" elif any(word in inst_lower for word in ["transform", "convert", "change to"]): intent = "transform" elif any(word in inst_lower for word in ["generate", "create", "make"]): intent = "generate" tone = "neutral" if any(word in inst_lower for word in ["formal", "professional", "business"]): tone = "formal" elif any(word in inst_lower for word in ["casual", "friendly", "conversational"]): tone = "casual" elif any(word in inst_lower for word in ["academic", "scholarly", "research"]): tone = "academic" elif any(word in inst_lower for word in ["persuasive", "convincing", "compelling"]): tone = "persuasive" target_audience = None audience_patterns = [ (r"for\s+a\s+(\d+[\s-]*year[\s-]*old)", "child"), (r"for\s+(executives|leaders|managers)", "executive"), (r"for\s+(beginners|novices)", "beginner"), (r"for\s+(experts|professionals)", "expert"), ] for pattern, audience_type in audience_patterns: match = re.search(pattern, inst_lower) if match: target_audience = match.group(1) if match.groups() else audience_type break scope = "full" if "introduction" in inst_lower: scope = "introduction" elif "conclusion" in inst_lower: scope = "conclusion" elif re.search(r'section\s+(\d+)', inst_lower): scope = f"section:{re.search(r'section\s+(\d+)', inst_lower).group(1)}" constraints = [] word_match = re.search(r'under\s+(\d+)\s+words', inst_lower) if word_match: constraints.append(f"keep under {word_match.group(1)} words") if "keep accuracy" in inst_lower: constraints.append("preserve technical accuracy") confidence = 0.5 if intent != "edit": confidence += 0.1 if tone != "neutral": confidence += 0.1 if constraints: confidence += 0.1 confidence = min(confidence, 0.9) # FIX #9: Generate clarification questions when ambiguous clarification_questions = [] needs_clarification = confidence < self.confidence_threshold if needs_clarification: if tone == "neutral": clarification_questions.append("What tone should I use? (formal, casual, academic, persuasive)") if target_audience is None: clarification_questions.append("Who is the target audience for this document?") if scope == "full" and len(instruction.split()) < 5: clarification_questions.append("Should I edit the full document or a specific section?") if not constraints: clarification_questions.append("Are there any length or style constraints I should follow?") return ParsedInstruction( intent=intent, target_audience=target_audience, tone=tone, scope=scope, domain="general", constraints=constraints, reasoning="Parsed using pattern matching", confidence=confidence, needs_clarification=needs_clarification, clarification_questions=clarification_questions[:3] ) # ============================================================================ # 3. FILE CONTEXT ACCUMULATOR # ============================================================================ class FileContextAccumulator: """Remember all uploaded files and cross-reference them""" def __init__(self, llm_client=None): self.llm_client = llm_client self.files: Dict[str, Dict] = {} self.file_summaries: Dict[str, str] = {} self.semantic_index: Dict[str, List[str]] = defaultdict(list) def add_file(self, filename: str, file_type: str, content: str, metadata: dict) -> None: summary = self._generate_file_summary(filename, file_type, content, metadata) keywords = self._extract_keywords(content, metadata) self.files[filename] = { "filename": filename, "type": file_type, "content": content[:3000], "metadata": metadata, "summary": summary, "keywords": keywords, "timestamp": datetime.now().isoformat() } self.file_summaries[filename] = summary for keyword in keywords: self.semantic_index[keyword].append(filename) def _generate_file_summary(self, filename: str, file_type: str, content: str, metadata: dict) -> str: if file_type == "csv": lines = content.strip().split('\n') if len(lines) > 1: headers = lines[0].split(',') return f"CSV with {len(lines)-1} data rows, columns: {', '.join(headers[:5])}" elif file_type == "json": try: data = json.loads(content[:1000]) if isinstance(data, dict): return f"JSON object with keys: {', '.join(list(data.keys())[:5])}" elif isinstance(data, list): return f"JSON array with {len(data)} items" except: pass words = len(content.split()) return f"File with {words} words. Type: {file_type}" def _extract_keywords(self, content: str, metadata: dict) -> List[str]: keywords = set() if "columns" in metadata: keywords.update(metadata["columns"]) words = content.lower().split()[:200] common_words = {"the", "a", "an", "and", "or", "but", "in", "on", "at", "to", "for"} for word in words: if len(word) > 3 and word not in common_words: keywords.add(word) return list(keywords)[:20] def get_file_context(self) -> str: if not self.files: return "No files uploaded." context_parts = ["## Uploaded Files Context\n"] for filename, file_info in self.files.items(): context_parts.append(f"**File:** {filename}") context_parts.append(f"Type: {file_info['type']}") context_parts.append(f"Summary: {file_info['summary']}") context_parts.append("") return "\n".join(context_parts) def get_detailed_file_context(self, filename: str = None) -> str: if filename and filename in self.files: file_info = self.files[filename] return f"""## File: {filename} Type: {file_info['type']} Summary: {file_info['summary']} Content Preview: {file_info['content'][:500]} """ result = "" for filename, file_info in self.files.items(): result += f"\n### {filename}\n{file_info['summary']}\n" return result or "No files uploaded." def find_relevant_file(self, query: str) -> Optional[str]: query_lower = query.lower() best_match = None best_score = 0 for filename, file_info in self.files.items(): score = 0 for keyword in file_info["keywords"]: if keyword in query_lower: score += 1 if filename.lower() in query_lower: score += 2 if any(word in query_lower for word in file_info["summary"].lower().split()[:10]): score += 1 if score > best_score and score > 0: best_score = score best_match = filename return best_match def get_cross_references(self, document_content: str) -> List[Dict]: suggestions = [] for filename, file_info in self.files.items(): doc_lower = document_content.lower() file_keywords = file_info["keywords"][:5] matched_keywords = [kw for kw in file_keywords if kw in doc_lower] if matched_keywords: suggestions.append({ "file": filename, "type": file_info["type"], "matched_terms": matched_keywords, "suggestion": f"Reference data from {filename} regarding {', '.join(matched_keywords[:3])}" }) return suggestions def clear(self) -> None: self.files.clear() self.file_summaries.clear() self.semantic_index.clear() # ============================================================================ # 4. EDIT PLANNER # ============================================================================ class EditPlanner: """Plan document transformations before executing them""" def __init__(self, llm_client=None): self.llm_client = llm_client def plan(self, parsed: ParsedInstruction, document: dict, conversation=None) -> EditPlan: if self.llm_client: try: return self._plan_with_ai(parsed, document, conversation) except Exception as e: print(f"AI planning failed: {e}") return self._plan_with_templates(parsed, document, conversation) def _plan_with_ai(self, parsed: ParsedInstruction, document: dict, conversation) -> EditPlan: conv_context = "" if conversation: conv_context = conversation.get_conversation_context() system_prompt = """You are an edit planner for a document AI. Create a detailed execution plan. Output JSON: { "strategy": "Overall approach description", "steps": ["Step 1", "Step 2", "Step 3"], "constraints": ["Constraint 1", "Constraint 2"], "target_metrics": {"metric": "value"}, "rationale": "Why this approach" }""" user_prompt = f"""Parsed Instruction: - Intent: {parsed.intent} - Audience: {parsed.target_audience} - Tone: {parsed.tone} - Scope: {parsed.scope} Document: {document.get('title', 'Untitled')} ({document.get('word_count', 0)} words) {conv_context} Create edit plan:""" try: messages = [ {"role": "system", "content": system_prompt}, {"role": "user", "content": user_prompt} ] response = self.llm_client.chat.completions.create( model="llama-3.3-70b-versatile", messages=messages, max_tokens=800, temperature=0.4 ) result_text = response.choices[0].message.content.strip() json_match = re.search(r'\{.*\}', result_text, re.DOTALL) if json_match: data = json.loads(json_match.group()) return EditPlan( strategy=data.get("strategy", "Apply requested edits"), steps=data.get("steps", ["Analyze document", "Apply changes", "Verify result"]), constraints=data.get("constraints", parsed.constraints), target_metrics=data.get("target_metrics", {}), rationale=data.get("rationale", "Based on user instruction") ) except Exception as e: print(f"AI planning error: {e}") return self._plan_with_templates(parsed, document, conversation) def _plan_with_templates(self, parsed: ParsedInstruction, document: dict, conversation) -> EditPlan: intent_plans = { "improve": { "strategy": "Enhance document quality by improving clarity, flow, and engagement", "steps": ["Identify areas needing improvement", "Rewrite for better clarity", "Enhance vocabulary", "Ensure consistent tone"] }, "analyze": { "strategy": "Perform comprehensive document analysis without modifying content", "steps": ["Analyze document structure", "Evaluate content quality", "Check for grammar issues", "Generate recommendations"] }, "transform": { "strategy": "Transform document style and tone according to requirements", "steps": ["Understand target style", "Rewrite to match desired tone", "Adjust vocabulary", "Preserve core meaning"] }, "generate": { "strategy": "Generate new content based on document context", "steps": ["Analyze existing content", "Identify gaps", "Generate relevant content", "Integrate smoothly"] } } plan_template = intent_plans.get(parsed.intent, intent_plans["improve"]) steps = plan_template["steps"].copy() if parsed.tone != "neutral": steps.append(f"Adjust content to {parsed.tone} tone") if parsed.scope != "full": steps.insert(1, f"Focus exclusively on {parsed.scope} section") constraints = parsed.constraints.copy() if parsed.target_audience: constraints.append(f"Target audience: {parsed.target_audience}") return EditPlan( strategy=plan_template["strategy"], steps=steps, constraints=constraints, target_metrics={"preserve_facts": True}, rationale=f"Template-based plan for {parsed.intent} operation" ) # ============================================================================ # 5. STREAMING RESPONSE HANDLER # ============================================================================ class StreamingResponseHandler: """Stream responses token-by-token instead of blocking""" def __init__(self, client): self.client = client # Item 1: Expose last-run speed metrics for UI display self.last_metrics: Dict[str, Any] = {} def stream_completion( self, messages: List[Dict], on_token: Optional[Callable[[str], None]] = None, model: str = "llama-3.3-70b-versatile", temperature: float = 0.3, max_tokens: int = 4000 ) -> str: full_content = "" token_count = 0 start_time = time.time() first_token_time: Optional[float] = None try: stream = self.client.chat.completions.create( model=model, messages=messages, temperature=temperature, max_tokens=max_tokens, stream=True, timeout=60 ) for chunk in stream: if chunk.choices and chunk.choices[0].delta.content: token = chunk.choices[0].delta.content if first_token_time is None: first_token_time = time.time() full_content += token token_count += 1 if on_token: on_token(token) elapsed_ms = int((time.time() - start_time) * 1000) ttft_ms = int((first_token_time - start_time) * 1000) if first_token_time else 0 tokens_per_sec = round(token_count / max(elapsed_ms / 1000, 0.001), 1) # Item 1: Store metrics for caller to display self.last_metrics = { "elapsed_ms": elapsed_ms, "ttft_ms": ttft_ms, "token_count": token_count, "tokens_per_sec": tokens_per_sec, "char_count": len(full_content), } print(f"Streaming: {token_count} tokens, {tokens_per_sec} tok/s, TTFT {ttft_ms}ms") return full_content except Exception as e: print(f"Streaming error: {e}") try: response = self.client.chat.completions.create( model=model, messages=messages, temperature=temperature, max_tokens=max_tokens, stream=False ) full_content = response.choices[0].message.content.strip() self.last_metrics = {"elapsed_ms": 0, "ttft_ms": 0, "token_count": 0, "tokens_per_sec": 0, "char_count": len(full_content)} if on_token: on_token(full_content) return full_content except Exception as e2: print(f"Fallback failed: {e2}") return f"Error: {str(e)}" # ============================================================================ # 6. DOCUMENT PROFILER # ============================================================================ class DocumentProfiler: """Deep AI-driven analysis of document (replaces regex analysis)""" def __init__(self, llm_client=None): self.llm_client = llm_client def profile(self, document: dict) -> DocumentProfile: content = document.get("content", "") if not content or len(content.strip()) < 50: return self._empty_profile("Document too short for analysis") if self.llm_client: try: return self._profile_with_ai(document) except Exception as e: print(f"AI profiling failed: {e}") return self._profile_with_fallback(document) def _profile_with_ai(self, document: dict) -> DocumentProfile: content = document.get("content", "") title = document.get("title", "Untitled") system_prompt = """You are a document profiler. Analyze the document and output JSON. Output format: { "structure": { "has_clear_intro": true/false, "has_body_paragraphs": true/false, "has_conclusion": true/false, "logical_flow": "good|fair|poor", "issues": ["specific structural issues"] }, "content": { "primary_purpose": "inform|persuade|entertain|instruct", "target_audience": "inferred audience description", "tone": "formal|casual|academic|persuasive", "reading_level": "elementary|high_school|college|expert" }, "quality": { "grammar_issues": ["specific grammar issues"], "clarity_problems": ["unclear sections"], "engagement_score": 0.0-1.0 }, "suggestions": ["specific, actionable suggestion 1", "suggestion 2"], "strengths": ["strength 1", "strength 2"] }""" user_prompt = f"""Title: {title} Content: {content[:3000]} Analyze this document:""" try: messages = [ {"role": "system", "content": system_prompt}, {"role": "user", "content": user_prompt} ] response = self.llm_client.chat.completions.create( model="llama-3.3-70b-versatile", messages=messages, max_tokens=1000, temperature=0.3 ) result_text = response.choices[0].message.content.strip() json_match = re.search(r'\{.*\}', result_text, re.DOTALL) if json_match: data = json.loads(json_match.group()) return DocumentProfile( structure=data.get("structure", {}), content=data.get("content", {}), quality=data.get("quality", {}), suggestions=data.get("suggestions", []), strengths=data.get("strengths", []), metadata={"analyzed_by": "AI", "timestamp": datetime.now().isoformat()} ) except Exception as e: print(f"AI profile error: {e}") return self._profile_with_fallback(document) def _profile_with_fallback(self, document: dict) -> DocumentProfile: content = document.get("content", "") has_intro = any(word in content[:500].lower() for word in ["introduction", "overview"]) has_conclusion = any(word in content[-500:].lower() for word in ["conclusion", "summary"]) paragraphs = [p for p in content.split('\n\n') if len(p.strip()) > 50] has_body = len(paragraphs) >= 2 grammar_issues = [] passive_matches = re.findall(r'\b(?:is|are|was|were|be|been|being)\s+\w+ed\b', content, re.I) if len(passive_matches) > 5: grammar_issues.append("Excessive passive voice usage") long_sentences = [s for s in re.split(r'[.!?]+', content) if len(s.split()) > 30] if len(long_sentences) > 3: grammar_issues.append(f"{len(long_sentences)} very long sentences") suggestions = [] if not has_intro: suggestions.append("Add a clear introduction") if not has_conclusion: suggestions.append("Add a conclusion to summarize key points") if len(content.split()) < 200: suggestions.append("Consider expanding with more details") tone = "neutral" formal_count = sum(content.lower().count(w) for w in ["therefore", "consequently"]) casual_count = sum(content.lower().count(w) for w in ["basically", "actually"]) if formal_count > casual_count * 2: tone = "formal" elif casual_count > formal_count * 2: tone = "casual" return DocumentProfile( structure={ "has_clear_intro": has_intro, "has_body_paragraphs": has_body, "has_conclusion": has_conclusion, "logical_flow": "fair" if has_intro and has_body else "poor", "issues": [] }, content={ "primary_purpose": "inform", "target_audience": "general audience", "tone": tone, "reading_level": "college" if len(content.split()) > 500 else "high_school" }, quality={ "grammar_issues": grammar_issues[:3], "clarity_problems": [], "engagement_score": 0.5 }, suggestions=suggestions[:5], strengths=[], metadata={"analyzed_by": "fallback"} ) def _empty_profile(self, reason: str) -> DocumentProfile: return DocumentProfile( structure={"has_clear_intro": False, "has_conclusion": False, "logical_flow": "poor", "issues": [reason]}, content={"primary_purpose": "unknown", "target_audience": "unknown", "tone": "neutral", "reading_level": "unknown"}, quality={"grammar_issues": [], "clarity_problems": [], "engagement_score": 0.0}, suggestions=["Add more content for proper analysis"], strengths=[], metadata={"error": reason} ) # ============================================================================ # 7. CONTEXTUAL EDITOR (Core Edit Engine) # ============================================================================ class ContextualEditor: """Execute document edits using full context - orchestrates all components""" def __init__(self, llm_client, streaming_handler: StreamingResponseHandler): self.llm_client = llm_client self.streaming_handler = streaming_handler self.instruction_parser = InstructionParser(llm_client) self.edit_planner = EditPlanner(llm_client) def edit( self, instruction: str, document: dict, conversation: Optional[ConversationManager] = None, file_context: Optional[FileContextAccumulator] = None, stream_callback: Optional[Callable[[str], None]] = None, selected_files: Optional[List[str]] = None, # Item 3: multi-file selection ) -> EditResult: start_time = time.time() parsed = self.instruction_parser.parse(instruction, document) if parsed.needs_clarification and parsed.confidence < 0.6: return EditResult( edited_document=document.get("content", ""), changes_made={}, reasoning=f"Need clarification", successful=False, execution_time_ms=int((time.time() - start_time) * 1000) ) conv_context = "" if conversation: conv_context = conversation.get_conversation_context() file_context_str = "" if file_context and file_context.files: # Item 2: Auto-reference files — inject content of matched/selected files if selected_files: # User explicitly chose files (Item 3) for fname in selected_files: detail = file_context.get_detailed_file_context(fname) file_context_str += detail + "\n" else: # Auto-detect the most relevant file relevant_file = file_context.find_relevant_file(instruction) if relevant_file: # Inject summary + content preview so AI can actually use the data file_context_str = file_context.get_detailed_file_context(relevant_file) else: file_context_str = file_context.get_file_context() plan = self.edit_planner.plan(parsed, document, conversation) prompt = self._build_edit_prompt( instruction=instruction, parsed=parsed, plan=plan, document=document, conv_context=conv_context, file_context=file_context_str ) messages = [ {"role": "system", "content": self._get_system_prompt(conversation)}, {"role": "user", "content": prompt} ] try: if stream_callback: edited_content = self.streaming_handler.stream_completion( messages=messages, on_token=stream_callback, model="llama-3.3-70b-versatile", max_tokens=4000 ) else: response = self.llm_client.chat.completions.create( model="llama-3.3-70b-versatile", messages=messages, temperature=0.3, max_tokens=4000 ) edited_content = response.choices[0].message.content.strip() changes_made = self._calculate_changes( document.get("content", ""), edited_content, plan ) execution_ms = int((time.time() - start_time) * 1000) return EditResult( edited_document=edited_content, changes_made=changes_made, reasoning=plan.rationale, successful=True, execution_time_ms=execution_ms ) except Exception as e: return EditResult( edited_document=document.get("content", ""), changes_made={}, reasoning=f"Edit failed: {str(e)}", successful=False, execution_time_ms=int((time.time() - start_time) * 1000) ) def _build_edit_prompt(self, instruction, parsed, plan, document, conv_context, file_context): content_preview = document.get("content", "") if len(content_preview) > 4000: content_preview = content_preview[:4000] + "\n...[truncated]..." prompt_parts = [ "## EDIT INSTRUCTION", f"User: {instruction}", "", "## PARSED INTENT", f"- Intent: {parsed.intent}", f"- Tone: {parsed.tone}", f"- Audience: {parsed.target_audience or 'Not specified'}", f"- Scope: {parsed.scope}", f"- Constraints: {', '.join(parsed.constraints) if parsed.constraints else 'None'}", "", "## EDIT PLAN", f"Strategy: {plan.strategy}", f"Steps:", ] for step in plan.steps: prompt_parts.append(f" {step}") # FIX #5: Make conversation context PROMINENT at top of prompt if conv_context and conv_context != "No previous conversation.": prompt_parts.extend([ "", "## ⚠️ CRITICAL - CONVERSATION CONTEXT (MUST FOLLOW)", conv_context, "IMPORTANT: The above history shows what was done previously.", "You MUST maintain any tone/style/constraints established in previous turns.", ]) if file_context and file_context != "No files uploaded.": prompt_parts.extend(["", file_context]) prompt_parts.extend([ "", "## DOCUMENT TO EDIT", "```", content_preview, "```", "", "Return ONLY the edited document content." ]) return "\n".join(prompt_parts) def _get_system_prompt(self, conversation=None) -> str: # FIX #5: Inject cumulative intent into system prompt intent_note = "" if conversation and conversation.intent_summary: intent_note = f"\nUser's overarching goal: {conversation.intent_summary}\nMaintain this goal across all edits." return f"""You are MozeAI Document Editor, a precise document editing AI.{intent_note} Return ONLY the edited document content - no explanations, no chat responses. Preserve the original meaning unless instructed otherwise. Apply changes exactly as described. ALWAYS maintain any tone, style, or constraints established in previous conversation turns.""" def _calculate_changes(self, old_content: str, new_content: str, plan: EditPlan) -> Dict: old_words = len(old_content.split()) new_words = len(new_content.split()) word_diff = new_words - old_words return { "additions": max(0, word_diff), "deletions": max(0, -word_diff), "net_change": word_diff, "old_word_count": old_words, "new_word_count": new_words, "sections_affected": ["content"], "key_changes": [f"Word count: {word_diff:+d} words ({old_words} → {new_words})"] } # ============================================================================ # FILE PROCESSING FUNCTIONS # ============================================================================ def extract_text_from_pdf(file): try: file.seek(0) pdf_reader = PyPDF2.PdfReader(file) text = "" for page_num, page in enumerate(pdf_reader.pages): page_text = page.extract_text() if page_text and page_text.strip(): text += f"\n--- Page {page_num + 1} ---\n" text += page_text.strip() + "\n" return text[:5000] if text.strip() else "No extractable text in PDF" except Exception as e: return f"Error reading PDF: {str(e)}" def extract_text_from_docx(file): try: file.seek(0) doc = docx.Document(file) text = "" for para in doc.paragraphs: if para.text and para.text.strip(): text += para.text.strip() + "\n\n" return text[:5000] if text.strip() else "No extractable text in document" except Exception as e: return f"Error reading Word document: {str(e)}" def extract_text_from_txt(file): try: file.seek(0) content = file.read().decode('utf-8') return content[:5000] if content.strip() else "File is empty" except UnicodeDecodeError: try: file.seek(0) content = file.read().decode('latin-1') return content[:5000] except: return "Error decoding text file" except Exception as e: return f"Error reading text file: {str(e)}" def extract_text_from_csv(file): try: file.seek(0) content = file.read().decode('utf-8') csv_reader = csv.reader(StringIO(content)) text = "CSV Data:\n\n" rows = list(csv_reader) if rows: text += "Headers: " + " | ".join(rows[0]) + "\n\n" for i, row in enumerate(rows[1:11], 1): text += f"Row {i}: " + " | ".join(row) + "\n" return text[:5000] if text.strip() else "CSV file appears empty" except Exception as e: return f"Error reading CSV: {str(e)}" def extract_text_from_json(file): try: file.seek(0) content = file.read().decode('utf-8') data = json.loads(content) formatted = json.dumps(data, indent=2) return formatted[:5000] if formatted else "JSON file is empty" except Exception as e: return f"Error reading JSON: {str(e)}" def process_uploaded_file(uploaded_file): file_type = uploaded_file.type file_name = uploaded_file.name.lower() if file_type == "application/pdf" or file_name.endswith('.pdf'): return extract_text_from_pdf(uploaded_file) elif file_type == "application/vnd.openxmlformats-officedocument.wordprocessingml.document" or file_name.endswith('.docx'): return extract_text_from_docx(uploaded_file) elif file_type == "text/plain" or file_name.endswith('.txt'): return extract_text_from_txt(uploaded_file) elif file_type == "text/csv" or file_name.endswith('.csv'): return extract_text_from_csv(uploaded_file) elif file_type == "application/json" or file_name.endswith('.json'): return extract_text_from_json(uploaded_file) else: return f"Unsupported file type: {file_type}" # ============================================================================ # DOCUMENT GENERATION FUNCTIONS # ============================================================================ def create_ppt_from_content(title, content, filename="presentation"): try: prs = Presentation() title_slide_layout = prs.slide_layouts[0] slide = prs.slides.add_slide(title_slide_layout) slide.shapes.title.text = title[:100] slide.placeholders[1].text = f"Created by MozeAI\n{datetime.now().strftime('%Y-%m-%d %H:%M:%S')}" content_slide_layout = prs.slide_layouts[1] lines = content.split('\n') current_slide = None current_text_frame = None for line in lines: line = line.strip() if not line: continue if len(line) < 60 and (line.endswith(':') or line.isupper() or re.match(r'^\d+\.', line)): current_slide = prs.slides.add_slide(content_slide_layout) current_slide.shapes.title.text = line.rstrip(':')[:100] content_box = current_slide.placeholders[1] current_text_frame = content_box.text_frame current_text_frame.text = "" else: if current_slide is None: current_slide = prs.slides.add_slide(content_slide_layout) current_slide.shapes.title.text = "Content" content_box = current_slide.placeholders[1] current_text_frame = content_box.text_frame current_text_frame.text = "" if current_text_frame: p = current_text_frame.add_paragraph() p.text = line[:150] p.font.size = Pt(18) ppt_bytes = BytesIO() prs.save(ppt_bytes) ppt_bytes.seek(0) return ppt_bytes except Exception as e: print(f"PPT error: {e}") return None def create_word_from_content(title, content, filename="document"): try: doc = WordDocument() title_heading = doc.add_heading(title, 0) title_heading.alignment = WD_ALIGN_PARAGRAPH.CENTER doc.add_paragraph(f"Generated by MozeAI on {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}") doc.add_paragraph() paragraphs = content.split('\n\n') for para in paragraphs: if para.strip(): doc.add_paragraph(para.strip()) word_bytes = BytesIO() doc.save(word_bytes) word_bytes.seek(0) return word_bytes except Exception as e: return None def create_real_excel_file(title, data_rows): try: from openpyxl import Workbook from openpyxl.styles import Font, PatternFill, Alignment from openpyxl.utils import get_column_letter wb = Workbook() ws = wb.active ws.title = title[:31].replace('/', '_') for row_idx, row in enumerate(data_rows, 1): for col_idx, value in enumerate(row, 1): cell = ws.cell(row=row_idx, column=col_idx, value=value) if row_idx == 1: cell.font = Font(bold=True, color="FFFFFF") cell.fill = PatternFill(start_color="4472C4", end_color="4472C4", fill_type="solid") for col in ws.columns: max_length = 0 for cell in col: try: if len(str(cell.value)) > max_length: max_length = len(str(cell.value)) except: pass ws.column_dimensions[get_column_letter(col[0].column)].width = min(max_length + 2, 50) output = BytesIO() wb.save(output) output.seek(0) return output except Exception as e: print(f"Excel error: {e}") return None def create_csv_from_data(title, data_rows): try: output = BytesIO() output.write('\ufeff'.encode('utf-8')) writer = csv.writer(output) for row in data_rows: writer.writerow(row) output.seek(0) return output except Exception as e: return None def export_chat_history(): """FIX #10: Enhanced export includes intent + timeline""" if not st.session_state.chat_history: return None export_content = "=" * 70 + "\n" export_content += "CHAT HISTORY WITH MOZEAI\n" export_content += f"Exported on: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n" export_content += "=" * 70 + "\n\n" # Include intent summary core = st.session_state.get("intelligence_core") if core: conv_manager = core["conversation_manager"] if conv_manager.intent_summary: export_content += f"=== SESSION INTENT ===\n{conv_manager.intent_summary}\n\n" evolution = conv_manager.get_document_evolution() if evolution: export_content += "=== EDIT TIMELINE ===\n" for i, turn in enumerate(evolution, 1): changes = turn.get("changes", {}) net = changes.get("net_change", 0) export_content += f"{i}. {turn['query'][:60]}... → {net:+d} words\n" export_content += "\n" export_content += "=== CONVERSATION ===\n\n" for idx, (role, msg) in enumerate(st.session_state.chat_history, 1): if role == "user": export_content += f"[{idx}] USER:\n{msg}\n\n" else: export_content += f"[{idx}] MOZEAI:\n{msg}\n\n" return export_content # ============================================================================ # WEB & UTILITY FUNCTIONS # ============================================================================ def get_current_datetime(): tz = pytz.timezone('Asia/Seoul') now = datetime.now(tz) return f"Date: {now.strftime('%B %d, %Y')}\nTime: {now.strftime('%I:%M %p')}\nTimezone: Asia/Seoul" def internet_search(query): try: clean_query = query.strip() url = "https://html.duckduckgo.com/html/" params = {"q": clean_query} headers = {"User-Agent": "Mozilla/5.0"} response = requests.post(url, data=params, headers=headers, timeout=10) if response.status_code == 200: results = re.findall(r'([^<]+)', response.text) snippets = re.findall(r']*>([^<]+)', response.text) if results: context = f"SEARCH RESULTS for '{clean_query}':\n\n" for i in range(min(3, len(results))): context += f"- {results[i]}\n" if i < len(snippets): snippet = re.sub(r'<[^>]+>', '', snippets[i]) context += f" {snippet[:300]}...\n\n" return context[:2000] return "" except: return "" def generate_image_with_quality(prompt, quality="high", style="realistic"): try: enhanced_prompt = f"{prompt}, high quality, detailed" encoded_prompt = requests.utils.quote(enhanced_prompt) timestamp = int(time.time()) image_url = f"https://image.pollinations.ai/prompt/{encoded_prompt}?width=1024&height=1024&seed={timestamp}" return image_url except Exception as e: return None def generate_and_display_image(prompt, is_edit=False): image_url = generate_image_with_quality(prompt) if image_url: return f"Generated Image for: '{prompt}'\n\n![Image]({image_url})" else: return "Sorry, I couldn't generate an image." def llm_with_fallback(messages, max_retries=2): models_to_try = [ "llama-3.3-70b-versatile", "llama-3.1-70b-versatile", "mixtral-8x7b-32768" ] for model in models_to_try: for attempt in range(max_retries): try: completion = client.chat.completions.create( model=model, temperature=0.3, max_tokens=800, messages=messages, timeout=30 ) st.session_state.last_model_used = model return completion.choices[0].message.content.strip() except Exception as e: if attempt < max_retries - 1: time.sleep(2 ** attempt) continue return "AI service temporarily unavailable." def reason(question, context): messages = [ {"role": "system", "content": "You are MozeAI, a helpful AI assistant."}, {"role": "user", "content": f"{context}\n\nUSER QUESTION: {question}\n\nANSWER:"} ] return llm_with_fallback(messages) def clean_answer(text): text = text.split("🧠")[0] text = text.split("Plan:")[0] return text.strip() # ============================================================================ # ENHANCED RUN AGENT WITH INTELLIGENCE CORE # ============================================================================ def handle_intelligent_edit(instruction: str, stream_callback=None, selected_files=None): """Use contextual editor for intelligent document editing""" core = st.session_state.intelligence_core document_state = { "content": st.session_state.workspace.current_document["content"], "title": st.session_state.workspace.current_document["title"], "word_count": len(st.session_state.workspace.current_document["content"].split()), "char_count": len(st.session_state.workspace.current_document["content"]) } # Check clarification BEFORE editing parser = core["instruction_parser"] parsed = parser.parse(instruction, document_state) if parsed.needs_clarification and parsed.confidence < 0.6 and parsed.clarification_questions: st.session_state.pending_clarification = { "instruction": instruction, "questions": parsed.clarification_questions, "parsed": parsed } return None # Signal that clarification is needed result = core["contextual_editor"].edit( instruction=instruction, document=document_state, conversation=core["conversation_manager"], file_context=core["file_accumulator"], stream_callback=stream_callback, selected_files=selected_files, # Item 2 & 3: pass through ) if result.successful: core["conversation_manager"].add_user_message(instruction, document_state) core["conversation_manager"].add_assistant_message("Document edited", result.changes_made) st.session_state.workspace.update_document(result.edited_document, f"AI Edit: {instruction[:100]}") return result def profile_current_document(): """Profile current document using AI""" core = st.session_state.intelligence_core document_state = { "content": st.session_state.workspace.current_document["content"], "title": st.session_state.workspace.current_document["title"] } return core["document_profiler"].profile(document_state) def run_agent(query: str, stream_callback=None, selected_files=None): """Enhanced agent with intelligence core integration""" q = query.lower().strip() # Document workspace commands if q.startswith("/"): return handle_document_command(q) # Check for document editing commands edit_keywords = ["improve", "rewrite", "summarize", "expand", "shorten", "fix grammar", "make formal", "make academic", "translate"] if any(keyword in q for keyword in edit_keywords) and st.session_state.workspace.current_document["content"]: # Item 6: Edit validation — snapshot word count before edit pre_word_count = len(st.session_state.workspace.current_document["content"].split()) result = handle_intelligent_edit(query, stream_callback=stream_callback, selected_files=selected_files) # Clarification needed if result is None: return "__CLARIFICATION_NEEDED__" if result.successful: # Item 6: Validate edit didn't produce empty/trivially-short output post_word_count = len(result.edited_document.split()) if post_word_count < max(10, pre_word_count * 0.1): result.successful = False result.reasoning = ( f"Edit validation failed: output was only {post_word_count} words " f"(original was {pre_word_count}). Original preserved." ) st.session_state.workspace.update_document( st.session_state.workspace.version_history[-1]["content"], "Rollback: edit validation failed" ) return result # Compute file cross-references doc_content = st.session_state.workspace.current_document["content"] file_acc = st.session_state.intelligence_core["file_accumulator"] cross_refs = file_acc.get_cross_references(doc_content) st.session_state.pending_file_suggestions = cross_refs return result else: return f"⚠️ {result.reasoning}" # Analysis command if "analyze document" in q or "profile document" in q: with st.spinner("Analyzing document..."): profile = profile_current_document() st.session_state.doc_profile_cache = profile result = f"## Document Analysis\n\n" result += f"**Tone:** {profile.content.get('tone', 'unknown').title()}\n" result += f"**Purpose:** {profile.content.get('primary_purpose', 'unknown').title()}\n" result += f"**Reading Level:** {profile.content.get('reading_level', 'unknown')}\n\n" if profile.strengths: result += "**Strengths:**\n" for s in profile.strengths[:3]: result += f"- {s}\n" result += "\n" if profile.suggestions: result += "**Suggestions:**\n" for s in profile.suggestions[:3]: result += f"- {s}\n" return result # Clear context if any(phrase in q for phrase in ["clear context", "new chat", "start fresh"]): st.session_state.workspace = DocumentWorkspace() st.session_state.intelligence_core["conversation_manager"].clear() st.session_state.intelligence_core["file_accumulator"].clear() st.session_state.chat_history = [] st.session_state.uploaded_files = {} st.session_state.pending_clarification = None st.session_state.pending_file_suggestions = [] return "✨ Everything cleared! Ready for a new session." # What is a word? if q == "what is a word": return "A **Word document** (.docx) is created by Microsoft Word. Try 'make a word about dogs'" # Excel generation if "make an excel" in q or "create an excel" in q or "generate an excel" in q: topic = q.replace("make an excel", "").replace("create an excel", "").replace("generate an excel", "").strip() topic = topic or "Sample_Data" data_rows = [ ["Item", "Category", "Quantity", "Price", "Total"], ["Product A", "Electronics", 10, 99.99, 999.90], ["Product B", "Clothing", 25, 49.99, 1249.75], ["Product C", "Food", 50, 9.99, 499.50] ] excel_data = create_real_excel_file(topic, data_rows) if excel_data: st.session_state.excel_data = excel_data st.session_state.excel_topic = topic st.session_state.show_excel_download = True return f"📊 Created Excel file: {topic}. Scroll down to download!" # PowerPoint generation if any(phrase in q for phrase in ["make a ppt", "create a powerpoint"]): topic = q.replace("make a ppt", "").replace("create a powerpoint", "").strip() or "Presentation" content = f"Introduction to {topic}\n- Key point 1\n- Key point 2\n\nConclusion\n- Summary" ppt_bytes = create_ppt_from_content(topic, content) if ppt_bytes: st.session_state.ppt_data = ppt_bytes st.session_state.ppt_topic = topic st.session_state.show_ppt_download = True return f"📊 Created PowerPoint: {topic}. Scroll down to download!" # Word generation if any(phrase in q for phrase in ["make a word", "create a document"]): topic = q.replace("make a word", "").replace("create a document", "").strip() or "Document" content = f"# {topic}\n\nThis document covers important information about {topic}.\n\n## Introduction\n\nContent here.\n\n## Conclusion\n\nSummary." word_bytes = create_word_from_content(topic, content) if word_bytes: st.session_state.word_data = word_bytes st.session_state.word_topic = topic st.session_state.show_word_download = True return f"📄 Created Word document: {topic}. Scroll down to download!" # Image generation if any(phrase in q for phrase in ["generate image", "create image"]): image_prompt = q.replace("generate image", "").replace("create image", "").strip() if not image_prompt: image_prompt = "a beautiful landscape" return generate_and_display_image(image_prompt) # Item 9: Streaming for non-edit queries — use stream_callback if provided search_result = internet_search(query) context = get_current_datetime() if search_result: context += "\n" + search_result if stream_callback: messages = [ {"role": "system", "content": "You are MozeAI, a helpful AI assistant."}, {"role": "user", "content": f"{context}\n\nUSER QUESTION: {query}\n\nANSWER:"} ] handler = st.session_state.intelligence_core["streaming_handler"] return handler.stream_completion(messages, on_token=stream_callback) else: answer = reason(query, context) return answer def handle_document_command(command: str) -> str: """Handle slash commands""" cmd = command.lower().strip() workspace = st.session_state.workspace core = st.session_state.intelligence_core if cmd == "/analyze": analysis = workspace.analyze_document() return f"""## Document Analysis **Structure:** {len(analysis['structure']['headings'])} headings, {analysis['structure']['paragraph_count']} paragraphs **Readability:** {analysis['readability']['level']} **Style:** {analysis['style_analysis']['detected_style']} **Suggestions:** {chr(10).join(f'- {s}' for s in analysis['suggestions'])}""" elif cmd == "/stats": meta = workspace.current_document["metadata"] return f"""## Document Stats **Title:** {workspace.current_document['title']} **Words:** {meta['word_count']} **Characters:** {meta['char_count']} **Reading Time:** {meta['reading_time']} min **Versions:** {len(workspace.version_history)}""" elif cmd.startswith("/version"): parts = cmd.split() if len(parts) > 1 and parts[1].isdigit(): if workspace.restore_version(int(parts[1])): return f"✅ Restored version {parts[1]}" return f"Versions: {len(workspace.version_history)} saved" elif cmd == "/conversation": summary = core["conversation_manager"].summarize_intent(client) return f"**Conversation Intent:** {summary}\n**Turns:** {len(core['conversation_manager'].turns)}" elif cmd == "/help": return """## Commands **Document:** `/analyze`, `/stats`, `/version N` **Conversation:** `/conversation`, `/clear` **Editing:** Just tell me what to do, like "make this formal" or "add a conclusion" """ else: return f"Unknown command. Type `/help` for available commands." # ============================================================================ # FIX #2: CLARIFICATION DIALOG COMPONENT # ============================================================================ def render_clarification_dialog(): """FIX #2: Show clarification questions when instruction is ambiguous""" pending = st.session_state.get("pending_clarification") if not pending: return st.warning("🤔 I need a bit more info to edit your document precisely:") with st.container(): st.markdown(f"**Your instruction:** _{pending['instruction']}_") answers = {} for i, question in enumerate(pending["questions"]): answer = st.text_input(f"Q{i+1}: {question}", key=f"clarif_q_{i}") answers[question] = answer col1, col2 = st.columns(2) with col1: if st.button("✅ Proceed with clarification", use_container_width=True, type="primary"): # Build enriched instruction clarifications = "; ".join([f"{q}: {a}" for q, a in answers.items() if a]) enriched = f"{pending['instruction']}. Clarifications: {clarifications}" st.session_state.pending_clarification = None # Now execute with enriched instruction with st.spinner("Applying edit..."): result = handle_intelligent_edit(enriched) if result and result.successful: st.success("✅ Done!") st.rerun() with col2: if st.button("⏭️ Skip & proceed anyway", use_container_width=True): st.session_state.pending_clarification = None with st.spinner("Applying edit..."): result = handle_intelligent_edit(pending["instruction"]) if result and result.successful: st.rerun() # ============================================================================ # FIX #4: FILE SUGGESTIONS BANNER # ============================================================================ def render_file_suggestions(): """FIX #4: Show cross-reference suggestions from uploaded files""" suggestions = st.session_state.get("pending_file_suggestions", []) if not suggestions: return with st.expander("💡 File Reference Opportunities", expanded=True): for suggestion in suggestions[:3]: st.info( f"📎 **{suggestion['file']}** — " f"matches terms in your document: `{'`, `'.join(suggestion['matched_terms'][:3])}`\n\n" f"{suggestion['suggestion']}" ) if st.button("✖ Dismiss", key="dismiss_file_suggestions"): st.session_state.pending_file_suggestions = [] st.rerun() # ============================================================================ # ITEM 4: SMART ROLLBACK / UNDO BY INTENT # ============================================================================ def smart_rollback(target_description: str = "") -> bool: """ Item 4: Rollback to the best matching version by intent keyword. If no keyword given, rolls back one version. """ workspace = st.session_state.workspace versions = workspace.version_history if not versions: return False if not target_description: # Simple one-step undo: restore second-to-last "After:" version after_versions = [v for v in versions if v["description"].startswith("After:")] if len(after_versions) >= 2: workspace.restore_version(after_versions[-2]["id"]) return True return False # Keyword search across version descriptions keyword = target_description.lower() best = None for v in reversed(versions): if keyword in v["description"].lower(): best = v break if best: workspace.restore_version(best["id"]) return True return False # ============================================================================ # ITEM 5: CONVERSATION DASHBOARD # ============================================================================ def render_conversation_dashboard(): """Item 5: Expandable timeline of turns + word-count evolution""" core = st.session_state.get("intelligence_core") if not core: return conv = core["conversation_manager"] if not conv.turns: st.caption("No conversation yet — start editing to see the timeline.") return evolution = conv.get_document_evolution() total_edits = conv.cumulative_edits["total_edits"] intent = conv.intent_summary or "Not summarized yet" st.markdown(f"**Session intent:** _{intent}_") st.caption(f"Total edits: {total_edits} | Turns: {len(evolution)}") for i, turn in enumerate(evolution, 1): changes = turn.get("changes", {}) net = changes.get("net_change", 0) old_wc = turn["document_state"].get("word_count", 0) new_wc = changes.get("new_word_count", old_wc + net) arrow = "📈" if net > 0 else ("📉" if net < 0 else "➡️") label = turn["query"][:45] + ("…" if len(turn["query"]) > 45 else "") st.markdown( f"**{i}.** {arrow} _{label}_ \n" f"{old_wc} → {new_wc} words ({net:+d})", unsafe_allow_html=True ) # Smart rollback controls st.markdown("---") st.markdown("**↩ Undo / Rollback**") col_a, col_b = st.columns([2, 1]) with col_a: rollback_kw = st.text_input("Roll back to edit containing…", placeholder="e.g. 'formal'", key="rollback_kw", label_visibility="collapsed") with col_b: if st.button("↩ Undo", use_container_width=True): keyword = rollback_kw.strip() if rollback_kw.strip() else "" if smart_rollback(keyword): st.success("✅ Rolled back!") st.rerun() else: st.warning("No matching version found.") # ============================================================================ # ITEM 7: EDIT PLAN DISPLAY # ============================================================================ def render_edit_plan(plan) -> None: """Item 7: Show the AI's edit plan as an expandable checklist in chat""" if not plan: return with st.expander("🗺️ Edit Plan", expanded=False): st.markdown(f"**Strategy:** {plan.strategy}") for step in plan.steps: st.markdown(f"- ☑ {step}") if plan.constraints: st.markdown("**Constraints:** " + " · ".join(plan.constraints)) if plan.rationale: st.caption(f"Rationale: {plan.rationale}") # ============================================================================ # ITEM 8: CUMULATIVE TIMELINE IN SIDEBAR # ============================================================================ def render_cumulative_timeline(): """Item 8: Compact edit timeline for sidebar display""" core = st.session_state.get("intelligence_core") if not core: return conv = core["conversation_manager"] if not conv.turns: return st.markdown("**📅 Edit Timeline**") for i, turn in enumerate(conv.turns[-5:], 1): changes = turn.edits_made net = changes.get("net_change", 0) label = turn.user_query[:30] + ("…" if len(turn.user_query) > 30 else "") color = "#4CAF50" if net >= 0 else "#f44336" st.markdown( f"{i}. {label} " f"{net:+d}w", unsafe_allow_html=True ) # ============================================================================ # UI COMPONENTS # ============================================================================ def render_document_explorer(): with st.sidebar: st.markdown("### 📁 Document Explorer") col1, col2 = st.columns(2) with col1: if st.button("📄 New", use_container_width=True): st.session_state.workspace.current_document["content"] = "" st.session_state.workspace.current_document["title"] = "Untitled Document" st.session_state.workspace.save_version("New document") st.rerun() with col2: if st.button("💾 Save", use_container_width=True): st.session_state.workspace.save_version("Manual save") st.success("Saved!") st.markdown("---") uploaded_files = st.file_uploader( "Upload files", type=['pdf', 'docx', 'txt', 'csv', 'json'], accept_multiple_files=True, key="file_uploader" ) if uploaded_files: for file in uploaded_files: if file.name not in st.session_state.uploaded_files: content = process_uploaded_file(file) if content and not content.startswith("Error"): st.session_state.uploaded_files[file.name] = content st.session_state.intelligence_core["file_accumulator"].add_file( file.name, file.type, content, {} ) st.success(f"✅ {file.name}") # Item 3: Multi-file selection — let user choose which files to reference in next edit all_files = list(st.session_state.uploaded_files.keys()) if all_files: st.markdown("**📎 Reference in next edit**") selected = st.multiselect( "Select files to inject", options=all_files, default=[], key="selected_ref_files", label_visibility="collapsed" ) st.session_state.selected_ref_files = selected if selected: st.caption(f"✅ {len(selected)} file(s) will be injected into the edit prompt") st.markdown("---") st.markdown("**Version History**") if st.button("📜 View Versions", use_container_width=True): versions = st.session_state.workspace.version_history if versions: for v in versions[-3:]: st.caption(f"v{v['id']}: {v['description'][:30]}") def render_document_editor(): st.markdown("### 📝 Document Editor") new_title = st.text_input( "Title", value=st.session_state.workspace.current_document["title"], key="doc_title" ) if new_title != st.session_state.workspace.current_document["title"]: st.session_state.workspace.current_document["title"] = new_title col1, col2, col3, col4 = st.columns(4) with col1: if st.button("🔍 Analyze", use_container_width=True): analysis = st.session_state.workspace.analyze_document() st.session_state.last_analysis = analysis st.info(f"Readability: {analysis['readability']['level']}") with col2: track_status = "✅ Track ON" if st.session_state.workspace.track_changes else "⭕ Track OFF" if st.button(track_status, use_container_width=True): st.session_state.workspace.track_changes = not st.session_state.workspace.track_changes st.rerun() with col3: if st.button("📊 Stats", use_container_width=True): meta = st.session_state.workspace.current_document["metadata"] st.info(f"{meta['word_count']} words, {meta['reading_time']} min read") with col4: if st.button("🧹 Clear", use_container_width=True): st.session_state.workspace.current_document["content"] = "" st.rerun() st.markdown("---") content = st.text_area( "Content", value=st.session_state.workspace.current_document["content"], height=400, key="doc_editor", label_visibility="collapsed" ) if content != st.session_state.workspace.current_document["content"]: st.session_state.workspace.update_document(content, "Manual edit") def render_ai_copilot(): with st.sidebar: st.markdown("### 🤖 AI Copilot") # Intent summary core = st.session_state.get("intelligence_core") if core: conv_manager = core["conversation_manager"] if conv_manager.turns: intent = conv_manager.intent_summary or conv_manager.summarize_intent(client) if intent and "No conversation" not in intent: st.info(f"📍 **Goal:** {intent}") total = conv_manager.cumulative_edits["total_edits"] if total > 0: st.caption(f"🔄 {total} edit{'s' if total != 1 else ''} this session") st.markdown("---") quick_actions = [ ("✨ Improve", "improve this document"), ("🎓 Academic", "make this academic"), ("📝 Summarize", "summarize this document"), ("🔧 Fix Grammar", "fix grammar"), ] for label, instruction in quick_actions: if st.button(label, use_container_width=True): with st.spinner("AI editing..."): result = handle_intelligent_edit(instruction) if result and result.successful: st.success("Done!") st.rerun() st.markdown("---") custom = st.text_area("Custom instruction", placeholder="e.g., 'Rewrite for a 12-year-old'", height=80) if st.button("Apply", use_container_width=True, type="primary"): if custom: with st.spinner("AI working..."): result = handle_intelligent_edit(custom) if result and result.successful: st.success("Updated!") st.rerun() st.markdown("---") # Item 8: Cumulative timeline render_cumulative_timeline() st.markdown("---") # Item 5: Conversation dashboard in expander with st.expander("📊 Conversation Dashboard", expanded=False): render_conversation_dashboard() st.markdown("---") st.markdown("**Quick Stats**") meta = st.session_state.workspace.current_document["metadata"] st.caption(f"Words: {meta['word_count']}") st.caption(f"Versions: {len(st.session_state.workspace.version_history)}") # ============================================================================ # FULLY UPGRADED CHAT INTERFACE # Items: 1 (token metrics), 3 (selected_files), 7 (plan display), # 9 (streaming non-edit), 10 (keyboard shortcut hint) # ============================================================================ def render_chat_interface(): st.markdown("---") st.markdown("### 💬 Chat") # Show clarification dialog if pending render_clarification_dialog() # Show file suggestions if any render_file_suggestions() # Show chat history for role, msg in st.session_state.chat_history[-10:]: with st.chat_message(role): st.markdown(msg) # Item 10: Keyboard shortcut hint st.caption("💡 Tip: Press **Enter** to send · Use `/help` for commands · `Ctrl+Z` style undo: type **undo**") query = st.chat_input("Ask me to edit, analyze, or generate…") # Item 10: "undo" as a text shortcut for smart rollback if query and query.strip().lower() in ("undo", "undo last"): with st.chat_message("user"): st.markdown(query) st.session_state.chat_history.append(("user", query)) with st.chat_message("assistant"): if smart_rollback(): msg = "↩️ Undone — restored the previous version." else: msg = "⚠️ Nothing to undo." st.markdown(msg) st.session_state.chat_history.append(("assistant", msg)) st.rerun() return if query: st.session_state.chat_history.append(("user", query)) with st.chat_message("user"): st.markdown(query) with st.chat_message("assistant"): q_lower = query.lower().strip() edit_keywords = ["improve", "rewrite", "summarize", "expand", "shorten", "fix grammar", "make formal", "make academic", "translate"] is_edit = (any(kw in q_lower for kw in edit_keywords) and st.session_state.workspace.current_document["content"] and not q_lower.startswith("/")) if is_edit: # Item 7: Show plan before streaming starts plan_placeholder = st.empty() # Item 1: Prepare live token-speed display response_placeholder = st.empty() streaming_text = "" metrics_placeholder = st.empty() stream_start = time.time() token_count_ref = [0] def display_token(token: str): nonlocal streaming_text streaming_text += token token_count_ref[0] += 1 elapsed = max(time.time() - stream_start, 0.001) tps = round(token_count_ref[0] / elapsed, 1) wc = len(streaming_text.split()) response_placeholder.markdown(streaming_text + "▌") # Item 1: live token speed metrics_placeholder.caption( f"✍️ {wc} words · {token_count_ref[0]} tokens · **{tps} tok/s**" ) # Item 3: pick up selected files from sidebar selected_files = st.session_state.get("selected_ref_files", []) or None with st.spinner(""): response = run_agent(query, stream_callback=display_token, selected_files=selected_files) metrics_placeholder.empty() if response == "__CLARIFICATION_NEEDED__": response_placeholder.empty() plan_placeholder.empty() st.rerun() return elif isinstance(response, EditResult) and response.successful: response_placeholder.empty() plan_placeholder.empty() # Item 7: Retrieve and display the edit plan used core = st.session_state.intelligence_core last_plan = None try: doc_snap = { "content": st.session_state.workspace.current_document["content"], "title": st.session_state.workspace.current_document["title"], "word_count": len(st.session_state.workspace.current_document["content"].split()) } parsed = core["instruction_parser"].parse(query, doc_snap) last_plan = core["edit_planner"].plan(parsed, doc_snap, core["conversation_manager"]) except Exception: pass if last_plan: render_edit_plan(last_plan) # Metrics card changes = response.changes_made old_wc = changes.get("old_word_count", 0) new_wc = changes.get("new_word_count", 0) net = changes.get("net_change", 0) additions = changes.get("additions", 0) deletions = changes.get("deletions", 0) exec_ms = response.execution_time_ms # Item 1: Pull token speed from streaming handler sh = core["streaming_handler"] tps = sh.last_metrics.get("tokens_per_sec", 0) ttft = sh.last_metrics.get("ttft_ms", 0) token_total = sh.last_metrics.get("token_count", 0) st.success("✅ Edit Complete") col1, col2, col3, col4, col5 = st.columns(5) with col1: st.metric("Words Added", f"+{additions}" if additions else "0") with col2: st.metric("Words Removed", f"-{deletions}" if deletions else "0") with col3: st.metric("Net Change", f"{net:+d}") with col4: st.metric("Time", f"{exec_ms}ms") with col5: # Item 1: token speed metric st.metric("Speed", f"{tps} tok/s") st.caption( f"📝 {old_wc} → {new_wc} words · " f"{token_total} tokens · TTFT {ttft}ms" ) if response.reasoning: st.info(f"**Why:** {response.reasoning}") summary_msg = ( f"✅ **Edit complete** in {exec_ms}ms · {tps} tok/s\n\n" f"Words: {old_wc} → {new_wc} ({net:+d})\n\n" f"Reasoning: {response.reasoning}" ) st.session_state.chat_history.append(("assistant", summary_msg)) elif isinstance(response, EditResult) and not response.successful: response_placeholder.warning(f"⚠️ {response.reasoning}") st.session_state.chat_history.append(("assistant", f"⚠️ {response.reasoning}")) elif isinstance(response, str): response_placeholder.markdown(response) st.session_state.chat_history.append(("assistant", response)) else: # Item 9: Streaming for non-edit queries response_placeholder = st.empty() metrics_placeholder = st.empty() streaming_text = "" stream_start = time.time() token_count_ref = [0] def display_token_general(token: str): nonlocal streaming_text streaming_text += token token_count_ref[0] += 1 elapsed = max(time.time() - stream_start, 0.001) tps = round(token_count_ref[0] / elapsed, 1) response_placeholder.markdown(streaming_text + "▌") metrics_placeholder.caption(f"⚡ {tps} tok/s") with st.spinner(""): response = run_agent(query, stream_callback=display_token_general) metrics_placeholder.empty() response_placeholder.empty() if isinstance(response, str): st.markdown(response) st.session_state.chat_history.append(("assistant", response)) st.rerun() def render_download_buttons(): if st.session_state.get("show_ppt_download", False) and st.session_state.get("ppt_data"): st.download_button( label="📥 Download PowerPoint", data=st.session_state.ppt_data, file_name=f"{st.session_state.ppt_topic}.pptx", mime="application/vnd.openxmlformats-officedocument.presentationml.presentation" ) st.session_state.show_ppt_download = False if st.session_state.get("show_word_download", False) and st.session_state.get("word_data"): st.download_button( label="📥 Download Word Document", data=st.session_state.word_data, file_name=f"{st.session_state.word_topic}.docx", mime="application/vnd.openxmlformats-officedocument.wordprocessingml.document" ) st.session_state.show_word_download = False if st.session_state.get("show_excel_download", False) and st.session_state.get("excel_data"): st.download_button( label="📥 Download Excel File", data=st.session_state.excel_data, file_name=f"{st.session_state.excel_topic}.xlsx", mime="application/vnd.openxmlformats-officedocument.spreadsheetml.sheet" ) st.session_state.show_excel_download = False # ============================================================================ # MAIN APPLICATION # ============================================================================ def init_session_state(): """Initialize all session state variables""" if "workspace" not in st.session_state: st.session_state.workspace = DocumentWorkspace() if "chat_history" not in st.session_state: st.session_state.chat_history = [] if "uploaded_files" not in st.session_state: st.session_state.uploaded_files = {} if "show_ppt_download" not in st.session_state: st.session_state.show_ppt_download = False st.session_state.ppt_data = None st.session_state.ppt_topic = "" if "show_word_download" not in st.session_state: st.session_state.show_word_download = False st.session_state.word_data = None st.session_state.word_topic = "" if "show_excel_download" not in st.session_state: st.session_state.show_excel_download = False st.session_state.excel_data = None st.session_state.excel_topic = "" if "intelligence_core" not in st.session_state: st.session_state.intelligence_core = None if "doc_profile_cache" not in st.session_state: st.session_state.doc_profile_cache = None if "last_model_used" not in st.session_state: st.session_state.last_model_used = None if "last_analysis" not in st.session_state: st.session_state.last_analysis = None # FIX #2: Clarification state if "pending_clarification" not in st.session_state: st.session_state.pending_clarification = None # FIX #4: File suggestions state if "pending_file_suggestions" not in st.session_state: st.session_state.pending_file_suggestions = [] # Item 3: Multi-file selection state if "selected_ref_files" not in st.session_state: st.session_state.selected_ref_files = [] def apply_custom_css(): st.markdown(""" """, unsafe_allow_html=True) def main(): st.set_page_config( page_title="MozeAI Document Studio", page_icon="📝", layout="wide" ) init_session_state() apply_custom_css() # Initialize Groq client groq_api_key = None try: if "GROQ_API_KEY" in st.secrets: groq_api_key = st.secrets["GROQ_API_KEY"] except: pass if not groq_api_key: groq_api_key = os.environ.get("GROQ_API_KEY") if not groq_api_key: st.error("GROQ_API_KEY not found. Please set it in secrets or environment.") st.stop() global client client = Groq(api_key=groq_api_key) # Initialize intelligence core if not exists if st.session_state.intelligence_core is None: streaming_handler = StreamingResponseHandler(client) contextual_editor = ContextualEditor(client, streaming_handler) st.session_state.intelligence_core = { "conversation_manager": ConversationManager(), "instruction_parser": InstructionParser(client), "file_accumulator": FileContextAccumulator(client), "edit_planner": EditPlanner(client), "streaming_handler": streaming_handler, "document_profiler": DocumentProfiler(client), "contextual_editor": contextual_editor } # Header st.markdown('

📝 MozeAI Document Studio

', unsafe_allow_html=True) st.markdown('

Intelligent Document Workspace

', unsafe_allow_html=True) st.markdown("---") # Sidebar with tabs with st.sidebar: tab1, tab2 = st.tabs(["📁 Explorer", "🤖 Copilot"]) with tab1: render_document_explorer() with tab2: render_ai_copilot() # Main content render_document_editor() # Download buttons render_download_buttons() # Chat interface render_chat_interface() # Footer st.markdown("---") st.markdown( '

MozeAI Document Studio | Created by Mukiibi Moses

', unsafe_allow_html=True ) if __name__ == "__main__": main()