feat: implement remove document

Daniele Briggi · Daniele Briggi · commit 3ce34c681d2a · 2025-08-06T10:13:48.000Z
diff --git a/src/sqlite_rag/cli.py b/src/sqlite_rag/cli.py
@@ -77,7 +77,35 @@ def list_documents():
 @app.command()
 def remove(identifier: str):
     """Remove document by path or UUID"""
-    pass
+    rag = SQLiteRag()
+    
+    # Find the document first
+    document = rag.find_document(identifier)
+    if not document:
+        typer.echo(f"Document not found: {identifier}")
+        raise typer.Exit(1)
+    
+    # Show document details
+    typer.echo(f"Found document:")
+    typer.echo(f"ID: {document.id}")
+    typer.echo(f"URI: {document.uri or 'N/A'}")
+    typer.echo(f"Created: {document.created_at.strftime('%Y-%m-%d %H:%M:%S') if document.created_at else 'N/A'}")
+    typer.echo(f"Content preview: {document.content[:200]}{'...' if len(document.content) > 200 else ''}")
+    typer.echo()
+    
+    # Ask for confirmation
+    confirm = typer.confirm("Are you sure you want to delete this document?")
+    if not confirm:
+        typer.echo("Cancelled.")
+        return
+    
+    # Remove the document
+    success = rag.remove_document(identifier)
+    if success:
+        typer.echo("Document removed successfully.")
+    else:
+        typer.echo("Failed to remove document.")
+        raise typer.Exit(1)
 
 
 @app.command()
diff --git a/src/sqlite_rag/repository.py b/src/sqlite_rag/repository.py
@@ -57,3 +57,41 @@ def list_documents(self) -> list[Document]:
             )
 
         return documents
+
+    def find_document_by_id_or_uri(self, identifier: str) -> Document | None:
+        """Find document by ID or URI"""
+        cursor = self._conn.cursor()
+        cursor.execute(
+            "SELECT id, content, uri, metadata, created_at FROM documents WHERE id = ? OR uri = ?", 
+            (identifier, identifier)
+        )
+        row = cursor.fetchone()
+        
+        if row:
+            doc_id, content, uri, metadata, created_at = row
+            return Document(
+                id=doc_id,
+                content=content,
+                uri=uri,
+                metadata=json.loads(metadata),
+                created_at=created_at,
+            )
+        return None
+
+    def remove_document(self, document_id: str) -> bool:
+        """Remove document and its chunks by document ID"""
+        cursor = self._conn.cursor()
+        
+        # Check if document exists
+        cursor.execute("SELECT COUNT(*) FROM documents WHERE id = ?", (document_id,))
+        if cursor.fetchone()[0] == 0:
+            return False
+        
+        # Remove chunks first (foreign key constraint)
+        cursor.execute("DELETE FROM chunks WHERE document_id = ?", (document_id,))
+        
+        # Remove document
+        cursor.execute("DELETE FROM documents WHERE id = ?", (document_id,))
+        
+        self._conn.commit()
+        return True
diff --git a/src/sqlite_rag/sqliterag.py b/src/sqlite_rag/sqliterag.py
@@ -97,6 +97,22 @@ def list_documents(self) -> list[Document]:
 
         return self._repository.list_documents()
 
+    def find_document(self, identifier: str) -> Document | None:
+        """Find document by ID or URI"""
+        self._ensure_initialized()
+        return self._repository.find_document_by_id_or_uri(identifier)
+
+    def remove_document(self, identifier: str) -> bool:
+        """Remove document by ID or URI"""
+        self._ensure_initialized()
+        
+        # First find the document to get its ID
+        document = self._repository.find_document_by_id_or_uri(identifier)
+        if not document:
+            return False
+        
+        return self._repository.remove_document(document.id or "")
+
     # def search(
     #     self, query: str, top_k: int = 10
     # ) -> list[Document]:
diff --git a/tests/test_engine.py b/tests/test_engine.py
@@ -1,10 +1,10 @@
 import pytest
 
-from chunker import Chunker
-from engine import Engine
-from models.chunk import Chunk
-from models.document import Document
-from repository import Repository
+from sqlite_rag.chunker import Chunker
+from sqlite_rag.engine import Engine
+from sqlite_rag.models.chunk import Chunk
+from sqlite_rag.models.document import Document
+from sqlite_rag.repository import Repository
 
 
 @pytest.fixture
diff --git a/tests/test_repository.py b/tests/test_repository.py
@@ -113,3 +113,96 @@ def test_list_documents_empty(self, db_conn):
         documents = repo.list_documents()
 
         assert len(documents) == 0
+
+    def test_find_document_by_id_or_uri_by_id(self, db_conn):
+        conn, settings = db_conn
+        repo = Repository(conn, settings)
+
+        # Add a document
+        doc = Document(
+            content="Test document content.", 
+            uri="test.txt", 
+            metadata={"author": "test"}
+        )
+        doc_id = repo.add_document(doc)
+
+        # Find by ID
+        found_doc = repo.find_document_by_id_or_uri(doc_id)
+
+        assert found_doc is not None
+        assert found_doc.id == doc_id
+        assert found_doc.content == "Test document content."
+        assert found_doc.uri == "test.txt"
+        assert found_doc.metadata == {"author": "test"}
+
+    def test_find_document_by_id_or_uri_by_uri(self, db_conn):
+        conn, settings = db_conn
+        repo = Repository(conn, settings)
+
+        # Add a document
+        doc = Document(
+            content="Test document content.", 
+            uri="test.txt", 
+            metadata={"author": "test"}
+        )
+        repo.add_document(doc)
+
+        # Find by URI
+        found_doc = repo.find_document_by_id_or_uri("test.txt")
+
+        assert found_doc is not None
+        assert found_doc.content == "Test document content."
+        assert found_doc.uri == "test.txt"
+        assert found_doc.metadata == {"author": "test"}
+
+    def test_find_document_by_id_or_uri_not_found(self, db_conn):
+        conn, settings = db_conn
+        repo = Repository(conn, settings)
+
+        # Try to find non-existent document
+        found_doc = repo.find_document_by_id_or_uri("nonexistent")
+
+        assert found_doc is None
+
+    def test_remove_document_success(self, db_conn):
+        conn, settings = db_conn
+        repo = Repository(conn, settings)
+
+        # Add a document with chunks
+        doc = Document(
+            content="Test document content.", 
+            uri="test.txt", 
+            metadata={"author": "test"}
+        )
+        doc.chunks = [
+            Chunk(content="Chunk 1", embedding=b"\x00" * 384),
+            Chunk(content="Chunk 2", embedding=b"\x00" * 384),
+        ]
+        doc_id = repo.add_document(doc)
+
+        # Verify document and chunks exist
+        cursor = conn.cursor()
+        cursor.execute("SELECT COUNT(*) FROM documents WHERE id = ?", (doc_id,))
+        assert cursor.fetchone()[0] == 1
+        cursor.execute("SELECT COUNT(*) FROM chunks WHERE document_id = ?", (doc_id,))
+        assert cursor.fetchone()[0] == 2
+
+        # Remove document
+        success = repo.remove_document(doc_id)
+
+        assert success is True
+
+        # Verify document and chunks are removed
+        cursor.execute("SELECT COUNT(*) FROM documents WHERE id = ?", (doc_id,))
+        assert cursor.fetchone()[0] == 0
+        cursor.execute("SELECT COUNT(*) FROM chunks WHERE document_id = ?", (doc_id,))
+        assert cursor.fetchone()[0] == 0
+
+    def test_remove_document_not_found(self, db_conn):
+        conn, settings = db_conn
+        repo = Repository(conn, settings)
+
+        # Try to remove non-existent document
+        success = repo.remove_document("nonexistent-id")
+
+        assert success is False
diff --git a/tests/test_sqlite_rag.py b/tests/test_sqlite_rag.py
@@ -99,3 +99,115 @@ def test_list_documents(self, db_settings):
         assert len(documents) == 2
         assert documents[0].content == "Document 1 content."
         assert documents[1].content == "Document 2 content."
+
+    def test_find_document_by_id(self, db_settings):
+        rag = SQLiteRag(db_settings)
+        
+        rag.add_text("Test document content.", uri="test.txt", metadata={"author": "test"})
+        documents = rag.list_documents()
+        doc_id = documents[0].id
+        
+        # Find by ID
+        assert doc_id is not None
+        found_doc = rag.find_document(doc_id)
+        
+        assert found_doc is not None
+        assert found_doc.id == doc_id
+        assert found_doc.content == "Test document content."
+        assert found_doc.uri == "test.txt"
+        assert found_doc.metadata == {"author": "test"}
+
+    def test_find_document_by_uri(self, db_settings):
+        rag = SQLiteRag(db_settings)
+        
+        rag.add_text("Test document content.", uri="test.txt", metadata={"author": "test"})
+        
+        # Find by URI
+        found_doc = rag.find_document("test.txt")
+        
+        assert found_doc is not None
+        assert found_doc.content == "Test document content."
+        assert found_doc.uri == "test.txt"
+        assert found_doc.metadata == {"author": "test"}
+
+    def test_find_document_not_found(self, db_settings):
+        rag = SQLiteRag(db_settings)
+        
+        found_doc = rag.find_document("nonexistent")
+        
+        assert found_doc is None
+
+    def test_remove_document_by_id(self, db_settings):
+        rag = SQLiteRag(db_settings)
+        
+        rag.add_text("Test document content.", uri="test.txt", metadata={"author": "test"})
+        documents = rag.list_documents()
+        doc_id = documents[0].id
+        
+        # Verify document exists
+        assert len(documents) == 1
+        
+        # Remove by ID
+        assert doc_id is not None
+        success = rag.remove_document(doc_id)
+        
+        assert success is True
+        
+        # Verify document is removed
+        documents = rag.list_documents()
+        assert len(documents) == 0
+
+    def test_remove_document_by_uri(self, db_settings):
+        rag = SQLiteRag(db_settings)
+        
+        rag.add_text("Test document content.", uri="test.txt", metadata={"author": "test"})
+        
+        # Verify document exists
+        documents = rag.list_documents()
+        assert len(documents) == 1
+        
+        # Remove by URI
+        success = rag.remove_document("test.txt")
+        
+        assert success is True
+        
+        # Verify document is removed
+        documents = rag.list_documents()
+        assert len(documents) == 0
+
+    def test_remove_document_not_found(self, db_settings):
+        rag = SQLiteRag(db_settings)
+        
+        success = rag.remove_document("nonexistent")
+        
+        assert success is False
+
+    def test_remove_document_with_chunks(self, db_settings):
+        rag = SQLiteRag(db_settings)
+        
+        # Add document that will create chunks
+        rag.add_text("This is a longer document that should create multiple chunks when processed by the chunker.", uri="test.txt")
+        
+        # Verify document and chunks exist
+        documents = rag.list_documents()
+        assert len(documents) == 1
+        doc_id = documents[0].id
+        
+        cursor = rag._conn.cursor()
+        cursor.execute("SELECT COUNT(*) FROM chunks WHERE document_id = ?", (doc_id,))
+        chunk_count = cursor.fetchone()[0]
+        assert chunk_count > 0
+        
+        # Remove document
+        assert doc_id is not None
+        success = rag.remove_document(doc_id)
+        
+        assert success is True
+        
+        # Verify document and chunks are removed
+        documents = rag.list_documents()
+        assert len(documents) == 0
+        
+        cursor.execute("SELECT COUNT(*) FROM chunks WHERE document_id = ?", (doc_id,))
+        chunk_count = cursor.fetchone()[0]
+        assert chunk_count == 0
