llamaindex vector index, llmindex mangement command

src/paperless/ai/embedding.py

@@ -0,0 +1,67 @@
+from llama_index.embeddings.huggingface import HuggingFaceEmbedding
+from llama_index.embeddings.openai import OpenAIEmbedding
+
+from documents.models import Document
+from documents.models import Note
+from paperless.config import AIConfig
+
+EMBEDDING_DIMENSIONS = {
+    "text-embedding-3-small": 1536,
+    "sentence-transformers/all-MiniLM-L6-v2": 384,
+}
+
+
+def get_embedding_model():
+    config = AIConfig()
+
+    match config.llm_embedding_backend:
+        case "openai":
+            return OpenAIEmbedding(
+                model=config.llm_embedding_model or "text-embedding-3-small",
+                api_key=config.llm_api_key,
+            )
+        case "local":
+            return HuggingFaceEmbedding(
+                model_name=config.llm_embedding_model
+                or "sentence-transformers/all-MiniLM-L6-v2",
+            )
+        case _:
+            raise ValueError(
+                f"Unsupported embedding backend: {config.llm_embedding_backend}",
+            )
+
+
+def get_embedding_dim() -> int:
+    config = AIConfig()
+    model = config.llm_embedding_model or (
+        "text-embedding-3-small"
+        if config.llm_embedding_backend == "openai"
+        else "sentence-transformers/all-MiniLM-L6-v2"
+    )
+    if model not in EMBEDDING_DIMENSIONS:
+        raise ValueError(f"Unknown embedding model: {model}")
+    return EMBEDDING_DIMENSIONS[model]
+
+
+def build_llm_index_text(doc: Document) -> str:
+    lines = [
+        f"Title: {doc.title}",
+        f"Filename: {doc.filename}",
+        f"Created: {doc.created}",
+        f"Added: {doc.added}",
+        f"Modified: {doc.modified}",
+        f"Tags: {', '.join(tag.name for tag in doc.tags.all())}",
+        f"Document Type: {doc.document_type.name if doc.document_type else ''}",
+        f"Correspondent: {doc.correspondent.name if doc.correspondent else ''}",
+        f"Storage Path: {doc.storage_path.name if doc.storage_path else ''}",
+        f"Archive Serial Number: {doc.archive_serial_number or ''}",
+        f"Notes: {','.join([str(c.note) for c in Note.objects.filter(document=doc)])}",
+    ]
+
+    for instance in doc.custom_fields.all():
+        lines.append(f"Custom Field - {instance.field.name}: {instance}")
+
+    lines.append("\nContent:\n")
+    lines.append(doc.content or "")
+
+    return "\n".join(lines)

src/paperless/ai/indexing.py

@@ -0,0 +1,52 @@
+import logging
+
+import llama_index.core.settings as llama_settings
+from django.conf import settings
+from llama_index.core import StorageContext
+from llama_index.core import VectorStoreIndex
+from llama_index.core import load_index_from_storage
+from llama_index.core.retrievers import VectorIndexRetriever
+from llama_index.vector_stores.faiss import FaissVectorStore
+
+from documents.models import Document
+from paperless.ai.embedding import get_embedding_model
+
+logger = logging.getLogger("paperless.ai.indexing")
+
+
+def load_index() -> VectorStoreIndex:
+    """Loads the persisted LlamaIndex from disk."""
+    vector_store = FaissVectorStore.from_persist_dir(settings.LLM_INDEX_DIR)
+    embed_model = get_embedding_model()
+
+    llama_settings.Settings.embed_model = embed_model
+    llama_settings.Settings.chunk_size = 512
+
+    storage_context = StorageContext.from_defaults(
+        vector_store=vector_store,
+        persist_dir=settings.LLM_INDEX_DIR,
+    )
+    return load_index_from_storage(storage_context)
+
+
+def query_similar_documents(document: Document, top_k: int = 5) -> list[Document]:
+    """Runs a similarity query and returns top-k similar Document objects."""
+
+    # Load index
+    index = load_index()
+    retriever = VectorIndexRetriever(index=index, similarity_top_k=top_k)
+
+    # Build query from the document text
+    query_text = (document.title or "") + "\n" + (document.content or "")
+
+    # Query
+    results = retriever.retrieve(query_text)
+
+    # Each result.node.metadata["document_id"] should match our stored doc
+    document_ids = [
+        int(node.metadata["document_id"])
+        for node in results
+        if "document_id" in node.metadata
+    ]
+
+    return list(Document.objects.filter(pk__in=document_ids))