""" RasoSpeak v2 — Document Import Agent Import documents, notes, PDFs, and files into your AI memory. Supported formats: - Text files (.txt) - Markdown (.md) - PDF (.pdf) - via PyPDF2 - Web pages (URLs) - Notes and snippets """ import json import logging import re import time from datetime import datetime from pathlib import Path from typing import Optional import httpx from .base_agent import BaseAgent from .shared_memory_agent import SharedMemoryAgent from config.settings import settings log = logging.getLogger("rasospeak.document") class DocumentAgent(BaseAgent): """ Agent for importing documents and files into your AI memory. This allows you to: - Import notes, PDFs, text files - Add web page content - Paste snippets and articles - Import from URLs All imported content goes into shared memory and can be queried by your AI partner. """ name = "DocumentAgent" def __init__(self): self._storage_path = Path(settings.documents_path or "./memory/documents") self._shared_memory: Optional[SharedMemoryAgent] = None self._second_brain = None # Second Brain for enhanced memory storage self._ensure_storage() def _ensure_storage(self): """Create storage directories.""" self._storage_path.mkdir(parents=True, exist_ok=True) log.info(f"Document storage: {self._storage_path}") async def initialize(self): """Initialize document agent.""" log.info("✅ DocumentAgent initialized") def set_second_brain(self, second_brain): """Connect to Second Brain for enhanced memory storage.""" self._second_brain = second_brain log.info("DocumentAgent connected to SecondBrainAgent") def set_shared_memory(self, shared_memory: SharedMemoryAgent): """Connect to shared memory.""" self._shared_memory = shared_memory # ══════════════════════════════════════════════════════ # IMPORT METHODS # ══════════════════════════════════════════════════════ async def import_text( self, content: str, title: str = None, tags: list = None, category: str = "note", ) -> dict: """ Import text content directly. Args: content: The text content to import title: Optional title tags: Optional tags category: note | article | snippet | book | other Returns: Import result with ID """ doc_id = f"doc_{int(time.time())}" title = title or f"Imported Note {datetime.now().strftime('%Y-%m-%d %H:%M')}" document = { "id": doc_id, "title": title, "content": content, "category": category, "tags": tags or [], "imported_at": datetime.utcnow().isoformat(), "word_count": len(content.split()), "char_count": len(content), } # Save to disk await self._save_document(doc_id, document) # Store in Second Brain (handles dual-write to legacy automatically) if self._second_brain: await self._second_brain.add_document( content=content, title=title, doc_type=category, tags=tags, ) # Extract and store key points key_points = self._extract_key_points(content) for point in key_points: await self._second_brain.add_user_fact( fact=point, category="document_key_point", ) log.info(f"📄 Imported document: {title} ({document['word_count']} words)") return { "imported": True, "document_id": doc_id, "title": title, "word_count": document["word_count"], } async def import_file( self, file_path: str, tags: list = None, ) -> dict: """ Import a file from disk. Supports: .txt, .md, .json """ path = Path(file_path) if not path.exists(): return {"error": "File not found", "path": file_path} # Determine type ext = path.suffix.lower() if ext in [".txt", ".md"]: content = path.read_text(encoding="utf-8") return await self.import_text( content=content, title=path.stem, tags=tags, category="file" ) elif ext == ".json": try: data = json.loads(path.read_text()) content = json.dumps(data, indent=2) return await self.import_text( content=content, title=path.stem, tags=tags, category="data" ) except Exception as e: return {"error": f"Failed to parse JSON: {e}"} else: return {"error": f"Unsupported file type: {ext}"} async def import_url( self, url: str, title: str = None, tags: list = None, ) -> dict: """ Import content from a URL (web page). Fetches the page and extracts main content. """ log.info(f"🌐 Fetching URL: {url}") try: client = httpx.Client(timeout=30) resp = client.get(url, headers={ "User-Agent": "Mozilla/5.0 (compatible; RasoSpeak/1.0)" }) resp.raise_for_status() html = resp.text # Extract text content (simple approach) content = self._extract_text_from_html(html) # Get title if not provided if not title: title_match = re.search(r"(.*?)", html, re.IGNORECASE) title = title_match.group(1) if title_match else url return await self.import_text( content=content, title=title, tags=tags + ["web", "url"] if tags else ["web", "url"], category="article" ) except Exception as e: log.error(f"Failed to fetch URL: {e}") return {"error": f"Failed to fetch URL: {e}"} async def import_pdf( self, pdf_path: str, title: str = None, tags: list = None, ) -> dict: """ Import a PDF file. Note: Requires pypdf or PyPDF2 installed. """ try: from pypdf import PdfReader except ImportError: return {"error": "pypdf not installed. Run: pip install pypdf"} path = Path(pdf_path) if not path.exists(): return {"error": "PDF not found"} try: reader = PdfReader(pdf_path) text = "" for page in reader.pages: text += page.extract_text() + "\n" title = title or path.stem return await self.import_text( content=text, title=title, tags=tags + ["pdf"] if tags else ["pdf"], category="document" ) except Exception as e: return {"error": f"Failed to read PDF: {e}"} async def import_snippet( self, content: str, label: str = None, ) -> dict: """Import a quick snippet/clipboard content.""" return await self.import_text( content=content, title=label or f"Snippet {datetime.now().strftime('%Y-%m-%d %H:%M')}", tags=["snippet", "clipboard"], category="snippet" ) # ══════════════════════════════════════════════════════ # DOCUMENT MANAGEMENT # ══════════════════════════════════════════════════════ async def list_documents( self, category: str = None, limit: int = 20, ) -> dict: """List all imported documents.""" documents = [] for path in sorted(self._storage_path.glob("*.json"), reverse=True)[:limit]: try: doc = json.loads(path.read_text()) if category and doc.get("category") != category: continue documents.append({ "id": doc.get("id"), "title": doc.get("title"), "category": doc.get("category"), "tags": doc.get("tags", []), "imported_at": doc.get("imported_at"), "word_count": doc.get("word_count"), }) except Exception: continue return { "documents": documents, "total": len(documents), } async def get_document(self, doc_id: str) -> dict: """Get a specific document.""" path = self._storage_path / f"{doc_id}.json" if path.exists(): return json.loads(path.read_text()) return {"error": "Document not found"} async def delete_document(self, doc_id: str) -> dict: """Delete a document.""" path = self._storage_path / f"{doc_id}.json" if path.exists(): path.unlink() log.info(f"Deleted document: {doc_id}") return {"deleted": True} return {"error": "Document not found"} async def search_documents(self, query: str, limit: int = 10) -> dict: """Search within imported documents.""" results = [] for path in self._storage_path.glob("*.json"): try: doc = json.loads(path.read_text()) content = doc.get("content", "").lower() title = doc.get("title", "").lower() if query.lower() in content or query.lower() in title: results.append({ "id": doc.get("id"), "title": doc.get("title"), "snippet": content[content.find(query.lower()):content.find(query.lower()) + 200] if query.lower() in content else content[:200], }) except Exception: continue return { "query": query, "results": results[:limit], "total": len(results), } # ══════════════════════════════════════════════════════ # PRIVATE HELPERS # ══════════════════════════════════════════════════════ async def _save_document(self, doc_id: str, document: dict): """Save document to disk.""" path = self._storage_path / f"{doc_id}.json" path.write_text(json.dumps(document, indent=2)) def _extract_text_from_html(self, html: str) -> str: """Extract readable text from HTML.""" # Remove script and style text = re.sub(r'', '', html, flags=re.DOTALL | re.IGNORECASE) text = re.sub(r'', '', text, flags=re.DOTALL | re.IGNORECASE) # Replace common tags with newlines text = re.sub(r'', '\n', text, flags=re.IGNORECASE) text = re.sub(r'

', '\n', text, flags=re.IGNORECASE) text = re.sub(r'', '\n', text, flags=re.IGNORECASE) # Remove remaining tags text = re.sub(r'<[^>]+>', '', text) # Clean up whitespace text = re.sub(r'\n\s*\n', '\n\n', text) text = text.strip() return text[:50000] # Limit to 50k chars def _extract_key_points(self, content: str) -> list: """Extract key points from content for memory.""" # Simple extraction - split by sentences and find important ones sentences = re.split(r'[.!?]+', content) key_points = [] for sentence in sentences[:20]: # First 20 sentences sentence = sentence.strip() if len(sentence) > 20 and len(sentence) < 200: # Look for sentences with key words key_words = ["important", "key", "remember", "note", "main", "critical", "essential"] if any(word in sentence.lower() for word in key_words): key_points.append(sentence) return key_points[:5] # Return top 5 async def shutdown(self): log.info("DocumentAgent shut down")