db_manager.py

import sqlite3
from dataclasses import dataclass
from typing import List, Optional
import json
import datetime
import os

@dataclass
class ArxivDocument:
    arxiv_id: str
    title: str
    paragraphs: List[str]
    summaries: List[str]
    html_content: str
    bibliography: str
    processed_date: str

class DocumentStore:
    def __init__(self, db_path: str = "arxiv_docs.sqlite"):
        self.db_path = db_path
        self._init_schema()
    
    def _init_schema(self):
        with sqlite3.connect(self.db_path) as conn:
            conn.execute("""
                CREATE TABLE IF NOT EXISTS documents (
                    arxiv_id TEXT PRIMARY KEY,
                    title TEXT,
                    paragraphs TEXT,
                    summaries TEXT,
                    html_content TEXT,
                    bibliography TEXT,
                    processed_date TEXT
                )
            """)
    
    def get_document(self, arxiv_id: str) -> Optional[ArxivDocument]:
        with sqlite3.connect(self.db_path) as conn:
            cursor = conn.cursor()
            cursor.execute(
                "SELECT * FROM documents WHERE arxiv_id = ?",
                (arxiv_id,)
            )
            result = cursor.fetchone()
            
            if result:
                return ArxivDocument(
                    arxiv_id=result[0],
                    title=result[1],
                    paragraphs=json.loads(result[2]),
                    summaries=json.loads(result[3]),
                    html_content=result[4],
                    bibliography=result[5],
                    processed_date=result[6]
                )
        return None

    def save_document(self, doc: ArxivDocument):
        with sqlite3.connect(self.db_path) as conn:
            cursor = conn.cursor()
            cursor.execute("""
                INSERT OR REPLACE INTO documents 
                (arxiv_id, title, paragraphs, summaries, html_content, bibliography, processed_date)
                VALUES (?, ?, ?, ?, ?, ?, ?)
            """, (
                doc.arxiv_id,
                doc.title,
                json.dumps(doc.paragraphs),
                json.dumps(doc.summaries),
                doc.html_content,
                doc.bibliography,
                doc.processed_date
            ))
            conn.commit()