"""Per-user state management with plain classes and CRUD helpers.""" from datetime import datetime import uuid class Source: def __init__(self, id, filename, file_type, size_mb, source_url, chunk_count, status, error_message, created_at, file_path=None): self.id = id self.filename = filename self.file_type = file_type # "pdf", "pptx", "txt", "url", "youtube" self.size_mb = size_mb self.source_url = source_url self.chunk_count = chunk_count self.status = status # "ready", "processing", "failed" self.error_message = error_message self.created_at = created_at self.file_path = file_path def to_dict(self): return { "id": self.id, "filename": self.filename, "file_type": self.file_type, "size_mb": self.size_mb, "source_url": self.source_url, "chunk_count": self.chunk_count, "status": self.status, "error_message": self.error_message, "created_at": self.created_at, } @classmethod def from_dict(cls, d): return cls( id=d["id"], filename=d["filename"], file_type=d["file_type"], size_mb=d["size_mb"], source_url=d["source_url"], chunk_count=d["chunk_count"], status=d["status"], error_message=d["error_message"], created_at=d["created_at"], ) class Message: def __init__(self, id, role, content, citations, created_at): self.id = id self.role = role # "user" or "assistant" self.content = content self.citations = citations # [{source, page, text}] self.created_at = created_at def to_dict(self): return { "id": self.id, "role": self.role, "content": self.content, "citations": self.citations, "created_at": self.created_at, } @classmethod def from_dict(cls, d): return cls( id=d["id"], role=d["role"], content=d["content"], citations=d.get("citations", []), created_at=d["created_at"], ) class Artifact: def __init__(self, id, type, title, content, audio_path, created_at): self.id = id self.type = type # "conversation_summary", "document_summary", "podcast", "quiz" self.title = title self.content = content self.audio_path = audio_path self.created_at = created_at def to_dict(self): return { "id": self.id, "type": self.type, "title": self.title, "content": self.content, "audio_path": self.audio_path, "created_at": self.created_at, } @classmethod def from_dict(cls, d): return cls( id=d["id"], type=d["type"], title=d["title"], content=d.get("content", ""), audio_path=d.get("audio_path"), created_at=d["created_at"], ) class Notebook: def __init__(self, id, title, created_at, sources=None, messages=None, artifacts=None): self.id = id self.title = title self.created_at = created_at self.sources = sources if sources is not None else [] self.messages = messages if messages is not None else [] self.artifacts = artifacts if artifacts is not None else [] def to_dict(self): return { "id": self.id, "title": self.title, "created_at": self.created_at, "sources": [s.to_dict() for s in self.sources], "messages": [m.to_dict() for m in self.messages], "artifacts": [a.to_dict() for a in self.artifacts], } @classmethod def from_dict(cls, d): return cls( id=d["id"], title=d["title"], created_at=d["created_at"], sources=[Source.from_dict(s) for s in d.get("sources", [])], messages=[Message.from_dict(m) for m in d.get("messages", [])], artifacts=[Artifact.from_dict(a) for a in d.get("artifacts", [])], ) class UserData: def __init__(self, user_id, user_name, notebooks=None, active_notebook_id=None): self.user_id = user_id self.user_name = user_name self.notebooks = notebooks if notebooks is not None else {} self.active_notebook_id = active_notebook_id def to_dict(self): return { "user_id": self.user_id, "user_name": self.user_name, "notebooks": {nb_id: nb.to_dict() for nb_id, nb in self.notebooks.items()}, "active_notebook_id": self.active_notebook_id, } @classmethod def from_dict(cls, d): notebooks = { nb_id: Notebook.from_dict(nb_data) for nb_id, nb_data in d.get("notebooks", {}).items() } return cls( user_id=d["user_id"], user_name=d["user_name"], notebooks=notebooks, active_notebook_id=d.get("active_notebook_id"), ) def create_default_user_data(user_id, user_name): nb_id = str(uuid.uuid4()) default_nb = Notebook( id=nb_id, title="My First Notebook", created_at=datetime.now().isoformat(), ) return UserData( user_id=user_id, user_name=user_name, notebooks={nb_id: default_nb}, active_notebook_id=nb_id, ) def get_active_notebook(state): if state and state.active_notebook_id and state.active_notebook_id in state.notebooks: return state.notebooks[state.active_notebook_id] return None def create_notebook(state, title): nb_id = str(uuid.uuid4()) state.notebooks[nb_id] = Notebook( id=nb_id, title=title, created_at=datetime.now().isoformat(), ) state.active_notebook_id = nb_id return state def delete_notebook(state, nb_id): if nb_id in state.notebooks: # Clean up Pinecone vectors for this notebook try: from persistence.vector_store import VectorStore VectorStore().delete_namespace(nb_id) except Exception: pass # Best-effort cleanup del state.notebooks[nb_id] remaining = list(state.notebooks.keys()) state.active_notebook_id = remaining[0] if remaining else None return state def rename_notebook(state, nb_id, new_title): if nb_id in state.notebooks: state.notebooks[nb_id].title = new_title return state def get_notebook_choices(state): """Return list of (display_label, notebook_id) for gr.Radio.""" choices = [] for nb_id, nb in state.notebooks.items(): src_count = len(nb.sources) msg_count = len(nb.messages) label = nb.title if src_count > 0 or msg_count > 0: label += f" ({src_count}s, {msg_count}m)" choices.append((label, nb_id)) return choices def get_all_artifacts(notebook, artifact_type): return [a for a in reversed(notebook.artifacts) if a.type == artifact_type] def get_latest_artifact(notebook, artifact_type): for artifact in reversed(notebook.artifacts): if artifact.type == artifact_type: return artifact return None