import os from datetime import datetime from pathlib import Path from typing import Any, ClassVar, Dict, List, Literal, Optional, Union from loguru import logger from pydantic import BaseModel, ConfigDict, Field, field_validator from surreal_commands import submit_command from surrealdb import RecordID from open_notebook.database.repository import ensure_record_id, repo_query from open_notebook.domain.base import ObjectModel from open_notebook.exceptions import DatabaseOperationError, InvalidInputError class Notebook(ObjectModel): table_name: ClassVar[str] = "notebook" name: str description: str archived: Optional[bool] = False last_viewed_at: Optional[datetime] = None @field_validator("name") @classmethod def name_must_not_be_empty(cls, v): if not v.strip(): raise InvalidInputError("Notebook name cannot be empty") return v async def get_sources(self, include_full_text: bool = False) -> List["Source"]: try: source_projection = "" if include_full_text else " omit source.full_text" srcs = await repo_query( f""" select *{source_projection} from ( select in as source from reference where out=$id fetch source ) order by source.updated desc """, {"id": ensure_record_id(self.id)}, ) return [Source(**src["source"]) for src in srcs] if srcs else [] except Exception as e: logger.error(f"Error fetching sources for notebook {self.id}: {str(e)}") logger.exception(e) raise DatabaseOperationError(e) async def get_notes(self, include_content: bool = False) -> List["Note"]: try: note_projection = ( " omit note.embedding" if include_content else " omit note.content, note.embedding" ) srcs = await repo_query( f""" select *{note_projection} from ( select in as note from artifact where out=$id fetch note ) order by note.updated desc """, {"id": ensure_record_id(self.id)}, ) return [Note(**src["note"]) for src in srcs] if srcs else [] except Exception as e: logger.error(f"Error fetching notes for notebook {self.id}: {str(e)}") logger.exception(e) raise DatabaseOperationError(e) async def get_context(self) -> str: """ Build long-form notebook context for podcast and LLM workflows. Normal list retrieval omits large source/note bodies, so this method uses opt-in full-content fetches and formats only substantive context blocks. """ sources = await self.get_sources(include_full_text=True) notes = await self.get_notes(include_content=True) context_blocks = [] insights_by_source = await SourceInsight.get_for_sources( [source.id for source in sources if source.id] ) for source in sources: source_context = await source.get_context( context_size="long", insights=insights_by_source.get(source.id or "", []), ) if isinstance(source_context, dict): title = source_context.get("title") or source.title or "Untitled source" full_text = source_context.get("full_text") insights = source_context.get("insights") or [] content_parts = [] if full_text: content_parts.append(str(full_text)) insight_lines = [] for insight in insights: if not isinstance(insight, dict): continue insight_content = insight.get("content") if not insight_content: continue insight_type = insight.get("insight_type") or "Insight" insight_lines.append(f"- {insight_type}: {insight_content}") if insight_lines: content_parts.append("Insights:\n" + "\n".join(insight_lines)) content = "\n\n".join(content_parts).strip() else: title = source.title or "Untitled source" content = str(source_context).strip() if content: context_blocks.append(f"## Source: {title}\n\n{content}") for note in notes: note_context = note.get_context(context_size="long") if isinstance(note_context, dict): title = note_context.get("title") or note.title or "Untitled note" content = note_context.get("content") content = str(content).strip() if content else "" else: title = note.title or "Untitled note" content = str(note_context).strip() if content: context_blocks.append(f"## Note: {title}\n\n{content}") return "\n\n".join(context_blocks) async def get_chat_sessions(self) -> List["ChatSession"]: try: srcs = await repo_query( """ select * from ( select <- chat_session as chat_session from refers_to where out=$id fetch chat_session ) order by chat_session.updated desc """, {"id": ensure_record_id(self.id)}, ) return ( [ChatSession(**src["chat_session"][0]) for src in srcs] if srcs else [] ) except Exception as e: logger.error( f"Error fetching chat sessions for notebook {self.id}: {str(e)}" ) logger.exception(e) raise DatabaseOperationError(e) async def get_delete_preview(self) -> Dict[str, Any]: """ Get counts of items that would be affected by deleting this notebook. Returns a dict with: - note_count: Number of notes that will be deleted - exclusive_source_count: Sources only in this notebook (can be deleted) - shared_source_count: Sources in other notebooks (will be unlinked only) """ try: notebook_id = ensure_record_id(self.id) # Count notes note_result = await repo_query( "SELECT count() as count FROM artifact WHERE out = $notebook_id GROUP ALL", {"notebook_id": notebook_id}, ) note_count = note_result[0]["count"] if note_result else 0 # Get sources with count of references to OTHER notebooks # If assigned_others = 0, source is exclusive to this notebook # If assigned_others > 0, source is shared with other notebooks source_counts = await repo_query( """ SELECT id, count(->reference[WHERE out != $notebook_id].out) as assigned_others FROM (SELECT VALUE <-reference.in AS sources FROM $notebook_id)[0] """, {"notebook_id": notebook_id}, ) exclusive_count = 0 shared_count = 0 for src in source_counts: if src.get("assigned_others", 0) == 0: exclusive_count += 1 else: shared_count += 1 return { "note_count": note_count, "exclusive_source_count": exclusive_count, "shared_source_count": shared_count, } except Exception as e: logger.error(f"Error getting delete preview for notebook {self.id}: {e}") logger.exception(e) raise DatabaseOperationError(e) async def delete(self, delete_exclusive_sources: bool = False) -> Dict[str, int]: """ Delete notebook with cascade deletion of notes and optional source deletion. Args: delete_exclusive_sources: If True, also delete sources that belong only to this notebook. Default is False. Returns: Dict with counts: deleted_notes, deleted_sources, unlinked_sources """ if self.id is None: raise InvalidInputError("Cannot delete notebook without an ID") try: notebook_id = ensure_record_id(self.id) deleted_notes = 0 deleted_sources = 0 unlinked_sources = 0 # 1. Get and delete all notes linked to this notebook notes = await self.get_notes() for note in notes: await note.delete() deleted_notes += 1 logger.info(f"Deleted {deleted_notes} notes for notebook {self.id}") # Delete artifact relationships await repo_query( "DELETE artifact WHERE out = $notebook_id", {"notebook_id": notebook_id}, ) # 2. Handle sources if delete_exclusive_sources: # Find sources with count of references to OTHER notebooks # If assigned_others = 0, source is exclusive to this notebook source_counts = await repo_query( """ SELECT id, count(->reference[WHERE out != $notebook_id].out) as assigned_others FROM (SELECT VALUE <-reference.in AS sources FROM $notebook_id)[0] """, {"notebook_id": notebook_id}, ) for src in source_counts: source_id = src.get("id") if source_id and src.get("assigned_others", 0) == 0: # Exclusive source - delete it try: source = await Source.get(str(source_id)) await source.delete() deleted_sources += 1 except Exception as e: logger.warning( f"Failed to delete exclusive source {source_id}: {e}" ) else: unlinked_sources += 1 else: # Just count sources that will be unlinked source_result = await repo_query( "SELECT count() as count FROM reference WHERE out = $notebook_id GROUP ALL", {"notebook_id": notebook_id}, ) unlinked_sources = source_result[0]["count"] if source_result else 0 # Delete reference relationships (unlink all sources) await repo_query( "DELETE reference WHERE out = $notebook_id", {"notebook_id": notebook_id}, ) logger.info( f"Unlinked {unlinked_sources} sources, deleted {deleted_sources} " f"exclusive sources for notebook {self.id}" ) # 3. Delete the notebook record itself await super().delete() logger.info(f"Deleted notebook {self.id}") return { "deleted_notes": deleted_notes, "deleted_sources": deleted_sources, "unlinked_sources": unlinked_sources, } except Exception as e: logger.error(f"Error deleting notebook {self.id}: {e}") logger.exception(e) raise DatabaseOperationError(f"Failed to delete notebook: {e}") class Asset(BaseModel): file_path: Optional[str] = None url: Optional[str] = None class SourceEmbedding(ObjectModel): table_name: ClassVar[str] = "source_embedding" content: str async def get_source(self) -> "Source": try: src = await repo_query( """ select source.* from $id fetch source """, {"id": ensure_record_id(self.id)}, ) return Source(**src[0]["source"]) except Exception as e: logger.error(f"Error fetching source for embedding {self.id}: {str(e)}") logger.exception(e) raise DatabaseOperationError(e) class SourceInsight(ObjectModel): table_name: ClassVar[str] = "source_insight" insight_type: str content: str @classmethod async def get_for_sources( cls, source_ids: List[str] ) -> Dict[str, List["SourceInsight"]]: """ Batch-fetch insights for many sources in a single query. Building notebook/chat context otherwise calls get_insights() once per source - fine for one source, but O(n) round trips (each paying its own connection setup - no pooling in the repository layer) when a caller loops over every source in a notebook. """ grouped: Dict[str, List[SourceInsight]] = {sid: [] for sid in source_ids if sid} if not grouped: return grouped try: result = await repo_query( "SELECT * FROM source_insight WHERE source IN $source_ids", {"source_ids": [ensure_record_id(sid) for sid in grouped]}, ) except Exception as e: logger.error(f"Error batch-fetching insights for sources: {str(e)}") logger.exception(e) raise DatabaseOperationError("Failed to fetch insights for sources") for row in result: key = str(row.get("source")) grouped.setdefault(key, []).append(cls(**row)) return grouped async def get_source(self) -> "Source": try: src = await repo_query( """ select source.* from $id fetch source """, {"id": ensure_record_id(self.id)}, ) return Source(**src[0]["source"]) except Exception as e: logger.error(f"Error fetching source for insight {self.id}: {str(e)}") logger.exception(e) raise DatabaseOperationError(e) async def save_as_note(self, notebook_id: Optional[str] = None) -> Any: source = await self.get_source() note = Note( title=f"{self.insight_type} from source {source.title}", content=self.content, ) await note.save() if notebook_id: await note.add_to_notebook(notebook_id) return note class Source(ObjectModel): model_config = ConfigDict(arbitrary_types_allowed=True) table_name: ClassVar[str] = "source" asset: Optional[Asset] = None title: Optional[str] = None topics: Optional[List[str]] = Field(default_factory=list) full_text: Optional[str] = None last_viewed_at: Optional[datetime] = None command: Optional[Union[str, RecordID]] = Field( default=None, description="Link to surreal-commands processing job" ) @field_validator("command", mode="before") @classmethod def parse_command(cls, value): """Parse command field to ensure RecordID format""" if isinstance(value, str) and value: return ensure_record_id(value) return value @field_validator("id", mode="before") @classmethod def parse_id(cls, value): """Parse id field to handle both string and RecordID inputs""" if value is None: return None if isinstance(value, RecordID): return str(value) return str(value) if value else None async def get_status(self) -> Optional[str]: """Get the processing status of the associated command""" if not self.command: return None try: from surreal_commands import get_command_status status = await get_command_status(str(self.command)) return status.status if status else "unknown" except Exception as e: logger.warning(f"Failed to get command status for {self.command}: {e}") return "unknown" async def get_processing_progress(self) -> Optional[Dict[str, Any]]: """Get detailed processing information for the associated command""" if not self.command: return None try: from surreal_commands import get_command_status status_result = await get_command_status(str(self.command)) if not status_result: return None # Extract execution metadata if available result = getattr(status_result, "result", None) execution_metadata = ( result.get("execution_metadata", {}) if isinstance(result, dict) else {} ) return { "status": status_result.status, "started_at": execution_metadata.get("started_at"), "completed_at": execution_metadata.get("completed_at"), "error": getattr(status_result, "error_message", None), "result": result, } except Exception as e: logger.warning(f"Failed to get command progress for {self.command}: {e}") return None async def get_context( self, context_size: Literal["short", "long"] = "short", insights: Optional[List["SourceInsight"]] = None, ) -> Dict[str, Any]: # Callers looping over many sources can batch-fetch insights up front # via SourceInsight.get_for_sources() and pass them in here, instead # of paying a separate query per source. insight_objects = insights if insights is not None else await self.get_insights() insights = [insight.model_dump() for insight in insight_objects] if context_size == "long": return dict( id=self.id, title=self.title, insights=insights, full_text=self.full_text, ) else: return dict(id=self.id, title=self.title, insights=insights) async def get_embedded_chunks(self) -> int: try: result = await repo_query( """ select count() as chunks from source_embedding where source=$id GROUP ALL """, {"id": ensure_record_id(self.id)}, ) if len(result) == 0: return 0 return result[0]["chunks"] except Exception as e: logger.error(f"Error fetching chunks count for source {self.id}: {str(e)}") logger.exception(e) raise DatabaseOperationError(f"Failed to count chunks for source: {str(e)}") async def get_insights(self) -> List[SourceInsight]: try: result = await repo_query( """ SELECT * FROM source_insight WHERE source=$id """, {"id": ensure_record_id(self.id)}, ) return [SourceInsight(**insight) for insight in result] except Exception as e: logger.error(f"Error fetching insights for source {self.id}: {str(e)}") logger.exception(e) raise DatabaseOperationError("Failed to fetch insights for source") async def add_to_notebook(self, notebook_id: str) -> Any: if not notebook_id: raise InvalidInputError("Notebook ID must be provided") await Notebook.get(notebook_id) # raises NotFoundError if invalid/missing return await self.relate("reference", notebook_id) async def vectorize(self) -> str: """ Submit vectorization as a background job using the embed_source command. This method leverages the job-based architecture to prevent HTTP connection pool exhaustion when processing large documents. The embed_source command: 1. Detects content type from file path 2. Chunks text using content-type aware splitter 3. Generates all embeddings in batches 4. Bulk inserts source_embedding records Returns: str: The command/job ID that can be used to track progress via the commands API Raises: ValueError: If source has no text to vectorize DatabaseOperationError: If job submission fails """ logger.info(f"Submitting embed_source job for source {self.id}") try: if not self.full_text or not self.full_text.strip(): raise ValueError(f"Source {self.id} has no text to vectorize") # Submit the embed_source command command_id = submit_command( "open_notebook", "embed_source", {"source_id": str(self.id)}, ) command_id_str = str(command_id) logger.info( f"Embed source job submitted for source {self.id}: " f"command_id={command_id_str}" ) return command_id_str except ValueError: raise except Exception as e: logger.error(f"Failed to submit embed_source job for source {self.id}: {e}") logger.exception(e) raise DatabaseOperationError(e) async def add_insight(self, insight_type: str, content: str) -> str: """ Submit insight creation as an async command (fire-and-forget). Submits a create_insight command that handles database operations with automatic retry logic for transaction conflicts. The command also submits an embed_insight command for async embedding. This method returns immediately after submitting the command - it does NOT wait for the insight to be created. Use this for batch operations where throughput is more important than immediate confirmation. Args: insight_type: Type/category of the insight content: The insight content text Returns: command_id for optional tracking Raises: InvalidInputError: If insight_type or content is empty DatabaseOperationError: If submitting the command fails. Matches vectorize()'s contract - callers (transformation.py, source.py) run inside surreal-commands jobs whose outer exception handling already retries transient failures, so a swallowed submission failure here previously meant a transformation could report success while the insight was silently never persisted. """ if not insight_type or not content: raise InvalidInputError("Insight type and content must be provided") try: # Submit create_insight command (fire-and-forget) # Command handles retries internally for transaction conflicts command_id = submit_command( "open_notebook", "create_insight", { "source_id": str(self.id), "insight_type": insight_type, "content": content, }, ) logger.info( f"Submitted create_insight command {command_id} for source {self.id} " f"(type={insight_type})" ) return str(command_id) except Exception as e: logger.exception(f"Error submitting create_insight for source {self.id}: {e}") raise DatabaseOperationError(e) def _prepare_save_data(self) -> dict: """Override to ensure command field is always RecordID format for database""" data = super()._prepare_save_data() # Ensure command field is RecordID format if not None if data.get("command") is not None: data["command"] = ensure_record_id(data["command"]) return data async def delete(self) -> bool: """Delete source and clean up associated file, embeddings, and insights.""" # Clean up uploaded file if it exists if self.asset and self.asset.file_path: file_path = Path(self.asset.file_path) if file_path.exists(): try: os.unlink(file_path) logger.info(f"Deleted file for source {self.id}: {file_path}") except Exception as e: logger.warning( f"Failed to delete file {file_path} for source {self.id}: {e}. " "Continuing with database deletion." ) else: logger.debug( f"File {file_path} not found for source {self.id}, skipping cleanup" ) # Delete associated embeddings and insights to prevent orphaned records try: source_id = ensure_record_id(self.id) await repo_query( "DELETE source_embedding WHERE source = $source_id", {"source_id": source_id}, ) await repo_query( "DELETE source_insight WHERE source = $source_id", {"source_id": source_id}, ) logger.debug(f"Deleted embeddings and insights for source {self.id}") except Exception as e: logger.warning( f"Failed to delete embeddings/insights for source {self.id}: {e}. " "Continuing with source deletion." ) # Call parent delete to remove database record return await super().delete() class Note(ObjectModel): table_name: ClassVar[str] = "note" title: Optional[str] = None note_type: Optional[Literal["human", "ai"]] = None content: Optional[str] = None @field_validator("content") @classmethod def content_must_not_be_empty(cls, v): if v is not None and not v.strip(): raise InvalidInputError("Note content cannot be empty") return v async def save(self) -> Optional[str]: """ Save the note and submit embedding command. Overrides ObjectModel.save() to submit an async embed_note command after saving, instead of inline embedding. Returns: Optional[str]: The command_id if embedding was submitted, None otherwise (either no content to embed, or submission failed) """ # Call parent save (without embedding) await super().save() # Submit embedding command (fire-and-forget) if note has content. # Unlike Source.vectorize()/add_insight() (explicit, dedicated calls # whose whole point is the submission), this runs automatically # inside save() - the note itself is already durably saved above, # so a submission hiccup here shouldn't fail an otherwise-successful # save with a 500. Best-effort: log and move on. if self.id and self.content and self.content.strip(): try: command_id = submit_command( "open_notebook", "embed_note", {"note_id": str(self.id)}, ) logger.debug(f"Submitted embed_note command {command_id} for {self.id}") return command_id except Exception as e: logger.error(f"Failed to submit embed_note command for {self.id}: {e}") return None return None async def add_to_notebook(self, notebook_id: str) -> Any: if not notebook_id: raise InvalidInputError("Notebook ID must be provided") await Notebook.get(notebook_id) # raises NotFoundError if invalid/missing return await self.relate("artifact", notebook_id) def get_context( self, context_size: Literal["short", "long"] = "short" ) -> Dict[str, Any]: if context_size == "long": return dict(id=self.id, title=self.title, content=self.content) else: return dict( id=self.id, title=self.title, content=self.content[:100] if self.content else None, ) class ChatSession(ObjectModel): table_name: ClassVar[str] = "chat_session" nullable_fields: ClassVar[set[str]] = {"model_override"} title: Optional[str] = None model_override: Optional[str] = None async def relate_to_notebook(self, notebook_id: str) -> Any: if not notebook_id: raise InvalidInputError("Notebook ID must be provided") return await self.relate("refers_to", notebook_id) async def relate_to_source(self, source_id: str) -> Any: if not source_id: raise InvalidInputError("Source ID must be provided") return await self.relate("refers_to", source_id) async def text_search( keyword: str, results: int, source: bool = True, note: bool = True ): if not keyword: raise InvalidInputError("Search keyword cannot be empty") try: search_results = await repo_query( """ select * from fn::text_search($keyword, $results, $source, $note) """, {"keyword": keyword, "results": results, "source": source, "note": note}, ) return search_results except RuntimeError as e: # SurrealDB's search::highlight can compute a byte position that exceeds the # stored string length on large or multi-byte chunks, aborting the whole query # ("position overflow"). Fall back to vector search so the user still gets # results instead of a 500. See issue #648. if "position overflow" in str(e): logger.warning( f"Highlight position overflow, falling back to vector search: {str(e)}" ) try: return await vector_search(keyword, results, source, note) except Exception as ve: # Both search paths failed (e.g. no embedding model configured). # Surface the failure instead of returning [] — an empty list would # be indistinguishable from a legitimate "no matches" and mask a # total search outage from callers. logger.error(f"Vector search fallback also failed: {str(ve)}") logger.exception(ve) raise DatabaseOperationError(ve) logger.error(f"Error performing text search: {str(e)}") logger.exception(e) raise DatabaseOperationError(e) except Exception as e: logger.error(f"Error performing text search: {str(e)}") logger.exception(e) raise DatabaseOperationError(e) async def vector_search( keyword: str, results: int, source: bool = True, note: bool = True, minimum_score=0.2, ): if not keyword: raise InvalidInputError("Search keyword cannot be empty") try: from open_notebook.utils.embedding import generate_embedding # Use unified embedding function (handles chunking if query is very long) embed = await generate_embedding(keyword) search_results = await repo_query( """ SELECT * FROM fn::vector_search($embed, $results, $source, $note, $minimum_score); """, { "embed": embed, "results": results, "source": source, "note": note, "minimum_score": minimum_score, }, ) return search_results except Exception as e: logger.error(f"Error performing vector search: {str(e)}") logger.exception(e) raise DatabaseOperationError(e)