Add Themis application with custom widgets, views, and utilities
- Implemented custom form widgets for date, time, and datetime fields with DaisyUI styling. - Created utility functions for formatting dates, times, and numbers according to user preferences. - Developed views for profile settings, API key management, and notifications, including health check endpoints. - Added URL configurations for Themis tests and main application routes. - Established test cases for custom widgets to ensure proper functionality and integration. - Defined project metadata and dependencies in pyproject.toml for package management.
This commit is contained in:
1
mnemosyne/library/__init__.py
Normal file
1
mnemosyne/library/__init__.py
Normal file
@@ -0,0 +1 @@
|
||||
default_app_config = "library.apps.LibraryConfig"
|
||||
5
mnemosyne/library/admin.py
Normal file
5
mnemosyne/library/admin.py
Normal file
@@ -0,0 +1,5 @@
|
||||
# Library app does not use standard Django admin (neomodel StructuredNodes
|
||||
# are not Django ORM models). Custom admin views are provided as regular
|
||||
# app views in library/views.py, rendered within Themis's template structure.
|
||||
#
|
||||
# The embedding pipeline dashboard is at /library/embedding/
|
||||
0
mnemosyne/library/api/__init__.py
Normal file
0
mnemosyne/library/api/__init__.py
Normal file
83
mnemosyne/library/api/serializers.py
Normal file
83
mnemosyne/library/api/serializers.py
Normal file
@@ -0,0 +1,83 @@
|
||||
"""
|
||||
DRF serializers for the library app.
|
||||
|
||||
Serialize Neo4j neomodel nodes into JSON for the REST API.
|
||||
"""
|
||||
|
||||
from rest_framework import serializers
|
||||
|
||||
|
||||
class LibrarySerializer(serializers.Serializer):
|
||||
uid = serializers.CharField(read_only=True)
|
||||
name = serializers.CharField(max_length=200)
|
||||
library_type = serializers.ChoiceField(
|
||||
choices=["fiction", "technical", "music", "film", "art", "journal"]
|
||||
)
|
||||
description = serializers.CharField(required=False, allow_blank=True, default="")
|
||||
chunking_config = serializers.JSONField(required=False, default=dict)
|
||||
embedding_instruction = serializers.CharField(
|
||||
required=False, allow_blank=True, default=""
|
||||
)
|
||||
reranker_instruction = serializers.CharField(
|
||||
required=False, allow_blank=True, default=""
|
||||
)
|
||||
llm_context_prompt = serializers.CharField(
|
||||
required=False, allow_blank=True, default=""
|
||||
)
|
||||
created_at = serializers.DateTimeField(read_only=True)
|
||||
|
||||
|
||||
class CollectionSerializer(serializers.Serializer):
|
||||
uid = serializers.CharField(read_only=True)
|
||||
name = serializers.CharField(max_length=200)
|
||||
description = serializers.CharField(required=False, allow_blank=True, default="")
|
||||
metadata = serializers.JSONField(required=False, default=dict)
|
||||
created_at = serializers.DateTimeField(read_only=True)
|
||||
library_uid = serializers.CharField(
|
||||
required=False, write_only=True, help_text="UID of the parent library"
|
||||
)
|
||||
|
||||
|
||||
class ItemSerializer(serializers.Serializer):
|
||||
uid = serializers.CharField(read_only=True)
|
||||
title = serializers.CharField(max_length=500)
|
||||
item_type = serializers.CharField(required=False, allow_blank=True, default="")
|
||||
s3_key = serializers.CharField(read_only=True)
|
||||
content_hash = serializers.CharField(read_only=True)
|
||||
file_type = serializers.CharField(required=False, allow_blank=True, default="")
|
||||
file_size = serializers.IntegerField(read_only=True)
|
||||
metadata = serializers.JSONField(required=False, default=dict)
|
||||
created_at = serializers.DateTimeField(read_only=True)
|
||||
updated_at = serializers.DateTimeField(read_only=True)
|
||||
collection_uid = serializers.CharField(
|
||||
required=False, write_only=True, help_text="UID of the parent collection"
|
||||
)
|
||||
# Phase 2: Embedding pipeline fields
|
||||
embedding_status = serializers.CharField(read_only=True)
|
||||
embedding_model_name = serializers.CharField(read_only=True)
|
||||
chunk_count = serializers.IntegerField(read_only=True)
|
||||
image_count = serializers.IntegerField(read_only=True)
|
||||
|
||||
|
||||
class ChunkSerializer(serializers.Serializer):
|
||||
uid = serializers.CharField(read_only=True)
|
||||
chunk_index = serializers.IntegerField()
|
||||
chunk_s3_key = serializers.CharField()
|
||||
chunk_size = serializers.IntegerField(required=False, default=0)
|
||||
text_preview = serializers.CharField(required=False, allow_blank=True, default="")
|
||||
created_at = serializers.DateTimeField(read_only=True)
|
||||
|
||||
|
||||
class ConceptSerializer(serializers.Serializer):
|
||||
uid = serializers.CharField(read_only=True)
|
||||
name = serializers.CharField(max_length=200)
|
||||
concept_type = serializers.CharField(required=False, allow_blank=True, default="")
|
||||
|
||||
|
||||
class ImageSerializer(serializers.Serializer):
|
||||
uid = serializers.CharField(read_only=True)
|
||||
s3_key = serializers.CharField()
|
||||
image_type = serializers.CharField(required=False, allow_blank=True, default="")
|
||||
description = serializers.CharField(required=False, allow_blank=True, default="")
|
||||
metadata = serializers.JSONField(required=False, default=dict)
|
||||
created_at = serializers.DateTimeField(read_only=True)
|
||||
24
mnemosyne/library/api/urls.py
Normal file
24
mnemosyne/library/api/urls.py
Normal file
@@ -0,0 +1,24 @@
|
||||
"""
|
||||
URL patterns for the library DRF API.
|
||||
"""
|
||||
|
||||
from django.urls import path
|
||||
|
||||
from . import views
|
||||
|
||||
app_name = "library-api"
|
||||
|
||||
urlpatterns = [
|
||||
# Libraries
|
||||
path("libraries/", views.library_list_create, name="library-list"),
|
||||
path("libraries/<str:uid>/", views.library_detail, name="library-detail"),
|
||||
# Collections
|
||||
path("collections/", views.collection_list_create, name="collection-list"),
|
||||
path("collections/<str:uid>/", views.collection_detail, name="collection-detail"),
|
||||
# Items
|
||||
path("items/", views.item_list_create, name="item-list"),
|
||||
path("items/upload/", views.item_upload, name="item-upload"),
|
||||
path("items/<str:uid>/", views.item_detail, name="item-detail"),
|
||||
path("items/<str:uid>/reembed/", views.item_reembed, name="item-reembed"),
|
||||
path("items/<str:uid>/status/", views.item_status, name="item-status"),
|
||||
]
|
||||
426
mnemosyne/library/api/views.py
Normal file
426
mnemosyne/library/api/views.py
Normal file
@@ -0,0 +1,426 @@
|
||||
"""
|
||||
DRF API views for the library app.
|
||||
|
||||
All views are function-based per Red Panda Standards.
|
||||
"""
|
||||
|
||||
import hashlib
|
||||
import logging
|
||||
import os
|
||||
|
||||
from django.core.files.base import ContentFile
|
||||
from django.core.files.storage import default_storage
|
||||
from rest_framework import status
|
||||
from rest_framework.decorators import api_view, parser_classes, permission_classes
|
||||
from rest_framework.parsers import FormParser, JSONParser, MultiPartParser
|
||||
from rest_framework.permissions import IsAuthenticated
|
||||
from rest_framework.response import Response
|
||||
|
||||
from library.content_types import get_library_type_config
|
||||
|
||||
from .serializers import (
|
||||
CollectionSerializer,
|
||||
ItemSerializer,
|
||||
LibrarySerializer,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Library API
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@api_view(["GET", "POST"])
|
||||
@permission_classes([IsAuthenticated])
|
||||
def library_list_create(request):
|
||||
"""List all libraries or create a new one."""
|
||||
from library.models import Library
|
||||
|
||||
if request.method == "GET":
|
||||
libraries = Library.nodes.order_by("name")
|
||||
serializer = LibrarySerializer(libraries, many=True)
|
||||
return Response(serializer.data)
|
||||
|
||||
# POST — create
|
||||
serializer = LibrarySerializer(data=request.data)
|
||||
serializer.is_valid(raise_exception=True)
|
||||
data = serializer.validated_data
|
||||
|
||||
# Populate defaults from content-type config if not provided
|
||||
library_type = data["library_type"]
|
||||
defaults = get_library_type_config(library_type)
|
||||
|
||||
lib = Library(
|
||||
name=data["name"],
|
||||
library_type=library_type,
|
||||
description=data.get("description", ""),
|
||||
chunking_config=data.get("chunking_config") or defaults["chunking_config"],
|
||||
embedding_instruction=(
|
||||
data.get("embedding_instruction") or defaults["embedding_instruction"]
|
||||
),
|
||||
reranker_instruction=(
|
||||
data.get("reranker_instruction") or defaults["reranker_instruction"]
|
||||
),
|
||||
llm_context_prompt=(
|
||||
data.get("llm_context_prompt") or defaults["llm_context_prompt"]
|
||||
),
|
||||
)
|
||||
lib.save()
|
||||
return Response(LibrarySerializer(lib).data, status=status.HTTP_201_CREATED)
|
||||
|
||||
|
||||
@api_view(["GET", "PUT", "DELETE"])
|
||||
@permission_classes([IsAuthenticated])
|
||||
def library_detail(request, uid):
|
||||
"""Retrieve, update, or delete a library."""
|
||||
from library.models import Library
|
||||
|
||||
try:
|
||||
lib = Library.nodes.get(uid=uid)
|
||||
except Library.DoesNotExist:
|
||||
return Response(
|
||||
{"detail": "Library not found."}, status=status.HTTP_404_NOT_FOUND
|
||||
)
|
||||
|
||||
if request.method == "GET":
|
||||
return Response(LibrarySerializer(lib).data)
|
||||
|
||||
if request.method == "PUT":
|
||||
serializer = LibrarySerializer(data=request.data, partial=True)
|
||||
serializer.is_valid(raise_exception=True)
|
||||
data = serializer.validated_data
|
||||
for field in [
|
||||
"name",
|
||||
"library_type",
|
||||
"description",
|
||||
"chunking_config",
|
||||
"embedding_instruction",
|
||||
"reranker_instruction",
|
||||
"llm_context_prompt",
|
||||
]:
|
||||
if field in data:
|
||||
setattr(lib, field, data[field])
|
||||
lib.save()
|
||||
return Response(LibrarySerializer(lib).data)
|
||||
|
||||
# DELETE
|
||||
lib.delete()
|
||||
return Response(status=status.HTTP_204_NO_CONTENT)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Collection API
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@api_view(["GET", "POST"])
|
||||
@permission_classes([IsAuthenticated])
|
||||
def collection_list_create(request):
|
||||
"""List all collections or create a new one."""
|
||||
from library.models import Collection, Library
|
||||
|
||||
if request.method == "GET":
|
||||
# Optionally filter by library_uid query param
|
||||
library_uid = request.query_params.get("library_uid")
|
||||
if library_uid:
|
||||
try:
|
||||
lib = Library.nodes.get(uid=library_uid)
|
||||
collections = lib.collections.all()
|
||||
except Library.DoesNotExist:
|
||||
return Response(
|
||||
{"detail": "Library not found."}, status=status.HTTP_404_NOT_FOUND
|
||||
)
|
||||
else:
|
||||
collections = Collection.nodes.all()
|
||||
serializer = CollectionSerializer(collections, many=True)
|
||||
return Response(serializer.data)
|
||||
|
||||
# POST
|
||||
serializer = CollectionSerializer(data=request.data)
|
||||
serializer.is_valid(raise_exception=True)
|
||||
data = serializer.validated_data
|
||||
|
||||
col = Collection(
|
||||
name=data["name"],
|
||||
description=data.get("description", ""),
|
||||
metadata=data.get("metadata", {}),
|
||||
)
|
||||
col.save()
|
||||
|
||||
# Connect to library if library_uid provided
|
||||
library_uid = data.get("library_uid")
|
||||
if library_uid:
|
||||
try:
|
||||
lib = Library.nodes.get(uid=library_uid)
|
||||
lib.collections.connect(col)
|
||||
col.library.connect(lib)
|
||||
except Library.DoesNotExist:
|
||||
pass
|
||||
|
||||
return Response(CollectionSerializer(col).data, status=status.HTTP_201_CREATED)
|
||||
|
||||
|
||||
@api_view(["GET", "PUT", "DELETE"])
|
||||
@permission_classes([IsAuthenticated])
|
||||
def collection_detail(request, uid):
|
||||
"""Retrieve, update, or delete a collection."""
|
||||
from library.models import Collection
|
||||
|
||||
try:
|
||||
col = Collection.nodes.get(uid=uid)
|
||||
except Collection.DoesNotExist:
|
||||
return Response(
|
||||
{"detail": "Collection not found."}, status=status.HTTP_404_NOT_FOUND
|
||||
)
|
||||
|
||||
if request.method == "GET":
|
||||
return Response(CollectionSerializer(col).data)
|
||||
|
||||
if request.method == "PUT":
|
||||
serializer = CollectionSerializer(data=request.data, partial=True)
|
||||
serializer.is_valid(raise_exception=True)
|
||||
data = serializer.validated_data
|
||||
for field in ["name", "description", "metadata"]:
|
||||
if field in data:
|
||||
setattr(col, field, data[field])
|
||||
col.save()
|
||||
return Response(CollectionSerializer(col).data)
|
||||
|
||||
col.delete()
|
||||
return Response(status=status.HTTP_204_NO_CONTENT)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Item API
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@api_view(["GET", "POST"])
|
||||
@permission_classes([IsAuthenticated])
|
||||
def item_list_create(request):
|
||||
"""List all items or create a new one."""
|
||||
from library.models import Collection, Item
|
||||
|
||||
if request.method == "GET":
|
||||
collection_uid = request.query_params.get("collection_uid")
|
||||
if collection_uid:
|
||||
try:
|
||||
col = Collection.nodes.get(uid=collection_uid)
|
||||
items = col.items.all()
|
||||
except Collection.DoesNotExist:
|
||||
return Response(
|
||||
{"detail": "Collection not found."},
|
||||
status=status.HTTP_404_NOT_FOUND,
|
||||
)
|
||||
else:
|
||||
items = Item.nodes.all()
|
||||
serializer = ItemSerializer(items, many=True)
|
||||
return Response(serializer.data)
|
||||
|
||||
# POST
|
||||
serializer = ItemSerializer(data=request.data)
|
||||
serializer.is_valid(raise_exception=True)
|
||||
data = serializer.validated_data
|
||||
|
||||
item = Item(
|
||||
title=data["title"],
|
||||
item_type=data.get("item_type", ""),
|
||||
file_type=data.get("file_type", ""),
|
||||
metadata=data.get("metadata", {}),
|
||||
)
|
||||
item.save()
|
||||
|
||||
collection_uid = data.get("collection_uid")
|
||||
if collection_uid:
|
||||
try:
|
||||
col = Collection.nodes.get(uid=collection_uid)
|
||||
col.items.connect(item)
|
||||
except Collection.DoesNotExist:
|
||||
pass
|
||||
|
||||
return Response(ItemSerializer(item).data, status=status.HTTP_201_CREATED)
|
||||
|
||||
|
||||
@api_view(["GET", "PUT", "DELETE"])
|
||||
@permission_classes([IsAuthenticated])
|
||||
def item_detail(request, uid):
|
||||
"""Retrieve, update, or delete an item."""
|
||||
from library.models import Item
|
||||
|
||||
try:
|
||||
item = Item.nodes.get(uid=uid)
|
||||
except Item.DoesNotExist:
|
||||
return Response(
|
||||
{"detail": "Item not found."}, status=status.HTTP_404_NOT_FOUND
|
||||
)
|
||||
|
||||
if request.method == "GET":
|
||||
return Response(ItemSerializer(item).data)
|
||||
|
||||
if request.method == "PUT":
|
||||
serializer = ItemSerializer(data=request.data, partial=True)
|
||||
serializer.is_valid(raise_exception=True)
|
||||
data = serializer.validated_data
|
||||
for field in ["title", "item_type", "file_type", "metadata"]:
|
||||
if field in data:
|
||||
setattr(item, field, data[field])
|
||||
item.save()
|
||||
return Response(ItemSerializer(item).data)
|
||||
|
||||
item.delete()
|
||||
return Response(status=status.HTTP_204_NO_CONTENT)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Item Upload (Phase 2)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@api_view(["POST"])
|
||||
@permission_classes([IsAuthenticated])
|
||||
@parser_classes([MultiPartParser, FormParser])
|
||||
def item_upload(request):
|
||||
"""
|
||||
Upload a file to create a new Item and trigger embedding.
|
||||
|
||||
Expects multipart form data with:
|
||||
- file: The document file
|
||||
- title: Item title
|
||||
- collection_uid: (optional) UID of parent collection
|
||||
- auto_embed: (optional) Whether to auto-trigger embedding (default: true)
|
||||
"""
|
||||
from library.models import Collection, Item
|
||||
|
||||
uploaded_file = request.FILES.get("file")
|
||||
if not uploaded_file:
|
||||
return Response(
|
||||
{"detail": "No file provided."}, status=status.HTTP_400_BAD_REQUEST
|
||||
)
|
||||
|
||||
title = request.data.get("title", uploaded_file.name)
|
||||
collection_uid = request.data.get("collection_uid", "")
|
||||
auto_embed = request.data.get("auto_embed", "true").lower() in ("true", "1", "yes")
|
||||
|
||||
# Determine file type from extension
|
||||
_, ext = os.path.splitext(uploaded_file.name)
|
||||
file_type = ext.lstrip(".").lower()
|
||||
|
||||
# Read file data
|
||||
file_data = uploaded_file.read()
|
||||
content_hash = hashlib.sha256(file_data).hexdigest()
|
||||
|
||||
# Create Item node
|
||||
item = Item(
|
||||
title=title,
|
||||
file_type=file_type,
|
||||
file_size=len(file_data),
|
||||
content_hash=content_hash,
|
||||
embedding_status="pending",
|
||||
)
|
||||
item.save()
|
||||
|
||||
# Store file in S3
|
||||
s3_key = f"items/{item.uid}/original.{file_type}"
|
||||
try:
|
||||
default_storage.save(s3_key, ContentFile(file_data))
|
||||
item.s3_key = s3_key
|
||||
item.save()
|
||||
except Exception as exc:
|
||||
logger.error("Failed to store file to S3: %s", exc)
|
||||
item.delete()
|
||||
return Response(
|
||||
{"detail": f"File storage failed: {exc}"},
|
||||
status=status.HTTP_500_INTERNAL_SERVER_ERROR,
|
||||
)
|
||||
|
||||
# Connect to collection if specified
|
||||
if collection_uid:
|
||||
try:
|
||||
col = Collection.nodes.get(uid=collection_uid)
|
||||
col.items.connect(item)
|
||||
except Exception:
|
||||
logger.warning("Collection not found: %s", collection_uid)
|
||||
|
||||
# Auto-trigger embedding
|
||||
task_id = None
|
||||
if auto_embed:
|
||||
try:
|
||||
from library.tasks import embed_item
|
||||
|
||||
task = embed_item.delay(item.uid, request.user.id)
|
||||
task_id = task.id
|
||||
logger.info(
|
||||
"Auto-triggered embedding item_uid=%s task_id=%s",
|
||||
item.uid,
|
||||
task_id,
|
||||
)
|
||||
except Exception as exc:
|
||||
logger.warning("Failed to queue embedding task: %s", exc)
|
||||
|
||||
return Response(
|
||||
{
|
||||
**ItemSerializer(item).data,
|
||||
"task_id": task_id,
|
||||
},
|
||||
status=status.HTTP_201_CREATED,
|
||||
)
|
||||
|
||||
|
||||
@api_view(["POST"])
|
||||
@permission_classes([IsAuthenticated])
|
||||
def item_reembed(request, uid):
|
||||
"""Trigger re-embedding for an existing Item."""
|
||||
from library.models import Item
|
||||
|
||||
try:
|
||||
item = Item.nodes.get(uid=uid)
|
||||
except Item.DoesNotExist:
|
||||
return Response(
|
||||
{"detail": "Item not found."}, status=status.HTTP_404_NOT_FOUND
|
||||
)
|
||||
|
||||
try:
|
||||
from library.tasks import reembed_item
|
||||
|
||||
task = reembed_item.delay(uid, request.user.id)
|
||||
return Response(
|
||||
{
|
||||
"detail": "Re-embedding queued.",
|
||||
"item_uid": uid,
|
||||
"task_id": task.id,
|
||||
}
|
||||
)
|
||||
except Exception as exc:
|
||||
logger.error("Failed to queue reembed task: %s", exc)
|
||||
return Response(
|
||||
{"detail": f"Failed to queue task: {exc}"},
|
||||
status=status.HTTP_500_INTERNAL_SERVER_ERROR,
|
||||
)
|
||||
|
||||
|
||||
@api_view(["GET"])
|
||||
@permission_classes([IsAuthenticated])
|
||||
def item_status(request, uid):
|
||||
"""Get embedding status for an Item."""
|
||||
from library.models import Item
|
||||
|
||||
try:
|
||||
item = Item.nodes.get(uid=uid)
|
||||
except Item.DoesNotExist:
|
||||
return Response(
|
||||
{"detail": "Item not found."}, status=status.HTTP_404_NOT_FOUND
|
||||
)
|
||||
|
||||
return Response(
|
||||
{
|
||||
"uid": item.uid,
|
||||
"title": item.title,
|
||||
"embedding_status": item.embedding_status,
|
||||
"embedding_model_name": item.embedding_model_name,
|
||||
"chunk_count": item.chunk_count,
|
||||
"image_count": item.image_count,
|
||||
"error_message": item.error_message,
|
||||
}
|
||||
)
|
||||
7
mnemosyne/library/apps.py
Normal file
7
mnemosyne/library/apps.py
Normal file
@@ -0,0 +1,7 @@
|
||||
from django.apps import AppConfig
|
||||
|
||||
|
||||
class LibraryConfig(AppConfig):
|
||||
default_auto_field = "django.db.models.BigAutoField"
|
||||
name = "library"
|
||||
verbose_name = "Library"
|
||||
163
mnemosyne/library/content_types.py
Normal file
163
mnemosyne/library/content_types.py
Normal file
@@ -0,0 +1,163 @@
|
||||
"""
|
||||
Content-type system configuration for Mnemosyne library types.
|
||||
|
||||
Each library type has a default configuration that governs chunking,
|
||||
embedding, re-ranking, and LLM context injection.
|
||||
"""
|
||||
|
||||
# Default configurations per library type.
|
||||
# These are loaded into Library nodes via the load_library_types management command.
|
||||
LIBRARY_TYPE_DEFAULTS = {
|
||||
"fiction": {
|
||||
"chunking_config": {
|
||||
"strategy": "chapter_aware",
|
||||
"chunk_size": 1024,
|
||||
"chunk_overlap": 128,
|
||||
"respect_boundaries": ["chapter", "scene", "paragraph"],
|
||||
},
|
||||
"embedding_instruction": (
|
||||
"Represent this passage from a work of fiction for retrieval. "
|
||||
"Focus on narrative elements: characters, plot events, themes, "
|
||||
"setting, and emotional tone."
|
||||
),
|
||||
"reranker_instruction": (
|
||||
"Re-rank passages from fiction based on narrative relevance to the query. "
|
||||
"Prioritize character actions, dialogue, plot developments, and thematic elements."
|
||||
),
|
||||
"llm_context_prompt": (
|
||||
"The following excerpts are from fiction (novels, short stories, etc.). "
|
||||
"Treat this as creative/narrative content. Respect the literary context — "
|
||||
"characters, settings, and events are fictional. Cite specific passages "
|
||||
"when answering."
|
||||
),
|
||||
},
|
||||
"technical": {
|
||||
"chunking_config": {
|
||||
"strategy": "section_aware",
|
||||
"chunk_size": 512,
|
||||
"chunk_overlap": 64,
|
||||
"respect_boundaries": ["section", "subsection", "code_block", "list"],
|
||||
},
|
||||
"embedding_instruction": (
|
||||
"Represent this passage from technical documentation for retrieval. "
|
||||
"Focus on procedures, configurations, API references, code examples, "
|
||||
"and technical concepts."
|
||||
),
|
||||
"reranker_instruction": (
|
||||
"Re-rank passages from technical documentation based on procedural relevance. "
|
||||
"Prioritize step-by-step instructions, code examples, and specific configurations."
|
||||
),
|
||||
"llm_context_prompt": (
|
||||
"The following excerpts are from technical documentation (manuals, guides, "
|
||||
"reference material). Provide precise, actionable answers. Include code "
|
||||
"examples and exact configurations when available. Cite source sections."
|
||||
),
|
||||
},
|
||||
"music": {
|
||||
"chunking_config": {
|
||||
"strategy": "song_level",
|
||||
"chunk_size": 512,
|
||||
"chunk_overlap": 32,
|
||||
"respect_boundaries": ["song", "verse", "chorus"],
|
||||
},
|
||||
"embedding_instruction": (
|
||||
"Represent this music content (lyrics, liner notes, metadata) for retrieval. "
|
||||
"Focus on artist, album, genre, lyrical themes, and musical elements."
|
||||
),
|
||||
"reranker_instruction": (
|
||||
"Re-rank music content based on relevance to the query. "
|
||||
"Consider artist, genre, lyrical themes, and musical characteristics."
|
||||
),
|
||||
"llm_context_prompt": (
|
||||
"The following excerpts are song lyrics and music metadata. "
|
||||
"Consider the artistic and cultural context. Reference specific "
|
||||
"songs, albums, and artists when answering."
|
||||
),
|
||||
},
|
||||
"film": {
|
||||
"chunking_config": {
|
||||
"strategy": "scene_level",
|
||||
"chunk_size": 768,
|
||||
"chunk_overlap": 64,
|
||||
"respect_boundaries": ["scene", "act", "sequence"],
|
||||
},
|
||||
"embedding_instruction": (
|
||||
"Represent this film content (scripts, synopses, reviews) for retrieval. "
|
||||
"Focus on scenes, characters, visual elements, dialogue, and narrative structure."
|
||||
),
|
||||
"reranker_instruction": (
|
||||
"Re-rank film content based on cinematic relevance. "
|
||||
"Prioritize scene descriptions, character interactions, and visual elements."
|
||||
),
|
||||
"llm_context_prompt": (
|
||||
"The following excerpts are from film-related content (scripts, synopses, "
|
||||
"reviews). Consider the cinematic context — visual storytelling, "
|
||||
"direction, and performance. Cite specific scenes and films."
|
||||
),
|
||||
},
|
||||
"art": {
|
||||
"chunking_config": {
|
||||
"strategy": "description_level",
|
||||
"chunk_size": 512,
|
||||
"chunk_overlap": 32,
|
||||
"respect_boundaries": ["artwork", "description", "analysis"],
|
||||
},
|
||||
"embedding_instruction": (
|
||||
"Represent this art content (descriptions, catalogs, analysis) for retrieval. "
|
||||
"Focus on visual elements, style, medium, artist, period, and artistic movements."
|
||||
),
|
||||
"reranker_instruction": (
|
||||
"Re-rank art content based on visual and stylistic relevance. "
|
||||
"Prioritize descriptions of artwork, technique, composition, and artistic context."
|
||||
),
|
||||
"llm_context_prompt": (
|
||||
"The following excerpts describe artworks and artistic content. "
|
||||
"Consider visual elements, artistic technique, historical context, "
|
||||
"and the artist's intent. Reference specific works and movements."
|
||||
),
|
||||
},
|
||||
"journal": {
|
||||
"chunking_config": {
|
||||
"strategy": "entry_level",
|
||||
"chunk_size": 512,
|
||||
"chunk_overlap": 32,
|
||||
"respect_boundaries": ["entry", "date", "paragraph"],
|
||||
},
|
||||
"embedding_instruction": (
|
||||
"Represent this personal journal entry for retrieval. "
|
||||
"Focus on temporal context, personal reflections, mentioned people "
|
||||
"and places, and emotional content."
|
||||
),
|
||||
"reranker_instruction": (
|
||||
"Re-rank journal entries based on temporal and thematic relevance. "
|
||||
"Prioritize entries matching the time period, people, places, or topics in the query."
|
||||
),
|
||||
"llm_context_prompt": (
|
||||
"The following excerpts are from personal journal entries. "
|
||||
"This is private, reflective content. Respect the personal nature — "
|
||||
"answer with sensitivity. Note dates and temporal context when relevant."
|
||||
),
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def get_library_type_config(library_type):
|
||||
"""
|
||||
Get the default configuration for a library type.
|
||||
|
||||
Args:
|
||||
library_type: One of 'fiction', 'technical', 'music', 'film', 'art', 'journal'
|
||||
|
||||
Returns:
|
||||
dict with keys: chunking_config, embedding_instruction,
|
||||
reranker_instruction, llm_context_prompt
|
||||
|
||||
Raises:
|
||||
ValueError: If library_type is not recognized
|
||||
"""
|
||||
if library_type not in LIBRARY_TYPE_DEFAULTS:
|
||||
raise ValueError(
|
||||
f"Unknown library type '{library_type}'. "
|
||||
f"Valid types: {', '.join(LIBRARY_TYPE_DEFAULTS.keys())}"
|
||||
)
|
||||
return LIBRARY_TYPE_DEFAULTS[library_type]
|
||||
100
mnemosyne/library/forms.py
Normal file
100
mnemosyne/library/forms.py
Normal file
@@ -0,0 +1,100 @@
|
||||
"""
|
||||
Django forms for Library admin views.
|
||||
|
||||
These forms are used by the custom admin views for Library, Collection,
|
||||
and Item CRUD. They are plain Django forms (not ModelForms) because
|
||||
neomodel StructuredNodes are not Django ORM models.
|
||||
"""
|
||||
|
||||
from django import forms
|
||||
|
||||
from .content_types import LIBRARY_TYPE_DEFAULTS
|
||||
|
||||
|
||||
LIBRARY_TYPE_CHOICES = [
|
||||
(key, key.capitalize()) for key in LIBRARY_TYPE_DEFAULTS.keys()
|
||||
]
|
||||
|
||||
|
||||
class LibraryForm(forms.Form):
|
||||
"""Form for creating/editing a Library node."""
|
||||
|
||||
name = forms.CharField(
|
||||
max_length=200,
|
||||
widget=forms.TextInput(attrs={"class": "input input-bordered w-full"}),
|
||||
)
|
||||
library_type = forms.ChoiceField(
|
||||
choices=LIBRARY_TYPE_CHOICES,
|
||||
widget=forms.Select(attrs={"class": "select select-bordered w-full"}),
|
||||
)
|
||||
description = forms.CharField(
|
||||
required=False,
|
||||
widget=forms.Textarea(
|
||||
attrs={"class": "textarea textarea-bordered w-full", "rows": 3}
|
||||
),
|
||||
)
|
||||
embedding_instruction = forms.CharField(
|
||||
required=False,
|
||||
widget=forms.Textarea(
|
||||
attrs={"class": "textarea textarea-bordered w-full", "rows": 3}
|
||||
),
|
||||
)
|
||||
reranker_instruction = forms.CharField(
|
||||
required=False,
|
||||
widget=forms.Textarea(
|
||||
attrs={"class": "textarea textarea-bordered w-full", "rows": 3}
|
||||
),
|
||||
)
|
||||
llm_context_prompt = forms.CharField(
|
||||
required=False,
|
||||
widget=forms.Textarea(
|
||||
attrs={"class": "textarea textarea-bordered w-full", "rows": 3}
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
class CollectionForm(forms.Form):
|
||||
"""Form for creating/editing a Collection node."""
|
||||
|
||||
name = forms.CharField(
|
||||
max_length=200,
|
||||
widget=forms.TextInput(attrs={"class": "input input-bordered w-full"}),
|
||||
)
|
||||
description = forms.CharField(
|
||||
required=False,
|
||||
widget=forms.Textarea(
|
||||
attrs={"class": "textarea textarea-bordered w-full", "rows": 3}
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
class ItemForm(forms.Form):
|
||||
"""Form for creating/editing an Item node."""
|
||||
|
||||
title = forms.CharField(
|
||||
max_length=500,
|
||||
widget=forms.TextInput(attrs={"class": "input input-bordered w-full"}),
|
||||
)
|
||||
item_type = forms.CharField(
|
||||
required=False,
|
||||
max_length=100,
|
||||
widget=forms.TextInput(attrs={"class": "input input-bordered w-full"}),
|
||||
)
|
||||
file_type = forms.CharField(
|
||||
required=False,
|
||||
max_length=50,
|
||||
widget=forms.TextInput(attrs={"class": "input input-bordered w-full"}),
|
||||
)
|
||||
file = forms.FileField(
|
||||
required=False,
|
||||
widget=forms.ClearableFileInput(
|
||||
attrs={"class": "file-input file-input-bordered w-full"}
|
||||
),
|
||||
help_text="Upload a document (PDF, EPUB, DOCX, PPTX, TXT, etc.)",
|
||||
)
|
||||
auto_embed = forms.BooleanField(
|
||||
required=False,
|
||||
initial=True,
|
||||
widget=forms.CheckboxInput(attrs={"class": "checkbox checkbox-primary"}),
|
||||
help_text="Automatically start embedding after upload",
|
||||
)
|
||||
0
mnemosyne/library/management/__init__.py
Normal file
0
mnemosyne/library/management/__init__.py
Normal file
0
mnemosyne/library/management/commands/__init__.py
Normal file
0
mnemosyne/library/management/commands/__init__.py
Normal file
43
mnemosyne/library/management/commands/embed_collection.py
Normal file
43
mnemosyne/library/management/commands/embed_collection.py
Normal file
@@ -0,0 +1,43 @@
|
||||
"""
|
||||
Management command to embed all items in a Collection.
|
||||
|
||||
Usage:
|
||||
python manage.py embed_collection <collection_uid>
|
||||
"""
|
||||
|
||||
import logging
|
||||
|
||||
from django.core.management.base import BaseCommand, CommandError
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class Command(BaseCommand):
|
||||
help = "Queue embedding tasks for all items in a Collection."
|
||||
|
||||
def add_arguments(self, parser):
|
||||
parser.add_argument(
|
||||
"collection_uid", type=str, help="UID of the Collection to embed"
|
||||
)
|
||||
|
||||
def handle(self, *args, **options):
|
||||
collection_uid = options["collection_uid"]
|
||||
|
||||
try:
|
||||
from library.models import Collection
|
||||
|
||||
col = Collection.nodes.get(uid=collection_uid)
|
||||
except Exception as exc:
|
||||
raise CommandError(f"Collection not found: {collection_uid} ({exc})")
|
||||
|
||||
items = col.items.all()
|
||||
self.stdout.write(f"Collection: {col.name} ({len(items)} items)")
|
||||
|
||||
from library.tasks import embed_collection
|
||||
|
||||
task = embed_collection.delay(collection_uid)
|
||||
self.stdout.write(
|
||||
self.style.SUCCESS(
|
||||
f"Batch task queued: {task.id} ({len(items)} items)"
|
||||
)
|
||||
)
|
||||
68
mnemosyne/library/management/commands/embed_item.py
Normal file
68
mnemosyne/library/management/commands/embed_item.py
Normal file
@@ -0,0 +1,68 @@
|
||||
"""
|
||||
Management command to embed a single Item via the CLI.
|
||||
|
||||
Usage:
|
||||
python manage.py embed_item <item_uid>
|
||||
python manage.py embed_item <item_uid> --sync
|
||||
"""
|
||||
|
||||
import logging
|
||||
|
||||
from django.core.management.base import BaseCommand, CommandError
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class Command(BaseCommand):
|
||||
help = "Run the embedding pipeline for a single Item."
|
||||
|
||||
def add_arguments(self, parser):
|
||||
parser.add_argument("item_uid", type=str, help="UID of the Item to embed")
|
||||
parser.add_argument(
|
||||
"--sync",
|
||||
action="store_true",
|
||||
help="Run synchronously instead of queueing a Celery task",
|
||||
)
|
||||
|
||||
def handle(self, *args, **options):
|
||||
item_uid = options["item_uid"]
|
||||
sync = options["sync"]
|
||||
|
||||
# Verify item exists
|
||||
try:
|
||||
from library.models import Item
|
||||
|
||||
item = Item.nodes.get(uid=item_uid)
|
||||
except Exception as exc:
|
||||
raise CommandError(f"Item not found: {item_uid} ({exc})")
|
||||
|
||||
self.stdout.write(f"Item: {item.title} (type={item.file_type}, status={item.embedding_status})")
|
||||
|
||||
if sync:
|
||||
self.stdout.write("Running embedding pipeline synchronously...")
|
||||
from library.services.pipeline import EmbeddingPipeline
|
||||
|
||||
pipeline = EmbeddingPipeline()
|
||||
|
||||
def progress_cb(percent, message):
|
||||
self.stdout.write(f" [{percent:3d}%] {message}")
|
||||
|
||||
try:
|
||||
result = pipeline.process_item(item_uid, progress_callback=progress_cb)
|
||||
self.stdout.write(
|
||||
self.style.SUCCESS(
|
||||
f"\nCompleted: {result.get('chunks_created', 0)} chunks, "
|
||||
f"{result.get('images_stored', 0)} images, "
|
||||
f"{result.get('concepts_extracted', 0)} concepts"
|
||||
)
|
||||
)
|
||||
except Exception as exc:
|
||||
raise CommandError(f"Embedding failed: {exc}")
|
||||
else:
|
||||
self.stdout.write("Queueing embedding task...")
|
||||
from library.tasks import embed_item
|
||||
|
||||
task = embed_item.delay(item_uid)
|
||||
self.stdout.write(
|
||||
self.style.SUCCESS(f"Task queued: {task.id}")
|
||||
)
|
||||
132
mnemosyne/library/management/commands/embedding_status.py
Normal file
132
mnemosyne/library/management/commands/embedding_status.py
Normal file
@@ -0,0 +1,132 @@
|
||||
"""
|
||||
Management command to display embedding pipeline status and statistics.
|
||||
|
||||
Usage:
|
||||
python manage.py embedding_status
|
||||
python manage.py embedding_status --library <uid>
|
||||
"""
|
||||
|
||||
import logging
|
||||
|
||||
from django.core.management.base import BaseCommand
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class Command(BaseCommand):
|
||||
help = "Display embedding pipeline status and statistics."
|
||||
|
||||
def add_arguments(self, parser):
|
||||
parser.add_argument(
|
||||
"--library",
|
||||
type=str,
|
||||
default="",
|
||||
help="Filter by library UID",
|
||||
)
|
||||
|
||||
def handle(self, *args, **options):
|
||||
library_uid = options["library"]
|
||||
|
||||
try:
|
||||
from neomodel import db
|
||||
except ImportError:
|
||||
self.stderr.write(self.style.ERROR("neomodel not available"))
|
||||
return
|
||||
|
||||
self.stdout.write(self.style.HTTP_INFO("\n=== Mnemosyne Embedding Pipeline Status ===\n"))
|
||||
|
||||
# System embedding model
|
||||
try:
|
||||
from llm_manager.models import LLMModel
|
||||
|
||||
embed_model = LLMModel.get_system_embedding_model()
|
||||
if embed_model:
|
||||
self.stdout.write(
|
||||
f"System Embedding Model: {embed_model.api.name}: {embed_model.name} "
|
||||
f"(dimensions={embed_model.vector_dimensions or '?'})"
|
||||
)
|
||||
else:
|
||||
self.stdout.write(
|
||||
self.style.WARNING("System Embedding Model: NOT CONFIGURED")
|
||||
)
|
||||
|
||||
chat_model = LLMModel.get_system_chat_model()
|
||||
if chat_model:
|
||||
self.stdout.write(f"System Chat Model: {chat_model.api.name}: {chat_model.name}")
|
||||
else:
|
||||
self.stdout.write(
|
||||
self.style.WARNING("System Chat Model: NOT CONFIGURED (concept extraction disabled)")
|
||||
)
|
||||
except Exception as exc:
|
||||
self.stdout.write(self.style.ERROR(f"Could not query LLM models: {exc}"))
|
||||
|
||||
self.stdout.write("")
|
||||
|
||||
# Item status counts
|
||||
try:
|
||||
statuses = ["pending", "processing", "completed", "failed"]
|
||||
self.stdout.write("Item Embedding Status:")
|
||||
|
||||
for status in statuses:
|
||||
if library_uid:
|
||||
query = (
|
||||
"MATCH (l:Library {uid: $lib_uid})-[:CONTAINS]->(c:Collection)"
|
||||
"-[:CONTAINS]->(i:Item {embedding_status: $status}) "
|
||||
"RETURN count(i)"
|
||||
)
|
||||
results, _ = db.cypher_query(
|
||||
query, {"lib_uid": library_uid, "status": status}
|
||||
)
|
||||
else:
|
||||
query = (
|
||||
"MATCH (i:Item {embedding_status: $status}) RETURN count(i)"
|
||||
)
|
||||
results, _ = db.cypher_query(query, {"status": status})
|
||||
|
||||
count = results[0][0] if results else 0
|
||||
style = {
|
||||
"completed": self.style.SUCCESS,
|
||||
"failed": self.style.ERROR,
|
||||
"processing": self.style.WARNING,
|
||||
"pending": self.style.NOTICE,
|
||||
}.get(status, str)
|
||||
self.stdout.write(f" {status:12s}: {style(str(count))}")
|
||||
except Exception as exc:
|
||||
self.stdout.write(self.style.ERROR(f"Could not query items: {exc}"))
|
||||
|
||||
self.stdout.write("")
|
||||
|
||||
# Node counts
|
||||
try:
|
||||
node_types = [
|
||||
("Library", "Library"),
|
||||
("Collection", "Collection"),
|
||||
("Item", "Item"),
|
||||
("Chunk", "Chunk"),
|
||||
("Concept", "Concept"),
|
||||
("Image", "Image"),
|
||||
("ImageEmbedding", "ImageEmbedding"),
|
||||
]
|
||||
self.stdout.write("Graph Node Counts:")
|
||||
for label, display in node_types:
|
||||
results, _ = db.cypher_query(f"MATCH (n:{label}) RETURN count(n)")
|
||||
count = results[0][0] if results else 0
|
||||
self.stdout.write(f" {display:20s}: {count}")
|
||||
except Exception as exc:
|
||||
self.stdout.write(self.style.ERROR(f"Could not query nodes: {exc}"))
|
||||
|
||||
# Chunks with embeddings
|
||||
try:
|
||||
results, _ = db.cypher_query(
|
||||
"MATCH (c:Chunk) WHERE c.embedding IS NOT NULL RETURN count(c)"
|
||||
)
|
||||
embedded = results[0][0] if results else 0
|
||||
results, _ = db.cypher_query("MATCH (c:Chunk) RETURN count(c)")
|
||||
total = results[0][0] if results else 0
|
||||
self.stdout.write(
|
||||
f"\nChunks with embeddings: {embedded}/{total}"
|
||||
)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
self.stdout.write("")
|
||||
94
mnemosyne/library/management/commands/load_library_types.py
Normal file
94
mnemosyne/library/management/commands/load_library_types.py
Normal file
@@ -0,0 +1,94 @@
|
||||
"""
|
||||
Management command to load default library type configurations.
|
||||
|
||||
Idempotent — safe to re-run. Creates Library nodes with default content-type
|
||||
configurations if they don't already exist. Does NOT overwrite existing
|
||||
libraries that have been customized.
|
||||
"""
|
||||
|
||||
import logging
|
||||
|
||||
from django.core.management.base import BaseCommand
|
||||
|
||||
from library.content_types import LIBRARY_TYPE_DEFAULTS
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class Command(BaseCommand):
|
||||
help = (
|
||||
"Load default library type configurations into Neo4j. "
|
||||
"Creates one Library node per type with default chunking, embedding, "
|
||||
"reranker, and LLM context settings. Safe to re-run."
|
||||
)
|
||||
|
||||
def add_arguments(self, parser):
|
||||
parser.add_argument(
|
||||
"--force",
|
||||
action="store_true",
|
||||
help="Update existing libraries with default configurations (overwrites customizations)",
|
||||
)
|
||||
|
||||
def handle(self, *args, **options):
|
||||
force = options["force"]
|
||||
|
||||
try:
|
||||
from library.models import Library
|
||||
except Exception as e:
|
||||
self.stderr.write(
|
||||
self.style.ERROR(f"Cannot import library models: {e}")
|
||||
)
|
||||
return
|
||||
|
||||
created_count = 0
|
||||
updated_count = 0
|
||||
skipped_count = 0
|
||||
|
||||
for library_type, config in LIBRARY_TYPE_DEFAULTS.items():
|
||||
display_name = library_type.capitalize()
|
||||
default_name = f"Default {display_name} Library"
|
||||
|
||||
# Check if a library of this type already exists
|
||||
existing = Library.nodes.filter(library_type=library_type)
|
||||
|
||||
if existing:
|
||||
if force:
|
||||
lib = existing[0]
|
||||
lib.chunking_config = config["chunking_config"]
|
||||
lib.embedding_instruction = config["embedding_instruction"]
|
||||
lib.reranker_instruction = config["reranker_instruction"]
|
||||
lib.llm_context_prompt = config["llm_context_prompt"]
|
||||
lib.save()
|
||||
updated_count += 1
|
||||
self.stdout.write(
|
||||
self.style.WARNING(f"Updated: {lib.name} ({library_type})")
|
||||
)
|
||||
else:
|
||||
skipped_count += 1
|
||||
self.stdout.write(
|
||||
self.style.NOTICE(
|
||||
f"Skipped: {existing[0].name} ({library_type}) — already exists"
|
||||
)
|
||||
)
|
||||
else:
|
||||
lib = Library(
|
||||
name=default_name,
|
||||
library_type=library_type,
|
||||
description=f"Default {display_name.lower()} library",
|
||||
chunking_config=config["chunking_config"],
|
||||
embedding_instruction=config["embedding_instruction"],
|
||||
reranker_instruction=config["reranker_instruction"],
|
||||
llm_context_prompt=config["llm_context_prompt"],
|
||||
)
|
||||
lib.save()
|
||||
created_count += 1
|
||||
self.stdout.write(
|
||||
self.style.SUCCESS(f"Created: {default_name} ({library_type})")
|
||||
)
|
||||
|
||||
self.stdout.write(
|
||||
self.style.SUCCESS(
|
||||
f"\nDone. Created: {created_count}, "
|
||||
f"Updated: {updated_count}, Skipped: {skipped_count}"
|
||||
)
|
||||
)
|
||||
189
mnemosyne/library/management/commands/setup_neo4j_indexes.py
Normal file
189
mnemosyne/library/management/commands/setup_neo4j_indexes.py
Normal file
@@ -0,0 +1,189 @@
|
||||
"""
|
||||
Management command to create Neo4j indexes for Mnemosyne content graph.
|
||||
|
||||
Creates:
|
||||
- Vector indexes (dynamic dimensions from system embedding model) for Chunk, Concept, and ImageEmbedding
|
||||
- Full-text indexes for text search on Chunk.text_preview and Concept.name
|
||||
- Constraint indexes enforced by neomodel (unique properties)
|
||||
"""
|
||||
|
||||
import logging
|
||||
|
||||
from django.core.management.base import BaseCommand
|
||||
from neomodel import db
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Default vector dimensions (used when no system embedding model is configured)
|
||||
DEFAULT_VECTOR_DIMENSIONS = 4096
|
||||
|
||||
# Full-text index definitions: (index_name, label, properties)
|
||||
FULLTEXT_INDEXES = [
|
||||
("chunk_text_fulltext", "Chunk", ["text_preview"]),
|
||||
("concept_name_fulltext", "Concept", ["name"]),
|
||||
("item_title_fulltext", "Item", ["title"]),
|
||||
("library_name_fulltext", "Library", ["name"]),
|
||||
]
|
||||
|
||||
|
||||
def _get_vector_dimensions():
|
||||
"""
|
||||
Get vector dimensions from the system embedding model.
|
||||
|
||||
Falls back to DEFAULT_VECTOR_DIMENSIONS if no model is configured
|
||||
or the model has no vector_dimensions set.
|
||||
|
||||
:returns: Tuple of (dimensions, source_description).
|
||||
"""
|
||||
try:
|
||||
from llm_manager.models import LLMModel
|
||||
|
||||
model = LLMModel.get_system_embedding_model()
|
||||
if model and model.vector_dimensions:
|
||||
return model.vector_dimensions, f"{model.api.name}: {model.name}"
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
return DEFAULT_VECTOR_DIMENSIONS, "default (no system embedding model)"
|
||||
|
||||
|
||||
class Command(BaseCommand):
|
||||
help = (
|
||||
"Create Neo4j vector, full-text, and constraint indexes "
|
||||
"for the Mnemosyne content graph. Vector dimensions are read "
|
||||
"from the system embedding model."
|
||||
)
|
||||
|
||||
def add_arguments(self, parser):
|
||||
parser.add_argument(
|
||||
"--drop",
|
||||
action="store_true",
|
||||
help="Drop existing indexes before recreating them",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--dimensions",
|
||||
type=int,
|
||||
default=0,
|
||||
help="Override vector dimensions (default: read from system embedding model)",
|
||||
)
|
||||
|
||||
def handle(self, *args, **options):
|
||||
drop = options["drop"]
|
||||
override_dims = options["dimensions"]
|
||||
|
||||
# Resolve vector dimensions
|
||||
if override_dims > 0:
|
||||
dimensions = override_dims
|
||||
source = f"CLI override ({override_dims})"
|
||||
else:
|
||||
dimensions, source = _get_vector_dimensions()
|
||||
|
||||
self.stdout.write(
|
||||
self.style.HTTP_INFO(
|
||||
f"Vector dimensions: {dimensions} (source: {source})"
|
||||
)
|
||||
)
|
||||
|
||||
# Vector index definitions (dynamic dimensions)
|
||||
vector_indexes = [
|
||||
("chunk_embedding_index", "Chunk", "embedding", dimensions, "cosine"),
|
||||
("concept_embedding_index", "Concept", "embedding", dimensions, "cosine"),
|
||||
("image_embedding_index", "ImageEmbedding", "embedding", dimensions, "cosine"),
|
||||
]
|
||||
|
||||
# Get existing indexes
|
||||
existing_indexes = self._get_existing_indexes()
|
||||
|
||||
if drop:
|
||||
self._drop_indexes(existing_indexes, vector_indexes)
|
||||
existing_indexes = self._get_existing_indexes()
|
||||
|
||||
# Create vector indexes
|
||||
for name, label, prop, dims, similarity in vector_indexes:
|
||||
if name in existing_indexes:
|
||||
self.stdout.write(
|
||||
self.style.NOTICE(f"Vector index '{name}' already exists, skipping")
|
||||
)
|
||||
continue
|
||||
try:
|
||||
cypher = (
|
||||
f"CREATE VECTOR INDEX {name} IF NOT EXISTS "
|
||||
f"FOR (n:{label}) ON (n.{prop}) "
|
||||
f"OPTIONS {{indexConfig: {{"
|
||||
f"`vector.dimensions`: {dims}, "
|
||||
f"`vector.similarity_function`: '{similarity}'"
|
||||
f"}}}}"
|
||||
)
|
||||
db.cypher_query(cypher)
|
||||
self.stdout.write(
|
||||
self.style.SUCCESS(
|
||||
f"Created vector index: {name} ({label}.{prop}, {dims}d {similarity})"
|
||||
)
|
||||
)
|
||||
except Exception as e:
|
||||
self.stderr.write(
|
||||
self.style.ERROR(f"Failed to create vector index '{name}': {e}")
|
||||
)
|
||||
|
||||
# Create full-text indexes
|
||||
for name, label, properties in FULLTEXT_INDEXES:
|
||||
if name in existing_indexes:
|
||||
self.stdout.write(
|
||||
self.style.NOTICE(
|
||||
f"Full-text index '{name}' already exists, skipping"
|
||||
)
|
||||
)
|
||||
continue
|
||||
try:
|
||||
props_str = ", ".join(f"n.{p}" for p in properties)
|
||||
cypher = (
|
||||
f"CREATE FULLTEXT INDEX {name} IF NOT EXISTS "
|
||||
f"FOR (n:{label}) ON EACH [{props_str}]"
|
||||
)
|
||||
db.cypher_query(cypher)
|
||||
self.stdout.write(
|
||||
self.style.SUCCESS(
|
||||
f"Created full-text index: {name} ({label}: {', '.join(properties)})"
|
||||
)
|
||||
)
|
||||
except Exception as e:
|
||||
self.stderr.write(
|
||||
self.style.ERROR(f"Failed to create full-text index '{name}': {e}")
|
||||
)
|
||||
|
||||
# Install neomodel constraints (unique indexes from model definitions)
|
||||
try:
|
||||
from neomodel import install_all_labels
|
||||
|
||||
install_all_labels()
|
||||
self.stdout.write(
|
||||
self.style.SUCCESS("Installed neomodel constraint indexes")
|
||||
)
|
||||
except Exception as e:
|
||||
self.stderr.write(
|
||||
self.style.ERROR(f"Failed to install neomodel labels: {e}")
|
||||
)
|
||||
|
||||
self.stdout.write(self.style.SUCCESS("\nNeo4j index setup complete."))
|
||||
|
||||
def _get_existing_indexes(self):
|
||||
"""Return set of existing index names."""
|
||||
try:
|
||||
results, _ = db.cypher_query("SHOW INDEXES YIELD name RETURN name")
|
||||
return {row[0] for row in results}
|
||||
except Exception:
|
||||
return set()
|
||||
|
||||
def _drop_indexes(self, existing_indexes, vector_indexes):
|
||||
"""Drop all Mnemosyne-managed indexes."""
|
||||
managed_names = {name for name, *_ in vector_indexes} | {
|
||||
name for name, *_ in FULLTEXT_INDEXES
|
||||
}
|
||||
for name in managed_names & existing_indexes:
|
||||
try:
|
||||
db.cypher_query(f"DROP INDEX {name} IF EXISTS")
|
||||
self.stdout.write(self.style.WARNING(f"Dropped index: {name}"))
|
||||
except Exception as e:
|
||||
self.stderr.write(
|
||||
self.style.ERROR(f"Failed to drop index '{name}': {e}")
|
||||
)
|
||||
96
mnemosyne/library/metrics.py
Normal file
96
mnemosyne/library/metrics.py
Normal file
@@ -0,0 +1,96 @@
|
||||
"""
|
||||
Prometheus metrics for the Mnemosyne embedding pipeline.
|
||||
|
||||
Exposes counters, histograms, and gauges for monitoring document parsing,
|
||||
chunking, embedding, and pipeline orchestration.
|
||||
"""
|
||||
|
||||
from prometheus_client import Counter, Gauge, Histogram
|
||||
|
||||
# --- Document Parsing ---
|
||||
|
||||
DOCUMENTS_PARSED_TOTAL = Counter(
|
||||
"mnemosyne_documents_parsed_total",
|
||||
"Total documents parsed",
|
||||
["file_type", "status"],
|
||||
)
|
||||
DOCUMENT_PARSE_DURATION = Histogram(
|
||||
"mnemosyne_document_parse_duration_seconds",
|
||||
"Time to parse a document",
|
||||
["file_type"],
|
||||
buckets=[0.1, 0.5, 1, 2, 5, 10, 30, 60, 120],
|
||||
)
|
||||
IMAGES_EXTRACTED_TOTAL = Counter(
|
||||
"mnemosyne_images_extracted_total",
|
||||
"Total images extracted from documents",
|
||||
["file_type"],
|
||||
)
|
||||
|
||||
# --- Chunking ---
|
||||
|
||||
CHUNKS_CREATED_TOTAL = Counter(
|
||||
"mnemosyne_chunks_created_total",
|
||||
"Total chunks created",
|
||||
["library_type", "strategy"],
|
||||
)
|
||||
CHUNK_SIZE_TOKENS = Histogram(
|
||||
"mnemosyne_chunk_size_tokens",
|
||||
"Distribution of chunk sizes in tokens",
|
||||
buckets=[32, 64, 128, 256, 512, 768, 1024, 2048],
|
||||
)
|
||||
|
||||
# --- Embedding ---
|
||||
|
||||
EMBEDDINGS_GENERATED_TOTAL = Counter(
|
||||
"mnemosyne_embeddings_generated_total",
|
||||
"Total embeddings generated",
|
||||
["model_name", "api_type", "content_type"],
|
||||
)
|
||||
EMBEDDING_BATCH_DURATION = Histogram(
|
||||
"mnemosyne_embedding_batch_duration_seconds",
|
||||
"Time per embedding batch request",
|
||||
["model_name", "api_type"],
|
||||
buckets=[0.1, 0.5, 1, 2, 5, 10, 30, 60],
|
||||
)
|
||||
EMBEDDING_API_ERRORS_TOTAL = Counter(
|
||||
"mnemosyne_embedding_api_errors_total",
|
||||
"Embedding API errors",
|
||||
["model_name", "api_type", "error_type"],
|
||||
)
|
||||
EMBEDDING_TOKENS_TOTAL = Counter(
|
||||
"mnemosyne_embedding_tokens_total",
|
||||
"Total tokens sent to embedding APIs",
|
||||
["model_name"],
|
||||
)
|
||||
|
||||
# --- Pipeline ---
|
||||
|
||||
PIPELINE_ITEMS_TOTAL = Counter(
|
||||
"mnemosyne_pipeline_items_total",
|
||||
"Total items processed by embedding pipeline",
|
||||
["status"],
|
||||
)
|
||||
PIPELINE_DURATION = Histogram(
|
||||
"mnemosyne_pipeline_item_duration_seconds",
|
||||
"Total time to process one item through the full pipeline",
|
||||
buckets=[1, 5, 10, 30, 60, 120, 300, 600],
|
||||
)
|
||||
PIPELINE_ITEMS_IN_PROGRESS = Gauge(
|
||||
"mnemosyne_pipeline_items_in_progress",
|
||||
"Items currently being processed",
|
||||
)
|
||||
|
||||
# --- Concept Extraction ---
|
||||
|
||||
CONCEPTS_EXTRACTED_TOTAL = Counter(
|
||||
"mnemosyne_concepts_extracted_total",
|
||||
"Total concepts extracted",
|
||||
["concept_type"],
|
||||
)
|
||||
|
||||
# --- System State ---
|
||||
|
||||
EMBEDDING_QUEUE_SIZE = Gauge(
|
||||
"mnemosyne_embedding_queue_size",
|
||||
"Items waiting in the embedding queue",
|
||||
)
|
||||
0
mnemosyne/library/migrations/__init__.py
Normal file
0
mnemosyne/library/migrations/__init__.py
Normal file
251
mnemosyne/library/models.py
Normal file
251
mnemosyne/library/models.py
Normal file
@@ -0,0 +1,251 @@
|
||||
"""
|
||||
Neo4j graph models for the Mnemosyne content library.
|
||||
|
||||
All content data (libraries, collections, items, chunks, concepts, images)
|
||||
lives in Neo4j as a knowledge graph. These models use neomodel's StructuredNode
|
||||
OGM — they do NOT participate in Django's ORM or migrations.
|
||||
"""
|
||||
|
||||
from neomodel import (
|
||||
ArrayProperty,
|
||||
DateTimeProperty,
|
||||
FloatProperty,
|
||||
IntegerProperty,
|
||||
JSONProperty,
|
||||
RelationshipTo,
|
||||
StringProperty,
|
||||
StructuredNode,
|
||||
StructuredRel,
|
||||
UniqueIdProperty,
|
||||
)
|
||||
|
||||
|
||||
# --- Relationship models ---
|
||||
|
||||
|
||||
class ReferencesRel(StructuredRel):
|
||||
"""Relationship properties for Item -> Concept REFERENCES edges."""
|
||||
|
||||
weight = FloatProperty(default=1.0)
|
||||
context = StringProperty(default="")
|
||||
|
||||
|
||||
class RelatedToRel(StructuredRel):
|
||||
"""Relationship properties for Item -> Item RELATED_TO edges."""
|
||||
|
||||
relationship_type = StringProperty(default="")
|
||||
weight = FloatProperty(default=1.0)
|
||||
|
||||
|
||||
class NearbyImageRel(StructuredRel):
|
||||
"""Relationship properties for Chunk -> Image HAS_NEARBY_IMAGE edges."""
|
||||
|
||||
proximity = StringProperty(default="same_page") # same_page, inline, same_slide, same_chapter
|
||||
|
||||
|
||||
# --- Node models ---
|
||||
|
||||
|
||||
class Library(StructuredNode):
|
||||
"""
|
||||
Top-level container representing a content library.
|
||||
|
||||
Each library has a type (fiction, technical, music, film, art, journal)
|
||||
that drives chunking strategy, embedding instructions, and LLM prompts.
|
||||
"""
|
||||
|
||||
uid = UniqueIdProperty()
|
||||
name = StringProperty(unique_index=True, required=True)
|
||||
library_type = StringProperty(
|
||||
required=True,
|
||||
choices={
|
||||
"fiction": "Fiction",
|
||||
"technical": "Technical",
|
||||
"music": "Music",
|
||||
"film": "Film",
|
||||
"art": "Art",
|
||||
"journal": "Journal",
|
||||
},
|
||||
)
|
||||
description = StringProperty(default="")
|
||||
|
||||
# Content-type configuration
|
||||
chunking_config = JSONProperty(default={})
|
||||
embedding_instruction = StringProperty(default="")
|
||||
reranker_instruction = StringProperty(default="")
|
||||
llm_context_prompt = StringProperty(default="")
|
||||
|
||||
created_at = DateTimeProperty(default_now=True)
|
||||
|
||||
# Relationships
|
||||
collections = RelationshipTo("Collection", "CONTAINS")
|
||||
|
||||
def __str__(self):
|
||||
return f"{self.name} ({self.library_type})"
|
||||
|
||||
|
||||
class Collection(StructuredNode):
|
||||
"""
|
||||
A grouping of items within a library.
|
||||
|
||||
Examples: a book series, an album discography, a project folder.
|
||||
"""
|
||||
|
||||
uid = UniqueIdProperty()
|
||||
name = StringProperty(required=True)
|
||||
description = StringProperty(default="")
|
||||
metadata = JSONProperty(default={})
|
||||
|
||||
created_at = DateTimeProperty(default_now=True)
|
||||
|
||||
# Relationships
|
||||
items = RelationshipTo("Item", "CONTAINS")
|
||||
library = RelationshipTo("Library", "BELONGS_TO")
|
||||
|
||||
def __str__(self):
|
||||
return self.name
|
||||
|
||||
|
||||
class Item(StructuredNode):
|
||||
"""
|
||||
An individual piece of content: a document, song, image set, journal entry, etc.
|
||||
|
||||
Items store their original file in S3 (via s3_key) and are chunked
|
||||
for embedding and retrieval.
|
||||
"""
|
||||
|
||||
uid = UniqueIdProperty()
|
||||
title = StringProperty(required=True)
|
||||
item_type = StringProperty(default="")
|
||||
s3_key = StringProperty(default="")
|
||||
content_hash = StringProperty(index=True)
|
||||
file_type = StringProperty(default="")
|
||||
file_size = IntegerProperty(default=0)
|
||||
metadata = JSONProperty(default={})
|
||||
|
||||
created_at = DateTimeProperty(default_now=True)
|
||||
updated_at = DateTimeProperty(default_now=True)
|
||||
|
||||
# Embedding pipeline fields (Phase 2)
|
||||
embedding_status = StringProperty(
|
||||
default="pending",
|
||||
choices={
|
||||
"pending": "Pending",
|
||||
"processing": "Processing",
|
||||
"completed": "Completed",
|
||||
"failed": "Failed",
|
||||
},
|
||||
)
|
||||
embedding_model_name = StringProperty(default="")
|
||||
chunk_count = IntegerProperty(default=0)
|
||||
image_count = IntegerProperty(default=0)
|
||||
error_message = StringProperty(default="")
|
||||
|
||||
# Relationships
|
||||
chunks = RelationshipTo("Chunk", "HAS_CHUNK")
|
||||
images = RelationshipTo("Image", "HAS_IMAGE")
|
||||
concepts = RelationshipTo("Concept", "REFERENCES", model=ReferencesRel)
|
||||
related_items = RelationshipTo("Item", "RELATED_TO", model=RelatedToRel)
|
||||
|
||||
def __str__(self):
|
||||
return self.title
|
||||
|
||||
|
||||
class Chunk(StructuredNode):
|
||||
"""
|
||||
A text chunk extracted from an Item for embedding and retrieval.
|
||||
|
||||
Chunk text is stored in S3; text_preview holds the first 500 chars
|
||||
for Neo4j full-text indexing.
|
||||
"""
|
||||
|
||||
uid = UniqueIdProperty()
|
||||
chunk_index = IntegerProperty(required=True)
|
||||
chunk_s3_key = StringProperty(required=True)
|
||||
chunk_size = IntegerProperty(default=0)
|
||||
text_preview = StringProperty(default="") # First 500 chars for full-text index
|
||||
embedding = ArrayProperty(FloatProperty()) # 4096d vector
|
||||
|
||||
created_at = DateTimeProperty(default_now=True)
|
||||
|
||||
# Relationships
|
||||
mentions = RelationshipTo("Concept", "MENTIONS")
|
||||
nearby_images = RelationshipTo("Image", "HAS_NEARBY_IMAGE", model=NearbyImageRel)
|
||||
|
||||
def __str__(self):
|
||||
return f"Chunk {self.chunk_index} ({self.uid})"
|
||||
|
||||
|
||||
class Concept(StructuredNode):
|
||||
"""
|
||||
A named entity or topic extracted from content.
|
||||
|
||||
Concepts form the backbone of the knowledge graph, linking items
|
||||
and chunks through shared references.
|
||||
"""
|
||||
|
||||
uid = UniqueIdProperty()
|
||||
name = StringProperty(unique_index=True, required=True)
|
||||
concept_type = StringProperty(
|
||||
default="",
|
||||
choices={
|
||||
"person": "Person",
|
||||
"place": "Place",
|
||||
"topic": "Topic",
|
||||
"technique": "Technique",
|
||||
"theme": "Theme",
|
||||
},
|
||||
)
|
||||
embedding = ArrayProperty(FloatProperty()) # 4096d vector
|
||||
|
||||
# Relationships
|
||||
related_concepts = RelationshipTo("Concept", "RELATED_TO")
|
||||
|
||||
def __str__(self):
|
||||
return self.name
|
||||
|
||||
|
||||
class Image(StructuredNode):
|
||||
"""
|
||||
An image associated with an Item (cover art, diagram, photo, etc.).
|
||||
|
||||
The image file is stored in S3; embeddings enable multimodal search.
|
||||
"""
|
||||
|
||||
uid = UniqueIdProperty()
|
||||
s3_key = StringProperty(required=True)
|
||||
image_type = StringProperty(
|
||||
default="",
|
||||
choices={
|
||||
"cover": "Cover",
|
||||
"diagram": "Diagram",
|
||||
"artwork": "Artwork",
|
||||
"still": "Still",
|
||||
"photo": "Photo",
|
||||
},
|
||||
)
|
||||
description = StringProperty(default="")
|
||||
metadata = JSONProperty(default={})
|
||||
|
||||
created_at = DateTimeProperty(default_now=True)
|
||||
|
||||
# Relationships
|
||||
embeddings = RelationshipTo("ImageEmbedding", "HAS_EMBEDDING")
|
||||
|
||||
def __str__(self):
|
||||
return f"Image {self.image_type} ({self.uid})"
|
||||
|
||||
|
||||
class ImageEmbedding(StructuredNode):
|
||||
"""
|
||||
A multimodal embedding vector for an Image node.
|
||||
|
||||
Generated by Qwen3-VL for unified text+image vector space.
|
||||
"""
|
||||
|
||||
uid = UniqueIdProperty()
|
||||
embedding = ArrayProperty(FloatProperty()) # 4096d multimodal vector
|
||||
created_at = DateTimeProperty(default_now=True)
|
||||
|
||||
def __str__(self):
|
||||
return f"ImageEmbedding ({self.uid})"
|
||||
11
mnemosyne/library/services/__init__.py
Normal file
11
mnemosyne/library/services/__init__.py
Normal file
@@ -0,0 +1,11 @@
|
||||
"""
|
||||
Library services for the Mnemosyne embedding pipeline.
|
||||
|
||||
Services:
|
||||
- parsers: Universal document parsing via PyMuPDF
|
||||
- text_utils: Text sanitization for embedding APIs
|
||||
- chunker: Content-type-aware chunking
|
||||
- embedding_client: Multi-backend embedding API client
|
||||
- pipeline: Orchestration of parse → chunk → embed → graph
|
||||
- concepts: LLM-based concept extraction
|
||||
"""
|
||||
250
mnemosyne/library/services/chunker.py
Normal file
250
mnemosyne/library/services/chunker.py
Normal file
@@ -0,0 +1,250 @@
|
||||
"""
|
||||
Content-type-aware chunking service.
|
||||
|
||||
Uses semantic-text-splitter with HuggingFace tokenizers to produce
|
||||
chunks that respect document structure boundaries per library type.
|
||||
"""
|
||||
|
||||
import logging
|
||||
from typing import Optional
|
||||
|
||||
from library.metrics import CHUNKS_CREATED_TOTAL, CHUNK_SIZE_TOKENS
|
||||
|
||||
from .parsers import ParseResult, TextBlock
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Default tokenizer when no model-specific tokenizer is available
|
||||
DEFAULT_TOKENIZER = "bert-base-uncased"
|
||||
|
||||
# Boundary markers used to detect structural elements in text
|
||||
_BOUNDARY_PATTERNS = {
|
||||
"chapter": [
|
||||
r"(?m)^chapter\s+\d+",
|
||||
r"(?m)^CHAPTER\s+\d+",
|
||||
r"(?m)^Chapter\s+\w+",
|
||||
],
|
||||
"scene": [r"(?m)^\*\s*\*\s*\*", r"(?m)^---+$", r"(?m)^###"],
|
||||
"section": [
|
||||
r"(?m)^#{1,3}\s+",
|
||||
r"(?m)^\d+\.\d*\s+\w",
|
||||
r"(?m)^Section\s+\d+",
|
||||
],
|
||||
"subsection": [r"(?m)^#{4,6}\s+", r"(?m)^\d+\.\d+\.\d+"],
|
||||
"entry": [
|
||||
r"(?m)^\d{4}-\d{2}-\d{2}",
|
||||
r"(?m)^(January|February|March|April|May|June|July|August|September|October|November|December)\s+\d",
|
||||
],
|
||||
"song": [r"(?m)^Track\s+\d+", r"(?m)^\[.*?\]$"],
|
||||
"verse": [r"(?m)^\[Verse", r"(?m)^\[Chorus"],
|
||||
}
|
||||
|
||||
|
||||
class ChunkResult:
|
||||
"""Result of chunking a document."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
chunks: list[str],
|
||||
chunk_page_map: dict[int, int],
|
||||
strategy: str,
|
||||
):
|
||||
"""
|
||||
:param chunks: List of chunk text strings.
|
||||
:param chunk_page_map: Mapping of chunk_index -> source page number.
|
||||
:param strategy: The chunking strategy used.
|
||||
"""
|
||||
self.chunks = chunks
|
||||
self.chunk_page_map = chunk_page_map
|
||||
self.strategy = strategy
|
||||
|
||||
def __len__(self) -> int:
|
||||
return len(self.chunks)
|
||||
|
||||
|
||||
class ContentTypeChunker:
|
||||
"""
|
||||
Content-type-aware document chunker.
|
||||
|
||||
Dispatches to different chunking strategies based on the library's
|
||||
chunking configuration, using semantic-text-splitter for
|
||||
token-aware splitting.
|
||||
"""
|
||||
|
||||
def __init__(self, tokenizer_name: Optional[str] = None):
|
||||
"""
|
||||
:param tokenizer_name: HuggingFace tokenizer name for token counting.
|
||||
"""
|
||||
self._tokenizer_name = tokenizer_name or DEFAULT_TOKENIZER
|
||||
self._splitter_cache: dict[tuple[int, int], object] = {}
|
||||
|
||||
def chunk(
|
||||
self,
|
||||
parse_result: ParseResult,
|
||||
chunking_config: dict,
|
||||
library_type: str = "",
|
||||
) -> ChunkResult:
|
||||
"""
|
||||
Chunk parsed document text using the library's chunking config.
|
||||
|
||||
:param parse_result: ParseResult from the document parser.
|
||||
:param chunking_config: Library chunking configuration dict.
|
||||
:param library_type: Library type for metrics labeling.
|
||||
:returns: ChunkResult with chunk texts and page mapping.
|
||||
"""
|
||||
strategy = chunking_config.get("strategy", "section_aware")
|
||||
chunk_size = chunking_config.get("chunk_size", 512)
|
||||
chunk_overlap = chunking_config.get("chunk_overlap", 64)
|
||||
|
||||
# Combine all text blocks into a single document text,
|
||||
# tracking page boundaries for chunk-page mapping
|
||||
full_text, page_offsets = self._combine_text_blocks(parse_result.text_blocks)
|
||||
|
||||
if not full_text.strip():
|
||||
logger.warning("No text to chunk strategy=%s", strategy)
|
||||
return ChunkResult(chunks=[], chunk_page_map={}, strategy=strategy)
|
||||
|
||||
logger.info(
|
||||
"Chunking text strategy=%s chunk_size=%d overlap=%d total_chars=%d",
|
||||
strategy,
|
||||
chunk_size,
|
||||
chunk_overlap,
|
||||
len(full_text),
|
||||
)
|
||||
|
||||
# Get or create the text splitter for this size/overlap
|
||||
splitter = self._get_splitter(chunk_size, chunk_overlap)
|
||||
|
||||
# Split into chunks
|
||||
try:
|
||||
chunks = splitter.chunks(full_text)
|
||||
except Exception as exc:
|
||||
logger.error("Chunking failed strategy=%s: %s", strategy, exc)
|
||||
raise
|
||||
|
||||
# Build chunk -> page mapping
|
||||
chunk_page_map = self._map_chunks_to_pages(chunks, full_text, page_offsets)
|
||||
|
||||
# Record metrics
|
||||
CHUNKS_CREATED_TOTAL.labels(
|
||||
library_type=library_type,
|
||||
strategy=strategy,
|
||||
).inc(len(chunks))
|
||||
|
||||
for chunk_text in chunks:
|
||||
CHUNK_SIZE_TOKENS.observe(len(chunk_text.split()))
|
||||
|
||||
logger.info(
|
||||
"Chunked document strategy=%s chunks=%d avg_size=%d",
|
||||
strategy,
|
||||
len(chunks),
|
||||
sum(len(c) for c in chunks) // max(len(chunks), 1),
|
||||
)
|
||||
|
||||
return ChunkResult(
|
||||
chunks=chunks,
|
||||
chunk_page_map=chunk_page_map,
|
||||
strategy=strategy,
|
||||
)
|
||||
|
||||
def _get_splitter(self, chunk_size: int, chunk_overlap: int):
|
||||
"""
|
||||
Get or create a semantic text splitter for the given parameters.
|
||||
|
||||
:param chunk_size: Maximum chunk size in tokens.
|
||||
:param chunk_overlap: Overlap between chunks in tokens.
|
||||
:returns: TextSplitter instance.
|
||||
"""
|
||||
cache_key = (chunk_size, chunk_overlap)
|
||||
if cache_key in self._splitter_cache:
|
||||
return self._splitter_cache[cache_key]
|
||||
|
||||
from semantic_text_splitter import TextSplitter
|
||||
from tokenizers import Tokenizer
|
||||
|
||||
try:
|
||||
tokenizer = Tokenizer.from_pretrained(self._tokenizer_name)
|
||||
splitter = TextSplitter.from_huggingface_tokenizer(
|
||||
tokenizer,
|
||||
capacity=chunk_size,
|
||||
overlap=chunk_overlap,
|
||||
)
|
||||
logger.debug(
|
||||
"Created text splitter tokenizer=%s capacity=%d overlap=%d",
|
||||
self._tokenizer_name,
|
||||
chunk_size,
|
||||
chunk_overlap,
|
||||
)
|
||||
except Exception as exc:
|
||||
logger.warning(
|
||||
"Failed to load tokenizer %s: %s, falling back to %s",
|
||||
self._tokenizer_name,
|
||||
exc,
|
||||
DEFAULT_TOKENIZER,
|
||||
)
|
||||
tokenizer = Tokenizer.from_pretrained(DEFAULT_TOKENIZER)
|
||||
splitter = TextSplitter.from_huggingface_tokenizer(
|
||||
tokenizer,
|
||||
capacity=chunk_size,
|
||||
overlap=chunk_overlap,
|
||||
)
|
||||
|
||||
self._splitter_cache[cache_key] = splitter
|
||||
return splitter
|
||||
|
||||
def _combine_text_blocks(
|
||||
self, text_blocks: list[TextBlock]
|
||||
) -> tuple[str, list[tuple[int, int]]]:
|
||||
"""
|
||||
Combine text blocks into a single string, tracking page offsets.
|
||||
|
||||
:param text_blocks: List of TextBlock from parser.
|
||||
:returns: Tuple of (combined_text, page_offsets) where page_offsets
|
||||
is a list of (char_offset, page_number).
|
||||
"""
|
||||
parts: list[str] = []
|
||||
page_offsets: list[tuple[int, int]] = []
|
||||
current_offset = 0
|
||||
|
||||
for block in text_blocks:
|
||||
page_offsets.append((current_offset, block.page))
|
||||
parts.append(block.text)
|
||||
current_offset += len(block.text) + 2 # +2 for paragraph separator
|
||||
|
||||
return "\n\n".join(parts), page_offsets
|
||||
|
||||
def _map_chunks_to_pages(
|
||||
self,
|
||||
chunks: list[str],
|
||||
full_text: str,
|
||||
page_offsets: list[tuple[int, int]],
|
||||
) -> dict[int, int]:
|
||||
"""
|
||||
Map each chunk index to its source page number.
|
||||
|
||||
:param chunks: List of chunk strings.
|
||||
:param full_text: Combined document text.
|
||||
:param page_offsets: List of (char_offset, page_number).
|
||||
:returns: Dict mapping chunk_index -> page_number.
|
||||
"""
|
||||
chunk_page_map: dict[int, int] = {}
|
||||
search_start = 0
|
||||
|
||||
for chunk_idx, chunk_text in enumerate(chunks):
|
||||
# Find where this chunk starts in the full text
|
||||
pos = full_text.find(chunk_text[:100], search_start)
|
||||
if pos == -1:
|
||||
pos = search_start
|
||||
|
||||
# Find which page this position belongs to
|
||||
page = 0
|
||||
for offset, page_num in page_offsets:
|
||||
if offset <= pos:
|
||||
page = page_num
|
||||
else:
|
||||
break
|
||||
|
||||
chunk_page_map[chunk_idx] = page
|
||||
search_start = max(search_start, pos)
|
||||
|
||||
return chunk_page_map
|
||||
267
mnemosyne/library/services/concepts.py
Normal file
267
mnemosyne/library/services/concepts.py
Normal file
@@ -0,0 +1,267 @@
|
||||
"""
|
||||
LLM-based concept extraction for the knowledge graph.
|
||||
|
||||
Uses the system chat model to extract named entities (people, places,
|
||||
topics, techniques, themes) from document chunks, then creates Concept
|
||||
nodes and MENTIONS/REFERENCES relationships in Neo4j.
|
||||
"""
|
||||
|
||||
import json
|
||||
import logging
|
||||
from typing import Optional
|
||||
|
||||
from library.metrics import CONCEPTS_EXTRACTED_TOTAL
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Prompt for concept extraction
|
||||
CONCEPT_EXTRACTION_PROMPT = """Extract named entities and key concepts from the following text.
|
||||
|
||||
Return a JSON array of objects, each with:
|
||||
- "name": the entity/concept name (lowercase, canonical form)
|
||||
- "type": one of "person", "place", "topic", "technique", "theme"
|
||||
|
||||
Only extract significant, specific concepts — not generic words.
|
||||
Return at most 20 concepts. Return ONLY the JSON array, no other text.
|
||||
|
||||
Text:
|
||||
{text}"""
|
||||
|
||||
|
||||
class ConceptExtractor:
|
||||
"""
|
||||
Extracts concepts from text using the system chat model.
|
||||
|
||||
Creates or updates Concept nodes in Neo4j and connects them
|
||||
to Chunk and Item nodes via MENTIONS and REFERENCES relationships.
|
||||
"""
|
||||
|
||||
def __init__(self, chat_model, user=None):
|
||||
"""
|
||||
:param chat_model: LLMModel instance for chat/completion.
|
||||
:param user: Optional Django user for usage tracking.
|
||||
"""
|
||||
self.chat_model = chat_model
|
||||
self.user = user
|
||||
|
||||
def extract_for_item(
|
||||
self,
|
||||
item,
|
||||
chunk_nodes: list,
|
||||
chunk_texts: list[str],
|
||||
) -> int:
|
||||
"""
|
||||
Extract concepts from all chunks of an item.
|
||||
|
||||
:param item: Item node.
|
||||
:param chunk_nodes: List of Chunk nodes.
|
||||
:param chunk_texts: List of chunk text strings.
|
||||
:returns: Total number of unique concepts extracted.
|
||||
"""
|
||||
all_concepts: dict[str, str] = {} # name -> type
|
||||
|
||||
# Sample chunks for extraction (don't process every chunk for large docs)
|
||||
sample_indices = self._select_sample_indices(len(chunk_texts), max_samples=10)
|
||||
|
||||
for idx in sample_indices:
|
||||
chunk_text = chunk_texts[idx]
|
||||
chunk_node = chunk_nodes[idx]
|
||||
|
||||
concepts = self._extract_from_text(chunk_text)
|
||||
if not concepts:
|
||||
continue
|
||||
|
||||
for concept_data in concepts:
|
||||
name = concept_data.get("name", "").strip().lower()
|
||||
concept_type = concept_data.get("type", "topic")
|
||||
|
||||
if not name or len(name) < 2:
|
||||
continue
|
||||
|
||||
all_concepts[name] = concept_type
|
||||
|
||||
# Connect chunk -> concept via MENTIONS
|
||||
concept_node = self._get_or_create_concept(name, concept_type)
|
||||
if concept_node:
|
||||
try:
|
||||
chunk_node.mentions.connect(concept_node)
|
||||
except Exception:
|
||||
pass # Already connected
|
||||
|
||||
# Connect item -> all concepts via REFERENCES
|
||||
for name, concept_type in all_concepts.items():
|
||||
concept_node = self._get_or_create_concept(name, concept_type)
|
||||
if concept_node:
|
||||
try:
|
||||
item.concepts.connect(concept_node, {"weight": 1.0})
|
||||
except Exception:
|
||||
pass # Already connected
|
||||
|
||||
CONCEPTS_EXTRACTED_TOTAL.labels(concept_type=concept_type).inc()
|
||||
|
||||
logger.info(
|
||||
"Extracted %d concepts for item_uid=%s",
|
||||
len(all_concepts),
|
||||
item.uid,
|
||||
)
|
||||
return len(all_concepts)
|
||||
|
||||
def _extract_from_text(self, text: str) -> list[dict]:
|
||||
"""
|
||||
Call the chat model to extract concepts from text.
|
||||
|
||||
:param text: Text to analyze.
|
||||
:returns: List of concept dicts with 'name' and 'type' keys.
|
||||
"""
|
||||
# Truncate very long text to avoid token limits
|
||||
if len(text) > 3000:
|
||||
text = text[:3000]
|
||||
|
||||
prompt = CONCEPT_EXTRACTION_PROMPT.format(text=text)
|
||||
|
||||
try:
|
||||
response_text = self._call_chat_model(prompt)
|
||||
concepts = self._parse_concept_response(response_text)
|
||||
logger.debug(
|
||||
"Extracted %d concepts from text chunk (len=%d)",
|
||||
len(concepts),
|
||||
len(text),
|
||||
)
|
||||
return concepts
|
||||
except Exception as exc:
|
||||
logger.warning("Concept extraction failed: %s", exc)
|
||||
return []
|
||||
|
||||
def _call_chat_model(self, prompt: str) -> str:
|
||||
"""
|
||||
Make a chat completion request to the system chat model.
|
||||
|
||||
:param prompt: User prompt text.
|
||||
:returns: Response text from the model.
|
||||
"""
|
||||
import requests
|
||||
|
||||
api = self.chat_model.api
|
||||
base_url = api.base_url.rstrip("/")
|
||||
|
||||
if api.api_type == "bedrock":
|
||||
# Bedrock Converse endpoint
|
||||
url = f"{base_url}/model/{self.chat_model.name}/converse"
|
||||
headers = {
|
||||
"Authorization": f"Bearer {api.api_key}",
|
||||
"Content-Type": "application/json",
|
||||
}
|
||||
body = {
|
||||
"messages": [{"role": "user", "content": [{"text": prompt}]}],
|
||||
}
|
||||
else:
|
||||
# OpenAI-compatible
|
||||
url = f"{base_url}/chat/completions"
|
||||
headers = {"Content-Type": "application/json"}
|
||||
if api.api_key:
|
||||
headers["Authorization"] = f"Bearer {api.api_key}"
|
||||
body = {
|
||||
"model": self.chat_model.name,
|
||||
"messages": [{"role": "user", "content": prompt}],
|
||||
"temperature": 0.1,
|
||||
"max_tokens": 1000,
|
||||
}
|
||||
|
||||
resp = requests.post(
|
||||
url, json=body, headers=headers, timeout=api.timeout_seconds or 60
|
||||
)
|
||||
resp.raise_for_status()
|
||||
data = resp.json()
|
||||
|
||||
# Parse response based on format
|
||||
if "output" in data:
|
||||
# Bedrock Converse format
|
||||
return data["output"]["message"]["content"][0]["text"]
|
||||
if "choices" in data:
|
||||
# OpenAI format
|
||||
return data["choices"][0]["message"]["content"]
|
||||
|
||||
raise ValueError(f"Unexpected chat response format: {list(data.keys())}")
|
||||
|
||||
def _parse_concept_response(self, response_text: str) -> list[dict]:
|
||||
"""
|
||||
Parse the LLM's concept extraction response into structured data.
|
||||
|
||||
:param response_text: Raw response text (expected JSON array).
|
||||
:returns: List of concept dicts.
|
||||
"""
|
||||
# Try to extract JSON from the response
|
||||
text = response_text.strip()
|
||||
|
||||
# Handle markdown code blocks
|
||||
if text.startswith("```"):
|
||||
lines = text.split("\n")
|
||||
text = "\n".join(lines[1:-1]) if len(lines) > 2 else text
|
||||
|
||||
try:
|
||||
concepts = json.loads(text)
|
||||
if isinstance(concepts, list):
|
||||
return [
|
||||
c for c in concepts
|
||||
if isinstance(c, dict) and "name" in c
|
||||
]
|
||||
except json.JSONDecodeError:
|
||||
pass
|
||||
|
||||
# Try to find JSON array in the response
|
||||
import re
|
||||
|
||||
match = re.search(r"\[.*\]", text, re.DOTALL)
|
||||
if match:
|
||||
try:
|
||||
concepts = json.loads(match.group())
|
||||
if isinstance(concepts, list):
|
||||
return [
|
||||
c for c in concepts
|
||||
if isinstance(c, dict) and "name" in c
|
||||
]
|
||||
except json.JSONDecodeError:
|
||||
pass
|
||||
|
||||
logger.debug("Could not parse concept response: %s", text[:200])
|
||||
return []
|
||||
|
||||
def _get_or_create_concept(self, name: str, concept_type: str):
|
||||
"""
|
||||
Get or create a Concept node by name.
|
||||
|
||||
:param name: Concept name (lowercase).
|
||||
:param concept_type: Concept type (person, place, topic, etc.).
|
||||
:returns: Concept node, or None on failure.
|
||||
"""
|
||||
from library.models import Concept
|
||||
|
||||
try:
|
||||
# Try to get existing
|
||||
existing = Concept.nodes.filter(name=name)
|
||||
if existing:
|
||||
return existing[0]
|
||||
|
||||
# Create new
|
||||
concept = Concept(name=name, concept_type=concept_type)
|
||||
concept.save()
|
||||
return concept
|
||||
except Exception as exc:
|
||||
logger.debug("Failed to get/create concept '%s': %s", name, exc)
|
||||
return None
|
||||
|
||||
def _select_sample_indices(
|
||||
self, total: int, max_samples: int = 10
|
||||
) -> list[int]:
|
||||
"""
|
||||
Select evenly-spaced sample indices for concept extraction.
|
||||
|
||||
:param total: Total number of chunks.
|
||||
:param max_samples: Maximum samples to take.
|
||||
:returns: List of chunk indices to process.
|
||||
"""
|
||||
if total <= max_samples:
|
||||
return list(range(total))
|
||||
|
||||
step = total / max_samples
|
||||
return [int(i * step) for i in range(max_samples)]
|
||||
396
mnemosyne/library/services/embedding_client.py
Normal file
396
mnemosyne/library/services/embedding_client.py
Normal file
@@ -0,0 +1,396 @@
|
||||
"""
|
||||
Multi-backend embedding client.
|
||||
|
||||
Dispatches embedding requests to OpenAI-compatible APIs (OpenAI, vLLM,
|
||||
llama-cpp, Ollama) or Amazon Bedrock via direct HTTP with Bearer token auth.
|
||||
"""
|
||||
|
||||
import base64
|
||||
import json
|
||||
import logging
|
||||
import time
|
||||
from typing import Optional
|
||||
|
||||
import requests
|
||||
|
||||
from library.metrics import (
|
||||
EMBEDDING_API_ERRORS_TOTAL,
|
||||
EMBEDDING_BATCH_DURATION,
|
||||
EMBEDDING_TOKENS_TOTAL,
|
||||
EMBEDDINGS_GENERATED_TOTAL,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class EmbeddingClient:
|
||||
"""
|
||||
Client for generating text and image embeddings via multiple backends.
|
||||
|
||||
Dispatches based on ``LLMApi.api_type``:
|
||||
|
||||
* ``openai``, ``vllm``, ``llama-cpp``, ``ollama`` — OpenAI-compatible
|
||||
``POST /embeddings``
|
||||
* ``bedrock`` — Amazon Bedrock Runtime ``POST /model/{id}/invoke``
|
||||
with Bearer token auth
|
||||
"""
|
||||
|
||||
def __init__(self, embedding_model, user=None):
|
||||
"""
|
||||
:param embedding_model: ``LLMModel`` instance for embeddings.
|
||||
:param user: Optional Django user for usage tracking.
|
||||
"""
|
||||
self.model = embedding_model
|
||||
self.api = embedding_model.api
|
||||
self.user = user
|
||||
self.base_url = self.api.base_url.rstrip("/")
|
||||
self.model_name = self.model.name
|
||||
self.api_type = self.api.api_type
|
||||
self.timeout = self.api.timeout_seconds or 120
|
||||
|
||||
logger.info(
|
||||
"EmbeddingClient initialized api=%s model=%s api_type=%s base_url=%s",
|
||||
self.api.name,
|
||||
self.model_name,
|
||||
self.api_type,
|
||||
self.base_url,
|
||||
)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Public API
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def embed_text(self, text: str) -> list[float]:
|
||||
"""
|
||||
Generate an embedding vector for a single text string.
|
||||
|
||||
:param text: Text to embed.
|
||||
:returns: Embedding vector as list of floats.
|
||||
"""
|
||||
if self.api_type == "bedrock":
|
||||
return self._embed_bedrock_single(text)
|
||||
return self._embed_openai_single(text)
|
||||
|
||||
def embed_texts(self, texts: list[str]) -> list[list[float]]:
|
||||
"""
|
||||
Generate embeddings for multiple texts.
|
||||
|
||||
:param texts: List of text strings.
|
||||
:returns: List of embedding vectors.
|
||||
"""
|
||||
if self.api_type == "bedrock":
|
||||
return self._embed_bedrock_batch(texts)
|
||||
return self._embed_openai_batch(texts)
|
||||
|
||||
def embed_image(self, image_data: bytes, image_ext: str = "png") -> Optional[list[float]]:
|
||||
"""
|
||||
Generate a multimodal embedding for an image.
|
||||
|
||||
Requires a model with ``supports_multimodal=True``.
|
||||
|
||||
:param image_data: Raw image bytes.
|
||||
:param image_ext: Image format extension.
|
||||
:returns: Embedding vector, or None if not supported.
|
||||
"""
|
||||
if not self.model.supports_multimodal:
|
||||
logger.debug(
|
||||
"Model %s does not support multimodal, skipping image embedding",
|
||||
self.model_name,
|
||||
)
|
||||
return None
|
||||
|
||||
b64 = base64.b64encode(image_data).decode("utf-8")
|
||||
mime_type = f"image/{image_ext}" if image_ext != "jpg" else "image/jpeg"
|
||||
|
||||
if self.api_type == "bedrock":
|
||||
return self._embed_bedrock_image(b64, mime_type)
|
||||
return self._embed_openai_image(b64, mime_type)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# OpenAI-compatible backend
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def _embed_openai_single(self, text: str) -> list[float]:
|
||||
"""Embed a single text via OpenAI-compatible /embeddings endpoint."""
|
||||
result = self._embed_openai_batch([text])
|
||||
return result[0]
|
||||
|
||||
def _embed_openai_batch(self, texts: list[str]) -> list[list[float]]:
|
||||
"""Embed a batch of texts via OpenAI-compatible /embeddings endpoint."""
|
||||
url = f"{self.base_url}/embeddings"
|
||||
payload = {"input": texts, "model": self.model_name}
|
||||
headers = self._openai_headers()
|
||||
|
||||
logger.debug(
|
||||
"OpenAI embedding request texts=%d model=%s",
|
||||
len(texts),
|
||||
self.model_name,
|
||||
)
|
||||
|
||||
with EMBEDDING_BATCH_DURATION.labels(
|
||||
model_name=self.model_name, api_type=self.api_type
|
||||
).time():
|
||||
try:
|
||||
resp = requests.post(
|
||||
url,
|
||||
json=payload,
|
||||
headers=headers,
|
||||
timeout=self.timeout * max(1, len(texts) // 10),
|
||||
)
|
||||
if resp.status_code != 200:
|
||||
logger.error(
|
||||
"OpenAI embedding failed status=%d body=%s",
|
||||
resp.status_code,
|
||||
resp.text[:500],
|
||||
)
|
||||
resp.raise_for_status()
|
||||
data = resp.json()
|
||||
except requests.RequestException as exc:
|
||||
EMBEDDING_API_ERRORS_TOTAL.labels(
|
||||
model_name=self.model_name,
|
||||
api_type=self.api_type,
|
||||
error_type=type(exc).__name__,
|
||||
).inc()
|
||||
logger.error("OpenAI embedding request failed: %s", exc)
|
||||
raise
|
||||
|
||||
embeddings = self._parse_openai_response(data)
|
||||
|
||||
# Metrics
|
||||
EMBEDDINGS_GENERATED_TOTAL.labels(
|
||||
model_name=self.model_name,
|
||||
api_type=self.api_type,
|
||||
content_type="text",
|
||||
).inc(len(embeddings))
|
||||
EMBEDDING_TOKENS_TOTAL.labels(model_name=self.model_name).inc(
|
||||
sum(len(t.split()) for t in texts)
|
||||
)
|
||||
|
||||
self._log_usage(len(texts), sum(len(t.split()) for t in texts))
|
||||
|
||||
logger.debug(
|
||||
"OpenAI embedding response texts=%d dimensions=%d",
|
||||
len(embeddings),
|
||||
len(embeddings[0]) if embeddings else 0,
|
||||
)
|
||||
return embeddings
|
||||
|
||||
def _embed_openai_image(self, b64_image: str, mime_type: str) -> Optional[list[float]]:
|
||||
"""Embed an image via OpenAI-compatible multimodal endpoint."""
|
||||
url = f"{self.base_url}/embeddings"
|
||||
payload = {
|
||||
"input": [{"type": "image_url", "image_url": {"url": f"data:{mime_type};base64,{b64_image}"}}],
|
||||
"model": self.model_name,
|
||||
}
|
||||
headers = self._openai_headers()
|
||||
|
||||
try:
|
||||
resp = requests.post(url, json=payload, headers=headers, timeout=self.timeout)
|
||||
resp.raise_for_status()
|
||||
data = resp.json()
|
||||
embeddings = self._parse_openai_response(data)
|
||||
if embeddings:
|
||||
EMBEDDINGS_GENERATED_TOTAL.labels(
|
||||
model_name=self.model_name,
|
||||
api_type=self.api_type,
|
||||
content_type="image",
|
||||
).inc()
|
||||
return embeddings[0]
|
||||
except Exception as exc:
|
||||
EMBEDDING_API_ERRORS_TOTAL.labels(
|
||||
model_name=self.model_name,
|
||||
api_type=self.api_type,
|
||||
error_type=type(exc).__name__,
|
||||
).inc()
|
||||
logger.warning("Image embedding failed: %s", exc)
|
||||
|
||||
return None
|
||||
|
||||
def _openai_headers(self) -> dict[str, str]:
|
||||
"""Build headers for OpenAI-compatible requests."""
|
||||
headers = {"Content-Type": "application/json"}
|
||||
if self.api.api_key:
|
||||
headers["Authorization"] = f"Bearer {self.api.api_key}"
|
||||
return headers
|
||||
|
||||
def _parse_openai_response(self, data) -> list[list[float]]:
|
||||
"""
|
||||
Parse embedding response from various OpenAI-compatible formats.
|
||||
|
||||
Handles:
|
||||
- OpenAI standard: ``{"data": [{"embedding": [...], "index": 0}]}``
|
||||
- Direct list of dicts: ``[{"embedding": [...]}]``
|
||||
- Direct list of vectors: ``[[0.1, 0.2, ...]]``
|
||||
- Dict with embeddings key: ``{"embeddings": [[...]]}``
|
||||
|
||||
:param data: Parsed JSON response.
|
||||
:returns: List of embedding vectors.
|
||||
"""
|
||||
if isinstance(data, list):
|
||||
if data and isinstance(data[0], dict) and "embedding" in data[0]:
|
||||
return [item["embedding"] for item in data]
|
||||
return data
|
||||
|
||||
if isinstance(data, dict):
|
||||
if "data" in data:
|
||||
return [
|
||||
item["embedding"]
|
||||
for item in sorted(data["data"], key=lambda x: x.get("index", 0))
|
||||
]
|
||||
if "embedding" in data:
|
||||
return [data["embedding"]]
|
||||
if "embeddings" in data:
|
||||
return data["embeddings"]
|
||||
|
||||
raise ValueError(f"Unexpected embedding response format: {type(data)}")
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Amazon Bedrock backend
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def _embed_bedrock_single(self, text: str) -> list[float]:
|
||||
"""Embed a single text via Bedrock Runtime InvokeModel."""
|
||||
url = f"{self.base_url}/model/{self.model_name}/invoke"
|
||||
headers = {
|
||||
"Authorization": f"Bearer {self.api.api_key}",
|
||||
"Content-Type": "application/json",
|
||||
"Accept": "application/json",
|
||||
}
|
||||
body = {"inputText": text, "normalize": True}
|
||||
|
||||
# Include dimensions if the model supports configurable output
|
||||
if self.model.vector_dimensions:
|
||||
body["dimensions"] = self.model.vector_dimensions
|
||||
|
||||
logger.debug(
|
||||
"Bedrock embedding request model=%s text_len=%d",
|
||||
self.model_name,
|
||||
len(text),
|
||||
)
|
||||
|
||||
with EMBEDDING_BATCH_DURATION.labels(
|
||||
model_name=self.model_name, api_type=self.api_type
|
||||
).time():
|
||||
try:
|
||||
resp = requests.post(
|
||||
url, json=body, headers=headers, timeout=self.timeout
|
||||
)
|
||||
if resp.status_code != 200:
|
||||
logger.error(
|
||||
"Bedrock embedding failed status=%d body=%s",
|
||||
resp.status_code,
|
||||
resp.text[:500],
|
||||
)
|
||||
resp.raise_for_status()
|
||||
data = resp.json()
|
||||
except requests.RequestException as exc:
|
||||
EMBEDDING_API_ERRORS_TOTAL.labels(
|
||||
model_name=self.model_name,
|
||||
api_type=self.api_type,
|
||||
error_type=type(exc).__name__,
|
||||
).inc()
|
||||
logger.error("Bedrock embedding request failed: %s", exc)
|
||||
raise
|
||||
|
||||
embedding = data.get("embedding")
|
||||
if not embedding:
|
||||
raise ValueError(f"Bedrock response missing 'embedding' key: {list(data.keys())}")
|
||||
|
||||
token_count = data.get("inputTextTokenCount", len(text.split()))
|
||||
|
||||
EMBEDDINGS_GENERATED_TOTAL.labels(
|
||||
model_name=self.model_name,
|
||||
api_type=self.api_type,
|
||||
content_type="text",
|
||||
).inc()
|
||||
EMBEDDING_TOKENS_TOTAL.labels(model_name=self.model_name).inc(token_count)
|
||||
self._log_usage(1, token_count)
|
||||
|
||||
logger.debug(
|
||||
"Bedrock embedding response dimensions=%d tokens=%d",
|
||||
len(embedding),
|
||||
token_count,
|
||||
)
|
||||
return embedding
|
||||
|
||||
def _embed_bedrock_batch(self, texts: list[str]) -> list[list[float]]:
|
||||
"""
|
||||
Embed multiple texts via Bedrock (client-side loop).
|
||||
|
||||
Bedrock InvokeModel accepts one input at a time, so we loop.
|
||||
|
||||
:param texts: List of text strings.
|
||||
:returns: List of embedding vectors.
|
||||
"""
|
||||
embeddings = []
|
||||
for i, text in enumerate(texts):
|
||||
embedding = self._embed_bedrock_single(text)
|
||||
embeddings.append(embedding)
|
||||
if (i + 1) % 10 == 0:
|
||||
logger.debug(
|
||||
"Bedrock batch progress %d/%d", i + 1, len(texts)
|
||||
)
|
||||
return embeddings
|
||||
|
||||
def _embed_bedrock_image(self, b64_image: str, mime_type: str) -> Optional[list[float]]:
|
||||
"""Embed an image via Bedrock multimodal endpoint."""
|
||||
url = f"{self.base_url}/model/{self.model_name}/invoke"
|
||||
headers = {
|
||||
"Authorization": f"Bearer {self.api.api_key}",
|
||||
"Content-Type": "application/json",
|
||||
"Accept": "application/json",
|
||||
}
|
||||
body = {
|
||||
"inputImage": b64_image,
|
||||
"normalize": True,
|
||||
}
|
||||
if self.model.vector_dimensions:
|
||||
body["dimensions"] = self.model.vector_dimensions
|
||||
|
||||
try:
|
||||
resp = requests.post(url, json=body, headers=headers, timeout=self.timeout)
|
||||
resp.raise_for_status()
|
||||
data = resp.json()
|
||||
embedding = data.get("embedding")
|
||||
if embedding:
|
||||
EMBEDDINGS_GENERATED_TOTAL.labels(
|
||||
model_name=self.model_name,
|
||||
api_type=self.api_type,
|
||||
content_type="image",
|
||||
).inc()
|
||||
return embedding
|
||||
except Exception as exc:
|
||||
EMBEDDING_API_ERRORS_TOTAL.labels(
|
||||
model_name=self.model_name,
|
||||
api_type=self.api_type,
|
||||
error_type=type(exc).__name__,
|
||||
).inc()
|
||||
logger.warning("Bedrock image embedding failed: %s", exc)
|
||||
|
||||
return None
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Usage tracking
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def _log_usage(self, text_count: int, token_count: int):
|
||||
"""
|
||||
Log embedding usage to LLMUsage model.
|
||||
|
||||
:param text_count: Number of texts embedded.
|
||||
:param token_count: Approximate token count.
|
||||
"""
|
||||
try:
|
||||
from llm_manager.models import LLMUsage
|
||||
|
||||
LLMUsage.objects.create(
|
||||
model=self.model,
|
||||
user=self.user,
|
||||
input_tokens=token_count,
|
||||
output_tokens=0,
|
||||
cached_tokens=0,
|
||||
total_cost=(token_count / 1000) * float(self.model.input_cost_per_1k),
|
||||
purpose="embeddings",
|
||||
)
|
||||
except Exception as exc:
|
||||
logger.warning("Failed to log embedding usage: %s", exc)
|
||||
360
mnemosyne/library/services/parsers.py
Normal file
360
mnemosyne/library/services/parsers.py
Normal file
@@ -0,0 +1,360 @@
|
||||
"""
|
||||
Universal document parsing service using PyMuPDF.
|
||||
|
||||
Handles text extraction and image extraction for all supported formats:
|
||||
PDF, EPUB, DOCX, PPTX, XLSX, XPS, MOBI, FB2, CBZ, TXT, HTML, and images.
|
||||
"""
|
||||
|
||||
import logging
|
||||
import os
|
||||
import tempfile
|
||||
from dataclasses import dataclass, field
|
||||
|
||||
import fitz # PyMuPDF
|
||||
|
||||
from library.metrics import (
|
||||
DOCUMENT_PARSE_DURATION,
|
||||
DOCUMENTS_PARSED_TOTAL,
|
||||
IMAGES_EXTRACTED_TOTAL,
|
||||
)
|
||||
|
||||
from .text_utils import remove_excessive_whitespace, sanitize_text
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# File extensions supported by PyMuPDF
|
||||
PYMUPDF_EXTENSIONS = {
|
||||
"pdf", "epub", "xps", "mobi", "fb2", "cbz", "svg",
|
||||
"docx", "pptx", "xlsx", "hwpx",
|
||||
}
|
||||
|
||||
# Plain text extensions — read directly, no PyMuPDF needed
|
||||
PLAINTEXT_EXTENSIONS = {"txt", "md", "csv", "tsv", "log", "json", "yaml", "yml", "xml"}
|
||||
|
||||
# Image extensions — store as Image nodes directly
|
||||
IMAGE_EXTENSIONS = {"jpg", "jpeg", "png", "gif", "bmp", "tiff", "tif", "webp", "svg"}
|
||||
|
||||
# Minimum image dimensions to extract (skip tiny icons/bullets)
|
||||
MIN_IMAGE_WIDTH = 50
|
||||
MIN_IMAGE_HEIGHT = 50
|
||||
|
||||
|
||||
@dataclass
|
||||
class TextBlock:
|
||||
"""A block of extracted text with page/section context."""
|
||||
|
||||
text: str
|
||||
page: int
|
||||
metadata: dict = field(default_factory=dict)
|
||||
|
||||
|
||||
@dataclass
|
||||
class ExtractedImage:
|
||||
"""An image extracted from a document."""
|
||||
|
||||
data: bytes
|
||||
ext: str
|
||||
width: int
|
||||
height: int
|
||||
source_page: int
|
||||
source_index: int
|
||||
|
||||
|
||||
@dataclass
|
||||
class ParseResult:
|
||||
"""Result of parsing a document: text blocks + images + metadata."""
|
||||
|
||||
text_blocks: list[TextBlock] = field(default_factory=list)
|
||||
images: list[ExtractedImage] = field(default_factory=list)
|
||||
metadata: dict = field(default_factory=dict)
|
||||
file_type: str = ""
|
||||
|
||||
|
||||
class DocumentParser:
|
||||
"""
|
||||
Universal document parser using PyMuPDF.
|
||||
|
||||
Extracts text and images from all supported document formats through
|
||||
a single unified interface.
|
||||
"""
|
||||
|
||||
def parse(self, file_path: str, file_type: str) -> ParseResult:
|
||||
"""
|
||||
Parse a document and extract text blocks and images.
|
||||
|
||||
:param file_path: Path to the document file.
|
||||
:param file_type: File extension (without dot), e.g. 'pdf', 'epub'.
|
||||
:returns: ParseResult with text blocks, images, and metadata.
|
||||
:raises ValueError: If the file type is not supported.
|
||||
"""
|
||||
file_type = file_type.lower().lstrip(".")
|
||||
|
||||
logger.info(
|
||||
"Parsing document file_type=%s path=%s",
|
||||
file_type,
|
||||
os.path.basename(file_path),
|
||||
)
|
||||
|
||||
if file_type in PLAINTEXT_EXTENSIONS:
|
||||
return self._parse_plaintext(file_path, file_type)
|
||||
|
||||
if file_type in IMAGE_EXTENSIONS:
|
||||
return self._parse_image_file(file_path, file_type)
|
||||
|
||||
if file_type in PYMUPDF_EXTENSIONS:
|
||||
return self._parse_with_pymupdf(file_path, file_type)
|
||||
|
||||
# HTML can be handled by PyMuPDF or direct read
|
||||
if file_type in ("html", "htm"):
|
||||
return self._parse_with_pymupdf(file_path, file_type)
|
||||
|
||||
raise ValueError(
|
||||
f"Unsupported file type '{file_type}'. "
|
||||
f"Supported: {sorted(PYMUPDF_EXTENSIONS | PLAINTEXT_EXTENSIONS | IMAGE_EXTENSIONS)}"
|
||||
)
|
||||
|
||||
def parse_bytes(self, data: bytes, file_type: str, filename: str = "") -> ParseResult:
|
||||
"""
|
||||
Parse document from bytes (e.g. from S3 download).
|
||||
|
||||
:param data: Raw file bytes.
|
||||
:param file_type: File extension (without dot).
|
||||
:param filename: Optional original filename for logging.
|
||||
:returns: ParseResult.
|
||||
"""
|
||||
suffix = f".{file_type.lower().lstrip('.')}"
|
||||
with tempfile.NamedTemporaryFile(suffix=suffix, delete=False) as tmp:
|
||||
tmp.write(data)
|
||||
tmp_path = tmp.name
|
||||
|
||||
try:
|
||||
logger.debug(
|
||||
"Parsing from bytes file_type=%s size=%d filename=%s",
|
||||
file_type,
|
||||
len(data),
|
||||
filename,
|
||||
)
|
||||
return self.parse(tmp_path, file_type)
|
||||
finally:
|
||||
os.unlink(tmp_path)
|
||||
|
||||
def _parse_with_pymupdf(self, file_path: str, file_type: str) -> ParseResult:
|
||||
"""
|
||||
Parse a document using PyMuPDF for text and image extraction.
|
||||
|
||||
:param file_path: Path to the document.
|
||||
:param file_type: Normalized file extension.
|
||||
:returns: ParseResult.
|
||||
"""
|
||||
with DOCUMENT_PARSE_DURATION.labels(file_type=file_type).time():
|
||||
try:
|
||||
doc = fitz.open(file_path)
|
||||
except Exception as exc:
|
||||
DOCUMENTS_PARSED_TOTAL.labels(file_type=file_type, status="error").inc()
|
||||
logger.error("Failed to open document file_type=%s: %s", file_type, exc)
|
||||
raise
|
||||
|
||||
text_blocks: list[TextBlock] = []
|
||||
images: list[ExtractedImage] = []
|
||||
image_global_index = 0
|
||||
|
||||
for page_num in range(len(doc)):
|
||||
page = doc[page_num]
|
||||
|
||||
# --- Text extraction ---
|
||||
try:
|
||||
text = page.get_text("text")
|
||||
if text and text.strip():
|
||||
cleaned = sanitize_text(text, log_changes=False)
|
||||
cleaned = remove_excessive_whitespace(cleaned)
|
||||
if cleaned.strip():
|
||||
text_blocks.append(
|
||||
TextBlock(text=cleaned, page=page_num)
|
||||
)
|
||||
logger.debug(
|
||||
"Extracted text page=%d chars=%d",
|
||||
page_num,
|
||||
len(cleaned),
|
||||
)
|
||||
except Exception as exc:
|
||||
logger.warning(
|
||||
"Text extraction failed page=%d: %s, continuing",
|
||||
page_num,
|
||||
exc,
|
||||
)
|
||||
|
||||
# --- Image extraction ---
|
||||
try:
|
||||
for img_info in page.get_images(full=True):
|
||||
xref = img_info[0]
|
||||
try:
|
||||
img_data = doc.extract_image(xref)
|
||||
if not img_data or not img_data.get("image"):
|
||||
continue
|
||||
|
||||
width = img_data.get("width", 0)
|
||||
height = img_data.get("height", 0)
|
||||
|
||||
# Skip tiny images (icons, bullets, etc.)
|
||||
if width < MIN_IMAGE_WIDTH or height < MIN_IMAGE_HEIGHT:
|
||||
logger.debug(
|
||||
"Skipping small image page=%d xref=%d size=%dx%d",
|
||||
page_num,
|
||||
xref,
|
||||
width,
|
||||
height,
|
||||
)
|
||||
continue
|
||||
|
||||
images.append(
|
||||
ExtractedImage(
|
||||
data=img_data["image"],
|
||||
ext=img_data.get("ext", "png"),
|
||||
width=width,
|
||||
height=height,
|
||||
source_page=page_num,
|
||||
source_index=image_global_index,
|
||||
)
|
||||
)
|
||||
image_global_index += 1
|
||||
|
||||
logger.debug(
|
||||
"Extracted image page=%d format=%s size=%dx%d bytes=%d",
|
||||
page_num,
|
||||
img_data.get("ext", "?"),
|
||||
width,
|
||||
height,
|
||||
len(img_data["image"]),
|
||||
)
|
||||
except Exception as exc:
|
||||
logger.warning(
|
||||
"Image extraction failed page=%d xref=%d: %s",
|
||||
page_num,
|
||||
xref,
|
||||
exc,
|
||||
)
|
||||
except Exception as exc:
|
||||
logger.warning(
|
||||
"Image listing failed page=%d: %s, continuing",
|
||||
page_num,
|
||||
exc,
|
||||
)
|
||||
|
||||
# Collect document metadata
|
||||
meta = doc.metadata or {}
|
||||
result_meta = {
|
||||
"page_count": len(doc),
|
||||
"title": meta.get("title", ""),
|
||||
"author": meta.get("author", ""),
|
||||
"subject": meta.get("subject", ""),
|
||||
"creator": meta.get("creator", ""),
|
||||
}
|
||||
|
||||
doc.close()
|
||||
|
||||
DOCUMENTS_PARSED_TOTAL.labels(file_type=file_type, status="success").inc()
|
||||
IMAGES_EXTRACTED_TOTAL.labels(file_type=file_type).inc(len(images))
|
||||
|
||||
logger.info(
|
||||
"Parsed document file_type=%s pages=%d text_blocks=%d images=%d",
|
||||
file_type,
|
||||
result_meta["page_count"],
|
||||
len(text_blocks),
|
||||
len(images),
|
||||
)
|
||||
|
||||
return ParseResult(
|
||||
text_blocks=text_blocks,
|
||||
images=images,
|
||||
metadata=result_meta,
|
||||
file_type=file_type,
|
||||
)
|
||||
|
||||
def _parse_plaintext(self, file_path: str, file_type: str) -> ParseResult:
|
||||
"""
|
||||
Parse a plain text file by direct read.
|
||||
|
||||
:param file_path: Path to the text file.
|
||||
:param file_type: Normalized file extension.
|
||||
:returns: ParseResult.
|
||||
"""
|
||||
with DOCUMENT_PARSE_DURATION.labels(file_type=file_type).time():
|
||||
try:
|
||||
with open(file_path, "r", encoding="utf-8", errors="replace") as f:
|
||||
content = f.read()
|
||||
except Exception as exc:
|
||||
DOCUMENTS_PARSED_TOTAL.labels(file_type=file_type, status="error").inc()
|
||||
logger.error("Failed to read text file file_type=%s: %s", file_type, exc)
|
||||
raise
|
||||
|
||||
cleaned = sanitize_text(content, log_changes=True)
|
||||
cleaned = remove_excessive_whitespace(cleaned)
|
||||
|
||||
text_blocks = []
|
||||
if cleaned.strip():
|
||||
text_blocks.append(TextBlock(text=cleaned, page=0))
|
||||
|
||||
DOCUMENTS_PARSED_TOTAL.labels(file_type=file_type, status="success").inc()
|
||||
|
||||
logger.info(
|
||||
"Parsed plaintext file_type=%s chars=%d",
|
||||
file_type,
|
||||
len(cleaned),
|
||||
)
|
||||
|
||||
return ParseResult(
|
||||
text_blocks=text_blocks,
|
||||
images=[],
|
||||
metadata={"page_count": 1},
|
||||
file_type=file_type,
|
||||
)
|
||||
|
||||
def _parse_image_file(self, file_path: str, file_type: str) -> ParseResult:
|
||||
"""
|
||||
Handle a standalone image file — store as a single ExtractedImage.
|
||||
|
||||
:param file_path: Path to the image file.
|
||||
:param file_type: Normalized file extension.
|
||||
:returns: ParseResult with one image and no text.
|
||||
"""
|
||||
with DOCUMENT_PARSE_DURATION.labels(file_type=file_type).time():
|
||||
try:
|
||||
from PIL import Image as PILImage
|
||||
|
||||
with open(file_path, "rb") as f:
|
||||
data = f.read()
|
||||
|
||||
img = PILImage.open(file_path)
|
||||
width, height = img.size
|
||||
img.close()
|
||||
except Exception as exc:
|
||||
DOCUMENTS_PARSED_TOTAL.labels(file_type=file_type, status="error").inc()
|
||||
logger.error("Failed to read image file_type=%s: %s", file_type, exc)
|
||||
raise
|
||||
|
||||
DOCUMENTS_PARSED_TOTAL.labels(file_type=file_type, status="success").inc()
|
||||
IMAGES_EXTRACTED_TOTAL.labels(file_type=file_type).inc(1)
|
||||
|
||||
logger.info(
|
||||
"Parsed image file file_type=%s size=%dx%d bytes=%d",
|
||||
file_type,
|
||||
width,
|
||||
height,
|
||||
len(data),
|
||||
)
|
||||
|
||||
return ParseResult(
|
||||
text_blocks=[],
|
||||
images=[
|
||||
ExtractedImage(
|
||||
data=data,
|
||||
ext=file_type,
|
||||
width=width,
|
||||
height=height,
|
||||
source_page=0,
|
||||
source_index=0,
|
||||
)
|
||||
],
|
||||
metadata={"page_count": 0, "width": width, "height": height},
|
||||
file_type=file_type,
|
||||
)
|
||||
581
mnemosyne/library/services/pipeline.py
Normal file
581
mnemosyne/library/services/pipeline.py
Normal file
@@ -0,0 +1,581 @@
|
||||
"""
|
||||
Embedding pipeline orchestrator.
|
||||
|
||||
Coordinates the full ingestion flow:
|
||||
parse → chunk → embed → store → graph construction.
|
||||
"""
|
||||
|
||||
import hashlib
|
||||
import logging
|
||||
import time
|
||||
from typing import Optional
|
||||
|
||||
from django.conf import settings
|
||||
from django.core.cache import cache
|
||||
from django.core.files.base import ContentFile
|
||||
from django.core.files.storage import default_storage
|
||||
|
||||
from library.metrics import (
|
||||
PIPELINE_DURATION,
|
||||
PIPELINE_ITEMS_IN_PROGRESS,
|
||||
PIPELINE_ITEMS_TOTAL,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# S3 key patterns
|
||||
ORIGINAL_S3_KEY = "items/{item_uid}/original.{ext}"
|
||||
CHUNK_S3_KEY = "chunks/{item_uid}/chunk_{index}.txt"
|
||||
IMAGE_S3_KEY = "images/{item_uid}/{index}.{ext}"
|
||||
|
||||
# Batch sizes
|
||||
EMBEDDING_BATCH_SIZE = getattr(settings, "EMBEDDING_BATCH_SIZE", 8)
|
||||
|
||||
|
||||
class EmbeddingPipeline:
|
||||
"""
|
||||
Orchestrates the complete embedding pipeline for a single Item.
|
||||
|
||||
Stages:
|
||||
1. Parse document (text + images)
|
||||
2. Chunk text (content-type-aware)
|
||||
3. Store chunks in S3 + Neo4j
|
||||
4. Embed text chunks
|
||||
5. Store images in S3 + Neo4j
|
||||
6. Embed images (multimodal, if available)
|
||||
7. Extract concepts (if system chat model available)
|
||||
"""
|
||||
|
||||
def __init__(self, user=None):
|
||||
"""
|
||||
:param user: Optional Django user for usage tracking.
|
||||
"""
|
||||
self.user = user
|
||||
|
||||
def process_item(
|
||||
self,
|
||||
item_uid: str,
|
||||
progress_callback: Optional[callable] = None,
|
||||
) -> dict:
|
||||
"""
|
||||
Run the full embedding pipeline for an Item.
|
||||
|
||||
:param item_uid: UID of the Item node to process.
|
||||
:param progress_callback: Optional callback(percent, message).
|
||||
:returns: Dict with processing results.
|
||||
:raises ValueError: If item not found or no embedding model configured.
|
||||
"""
|
||||
from library.models import Item
|
||||
|
||||
start_time = time.time()
|
||||
PIPELINE_ITEMS_IN_PROGRESS.inc()
|
||||
|
||||
try:
|
||||
item = Item.nodes.get(uid=item_uid)
|
||||
except Exception:
|
||||
PIPELINE_ITEMS_TOTAL.labels(status="failed").inc()
|
||||
PIPELINE_ITEMS_IN_PROGRESS.dec()
|
||||
raise ValueError(f"Item not found: {item_uid}")
|
||||
|
||||
logger.info(
|
||||
"Pipeline starting item_uid=%s title='%s' file_type=%s",
|
||||
item_uid,
|
||||
item.title,
|
||||
item.file_type,
|
||||
)
|
||||
|
||||
# Mark as processing
|
||||
item.embedding_status = "processing"
|
||||
item.error_message = ""
|
||||
item.save()
|
||||
|
||||
try:
|
||||
result = self._run_pipeline(item, progress_callback)
|
||||
|
||||
# Mark as completed
|
||||
item.embedding_status = "completed"
|
||||
item.chunk_count = result.get("chunks_created", 0)
|
||||
item.image_count = result.get("images_stored", 0)
|
||||
item.embedding_model_name = result.get("model_name", "")
|
||||
item.error_message = ""
|
||||
item.save()
|
||||
|
||||
elapsed = time.time() - start_time
|
||||
PIPELINE_ITEMS_TOTAL.labels(status="completed").inc()
|
||||
PIPELINE_DURATION.observe(elapsed)
|
||||
|
||||
logger.info(
|
||||
"Pipeline completed item_uid=%s chunks=%d images=%d concepts=%d elapsed=%.2fs",
|
||||
item_uid,
|
||||
result.get("chunks_created", 0),
|
||||
result.get("images_stored", 0),
|
||||
result.get("concepts_extracted", 0),
|
||||
elapsed,
|
||||
)
|
||||
|
||||
if progress_callback:
|
||||
progress_callback(100, "Completed")
|
||||
|
||||
return result
|
||||
|
||||
except Exception as exc:
|
||||
item.embedding_status = "failed"
|
||||
item.error_message = str(exc)[:500]
|
||||
item.save()
|
||||
|
||||
PIPELINE_ITEMS_TOTAL.labels(status="failed").inc()
|
||||
logger.error(
|
||||
"Pipeline failed item_uid=%s: %s",
|
||||
item_uid,
|
||||
exc,
|
||||
exc_info=True,
|
||||
)
|
||||
raise
|
||||
|
||||
finally:
|
||||
PIPELINE_ITEMS_IN_PROGRESS.dec()
|
||||
|
||||
def _run_pipeline(self, item, progress_callback) -> dict:
|
||||
"""
|
||||
Execute pipeline stages sequentially.
|
||||
|
||||
:param item: Item node instance.
|
||||
:param progress_callback: Optional progress callback.
|
||||
:returns: Results dict.
|
||||
"""
|
||||
from llm_manager.models import LLMModel
|
||||
|
||||
from .chunker import ContentTypeChunker
|
||||
from .concepts import ConceptExtractor
|
||||
from .embedding_client import EmbeddingClient
|
||||
from .parsers import DocumentParser
|
||||
|
||||
result = {
|
||||
"chunks_created": 0,
|
||||
"chunks_embedded": 0,
|
||||
"images_stored": 0,
|
||||
"images_embedded": 0,
|
||||
"concepts_extracted": 0,
|
||||
"model_name": "",
|
||||
}
|
||||
|
||||
# --- Resolve library context ---
|
||||
library = self._get_item_library(item)
|
||||
chunking_config = library.chunking_config if library else {}
|
||||
embedding_instruction = library.embedding_instruction if library else ""
|
||||
library_type = library.library_type if library else ""
|
||||
|
||||
# --- Get system embedding model ---
|
||||
embedding_model = LLMModel.get_system_embedding_model()
|
||||
if not embedding_model:
|
||||
raise ValueError(
|
||||
"No system embedding model configured. "
|
||||
"Set one via Django admin > LLM Models > Set as System Embedding Model."
|
||||
)
|
||||
|
||||
result["model_name"] = embedding_model.name
|
||||
embed_client = EmbeddingClient(embedding_model, user=self.user)
|
||||
|
||||
# --- Check dimension compatibility ---
|
||||
if embedding_model.vector_dimensions:
|
||||
self._check_dimension_compatibility(embedding_model.vector_dimensions)
|
||||
|
||||
if progress_callback:
|
||||
progress_callback(5, "Parsing document")
|
||||
|
||||
# --- Stage 1: Parse ---
|
||||
parser = DocumentParser()
|
||||
file_data = self._read_item_from_s3(item)
|
||||
if not file_data:
|
||||
logger.warning("No file data for item_uid=%s, skipping", item.uid)
|
||||
return result
|
||||
|
||||
parse_result = parser.parse_bytes(
|
||||
file_data,
|
||||
item.file_type,
|
||||
filename=item.title,
|
||||
)
|
||||
|
||||
if progress_callback:
|
||||
progress_callback(20, "Chunking text")
|
||||
|
||||
# --- Stage 2: Chunk ---
|
||||
chunker = ContentTypeChunker()
|
||||
chunk_result = chunker.chunk(parse_result, chunking_config, library_type)
|
||||
|
||||
if progress_callback:
|
||||
progress_callback(30, "Storing chunks")
|
||||
|
||||
# --- Stage 3: Store chunks in S3 + Neo4j ---
|
||||
chunk_nodes = self._store_chunks(item, chunk_result)
|
||||
result["chunks_created"] = len(chunk_nodes)
|
||||
|
||||
if progress_callback:
|
||||
progress_callback(40, "Embedding text chunks")
|
||||
|
||||
# --- Stage 4: Embed text chunks ---
|
||||
if chunk_result.chunks:
|
||||
self._embed_chunks(
|
||||
item,
|
||||
chunk_nodes,
|
||||
chunk_result.chunks,
|
||||
embed_client,
|
||||
embedding_instruction,
|
||||
progress_callback,
|
||||
)
|
||||
result["chunks_embedded"] = len(chunk_nodes)
|
||||
|
||||
if progress_callback:
|
||||
progress_callback(70, "Storing images")
|
||||
|
||||
# --- Stage 5: Store images ---
|
||||
image_nodes = self._store_images(item, parse_result.images)
|
||||
result["images_stored"] = len(image_nodes)
|
||||
|
||||
# Associate images with nearby chunks
|
||||
self._associate_images_with_chunks(
|
||||
chunk_nodes, image_nodes, chunk_result, parse_result
|
||||
)
|
||||
|
||||
if progress_callback:
|
||||
progress_callback(80, "Embedding images")
|
||||
|
||||
# --- Stage 6: Embed images (multimodal) ---
|
||||
if image_nodes and embedding_model.supports_multimodal:
|
||||
embedded_count = self._embed_images(image_nodes, embed_client)
|
||||
result["images_embedded"] = embedded_count
|
||||
|
||||
if progress_callback:
|
||||
progress_callback(90, "Extracting concepts")
|
||||
|
||||
# --- Stage 7: Concept extraction ---
|
||||
chat_model = LLMModel.get_system_chat_model()
|
||||
if chat_model and chunk_result.chunks:
|
||||
extractor = ConceptExtractor(chat_model, user=self.user)
|
||||
concepts_count = extractor.extract_for_item(
|
||||
item, chunk_nodes, chunk_result.chunks
|
||||
)
|
||||
result["concepts_extracted"] = concepts_count
|
||||
|
||||
# Update content hash to prevent redundant re-processing
|
||||
if file_data:
|
||||
item.content_hash = hashlib.sha256(file_data).hexdigest()
|
||||
item.save()
|
||||
|
||||
return result
|
||||
|
||||
def _get_item_library(self, item):
|
||||
"""
|
||||
Walk the graph to find the Library containing this Item.
|
||||
|
||||
:param item: Item node.
|
||||
:returns: Library node, or None.
|
||||
"""
|
||||
from library.models import Collection
|
||||
|
||||
try:
|
||||
# Item <- Collection <- Library
|
||||
from neomodel import db
|
||||
|
||||
results, _ = db.cypher_query(
|
||||
"MATCH (l:Library)-[:CONTAINS]->(c:Collection)-[:CONTAINS]->(i:Item {uid: $uid}) "
|
||||
"RETURN l.uid, l.library_type, l.chunking_config, l.embedding_instruction",
|
||||
{"uid": item.uid},
|
||||
)
|
||||
if results:
|
||||
from library.models import Library
|
||||
|
||||
return Library.nodes.get(uid=results[0][0])
|
||||
except Exception as exc:
|
||||
logger.warning("Could not resolve library for item_uid=%s: %s", item.uid, exc)
|
||||
|
||||
return None
|
||||
|
||||
def _read_item_from_s3(self, item) -> Optional[bytes]:
|
||||
"""
|
||||
Read the original file from S3 storage.
|
||||
|
||||
:param item: Item node with s3_key.
|
||||
:returns: File bytes, or None.
|
||||
"""
|
||||
if not item.s3_key:
|
||||
logger.warning("Item has no s3_key item_uid=%s", item.uid)
|
||||
return None
|
||||
|
||||
try:
|
||||
with default_storage.open(item.s3_key, "rb") as f:
|
||||
data = f.read()
|
||||
logger.debug(
|
||||
"Read item from S3 key=%s size=%d", item.s3_key, len(data)
|
||||
)
|
||||
return data
|
||||
except Exception as exc:
|
||||
logger.error(
|
||||
"Failed to read from S3 key=%s: %s", item.s3_key, exc
|
||||
)
|
||||
raise
|
||||
|
||||
def _store_chunks(self, item, chunk_result) -> list:
|
||||
"""
|
||||
Store chunk text in S3 and create Chunk nodes in Neo4j.
|
||||
|
||||
:param item: Item node.
|
||||
:param chunk_result: ChunkResult from chunker.
|
||||
:returns: List of Chunk node instances.
|
||||
"""
|
||||
from library.models import Chunk
|
||||
|
||||
# Delete existing chunks for this item
|
||||
for old_chunk in item.chunks.all():
|
||||
# Clean up S3
|
||||
try:
|
||||
default_storage.delete(old_chunk.chunk_s3_key)
|
||||
except Exception:
|
||||
pass
|
||||
old_chunk.delete()
|
||||
|
||||
chunk_nodes = []
|
||||
for idx, chunk_text in enumerate(chunk_result.chunks):
|
||||
s3_key = CHUNK_S3_KEY.format(item_uid=item.uid, index=idx)
|
||||
|
||||
# Store chunk text in S3
|
||||
try:
|
||||
default_storage.save(s3_key, ContentFile(chunk_text.encode("utf-8")))
|
||||
except Exception as exc:
|
||||
logger.error("Failed to store chunk %d to S3: %s", idx, exc)
|
||||
raise
|
||||
|
||||
# Create Chunk node
|
||||
chunk_node = Chunk(
|
||||
chunk_index=idx,
|
||||
chunk_s3_key=s3_key,
|
||||
chunk_size=len(chunk_text),
|
||||
text_preview=chunk_text[:500],
|
||||
)
|
||||
chunk_node.save()
|
||||
item.chunks.connect(chunk_node)
|
||||
chunk_nodes.append(chunk_node)
|
||||
|
||||
logger.info(
|
||||
"Stored %d chunks for item_uid=%s", len(chunk_nodes), item.uid
|
||||
)
|
||||
return chunk_nodes
|
||||
|
||||
def _embed_chunks(
|
||||
self,
|
||||
item,
|
||||
chunk_nodes: list,
|
||||
chunk_texts: list[str],
|
||||
embed_client,
|
||||
embedding_instruction: str,
|
||||
progress_callback: Optional[callable],
|
||||
):
|
||||
"""
|
||||
Generate embeddings for chunks and update Chunk nodes.
|
||||
|
||||
:param item: Item node.
|
||||
:param chunk_nodes: List of Chunk nodes.
|
||||
:param chunk_texts: List of chunk text strings.
|
||||
:param embed_client: EmbeddingClient instance.
|
||||
:param embedding_instruction: Instruction prefix for embedding.
|
||||
:param progress_callback: Optional progress callback.
|
||||
"""
|
||||
# Prepend embedding instruction if configured
|
||||
if embedding_instruction:
|
||||
texts_to_embed = [
|
||||
f"{embedding_instruction}\n\n{text}" for text in chunk_texts
|
||||
]
|
||||
else:
|
||||
texts_to_embed = chunk_texts
|
||||
|
||||
batch_size = EMBEDDING_BATCH_SIZE
|
||||
total_batches = (len(texts_to_embed) + batch_size - 1) // batch_size
|
||||
|
||||
for batch_idx in range(0, len(texts_to_embed), batch_size):
|
||||
batch_texts = texts_to_embed[batch_idx : batch_idx + batch_size]
|
||||
batch_nodes = chunk_nodes[batch_idx : batch_idx + batch_size]
|
||||
batch_num = batch_idx // batch_size + 1
|
||||
|
||||
logger.debug(
|
||||
"Embedding batch %d/%d size=%d item_uid=%s",
|
||||
batch_num,
|
||||
total_batches,
|
||||
len(batch_texts),
|
||||
item.uid,
|
||||
)
|
||||
|
||||
embeddings = embed_client.embed_texts(batch_texts)
|
||||
|
||||
for node, embedding in zip(batch_nodes, embeddings):
|
||||
node.embedding = embedding
|
||||
node.save()
|
||||
|
||||
if progress_callback:
|
||||
pct = 40 + (30 * (batch_idx + len(batch_texts)) / len(texts_to_embed))
|
||||
progress_callback(
|
||||
int(pct),
|
||||
f"Embedded {batch_idx + len(batch_texts)}/{len(texts_to_embed)} chunks",
|
||||
)
|
||||
|
||||
logger.info(
|
||||
"Embedded %d chunks for item_uid=%s", len(chunk_nodes), item.uid
|
||||
)
|
||||
|
||||
def _store_images(self, item, extracted_images) -> list:
|
||||
"""
|
||||
Store extracted images in S3 and create Image nodes in Neo4j.
|
||||
|
||||
:param item: Item node.
|
||||
:param extracted_images: List of ExtractedImage from parser.
|
||||
:returns: List of Image node instances.
|
||||
"""
|
||||
from library.models import Image
|
||||
|
||||
# Delete existing images for this item
|
||||
for old_image in item.images.all():
|
||||
try:
|
||||
default_storage.delete(old_image.s3_key)
|
||||
except Exception:
|
||||
pass
|
||||
old_image.delete()
|
||||
|
||||
image_nodes = []
|
||||
for img in extracted_images:
|
||||
s3_key = IMAGE_S3_KEY.format(
|
||||
item_uid=item.uid,
|
||||
index=img.source_index,
|
||||
ext=img.ext,
|
||||
)
|
||||
|
||||
try:
|
||||
default_storage.save(s3_key, ContentFile(img.data))
|
||||
except Exception as exc:
|
||||
logger.warning(
|
||||
"Failed to store image %d to S3: %s", img.source_index, exc
|
||||
)
|
||||
continue
|
||||
|
||||
image_node = Image(
|
||||
s3_key=s3_key,
|
||||
image_type="diagram", # Default; could be refined by content analysis
|
||||
metadata={
|
||||
"width": img.width,
|
||||
"height": img.height,
|
||||
"source_page": img.source_page,
|
||||
"content_type": f"image/{img.ext}",
|
||||
},
|
||||
)
|
||||
image_node.save()
|
||||
item.images.connect(image_node)
|
||||
image_nodes.append(image_node)
|
||||
|
||||
if image_nodes:
|
||||
logger.info(
|
||||
"Stored %d images for item_uid=%s", len(image_nodes), item.uid
|
||||
)
|
||||
|
||||
return image_nodes
|
||||
|
||||
def _associate_images_with_chunks(
|
||||
self, chunk_nodes, image_nodes, chunk_result, parse_result
|
||||
):
|
||||
"""
|
||||
Create HAS_NEARBY_IMAGE relationships between chunks and images.
|
||||
|
||||
Associates images with chunks from the same page/section.
|
||||
|
||||
:param chunk_nodes: List of Chunk nodes.
|
||||
:param image_nodes: List of Image nodes.
|
||||
:param chunk_result: ChunkResult with page mapping.
|
||||
:param parse_result: ParseResult with image source pages.
|
||||
"""
|
||||
if not chunk_nodes or not image_nodes:
|
||||
return
|
||||
|
||||
# Build page -> images mapping
|
||||
page_images: dict[int, list] = {}
|
||||
for img_node, ext_img in zip(image_nodes, parse_result.images):
|
||||
page_images.setdefault(ext_img.source_page, []).append(img_node)
|
||||
|
||||
# Connect chunks to images on the same page
|
||||
connected = 0
|
||||
for chunk_idx, chunk_node in enumerate(chunk_nodes):
|
||||
page = chunk_result.chunk_page_map.get(chunk_idx, -1)
|
||||
nearby = page_images.get(page, [])
|
||||
for img_node in nearby:
|
||||
chunk_node.nearby_images.connect(
|
||||
img_node, {"proximity": "same_page"}
|
||||
)
|
||||
connected += 1
|
||||
|
||||
if connected:
|
||||
logger.debug(
|
||||
"Created %d chunk-image associations", connected
|
||||
)
|
||||
|
||||
def _embed_images(self, image_nodes: list, embed_client) -> int:
|
||||
"""
|
||||
Generate multimodal embeddings for Image nodes.
|
||||
|
||||
:param image_nodes: List of Image nodes.
|
||||
:param embed_client: EmbeddingClient with multimodal support.
|
||||
:returns: Number of images successfully embedded.
|
||||
"""
|
||||
from library.models import ImageEmbedding
|
||||
|
||||
embedded_count = 0
|
||||
for img_node in image_nodes:
|
||||
try:
|
||||
img_data = default_storage.open(img_node.s3_key, "rb").read()
|
||||
ext = img_node.s3_key.rsplit(".", 1)[-1] if "." in img_node.s3_key else "png"
|
||||
|
||||
embedding = embed_client.embed_image(img_data, ext)
|
||||
if embedding:
|
||||
emb_node = ImageEmbedding(embedding=embedding)
|
||||
emb_node.save()
|
||||
img_node.embeddings.connect(emb_node)
|
||||
embedded_count += 1
|
||||
except Exception as exc:
|
||||
logger.warning(
|
||||
"Image embedding failed s3_key=%s: %s",
|
||||
img_node.s3_key,
|
||||
exc,
|
||||
)
|
||||
|
||||
if embedded_count:
|
||||
logger.info("Embedded %d images", embedded_count)
|
||||
|
||||
return embedded_count
|
||||
|
||||
def _check_dimension_compatibility(self, model_dimensions: int):
|
||||
"""
|
||||
Check if the model's vector dimensions match the Neo4j index.
|
||||
|
||||
:param model_dimensions: Expected embedding dimensions.
|
||||
"""
|
||||
# Log a warning — actual enforcement is in setup_neo4j_indexes
|
||||
logger.debug(
|
||||
"System embedding model dimensions=%d", model_dimensions
|
||||
)
|
||||
|
||||
def reprocess_item(self, item_uid: str, progress_callback=None) -> dict:
|
||||
"""
|
||||
Re-embed an item: delete existing chunks/images, then re-process.
|
||||
|
||||
:param item_uid: UID of the Item to re-embed.
|
||||
:param progress_callback: Optional progress callback.
|
||||
:returns: Processing results dict.
|
||||
"""
|
||||
from library.models import Item
|
||||
|
||||
try:
|
||||
item = Item.nodes.get(uid=item_uid)
|
||||
except Exception:
|
||||
raise ValueError(f"Item not found: {item_uid}")
|
||||
|
||||
# Clear content hash to force re-processing
|
||||
item.content_hash = ""
|
||||
item.save()
|
||||
|
||||
logger.info("Re-processing item_uid=%s title='%s'", item_uid, item.title)
|
||||
return self.process_item(item_uid, progress_callback)
|
||||
165
mnemosyne/library/services/text_utils.py
Normal file
165
mnemosyne/library/services/text_utils.py
Normal file
@@ -0,0 +1,165 @@
|
||||
"""
|
||||
Text sanitization utilities for the embedding pipeline.
|
||||
|
||||
Ported from Spelunker's text_utils.py — ensures text can be safely
|
||||
processed by embedding APIs and LLMs.
|
||||
"""
|
||||
|
||||
import logging
|
||||
import re
|
||||
import unicodedata
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
# Common PDF ligatures
|
||||
_LIGATURE_MAP = {
|
||||
"\ufb01": "fi",
|
||||
"\ufb02": "fl",
|
||||
"\ufb00": "ff",
|
||||
"\ufb03": "ffi",
|
||||
"\ufb04": "ffl",
|
||||
"\ufb05": "ft",
|
||||
"\ufb06": "st",
|
||||
}
|
||||
|
||||
# Common special characters from PDF extraction
|
||||
_SPECIAL_CHAR_MAP = {
|
||||
"\u2018": "'", # left single quotation
|
||||
"\u2019": "'", # right single quotation
|
||||
"\u201c": '"', # left double quotation
|
||||
"\u201d": '"', # right double quotation
|
||||
"\u2013": "-", # en dash
|
||||
"\u2014": "-", # em dash
|
||||
"\u2026": "...", # horizontal ellipsis
|
||||
"\u00a0": " ", # non-breaking space
|
||||
}
|
||||
|
||||
# Zero-width characters
|
||||
_ZERO_WIDTH_CHARS = [
|
||||
"\u200b", # zero-width space
|
||||
"\u200c", # zero-width non-joiner
|
||||
"\u200d", # zero-width joiner
|
||||
"\ufeff", # zero-width no-break space (BOM)
|
||||
]
|
||||
|
||||
# Control characters pattern (exclude newline, tab, carriage return)
|
||||
_CONTROL_CHAR_RE = re.compile(r"[\x00-\x08\x0b-\x0c\x0e-\x1f\x7f-\x9f]")
|
||||
|
||||
|
||||
def sanitize_text(text: str, log_changes: bool = True) -> str:
|
||||
"""
|
||||
Sanitize text for embedding APIs by removing problematic characters.
|
||||
|
||||
Addresses common issues that cause "invalid tokens" errors:
|
||||
null bytes, control characters, zero-width characters, invalid UTF-8,
|
||||
and non-normalized Unicode.
|
||||
|
||||
:param text: Text to sanitize.
|
||||
:param log_changes: Whether to log sanitization actions.
|
||||
:returns: Sanitized text safe for tokenization.
|
||||
"""
|
||||
if not text:
|
||||
return text
|
||||
|
||||
original_length = len(text)
|
||||
changes: list[str] = []
|
||||
|
||||
# 1. Remove null bytes
|
||||
if "\x00" in text:
|
||||
text = text.replace("\x00", "")
|
||||
changes.append("removed null bytes")
|
||||
|
||||
# 2. Remove control characters
|
||||
if _CONTROL_CHAR_RE.search(text):
|
||||
text = _CONTROL_CHAR_RE.sub("", text)
|
||||
changes.append("removed control characters")
|
||||
|
||||
# 3. Remove zero-width characters
|
||||
for char in _ZERO_WIDTH_CHARS:
|
||||
if char in text:
|
||||
text = text.replace(char, "")
|
||||
changes.append("removed zero-width characters")
|
||||
break
|
||||
|
||||
# 4. Normalize Unicode to NFC form
|
||||
normalized = unicodedata.normalize("NFC", text)
|
||||
if normalized != text:
|
||||
text = normalized
|
||||
changes.append("normalized Unicode to NFC")
|
||||
|
||||
# 5. Replace invalid UTF-8 sequences
|
||||
try:
|
||||
text = text.encode("utf-8", errors="replace").decode("utf-8")
|
||||
if "\ufffd" in text:
|
||||
changes.append("replaced invalid UTF-8 sequences")
|
||||
except Exception as exc:
|
||||
logger.warning("Error during UTF-8 validation: %s", exc)
|
||||
|
||||
# 6. Clean PDF artifacts (ligatures, special chars)
|
||||
cleaned = clean_pdf_artifacts(text)
|
||||
if cleaned != text:
|
||||
text = cleaned
|
||||
changes.append("cleaned PDF artifacts")
|
||||
|
||||
if log_changes and changes:
|
||||
chars_removed = original_length - len(text)
|
||||
logger.info(
|
||||
"Text sanitization: %s original_length=%d final_length=%d chars_removed=%d",
|
||||
", ".join(changes),
|
||||
original_length,
|
||||
len(text),
|
||||
chars_removed,
|
||||
)
|
||||
|
||||
return text
|
||||
|
||||
|
||||
def clean_pdf_artifacts(text: str) -> str:
|
||||
"""
|
||||
Clean common PDF extraction artifacts.
|
||||
|
||||
Replaces ligatures and special characters with standard equivalents.
|
||||
|
||||
:param text: Text to clean.
|
||||
:returns: Cleaned text.
|
||||
"""
|
||||
for ligature, replacement in _LIGATURE_MAP.items():
|
||||
text = text.replace(ligature, replacement)
|
||||
|
||||
for special, replacement in _SPECIAL_CHAR_MAP.items():
|
||||
text = text.replace(special, replacement)
|
||||
|
||||
return text
|
||||
|
||||
|
||||
def remove_excessive_whitespace(text: str) -> str:
|
||||
"""
|
||||
Remove excessive whitespace while preserving paragraph structure.
|
||||
|
||||
:param text: Text to clean.
|
||||
:returns: Text with normalized whitespace.
|
||||
"""
|
||||
text = re.sub(r" +", " ", text)
|
||||
text = re.sub(r"\n\n+", "\n\n", text)
|
||||
text = "\n".join(line.strip() for line in text.split("\n"))
|
||||
return text.strip()
|
||||
|
||||
|
||||
def truncate_text(text: str, max_chars: int, suffix: str = "...") -> str:
|
||||
"""
|
||||
Truncate text to a maximum length preserving word boundaries.
|
||||
|
||||
:param text: Text to truncate.
|
||||
:param max_chars: Maximum number of characters.
|
||||
:param suffix: Suffix to add if truncated.
|
||||
:returns: Truncated text.
|
||||
"""
|
||||
if len(text) <= max_chars:
|
||||
return text
|
||||
|
||||
truncate_at = text.rfind(" ", 0, max_chars - len(suffix))
|
||||
if truncate_at == -1:
|
||||
truncate_at = max_chars - len(suffix)
|
||||
|
||||
return text[:truncate_at] + suffix
|
||||
282
mnemosyne/library/tasks.py
Normal file
282
mnemosyne/library/tasks.py
Normal file
@@ -0,0 +1,282 @@
|
||||
"""
|
||||
Celery tasks for the embedding pipeline.
|
||||
|
||||
All tasks pass UIDs (not model instances) per Red Panda Standards.
|
||||
Tasks are idempotent, include retry logic, and track progress
|
||||
via Memcached: library:task:{task_id}:progress.
|
||||
"""
|
||||
|
||||
import logging
|
||||
|
||||
from celery import shared_task
|
||||
from django.core.cache import cache
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Cache key pattern for task progress
|
||||
PROGRESS_KEY = "library:task:{task_id}:progress"
|
||||
|
||||
|
||||
def _update_progress(task, percent: int, message: str):
|
||||
"""
|
||||
Update task progress in Memcached and Celery state.
|
||||
|
||||
:param task: Celery task instance (self).
|
||||
:param percent: Progress percentage (0-100).
|
||||
:param message: Human-readable status message.
|
||||
"""
|
||||
try:
|
||||
task.update_state(state="PROGRESS", meta={"percent": percent, "message": message})
|
||||
cache.set(
|
||||
PROGRESS_KEY.format(task_id=task.request.id),
|
||||
{"percent": percent, "message": message},
|
||||
timeout=3600,
|
||||
)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
@shared_task(
|
||||
name="library.tasks.embed_item",
|
||||
bind=True,
|
||||
queue="embedding",
|
||||
max_retries=3,
|
||||
default_retry_delay=60,
|
||||
acks_late=True,
|
||||
)
|
||||
def embed_item(self, item_uid: str, user_id: int = None):
|
||||
"""
|
||||
Run the full embedding pipeline for a single Item.
|
||||
|
||||
:param item_uid: UID of the Item node to process.
|
||||
:param user_id: Optional user ID for usage tracking.
|
||||
:returns: Dict with processing results.
|
||||
"""
|
||||
logger.info("Task embed_item starting item_uid=%s task_id=%s", item_uid, self.request.id)
|
||||
|
||||
try:
|
||||
from library.services.pipeline import EmbeddingPipeline
|
||||
|
||||
user = _resolve_user(user_id)
|
||||
pipeline = EmbeddingPipeline(user=user)
|
||||
|
||||
def progress_cb(percent, message):
|
||||
_update_progress(self, percent, message)
|
||||
|
||||
result = pipeline.process_item(item_uid, progress_callback=progress_cb)
|
||||
|
||||
logger.info(
|
||||
"Task embed_item completed item_uid=%s chunks=%d images=%d",
|
||||
item_uid,
|
||||
result.get("chunks_created", 0),
|
||||
result.get("images_stored", 0),
|
||||
)
|
||||
return {"success": True, "item_uid": item_uid, **result}
|
||||
|
||||
except Exception as exc:
|
||||
logger.error(
|
||||
"Task embed_item failed item_uid=%s: %s",
|
||||
item_uid,
|
||||
exc,
|
||||
exc_info=True,
|
||||
)
|
||||
# Retry on transient errors
|
||||
if self.request.retries < self.max_retries:
|
||||
raise self.retry(exc=exc)
|
||||
return {"success": False, "item_uid": item_uid, "error": str(exc)}
|
||||
|
||||
|
||||
@shared_task(
|
||||
name="library.tasks.reembed_item",
|
||||
bind=True,
|
||||
queue="embedding",
|
||||
max_retries=3,
|
||||
default_retry_delay=60,
|
||||
acks_late=True,
|
||||
)
|
||||
def reembed_item(self, item_uid: str, user_id: int = None):
|
||||
"""
|
||||
Delete existing embeddings and re-process an Item.
|
||||
|
||||
:param item_uid: UID of the Item node to re-embed.
|
||||
:param user_id: Optional user ID for usage tracking.
|
||||
:returns: Dict with processing results.
|
||||
"""
|
||||
logger.info("Task reembed_item starting item_uid=%s", item_uid)
|
||||
|
||||
try:
|
||||
from library.services.pipeline import EmbeddingPipeline
|
||||
|
||||
user = _resolve_user(user_id)
|
||||
pipeline = EmbeddingPipeline(user=user)
|
||||
|
||||
def progress_cb(percent, message):
|
||||
_update_progress(self, percent, message)
|
||||
|
||||
result = pipeline.reprocess_item(item_uid, progress_callback=progress_cb)
|
||||
|
||||
logger.info("Task reembed_item completed item_uid=%s", item_uid)
|
||||
return {"success": True, "item_uid": item_uid, **result}
|
||||
|
||||
except Exception as exc:
|
||||
logger.error("Task reembed_item failed item_uid=%s: %s", item_uid, exc, exc_info=True)
|
||||
if self.request.retries < self.max_retries:
|
||||
raise self.retry(exc=exc)
|
||||
return {"success": False, "item_uid": item_uid, "error": str(exc)}
|
||||
|
||||
|
||||
@shared_task(
|
||||
name="library.tasks.embed_collection",
|
||||
bind=True,
|
||||
queue="batch",
|
||||
acks_late=True,
|
||||
)
|
||||
def embed_collection(self, collection_uid: str, user_id: int = None):
|
||||
"""
|
||||
Embed all items in a collection.
|
||||
|
||||
:param collection_uid: UID of the Collection node.
|
||||
:param user_id: Optional user ID for usage tracking.
|
||||
:returns: Dict with summary results.
|
||||
"""
|
||||
logger.info("Task embed_collection starting collection_uid=%s", collection_uid)
|
||||
|
||||
try:
|
||||
from library.models import Collection
|
||||
|
||||
col = Collection.nodes.get(uid=collection_uid)
|
||||
items = col.items.all()
|
||||
|
||||
results = {"total": len(items), "successful": 0, "failed": 0, "skipped": 0}
|
||||
|
||||
for i, item in enumerate(items):
|
||||
# Skip already-completed items with unchanged content
|
||||
if item.embedding_status == "completed" and item.content_hash:
|
||||
results["skipped"] += 1
|
||||
logger.debug("Skipping already-embedded item_uid=%s", item.uid)
|
||||
continue
|
||||
|
||||
try:
|
||||
embed_item.delay(item.uid, user_id)
|
||||
results["successful"] += 1
|
||||
except Exception as exc:
|
||||
results["failed"] += 1
|
||||
logger.error(
|
||||
"Failed to queue embed for item_uid=%s: %s", item.uid, exc
|
||||
)
|
||||
|
||||
_update_progress(
|
||||
self,
|
||||
int((i + 1) / len(items) * 100),
|
||||
f"Queued {i + 1}/{len(items)} items",
|
||||
)
|
||||
|
||||
logger.info(
|
||||
"Task embed_collection completed collection_uid=%s queued=%d skipped=%d failed=%d",
|
||||
collection_uid,
|
||||
results["successful"],
|
||||
results["skipped"],
|
||||
results["failed"],
|
||||
)
|
||||
return {"success": True, "collection_uid": collection_uid, **results}
|
||||
|
||||
except Exception as exc:
|
||||
logger.error(
|
||||
"Task embed_collection failed collection_uid=%s: %s",
|
||||
collection_uid,
|
||||
exc,
|
||||
exc_info=True,
|
||||
)
|
||||
return {"success": False, "collection_uid": collection_uid, "error": str(exc)}
|
||||
|
||||
|
||||
@shared_task(
|
||||
name="library.tasks.embed_library",
|
||||
bind=True,
|
||||
queue="batch",
|
||||
acks_late=True,
|
||||
)
|
||||
def embed_library(self, library_uid: str, user_id: int = None):
|
||||
"""
|
||||
Embed all items across all collections in a library.
|
||||
|
||||
:param library_uid: UID of the Library node.
|
||||
:param user_id: Optional user ID for usage tracking.
|
||||
:returns: Dict with summary results.
|
||||
"""
|
||||
logger.info("Task embed_library starting library_uid=%s", library_uid)
|
||||
|
||||
try:
|
||||
from library.models import Library
|
||||
|
||||
lib = Library.nodes.get(uid=library_uid)
|
||||
collections = lib.collections.all()
|
||||
|
||||
results = {"total_collections": len(collections), "items_queued": 0}
|
||||
|
||||
for col in collections:
|
||||
embed_collection.delay(col.uid, user_id)
|
||||
results["items_queued"] += len(col.items.all())
|
||||
|
||||
logger.info(
|
||||
"Task embed_library completed library_uid=%s collections=%d items=%d",
|
||||
library_uid,
|
||||
results["total_collections"],
|
||||
results["items_queued"],
|
||||
)
|
||||
return {"success": True, "library_uid": library_uid, **results}
|
||||
|
||||
except Exception as exc:
|
||||
logger.error(
|
||||
"Task embed_library failed library_uid=%s: %s",
|
||||
library_uid,
|
||||
exc,
|
||||
exc_info=True,
|
||||
)
|
||||
return {"success": False, "library_uid": library_uid, "error": str(exc)}
|
||||
|
||||
|
||||
@shared_task(
|
||||
name="library.tasks.batch_embed_items",
|
||||
bind=True,
|
||||
queue="batch",
|
||||
acks_late=True,
|
||||
)
|
||||
def batch_embed_items(self, item_uids: list[str], user_id: int = None):
|
||||
"""
|
||||
Queue embedding tasks for a specific list of items.
|
||||
|
||||
:param item_uids: List of Item UIDs.
|
||||
:param user_id: Optional user ID for usage tracking.
|
||||
:returns: Dict with queuing results.
|
||||
"""
|
||||
logger.info("Task batch_embed_items starting count=%d", len(item_uids))
|
||||
|
||||
queued = 0
|
||||
for uid in item_uids:
|
||||
try:
|
||||
embed_item.delay(uid, user_id)
|
||||
queued += 1
|
||||
except Exception as exc:
|
||||
logger.error("Failed to queue item_uid=%s: %s", uid, exc)
|
||||
|
||||
logger.info("Task batch_embed_items completed queued=%d/%d", queued, len(item_uids))
|
||||
return {"success": True, "queued": queued, "total": len(item_uids)}
|
||||
|
||||
|
||||
def _resolve_user(user_id: int = None):
|
||||
"""
|
||||
Resolve a user ID to a User instance.
|
||||
|
||||
:param user_id: Optional user ID.
|
||||
:returns: User instance, or None.
|
||||
"""
|
||||
if not user_id:
|
||||
return None
|
||||
try:
|
||||
from django.contrib.auth import get_user_model
|
||||
|
||||
User = get_user_model()
|
||||
return User.objects.get(pk=user_id)
|
||||
except Exception:
|
||||
return None
|
||||
@@ -0,0 +1,19 @@
|
||||
{% extends "themis/base.html" %}
|
||||
|
||||
{% block title %}Delete {{ collection.name }} — {{ themis_app_name }}{% endblock %}
|
||||
|
||||
{% block content %}
|
||||
<div class="max-w-lg">
|
||||
<h1 class="text-3xl font-bold mb-4 text-error">Delete Collection</h1>
|
||||
<div class="alert alert-warning mb-6">
|
||||
<span>Are you sure you want to delete <strong>{{ collection.name }}</strong>? This action cannot be undone.</span>
|
||||
</div>
|
||||
<form method="post">
|
||||
{% csrf_token %}
|
||||
<div class="flex gap-2">
|
||||
<button type="submit" class="btn btn-error">Delete</button>
|
||||
<a href="{% url 'library:collection-detail' uid=collection.uid %}" class="btn btn-ghost">Cancel</a>
|
||||
</div>
|
||||
</form>
|
||||
</div>
|
||||
{% endblock %}
|
||||
71
mnemosyne/library/templates/library/collection_detail.html
Normal file
71
mnemosyne/library/templates/library/collection_detail.html
Normal file
@@ -0,0 +1,71 @@
|
||||
{% extends "themis/base.html" %}
|
||||
|
||||
{% block title %}{{ collection.name }} — {{ themis_app_name }}{% endblock %}
|
||||
|
||||
{% block content %}
|
||||
<div class="mb-4">
|
||||
{% if library %}
|
||||
<a href="{% url 'library:library-detail' uid=library.uid %}" class="btn btn-ghost btn-sm">← {{ library.name }}</a>
|
||||
{% else %}
|
||||
<a href="{% url 'library:library-list' %}" class="btn btn-ghost btn-sm">← Libraries</a>
|
||||
{% endif %}
|
||||
</div>
|
||||
|
||||
<div class="flex justify-between items-start mb-6">
|
||||
<div>
|
||||
<h1 class="text-3xl font-bold">{{ collection.name }}</h1>
|
||||
{% if library %}<p class="opacity-60 mt-1">In: {{ library.name }}</p>{% endif %}
|
||||
{% if collection.description %}
|
||||
<p class="mt-3 opacity-80">{{ collection.description }}</p>
|
||||
{% endif %}
|
||||
</div>
|
||||
<div class="flex gap-2">
|
||||
<a href="{% url 'library:collection-edit' uid=collection.uid %}" class="btn btn-sm btn-outline">Edit</a>
|
||||
<a href="{% url 'library:collection-delete' uid=collection.uid %}" class="btn btn-sm btn-error btn-outline">Delete</a>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<!-- Items -->
|
||||
<div class="flex justify-between items-center mb-4">
|
||||
<h2 class="text-xl font-bold">Items</h2>
|
||||
<a href="{% url 'library:item-create' collection_uid=collection.uid %}" class="btn btn-sm btn-primary">
|
||||
+ New Item
|
||||
</a>
|
||||
</div>
|
||||
|
||||
{% if items %}
|
||||
<div class="overflow-x-auto">
|
||||
<table class="table table-zebra w-full">
|
||||
<thead>
|
||||
<tr>
|
||||
<th>Title</th>
|
||||
<th>Type</th>
|
||||
<th>File Type</th>
|
||||
<th>Actions</th>
|
||||
</tr>
|
||||
</thead>
|
||||
<tbody>
|
||||
{% for item in items %}
|
||||
<tr>
|
||||
<td>
|
||||
<a href="{% url 'library:item-detail' uid=item.uid %}" class="link link-hover font-medium">
|
||||
{{ item.title }}
|
||||
</a>
|
||||
</td>
|
||||
<td>{{ item.item_type|default:"-" }}</td>
|
||||
<td>{{ item.file_type|default:"-" }}</td>
|
||||
<td>
|
||||
<a href="{% url 'library:item-detail' uid=item.uid %}" class="btn btn-xs btn-ghost">View</a>
|
||||
<a href="{% url 'library:item-edit' uid=item.uid %}" class="btn btn-xs btn-ghost">Edit</a>
|
||||
</td>
|
||||
</tr>
|
||||
{% endfor %}
|
||||
</tbody>
|
||||
</table>
|
||||
</div>
|
||||
{% else %}
|
||||
<div class="text-center py-8 opacity-60">
|
||||
<p>No items in this collection yet.</p>
|
||||
</div>
|
||||
{% endif %}
|
||||
{% endblock %}
|
||||
43
mnemosyne/library/templates/library/collection_form.html
Normal file
43
mnemosyne/library/templates/library/collection_form.html
Normal file
@@ -0,0 +1,43 @@
|
||||
{% extends "themis/base.html" %}
|
||||
|
||||
{% block title %}{% if editing %}Edit Collection{% else %}New Collection{% endif %} — {{ themis_app_name }}{% endblock %}
|
||||
|
||||
{% block content %}
|
||||
<div class="mb-4">
|
||||
{% if library %}
|
||||
<a href="{% url 'library:library-detail' uid=library.uid %}" class="btn btn-ghost btn-sm">← {{ library.name }}</a>
|
||||
{% else %}
|
||||
<a href="{% url 'library:library-list' %}" class="btn btn-ghost btn-sm">← Libraries</a>
|
||||
{% endif %}
|
||||
</div>
|
||||
|
||||
<h1 class="text-3xl font-bold mb-6">
|
||||
{% if editing %}Edit Collection: {{ collection.name }}{% else %}New Collection{% endif %}
|
||||
</h1>
|
||||
{% if library %}<p class="opacity-60 mb-4">In library: {{ library.name }}</p>{% endif %}
|
||||
|
||||
<form method="post" class="max-w-2xl">
|
||||
{% csrf_token %}
|
||||
<div class="space-y-4">
|
||||
<div class="form-control">
|
||||
<label class="label"><span class="label-text font-medium">Name</span></label>
|
||||
{{ form.name }}
|
||||
{% if form.name.errors %}<p class="text-error text-sm mt-1">{{ form.name.errors.0 }}</p>{% endif %}
|
||||
</div>
|
||||
<div class="form-control">
|
||||
<label class="label"><span class="label-text font-medium">Description</span></label>
|
||||
{{ form.description }}
|
||||
</div>
|
||||
</div>
|
||||
<div class="flex gap-2 mt-6">
|
||||
<button type="submit" class="btn btn-primary">
|
||||
{% if editing %}Save Changes{% else %}Create Collection{% endif %}
|
||||
</button>
|
||||
{% if library %}
|
||||
<a href="{% url 'library:library-detail' uid=library.uid %}" class="btn btn-ghost">Cancel</a>
|
||||
{% else %}
|
||||
<a href="{% url 'library:library-list' %}" class="btn btn-ghost">Cancel</a>
|
||||
{% endif %}
|
||||
</div>
|
||||
</form>
|
||||
{% endblock %}
|
||||
163
mnemosyne/library/templates/library/embedding_dashboard.html
Normal file
163
mnemosyne/library/templates/library/embedding_dashboard.html
Normal file
@@ -0,0 +1,163 @@
|
||||
{% extends "themis/base.html" %}
|
||||
{% load humanize %}
|
||||
|
||||
{% block title %}Embedding Pipeline — {{ themis_app_name }}{% endblock %}
|
||||
|
||||
{% block content %}
|
||||
<div class="mb-4">
|
||||
<a href="{% url 'library:library-list' %}" class="btn btn-ghost btn-sm">← Libraries</a>
|
||||
</div>
|
||||
|
||||
<h1 class="text-3xl font-bold mb-6">Embedding Pipeline Dashboard</h1>
|
||||
|
||||
<!-- System Models -->
|
||||
<div class="card bg-base-200 mb-6">
|
||||
<div class="card-body">
|
||||
<h2 class="card-title">System Models</h2>
|
||||
<div class="overflow-x-auto">
|
||||
<table class="table">
|
||||
<tbody>
|
||||
<tr>
|
||||
<th class="w-48">Embedding Model</th>
|
||||
<td>
|
||||
{% if system_embedding_model %}
|
||||
<span class="font-semibold">{{ system_embedding_model.api.name }}: {{ system_embedding_model.name }}</span>
|
||||
{% if system_embedding_model.vector_dimensions %}
|
||||
<span class="badge badge-info badge-sm ml-2">{{ system_embedding_model.vector_dimensions }}d</span>
|
||||
{% endif %}
|
||||
{% if system_embedding_model.supports_multimodal %}
|
||||
<span class="badge badge-accent badge-sm ml-1">Multimodal</span>
|
||||
{% endif %}
|
||||
{% else %}
|
||||
<div class="flex items-center gap-2">
|
||||
<span class="badge badge-error">NOT CONFIGURED</span>
|
||||
<span class="text-sm opacity-60">Set via <a href="/admin/llm_manager/llmmodel/" class="link link-primary">Admin → LLM Models</a> → Action: "Set as System Embedding Model"</span>
|
||||
</div>
|
||||
{% endif %}
|
||||
</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<th>Chat Model</th>
|
||||
<td>
|
||||
{% if system_chat_model %}
|
||||
<span class="font-semibold">{{ system_chat_model.api.name }}: {{ system_chat_model.name }}</span>
|
||||
{% else %}
|
||||
<span class="text-sm opacity-60">Not configured — concept extraction disabled</span>
|
||||
{% endif %}
|
||||
</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<th>Reranker Model</th>
|
||||
<td>
|
||||
{% if system_reranker_model %}
|
||||
<span class="font-semibold">{{ system_reranker_model.api.name }}: {{ system_reranker_model.name }}</span>
|
||||
{% else %}
|
||||
<span class="text-sm opacity-60">Not configured — Phase 3</span>
|
||||
{% endif %}
|
||||
</td>
|
||||
</tr>
|
||||
</tbody>
|
||||
</table>
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
{% if not neo4j_available %}
|
||||
<div class="alert alert-warning mb-6">
|
||||
<svg xmlns="http://www.w3.org/2000/svg" class="stroke-current shrink-0 h-6 w-6" fill="none" viewBox="0 0 24 24"><path stroke-linecap="round" stroke-linejoin="round" stroke-width="2" d="M12 9v2m0 4h.01m-6.938 4h13.856c1.54 0 2.502-1.667 1.732-2.5L13.732 4c-.77-.833-1.962-.833-2.732 0L4.082 16.5c-.77.833.192 2.5 1.732 2.5z" /></svg>
|
||||
<span>Neo4j is not available. Item counts and graph statistics cannot be loaded.</span>
|
||||
</div>
|
||||
{% endif %}
|
||||
|
||||
<!-- Embedding Status -->
|
||||
{% if neo4j_available %}
|
||||
<div class="grid grid-cols-2 md:grid-cols-5 gap-4 mb-6">
|
||||
<div class="stat bg-base-200 rounded-lg">
|
||||
<div class="stat-title">Total Items</div>
|
||||
<div class="stat-value text-lg">{{ total_items }}</div>
|
||||
</div>
|
||||
{% for status, count in status_counts.items %}
|
||||
<div class="stat bg-base-200 rounded-lg">
|
||||
<div class="stat-title">
|
||||
{% if status == "completed" %}✓ Completed
|
||||
{% elif status == "processing" %}⟳ Processing
|
||||
{% elif status == "failed" %}✗ Failed
|
||||
{% elif status == "pending" %}◦ Pending
|
||||
{% else %}{{ status }}{% endif %}
|
||||
</div>
|
||||
<div class="stat-value text-lg
|
||||
{% if status == 'completed' %} text-success
|
||||
{% elif status == 'processing' %} text-warning
|
||||
{% elif status == 'failed' %} text-error
|
||||
{% endif %}">
|
||||
{{ count }}
|
||||
</div>
|
||||
{% if total_items > 0 %}
|
||||
<div class="stat-desc">{% widthratio count total_items 100 %}% of items</div>
|
||||
{% endif %}
|
||||
</div>
|
||||
{% endfor %}
|
||||
</div>
|
||||
{% endif %}
|
||||
|
||||
<!-- Actions -->
|
||||
{% if status_counts.pending and status_counts.pending > 0 %}
|
||||
<div class="card bg-base-200 mb-6">
|
||||
<div class="card-body">
|
||||
<h2 class="card-title">Actions</h2>
|
||||
<form method="post" action="{% url 'library:embed-all-pending' %}">
|
||||
{% csrf_token %}
|
||||
<button type="submit" class="btn btn-primary"
|
||||
onclick="return confirm('Queue embedding for {{ status_counts.pending }} pending items?')">
|
||||
Embed All Pending Items ({{ status_counts.pending }})
|
||||
</button>
|
||||
<p class="text-sm opacity-60 mt-2">
|
||||
This will queue Celery tasks for all pending items that have uploaded files.
|
||||
</p>
|
||||
</form>
|
||||
</div>
|
||||
</div>
|
||||
{% endif %}
|
||||
|
||||
<!-- Knowledge Graph Nodes -->
|
||||
{% if neo4j_available %}
|
||||
<div class="card bg-base-200 mb-6">
|
||||
<div class="card-body">
|
||||
<h2 class="card-title">Knowledge Graph</h2>
|
||||
<div class="grid grid-cols-2 md:grid-cols-4 gap-3">
|
||||
{% for label, count in node_counts.items %}
|
||||
<div class="stat bg-base-100 rounded-lg p-3">
|
||||
<div class="stat-title text-xs">{{ label }}</div>
|
||||
<div class="stat-value text-base">{{ count|intcomma }}</div>
|
||||
</div>
|
||||
{% endfor %}
|
||||
</div>
|
||||
{% if total_chunks > 0 %}
|
||||
<div class="mt-4">
|
||||
<div class="flex items-center gap-2">
|
||||
<span class="font-medium">Chunks with embeddings:</span>
|
||||
<span>{{ embedded_chunks|intcomma }} / {{ total_chunks|intcomma }}</span>
|
||||
<progress class="progress progress-primary w-48"
|
||||
value="{{ embedded_chunks }}"
|
||||
max="{{ total_chunks }}"></progress>
|
||||
<span class="text-sm opacity-60">{% widthratio embedded_chunks total_chunks 100 %}%</span>
|
||||
</div>
|
||||
</div>
|
||||
{% endif %}
|
||||
</div>
|
||||
</div>
|
||||
{% endif %}
|
||||
|
||||
<!-- Quick Links -->
|
||||
<div class="card bg-base-200">
|
||||
<div class="card-body">
|
||||
<h2 class="card-title">Quick Links</h2>
|
||||
<div class="flex flex-wrap gap-2">
|
||||
<a href="{% url 'library:library-list' %}" class="btn btn-outline btn-sm">Libraries</a>
|
||||
<a href="/llm/" class="btn btn-outline btn-sm">LLM Manager</a>
|
||||
<a href="/admin/llm_manager/llmmodel/" class="btn btn-outline btn-sm">Admin: LLM Models</a>
|
||||
<a href="/admin/llm_manager/llmusage/" class="btn btn-outline btn-sm">Admin: Usage</a>
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
{% endblock %}
|
||||
19
mnemosyne/library/templates/library/item_confirm_delete.html
Normal file
19
mnemosyne/library/templates/library/item_confirm_delete.html
Normal file
@@ -0,0 +1,19 @@
|
||||
{% extends "themis/base.html" %}
|
||||
|
||||
{% block title %}Delete {{ item.title }} — {{ themis_app_name }}{% endblock %}
|
||||
|
||||
{% block content %}
|
||||
<div class="max-w-lg">
|
||||
<h1 class="text-3xl font-bold mb-4 text-error">Delete Item</h1>
|
||||
<div class="alert alert-warning mb-6">
|
||||
<span>Are you sure you want to delete <strong>{{ item.title }}</strong>? This action cannot be undone.</span>
|
||||
</div>
|
||||
<form method="post">
|
||||
{% csrf_token %}
|
||||
<div class="flex gap-2">
|
||||
<button type="submit" class="btn btn-error">Delete</button>
|
||||
<a href="{% url 'library:item-detail' uid=item.uid %}" class="btn btn-ghost">Cancel</a>
|
||||
</div>
|
||||
</form>
|
||||
</div>
|
||||
{% endblock %}
|
||||
138
mnemosyne/library/templates/library/item_detail.html
Normal file
138
mnemosyne/library/templates/library/item_detail.html
Normal file
@@ -0,0 +1,138 @@
|
||||
{% extends "themis/base.html" %}
|
||||
{% load humanize %}
|
||||
|
||||
{% block title %}{{ item.title }} — {{ themis_app_name }}{% endblock %}
|
||||
|
||||
{% block content %}
|
||||
<div class="mb-4">
|
||||
<a href="{% url 'library:library-list' %}" class="btn btn-ghost btn-sm">← Libraries</a>
|
||||
</div>
|
||||
|
||||
<div class="flex justify-between items-start mb-6">
|
||||
<div>
|
||||
<h1 class="text-3xl font-bold">{{ item.title }}</h1>
|
||||
{% if item.item_type %}<div class="badge badge-outline mt-2">{{ item.item_type }}</div>{% endif %}
|
||||
{% if item.file_type %}<div class="badge badge-ghost mt-2 ml-1">{{ item.file_type }}</div>{% endif %}
|
||||
</div>
|
||||
<div class="flex gap-2">
|
||||
<a href="{% url 'library:item-edit' uid=item.uid %}" class="btn btn-sm btn-outline">Edit</a>
|
||||
<form method="post" action="{% url 'library:item-reembed' uid=item.uid %}" class="inline">
|
||||
{% csrf_token %}
|
||||
<button type="submit" class="btn btn-sm btn-outline btn-secondary" title="Re-embed this item">
|
||||
↻ Re-embed
|
||||
</button>
|
||||
</form>
|
||||
<a href="{% url 'library:item-delete' uid=item.uid %}" class="btn btn-sm btn-error btn-outline">Delete</a>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<!-- Embedding Status -->
|
||||
<div class="mb-6">
|
||||
<div class="flex items-center gap-3">
|
||||
<span class="font-medium">Embedding Status:</span>
|
||||
{% if item.embedding_status == "completed" %}
|
||||
<span class="badge badge-success">Completed</span>
|
||||
{% elif item.embedding_status == "processing" %}
|
||||
<span class="badge badge-warning">Processing</span>
|
||||
{% elif item.embedding_status == "failed" %}
|
||||
<span class="badge badge-error">Failed</span>
|
||||
{% else %}
|
||||
<span class="badge badge-ghost">Pending</span>
|
||||
{% endif %}
|
||||
|
||||
{% if item.embedding_model_name %}
|
||||
<span class="text-sm opacity-60">Model: {{ item.embedding_model_name }}</span>
|
||||
{% endif %}
|
||||
</div>
|
||||
{% if item.error_message %}
|
||||
<div class="alert alert-error mt-2">
|
||||
<span>{{ item.error_message }}</span>
|
||||
</div>
|
||||
{% endif %}
|
||||
</div>
|
||||
|
||||
<!-- Item Metadata -->
|
||||
<div class="grid grid-cols-1 md:grid-cols-4 gap-4 mb-6">
|
||||
<div class="stat bg-base-200 rounded-lg">
|
||||
<div class="stat-title">File Size</div>
|
||||
<div class="stat-value text-lg">{{ item.file_size|default:0|intcomma }} bytes</div>
|
||||
</div>
|
||||
<div class="stat bg-base-200 rounded-lg">
|
||||
<div class="stat-title">Chunks</div>
|
||||
<div class="stat-value text-lg">{{ item.chunk_count|default:0 }}</div>
|
||||
</div>
|
||||
<div class="stat bg-base-200 rounded-lg">
|
||||
<div class="stat-title">Images</div>
|
||||
<div class="stat-value text-lg">{{ item.image_count|default:0 }}</div>
|
||||
</div>
|
||||
<div class="stat bg-base-200 rounded-lg">
|
||||
<div class="stat-title">Concepts</div>
|
||||
<div class="stat-value text-lg">{{ concepts|length }}</div>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<!-- Concepts -->
|
||||
{% if concepts %}
|
||||
<div class="mb-6">
|
||||
<h2 class="text-xl font-bold mb-3">Referenced Concepts</h2>
|
||||
<div class="flex flex-wrap gap-2">
|
||||
{% for concept in concepts %}
|
||||
<div class="badge badge-lg badge-primary badge-outline">
|
||||
{{ concept.name }}
|
||||
{% if concept.concept_type %}
|
||||
<span class="ml-1 opacity-60 text-xs">({{ concept.concept_type }})</span>
|
||||
{% endif %}
|
||||
</div>
|
||||
{% endfor %}
|
||||
</div>
|
||||
</div>
|
||||
{% endif %}
|
||||
|
||||
<!-- Images -->
|
||||
{% if images %}
|
||||
<div class="mb-6">
|
||||
<h2 class="text-xl font-bold mb-3">Images ({{ images|length }})</h2>
|
||||
<div class="grid grid-cols-2 md:grid-cols-4 gap-3">
|
||||
{% for img in images %}
|
||||
<div class="card bg-base-200">
|
||||
<div class="card-body p-3">
|
||||
<span class="badge badge-sm">{{ img.image_type|default:"image" }}</span>
|
||||
{% if img.description %}
|
||||
<p class="text-xs opacity-60 mt-1">{{ img.description|truncatewords:10 }}</p>
|
||||
{% endif %}
|
||||
<p class="text-xs opacity-40 mt-1">{{ img.s3_key }}</p>
|
||||
</div>
|
||||
</div>
|
||||
{% endfor %}
|
||||
</div>
|
||||
</div>
|
||||
{% endif %}
|
||||
|
||||
<!-- Chunks Preview -->
|
||||
{% if chunks %}
|
||||
<div class="mb-6">
|
||||
<h2 class="text-xl font-bold mb-3">Chunks ({{ chunks|length }})</h2>
|
||||
<div class="space-y-2">
|
||||
{% for chunk in chunks|slice:":10" %}
|
||||
<div class="collapse collapse-arrow bg-base-200">
|
||||
<input type="checkbox" />
|
||||
<div class="collapse-title font-medium">
|
||||
Chunk {{ chunk.chunk_index }} <span class="text-sm opacity-60">({{ chunk.chunk_size }} chars)</span>
|
||||
{% if chunk.embedding %}
|
||||
<span class="badge badge-success badge-xs ml-2">embedded</span>
|
||||
{% else %}
|
||||
<span class="badge badge-ghost badge-xs ml-2">no vector</span>
|
||||
{% endif %}
|
||||
</div>
|
||||
<div class="collapse-content">
|
||||
<p class="text-sm whitespace-pre-wrap">{{ chunk.text_preview }}</p>
|
||||
</div>
|
||||
</div>
|
||||
{% endfor %}
|
||||
{% if chunks|length > 10 %}
|
||||
<p class="text-sm opacity-60">… and {{ chunks|length|add:"-10" }} more chunks</p>
|
||||
{% endif %}
|
||||
</div>
|
||||
</div>
|
||||
{% endif %}
|
||||
{% endblock %}
|
||||
67
mnemosyne/library/templates/library/item_form.html
Normal file
67
mnemosyne/library/templates/library/item_form.html
Normal file
@@ -0,0 +1,67 @@
|
||||
{% extends "themis/base.html" %}
|
||||
|
||||
{% block title %}{% if editing %}Edit Item{% else %}New Item{% endif %} — {{ themis_app_name }}{% endblock %}
|
||||
|
||||
{% block content %}
|
||||
<div class="mb-4">
|
||||
{% if collection %}
|
||||
<a href="{% url 'library:collection-detail' uid=collection.uid %}" class="btn btn-ghost btn-sm">← {{ collection.name }}</a>
|
||||
{% else %}
|
||||
<a href="{% url 'library:library-list' %}" class="btn btn-ghost btn-sm">← Libraries</a>
|
||||
{% endif %}
|
||||
</div>
|
||||
|
||||
<h1 class="text-3xl font-bold mb-6">
|
||||
{% if editing %}Edit Item: {{ item.title }}{% else %}New Item{% endif %}
|
||||
</h1>
|
||||
|
||||
<form method="post" enctype="multipart/form-data" class="max-w-2xl">
|
||||
{% csrf_token %}
|
||||
<div class="space-y-4">
|
||||
<div class="form-control">
|
||||
<label class="label"><span class="label-text font-medium">Title</span></label>
|
||||
{{ form.title }}
|
||||
{% if form.title.errors %}<p class="text-error text-sm mt-1">{{ form.title.errors.0 }}</p>{% endif %}
|
||||
</div>
|
||||
<div class="form-control">
|
||||
<label class="label"><span class="label-text font-medium">Item Type</span></label>
|
||||
{{ form.item_type }}
|
||||
</div>
|
||||
<div class="form-control">
|
||||
<label class="label"><span class="label-text font-medium">File Type</span></label>
|
||||
{{ form.file_type }}
|
||||
<label class="label"><span class="label-text-alt">Auto-detected from uploaded file if left blank</span></label>
|
||||
</div>
|
||||
|
||||
{% if not editing %}
|
||||
<!-- File upload (only on create) -->
|
||||
<div class="form-control">
|
||||
<label class="label"><span class="label-text font-medium">Document File</span></label>
|
||||
{{ form.file }}
|
||||
<label class="label">
|
||||
<span class="label-text-alt">{{ form.file.help_text }}</span>
|
||||
</label>
|
||||
{% if form.file.errors %}<p class="text-error text-sm mt-1">{{ form.file.errors.0 }}</p>{% endif %}
|
||||
</div>
|
||||
|
||||
<div class="form-control">
|
||||
<label class="label cursor-pointer justify-start gap-3">
|
||||
{{ form.auto_embed }}
|
||||
<span class="label-text">Auto-embed after upload</span>
|
||||
</label>
|
||||
<label class="label">
|
||||
<span class="label-text-alt">{{ form.auto_embed.help_text }}</span>
|
||||
</label>
|
||||
</div>
|
||||
{% endif %}
|
||||
</div>
|
||||
<div class="flex gap-2 mt-6">
|
||||
<button type="submit" class="btn btn-primary">
|
||||
{% if editing %}Save Changes{% else %}Create Item{% endif %}
|
||||
</button>
|
||||
{% if collection %}
|
||||
<a href="{% url 'library:collection-detail' uid=collection.uid %}" class="btn btn-ghost">Cancel</a>
|
||||
{% endif %}
|
||||
</div>
|
||||
</form>
|
||||
{% endblock %}
|
||||
@@ -0,0 +1,23 @@
|
||||
{% extends "themis/base.html" %}
|
||||
|
||||
{% block title %}Delete {{ library.name }} — {{ themis_app_name }}{% endblock %}
|
||||
|
||||
{% block content %}
|
||||
<div class="mb-4">
|
||||
<a href="{% url 'library:library-detail' uid=library.uid %}" class="btn btn-ghost btn-sm">← {{ library.name }}</a>
|
||||
</div>
|
||||
|
||||
<div class="max-w-lg">
|
||||
<h1 class="text-3xl font-bold mb-4 text-error">Delete Library</h1>
|
||||
<div class="alert alert-warning mb-6">
|
||||
<span>Are you sure you want to delete <strong>{{ library.name }}</strong>? This action cannot be undone.</span>
|
||||
</div>
|
||||
<form method="post">
|
||||
{% csrf_token %}
|
||||
<div class="flex gap-2">
|
||||
<button type="submit" class="btn btn-error">Delete</button>
|
||||
<a href="{% url 'library:library-detail' uid=library.uid %}" class="btn btn-ghost">Cancel</a>
|
||||
</div>
|
||||
</form>
|
||||
</div>
|
||||
{% endblock %}
|
||||
93
mnemosyne/library/templates/library/library_detail.html
Normal file
93
mnemosyne/library/templates/library/library_detail.html
Normal file
@@ -0,0 +1,93 @@
|
||||
{% extends "themis/base.html" %}
|
||||
|
||||
{% block title %}{{ library.name }} — {{ themis_app_name }}{% endblock %}
|
||||
|
||||
{% block content %}
|
||||
<div class="mb-4">
|
||||
<a href="{% url 'library:library-list' %}" class="btn btn-ghost btn-sm">← Libraries</a>
|
||||
</div>
|
||||
|
||||
<div class="flex justify-between items-start mb-6">
|
||||
<div>
|
||||
<h1 class="text-3xl font-bold">{{ library.name }}</h1>
|
||||
<div class="badge badge-primary mt-2">{{ library.library_type }}</div>
|
||||
{% if library.description %}
|
||||
<p class="mt-3 opacity-80">{{ library.description }}</p>
|
||||
{% endif %}
|
||||
</div>
|
||||
<div class="flex gap-2">
|
||||
<a href="{% url 'library:library-edit' uid=library.uid %}" class="btn btn-sm btn-outline">Edit</a>
|
||||
<a href="{% url 'library:library-delete' uid=library.uid %}" class="btn btn-sm btn-error btn-outline">Delete</a>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<!-- Content-Type Configuration -->
|
||||
<div class="collapse collapse-arrow bg-base-200 mb-6">
|
||||
<input type="checkbox" />
|
||||
<div class="collapse-title font-medium">Content-Type Configuration</div>
|
||||
<div class="collapse-content">
|
||||
<div class="grid grid-cols-1 gap-4">
|
||||
{% if library.embedding_instruction %}
|
||||
<div>
|
||||
<h4 class="font-semibold text-sm opacity-60">Embedding Instruction</h4>
|
||||
<p class="text-sm mt-1">{{ library.embedding_instruction }}</p>
|
||||
</div>
|
||||
{% endif %}
|
||||
{% if library.reranker_instruction %}
|
||||
<div>
|
||||
<h4 class="font-semibold text-sm opacity-60">Reranker Instruction</h4>
|
||||
<p class="text-sm mt-1">{{ library.reranker_instruction }}</p>
|
||||
</div>
|
||||
{% endif %}
|
||||
{% if library.llm_context_prompt %}
|
||||
<div>
|
||||
<h4 class="font-semibold text-sm opacity-60">LLM Context Prompt</h4>
|
||||
<p class="text-sm mt-1">{{ library.llm_context_prompt }}</p>
|
||||
</div>
|
||||
{% endif %}
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<!-- Collections -->
|
||||
<div class="flex justify-between items-center mb-4">
|
||||
<h2 class="text-xl font-bold">Collections</h2>
|
||||
<a href="{% url 'library:collection-create' library_uid=library.uid %}" class="btn btn-sm btn-primary">
|
||||
+ New Collection
|
||||
</a>
|
||||
</div>
|
||||
|
||||
{% if collections %}
|
||||
<div class="overflow-x-auto">
|
||||
<table class="table table-zebra w-full">
|
||||
<thead>
|
||||
<tr>
|
||||
<th>Name</th>
|
||||
<th>Description</th>
|
||||
<th>Actions</th>
|
||||
</tr>
|
||||
</thead>
|
||||
<tbody>
|
||||
{% for col in collections %}
|
||||
<tr>
|
||||
<td>
|
||||
<a href="{% url 'library:collection-detail' uid=col.uid %}" class="link link-hover font-medium">
|
||||
{{ col.name }}
|
||||
</a>
|
||||
</td>
|
||||
<td class="opacity-70">{{ col.description|truncatewords:15 }}</td>
|
||||
<td>
|
||||
<a href="{% url 'library:collection-detail' uid=col.uid %}" class="btn btn-xs btn-ghost">View</a>
|
||||
<a href="{% url 'library:collection-edit' uid=col.uid %}" class="btn btn-xs btn-ghost">Edit</a>
|
||||
</td>
|
||||
</tr>
|
||||
{% endfor %}
|
||||
</tbody>
|
||||
</table>
|
||||
</div>
|
||||
{% else %}
|
||||
<div class="text-center py-8 opacity-60">
|
||||
<p>No collections in this library yet.</p>
|
||||
</div>
|
||||
{% endif %}
|
||||
{% endblock %}
|
||||
59
mnemosyne/library/templates/library/library_form.html
Normal file
59
mnemosyne/library/templates/library/library_form.html
Normal file
@@ -0,0 +1,59 @@
|
||||
{% extends "themis/base.html" %}
|
||||
|
||||
{% block title %}{% if editing %}Edit Library{% else %}New Library{% endif %} — {{ themis_app_name }}{% endblock %}
|
||||
|
||||
{% block content %}
|
||||
<div class="mb-4">
|
||||
<a href="{% url 'library:library-list' %}" class="btn btn-ghost btn-sm">← Libraries</a>
|
||||
</div>
|
||||
|
||||
<h1 class="text-3xl font-bold mb-6">
|
||||
{% if editing %}Edit Library: {{ library.name }}{% else %}New Library{% endif %}
|
||||
</h1>
|
||||
|
||||
<form method="post" class="max-w-2xl">
|
||||
{% csrf_token %}
|
||||
<div class="space-y-4">
|
||||
<div class="form-control">
|
||||
<label class="label"><span class="label-text font-medium">Name</span></label>
|
||||
{{ form.name }}
|
||||
{% if form.name.errors %}<p class="text-error text-sm mt-1">{{ form.name.errors.0 }}</p>{% endif %}
|
||||
</div>
|
||||
|
||||
<div class="form-control">
|
||||
<label class="label"><span class="label-text font-medium">Library Type</span></label>
|
||||
{{ form.library_type }}
|
||||
</div>
|
||||
|
||||
<div class="form-control">
|
||||
<label class="label"><span class="label-text font-medium">Description</span></label>
|
||||
{{ form.description }}
|
||||
</div>
|
||||
|
||||
<div class="divider">Content-Type Configuration</div>
|
||||
|
||||
<div class="form-control">
|
||||
<label class="label"><span class="label-text font-medium">Embedding Instruction</span></label>
|
||||
{{ form.embedding_instruction }}
|
||||
<label class="label"><span class="label-text-alt opacity-60">Leave blank to use default for the selected library type</span></label>
|
||||
</div>
|
||||
|
||||
<div class="form-control">
|
||||
<label class="label"><span class="label-text font-medium">Reranker Instruction</span></label>
|
||||
{{ form.reranker_instruction }}
|
||||
</div>
|
||||
|
||||
<div class="form-control">
|
||||
<label class="label"><span class="label-text font-medium">LLM Context Prompt</span></label>
|
||||
{{ form.llm_context_prompt }}
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<div class="flex gap-2 mt-6">
|
||||
<button type="submit" class="btn btn-primary">
|
||||
{% if editing %}Save Changes{% else %}Create Library{% endif %}
|
||||
</button>
|
||||
<a href="{% url 'library:library-list' %}" class="btn btn-ghost">Cancel</a>
|
||||
</div>
|
||||
</form>
|
||||
{% endblock %}
|
||||
54
mnemosyne/library/templates/library/library_list.html
Normal file
54
mnemosyne/library/templates/library/library_list.html
Normal file
@@ -0,0 +1,54 @@
|
||||
{% extends "themis/base.html" %}
|
||||
|
||||
{% block title %}Libraries — {{ themis_app_name }}{% endblock %}
|
||||
|
||||
{% block content %}
|
||||
<div class="flex justify-between items-center mb-6">
|
||||
<h1 class="text-3xl font-bold">Libraries</h1>
|
||||
<div class="flex gap-2">
|
||||
<a href="{% url 'library:embedding-dashboard' %}" class="btn btn-outline btn-secondary">
|
||||
Embedding Pipeline
|
||||
</a>
|
||||
<a href="{% url 'library:library-create' %}" class="btn btn-primary">
|
||||
+ New Library
|
||||
</a>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
{% if error %}
|
||||
<div class="alert alert-warning mb-4">
|
||||
<span>{{ error }}</span>
|
||||
</div>
|
||||
{% endif %}
|
||||
|
||||
{% if libraries %}
|
||||
<div class="grid grid-cols-1 md:grid-cols-2 lg:grid-cols-3 gap-4">
|
||||
{% for lib in libraries %}
|
||||
<div class="card bg-base-200 shadow-md">
|
||||
<div class="card-body">
|
||||
<h2 class="card-title">
|
||||
<a href="{% url 'library:library-detail' uid=lib.uid %}" class="link link-hover">
|
||||
{{ lib.name }}
|
||||
</a>
|
||||
</h2>
|
||||
<div class="badge badge-outline">{{ lib.library_type }}</div>
|
||||
{% if lib.description %}
|
||||
<p class="text-sm opacity-70 mt-2">{{ lib.description|truncatewords:20 }}</p>
|
||||
{% endif %}
|
||||
<div class="card-actions justify-end mt-3">
|
||||
<a href="{% url 'library:library-detail' uid=lib.uid %}" class="btn btn-sm btn-ghost">View</a>
|
||||
<a href="{% url 'library:library-edit' uid=lib.uid %}" class="btn btn-sm btn-ghost">Edit</a>
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
{% endfor %}
|
||||
</div>
|
||||
{% else %}
|
||||
{% if not error %}
|
||||
<div class="text-center py-12 opacity-60">
|
||||
<p class="text-lg">No libraries yet.</p>
|
||||
<p class="mt-2">Create your first library to get started.</p>
|
||||
</div>
|
||||
{% endif %}
|
||||
{% endif %}
|
||||
{% endblock %}
|
||||
0
mnemosyne/library/tests/__init__.py
Normal file
0
mnemosyne/library/tests/__init__.py
Normal file
108
mnemosyne/library/tests/test_chunker.py
Normal file
108
mnemosyne/library/tests/test_chunker.py
Normal file
@@ -0,0 +1,108 @@
|
||||
"""
|
||||
Tests for the content-type-aware chunking service.
|
||||
"""
|
||||
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
from django.test import TestCase
|
||||
|
||||
from library.services.chunker import ChunkResult, ContentTypeChunker
|
||||
from library.services.parsers import ParseResult, TextBlock
|
||||
|
||||
|
||||
class ChunkResultTests(TestCase):
|
||||
"""Tests for ChunkResult dataclass."""
|
||||
|
||||
def test_len(self):
|
||||
result = ChunkResult(chunks=["a", "b", "c"], chunk_page_map={}, strategy="test")
|
||||
self.assertEqual(len(result), 3)
|
||||
|
||||
def test_empty(self):
|
||||
result = ChunkResult(chunks=[], chunk_page_map={}, strategy="test")
|
||||
self.assertEqual(len(result), 0)
|
||||
|
||||
|
||||
class ContentTypeChunkerTests(TestCase):
|
||||
"""Tests for ContentTypeChunker."""
|
||||
|
||||
def _make_parse_result(self, text: str, pages: int = 1) -> ParseResult:
|
||||
"""Helper to create a ParseResult with text blocks."""
|
||||
blocks = []
|
||||
if pages == 1:
|
||||
blocks = [TextBlock(text=text, page=0)]
|
||||
else:
|
||||
chunk_size = len(text) // pages
|
||||
for i in range(pages):
|
||||
start = i * chunk_size
|
||||
end = start + chunk_size if i < pages - 1 else len(text)
|
||||
blocks.append(TextBlock(text=text[start:end], page=i))
|
||||
return ParseResult(text_blocks=blocks, images=[], metadata={}, file_type="txt")
|
||||
|
||||
@patch("library.services.chunker.ContentTypeChunker._get_splitter")
|
||||
def test_chunk_dispatches_strategy(self, mock_splitter):
|
||||
"""Chunker uses the strategy from config."""
|
||||
mock_instance = MagicMock()
|
||||
mock_instance.chunks.return_value = ["chunk1", "chunk2"]
|
||||
mock_splitter.return_value = mock_instance
|
||||
|
||||
chunker = ContentTypeChunker()
|
||||
parse_result = self._make_parse_result("Some text to chunk into pieces")
|
||||
config = {"strategy": "chapter_aware", "chunk_size": 512, "chunk_overlap": 64}
|
||||
|
||||
result = chunker.chunk(parse_result, config, library_type="fiction")
|
||||
|
||||
self.assertIsInstance(result, ChunkResult)
|
||||
self.assertEqual(result.strategy, "chapter_aware")
|
||||
self.assertEqual(len(result.chunks), 2)
|
||||
mock_splitter.assert_called_once_with(512, 64)
|
||||
|
||||
@patch("library.services.chunker.ContentTypeChunker._get_splitter")
|
||||
def test_empty_text_returns_empty(self, mock_splitter):
|
||||
"""Empty text produces no chunks."""
|
||||
chunker = ContentTypeChunker()
|
||||
parse_result = ParseResult(text_blocks=[], images=[], metadata={}, file_type="txt")
|
||||
config = {"strategy": "section_aware", "chunk_size": 512, "chunk_overlap": 64}
|
||||
|
||||
result = chunker.chunk(parse_result, config)
|
||||
|
||||
self.assertEqual(len(result), 0)
|
||||
mock_splitter.assert_not_called()
|
||||
|
||||
@patch("library.services.chunker.ContentTypeChunker._get_splitter")
|
||||
def test_default_config_values(self, mock_splitter):
|
||||
"""Missing config keys use defaults."""
|
||||
mock_instance = MagicMock()
|
||||
mock_instance.chunks.return_value = ["chunk"]
|
||||
mock_splitter.return_value = mock_instance
|
||||
|
||||
chunker = ContentTypeChunker()
|
||||
parse_result = self._make_parse_result("Text")
|
||||
|
||||
result = chunker.chunk(parse_result, {})
|
||||
|
||||
# Default: strategy=section_aware, chunk_size=512, overlap=64
|
||||
self.assertEqual(result.strategy, "section_aware")
|
||||
mock_splitter.assert_called_once_with(512, 64)
|
||||
|
||||
@patch("library.services.chunker.ContentTypeChunker._get_splitter")
|
||||
def test_page_mapping(self, mock_splitter):
|
||||
"""Chunks are mapped to their source pages."""
|
||||
mock_instance = MagicMock()
|
||||
mock_instance.chunks.return_value = ["Page 0 text", "Page 1 text"]
|
||||
mock_splitter.return_value = mock_instance
|
||||
|
||||
chunker = ContentTypeChunker()
|
||||
parse_result = ParseResult(
|
||||
text_blocks=[
|
||||
TextBlock(text="Page 0 text content", page=0),
|
||||
TextBlock(text="Page 1 text content", page=1),
|
||||
],
|
||||
images=[],
|
||||
metadata={},
|
||||
file_type="pdf",
|
||||
)
|
||||
config = {"strategy": "section_aware", "chunk_size": 512, "chunk_overlap": 64}
|
||||
|
||||
result = chunker.chunk(parse_result, config)
|
||||
|
||||
self.assertIn(0, result.chunk_page_map)
|
||||
77
mnemosyne/library/tests/test_concepts.py
Normal file
77
mnemosyne/library/tests/test_concepts.py
Normal file
@@ -0,0 +1,77 @@
|
||||
"""
|
||||
Tests for the concept extraction service.
|
||||
"""
|
||||
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
from django.test import TestCase
|
||||
|
||||
from library.services.concepts import ConceptExtractor
|
||||
|
||||
|
||||
class ConceptExtractionParsingTests(TestCase):
|
||||
"""Tests for concept response parsing."""
|
||||
|
||||
def setUp(self):
|
||||
self.mock_model = MagicMock()
|
||||
self.mock_model.api.api_type = "openai"
|
||||
self.mock_model.api.base_url = "http://localhost:8080/v1"
|
||||
self.mock_model.api.api_key = "test"
|
||||
self.mock_model.api.timeout_seconds = 30
|
||||
self.mock_model.name = "test-chat"
|
||||
self.extractor = ConceptExtractor(self.mock_model)
|
||||
|
||||
def test_parse_valid_json_array(self):
|
||||
response = '[{"name": "python", "type": "topic"}, {"name": "django", "type": "technique"}]'
|
||||
result = self.extractor._parse_concept_response(response)
|
||||
self.assertEqual(len(result), 2)
|
||||
self.assertEqual(result[0]["name"], "python")
|
||||
self.assertEqual(result[1]["type"], "technique")
|
||||
|
||||
def test_parse_json_in_markdown_code_block(self):
|
||||
response = '```json\n[{"name": "python", "type": "topic"}]\n```'
|
||||
result = self.extractor._parse_concept_response(response)
|
||||
self.assertEqual(len(result), 1)
|
||||
|
||||
def test_parse_json_embedded_in_text(self):
|
||||
response = 'Here are the concepts: [{"name": "neo4j", "type": "technique"}] found in the text.'
|
||||
result = self.extractor._parse_concept_response(response)
|
||||
self.assertEqual(len(result), 1)
|
||||
|
||||
def test_parse_invalid_json_returns_empty(self):
|
||||
response = "This is not JSON at all."
|
||||
result = self.extractor._parse_concept_response(response)
|
||||
self.assertEqual(result, [])
|
||||
|
||||
def test_parse_filters_invalid_entries(self):
|
||||
response = '[{"name": "valid", "type": "topic"}, {"invalid": "entry"}, "string"]'
|
||||
result = self.extractor._parse_concept_response(response)
|
||||
self.assertEqual(len(result), 1)
|
||||
self.assertEqual(result[0]["name"], "valid")
|
||||
|
||||
|
||||
class SampleIndexSelectionTests(TestCase):
|
||||
"""Tests for sample index selection."""
|
||||
|
||||
def setUp(self):
|
||||
self.extractor = ConceptExtractor(MagicMock())
|
||||
|
||||
def test_small_total_returns_all(self):
|
||||
indices = self.extractor._select_sample_indices(5, max_samples=10)
|
||||
self.assertEqual(indices, [0, 1, 2, 3, 4])
|
||||
|
||||
def test_equal_total_returns_all(self):
|
||||
indices = self.extractor._select_sample_indices(10, max_samples=10)
|
||||
self.assertEqual(indices, list(range(10)))
|
||||
|
||||
def test_large_total_returns_max_samples(self):
|
||||
indices = self.extractor._select_sample_indices(100, max_samples=10)
|
||||
self.assertEqual(len(indices), 10)
|
||||
# Should be evenly spaced
|
||||
self.assertEqual(indices[0], 0)
|
||||
self.assertEqual(indices[-1], 90)
|
||||
|
||||
def test_returns_integers(self):
|
||||
indices = self.extractor._select_sample_indices(50, max_samples=7)
|
||||
for idx in indices:
|
||||
self.assertIsInstance(idx, int)
|
||||
165
mnemosyne/library/tests/test_embedding_client.py
Normal file
165
mnemosyne/library/tests/test_embedding_client.py
Normal file
@@ -0,0 +1,165 @@
|
||||
"""
|
||||
Tests for the multi-backend embedding client.
|
||||
"""
|
||||
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
from django.test import TestCase
|
||||
|
||||
from library.services.embedding_client import EmbeddingClient
|
||||
|
||||
|
||||
class MockLLMModel:
|
||||
"""Mock LLMModel for testing."""
|
||||
|
||||
def __init__(self, api_type="openai", supports_multimodal=False, vector_dimensions=None):
|
||||
self.name = "test-embedding-model"
|
||||
self.supports_multimodal = supports_multimodal
|
||||
self.vector_dimensions = vector_dimensions
|
||||
self.input_cost_per_1k = "0.0001"
|
||||
self.api = MockLLMApi(api_type=api_type)
|
||||
|
||||
|
||||
class MockLLMApi:
|
||||
"""Mock LLMApi for testing."""
|
||||
|
||||
def __init__(self, api_type="openai"):
|
||||
self.name = "Test API"
|
||||
self.api_type = api_type
|
||||
self.base_url = "http://localhost:8080/v1"
|
||||
self.api_key = "test-key"
|
||||
self.timeout_seconds = 60
|
||||
|
||||
|
||||
class EmbeddingClientInitTests(TestCase):
|
||||
"""Tests for EmbeddingClient initialization."""
|
||||
|
||||
def test_init_openai(self):
|
||||
model = MockLLMModel(api_type="openai")
|
||||
client = EmbeddingClient(model)
|
||||
self.assertEqual(client.api_type, "openai")
|
||||
self.assertEqual(client.model_name, "test-embedding-model")
|
||||
|
||||
def test_init_bedrock(self):
|
||||
model = MockLLMModel(api_type="bedrock")
|
||||
model.api.base_url = "https://bedrock-runtime.us-east-1.amazonaws.com"
|
||||
client = EmbeddingClient(model)
|
||||
self.assertEqual(client.api_type, "bedrock")
|
||||
|
||||
def test_init_with_user(self):
|
||||
model = MockLLMModel()
|
||||
user = MagicMock()
|
||||
client = EmbeddingClient(model, user=user)
|
||||
self.assertEqual(client.user, user)
|
||||
|
||||
|
||||
class OpenAIResponseParsingTests(TestCase):
|
||||
"""Tests for OpenAI-compatible response parsing."""
|
||||
|
||||
def setUp(self):
|
||||
self.model = MockLLMModel()
|
||||
self.client = EmbeddingClient(self.model)
|
||||
|
||||
def test_parse_standard_openai_format(self):
|
||||
data = {"data": [{"embedding": [0.1, 0.2, 0.3], "index": 0}]}
|
||||
result = self.client._parse_openai_response(data)
|
||||
self.assertEqual(len(result), 1)
|
||||
self.assertEqual(result[0], [0.1, 0.2, 0.3])
|
||||
|
||||
def test_parse_multi_embedding_openai(self):
|
||||
data = {
|
||||
"data": [
|
||||
{"embedding": [0.1, 0.2], "index": 1},
|
||||
{"embedding": [0.3, 0.4], "index": 0},
|
||||
]
|
||||
}
|
||||
result = self.client._parse_openai_response(data)
|
||||
self.assertEqual(len(result), 2)
|
||||
# Should be sorted by index
|
||||
self.assertEqual(result[0], [0.3, 0.4])
|
||||
self.assertEqual(result[1], [0.1, 0.2])
|
||||
|
||||
def test_parse_list_of_dicts(self):
|
||||
data = [{"embedding": [0.1, 0.2]}, {"embedding": [0.3, 0.4]}]
|
||||
result = self.client._parse_openai_response(data)
|
||||
self.assertEqual(len(result), 2)
|
||||
|
||||
def test_parse_dict_with_embedding_key(self):
|
||||
data = {"embedding": [0.1, 0.2, 0.3]}
|
||||
result = self.client._parse_openai_response(data)
|
||||
self.assertEqual(len(result), 1)
|
||||
|
||||
def test_parse_dict_with_embeddings_key(self):
|
||||
data = {"embeddings": [[0.1, 0.2], [0.3, 0.4]]}
|
||||
result = self.client._parse_openai_response(data)
|
||||
self.assertEqual(len(result), 2)
|
||||
|
||||
def test_unexpected_format_raises(self):
|
||||
with self.assertRaises(ValueError):
|
||||
self.client._parse_openai_response({"unexpected": "data"})
|
||||
|
||||
|
||||
class EmbeddingClientDispatchTests(TestCase):
|
||||
"""Tests for API type dispatch."""
|
||||
|
||||
@patch("library.services.embedding_client.requests.post")
|
||||
def test_embed_text_openai_dispatch(self, mock_post):
|
||||
mock_response = MagicMock()
|
||||
mock_response.status_code = 200
|
||||
mock_response.json.return_value = {
|
||||
"data": [{"embedding": [0.1, 0.2, 0.3], "index": 0}]
|
||||
}
|
||||
mock_post.return_value = mock_response
|
||||
|
||||
model = MockLLMModel(api_type="openai")
|
||||
client = EmbeddingClient(model)
|
||||
result = client.embed_text("test text")
|
||||
|
||||
self.assertEqual(result, [0.1, 0.2, 0.3])
|
||||
mock_post.assert_called_once()
|
||||
call_url = mock_post.call_args[0][0]
|
||||
self.assertIn("/embeddings", call_url)
|
||||
|
||||
@patch("library.services.embedding_client.requests.post")
|
||||
def test_embed_text_bedrock_dispatch(self, mock_post):
|
||||
mock_response = MagicMock()
|
||||
mock_response.status_code = 200
|
||||
mock_response.json.return_value = {
|
||||
"embedding": [0.4, 0.5, 0.6],
|
||||
"inputTextTokenCount": 5,
|
||||
}
|
||||
mock_post.return_value = mock_response
|
||||
|
||||
model = MockLLMModel(api_type="bedrock", vector_dimensions=1024)
|
||||
model.api.base_url = "https://bedrock-runtime.us-east-1.amazonaws.com"
|
||||
client = EmbeddingClient(model)
|
||||
result = client.embed_text("test text")
|
||||
|
||||
self.assertEqual(result, [0.4, 0.5, 0.6])
|
||||
call_url = mock_post.call_args[0][0]
|
||||
self.assertIn("/model/", call_url)
|
||||
self.assertIn("/invoke", call_url)
|
||||
|
||||
def test_embed_image_not_multimodal_returns_none(self):
|
||||
model = MockLLMModel(supports_multimodal=False)
|
||||
client = EmbeddingClient(model)
|
||||
result = client.embed_image(b"fake image data", "png")
|
||||
self.assertIsNone(result)
|
||||
|
||||
@patch("library.services.embedding_client.requests.post")
|
||||
def test_embed_texts_batch(self, mock_post):
|
||||
mock_response = MagicMock()
|
||||
mock_response.status_code = 200
|
||||
mock_response.json.return_value = {
|
||||
"data": [
|
||||
{"embedding": [0.1, 0.2], "index": 0},
|
||||
{"embedding": [0.3, 0.4], "index": 1},
|
||||
]
|
||||
}
|
||||
mock_post.return_value = mock_response
|
||||
|
||||
model = MockLLMModel()
|
||||
client = EmbeddingClient(model)
|
||||
results = client.embed_texts(["text1", "text2"])
|
||||
|
||||
self.assertEqual(len(results), 2)
|
||||
129
mnemosyne/library/tests/test_parsers.py
Normal file
129
mnemosyne/library/tests/test_parsers.py
Normal file
@@ -0,0 +1,129 @@
|
||||
"""
|
||||
Tests for the document parser service.
|
||||
"""
|
||||
|
||||
import os
|
||||
import tempfile
|
||||
|
||||
from django.test import TestCase
|
||||
|
||||
from library.services.parsers import (
|
||||
IMAGE_EXTENSIONS,
|
||||
PLAINTEXT_EXTENSIONS,
|
||||
PYMUPDF_EXTENSIONS,
|
||||
DocumentParser,
|
||||
ParseResult,
|
||||
)
|
||||
|
||||
|
||||
class DocumentParserPlaintextTests(TestCase):
|
||||
"""Tests for plain text parsing."""
|
||||
|
||||
def setUp(self):
|
||||
self.parser = DocumentParser()
|
||||
|
||||
def test_parse_txt_file(self):
|
||||
with tempfile.NamedTemporaryFile(mode="w", suffix=".txt", delete=False) as f:
|
||||
f.write("Hello World\n\nThis is a test document.")
|
||||
f.flush()
|
||||
path = f.name
|
||||
|
||||
try:
|
||||
result = self.parser.parse(path, "txt")
|
||||
self.assertIsInstance(result, ParseResult)
|
||||
self.assertEqual(result.file_type, "txt")
|
||||
self.assertEqual(len(result.text_blocks), 1)
|
||||
self.assertIn("Hello World", result.text_blocks[0].text)
|
||||
self.assertEqual(len(result.images), 0)
|
||||
finally:
|
||||
os.unlink(path)
|
||||
|
||||
def test_parse_md_file(self):
|
||||
with tempfile.NamedTemporaryFile(mode="w", suffix=".md", delete=False) as f:
|
||||
f.write("# Heading\n\nSome markdown content.")
|
||||
f.flush()
|
||||
path = f.name
|
||||
|
||||
try:
|
||||
result = self.parser.parse(path, "md")
|
||||
self.assertEqual(result.file_type, "md")
|
||||
self.assertIn("Heading", result.text_blocks[0].text)
|
||||
finally:
|
||||
os.unlink(path)
|
||||
|
||||
def test_parse_empty_file(self):
|
||||
with tempfile.NamedTemporaryFile(mode="w", suffix=".txt", delete=False) as f:
|
||||
f.write("")
|
||||
f.flush()
|
||||
path = f.name
|
||||
|
||||
try:
|
||||
result = self.parser.parse(path, "txt")
|
||||
self.assertEqual(len(result.text_blocks), 0)
|
||||
finally:
|
||||
os.unlink(path)
|
||||
|
||||
def test_parse_bytes(self):
|
||||
data = b"Hello from bytes"
|
||||
result = self.parser.parse_bytes(data, "txt", filename="test.txt")
|
||||
self.assertEqual(len(result.text_blocks), 1)
|
||||
self.assertIn("Hello from bytes", result.text_blocks[0].text)
|
||||
|
||||
|
||||
class DocumentParserValidationTests(TestCase):
|
||||
"""Tests for parser input validation."""
|
||||
|
||||
def setUp(self):
|
||||
self.parser = DocumentParser()
|
||||
|
||||
def test_unsupported_format_raises(self):
|
||||
with tempfile.NamedTemporaryFile(suffix=".xyz", delete=False) as f:
|
||||
f.write(b"data")
|
||||
path = f.name
|
||||
|
||||
try:
|
||||
with self.assertRaises(ValueError) as ctx:
|
||||
self.parser.parse(path, "xyz")
|
||||
self.assertIn("Unsupported file type", str(ctx.exception))
|
||||
finally:
|
||||
os.unlink(path)
|
||||
|
||||
def test_file_type_normalization(self):
|
||||
"""File type should be normalized (lowercase, no dot)."""
|
||||
with tempfile.NamedTemporaryFile(mode="w", suffix=".txt", delete=False) as f:
|
||||
f.write("test")
|
||||
path = f.name
|
||||
|
||||
try:
|
||||
result = self.parser.parse(path, ".TXT")
|
||||
self.assertEqual(result.file_type, "txt")
|
||||
finally:
|
||||
os.unlink(path)
|
||||
|
||||
|
||||
class SupportedExtensionsTests(TestCase):
|
||||
"""Tests for supported extension sets."""
|
||||
|
||||
def test_pymupdf_includes_pdf(self):
|
||||
self.assertIn("pdf", PYMUPDF_EXTENSIONS)
|
||||
|
||||
def test_pymupdf_includes_epub(self):
|
||||
self.assertIn("epub", PYMUPDF_EXTENSIONS)
|
||||
|
||||
def test_pymupdf_includes_docx(self):
|
||||
self.assertIn("docx", PYMUPDF_EXTENSIONS)
|
||||
|
||||
def test_pymupdf_includes_pptx(self):
|
||||
self.assertIn("pptx", PYMUPDF_EXTENSIONS)
|
||||
|
||||
def test_plaintext_includes_txt(self):
|
||||
self.assertIn("txt", PLAINTEXT_EXTENSIONS)
|
||||
|
||||
def test_plaintext_includes_md(self):
|
||||
self.assertIn("md", PLAINTEXT_EXTENSIONS)
|
||||
|
||||
def test_image_includes_png(self):
|
||||
self.assertIn("png", IMAGE_EXTENSIONS)
|
||||
|
||||
def test_image_includes_jpg(self):
|
||||
self.assertIn("jpg", IMAGE_EXTENSIONS)
|
||||
103
mnemosyne/library/tests/test_pipeline.py
Normal file
103
mnemosyne/library/tests/test_pipeline.py
Normal file
@@ -0,0 +1,103 @@
|
||||
"""
|
||||
Tests for the embedding pipeline orchestrator.
|
||||
|
||||
Pipeline tests mock external dependencies (Neo4j, S3, LLM APIs).
|
||||
"""
|
||||
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
from django.test import TestCase
|
||||
|
||||
from library.services.pipeline import (
|
||||
CHUNK_S3_KEY,
|
||||
IMAGE_S3_KEY,
|
||||
ORIGINAL_S3_KEY,
|
||||
EmbeddingPipeline,
|
||||
)
|
||||
|
||||
|
||||
class S3KeyPatternTests(TestCase):
|
||||
"""Tests for S3 key pattern formatting."""
|
||||
|
||||
def test_original_key_format(self):
|
||||
key = ORIGINAL_S3_KEY.format(item_uid="abc123", ext="pdf")
|
||||
self.assertEqual(key, "items/abc123/original.pdf")
|
||||
|
||||
def test_chunk_key_format(self):
|
||||
key = CHUNK_S3_KEY.format(item_uid="abc123", index=5)
|
||||
self.assertEqual(key, "chunks/abc123/chunk_5.txt")
|
||||
|
||||
def test_image_key_format(self):
|
||||
key = IMAGE_S3_KEY.format(item_uid="abc123", index=2, ext="png")
|
||||
self.assertEqual(key, "images/abc123/2.png")
|
||||
|
||||
|
||||
class EmbeddingPipelineInitTests(TestCase):
|
||||
"""Tests for pipeline initialization."""
|
||||
|
||||
def test_init_without_user(self):
|
||||
pipeline = EmbeddingPipeline()
|
||||
self.assertIsNone(pipeline.user)
|
||||
|
||||
def test_init_with_user(self):
|
||||
user = MagicMock()
|
||||
pipeline = EmbeddingPipeline(user=user)
|
||||
self.assertEqual(pipeline.user, user)
|
||||
|
||||
|
||||
class PipelineItemNotFoundTests(TestCase):
|
||||
"""Tests for handling missing items."""
|
||||
|
||||
@patch("library.services.pipeline.Item")
|
||||
def test_process_nonexistent_item_raises(self, mock_item_cls):
|
||||
mock_item_cls.nodes.get.side_effect = Exception("Not found")
|
||||
|
||||
pipeline = EmbeddingPipeline()
|
||||
with self.assertRaises(ValueError) as ctx:
|
||||
pipeline.process_item("nonexistent-uid")
|
||||
self.assertIn("Item not found", str(ctx.exception))
|
||||
|
||||
@patch("library.services.pipeline.Item")
|
||||
def test_reprocess_nonexistent_item_raises(self, mock_item_cls):
|
||||
mock_item_cls.nodes.get.side_effect = Exception("Not found")
|
||||
|
||||
pipeline = EmbeddingPipeline()
|
||||
with self.assertRaises(ValueError):
|
||||
pipeline.reprocess_item("nonexistent-uid")
|
||||
|
||||
|
||||
class PipelineNoEmbeddingModelTests(TestCase):
|
||||
"""Tests for handling missing system embedding model."""
|
||||
|
||||
@patch("library.services.pipeline.LLMModel")
|
||||
@patch("library.services.pipeline.default_storage")
|
||||
@patch("library.services.pipeline.DocumentParser")
|
||||
def test_no_embedding_model_raises(self, mock_parser, mock_storage, mock_llm):
|
||||
"""Pipeline raises ValueError if no system embedding model is configured."""
|
||||
mock_llm.get_system_embedding_model.return_value = None
|
||||
|
||||
# Mock item
|
||||
mock_item = MagicMock()
|
||||
mock_item.uid = "test-uid"
|
||||
mock_item.title = "Test"
|
||||
mock_item.file_type = "txt"
|
||||
mock_item.s3_key = "items/test-uid/original.txt"
|
||||
mock_item.embedding_status = "pending"
|
||||
mock_item.chunks.all.return_value = []
|
||||
mock_item.images.all.return_value = []
|
||||
|
||||
with patch("library.services.pipeline.Item") as mock_item_cls:
|
||||
mock_item_cls.nodes.get.return_value = mock_item
|
||||
|
||||
# Mock S3 read
|
||||
mock_storage.open.return_value.__enter__ = MagicMock(
|
||||
return_value=MagicMock(read=MagicMock(return_value=b"test content"))
|
||||
)
|
||||
mock_storage.open.return_value.__exit__ = MagicMock(return_value=False)
|
||||
|
||||
pipeline = EmbeddingPipeline()
|
||||
|
||||
with self.assertRaises(ValueError) as ctx:
|
||||
pipeline.process_item("test-uid")
|
||||
|
||||
self.assertIn("No system embedding model", str(ctx.exception))
|
||||
86
mnemosyne/library/tests/test_tasks.py
Normal file
86
mnemosyne/library/tests/test_tasks.py
Normal file
@@ -0,0 +1,86 @@
|
||||
"""
|
||||
Tests for Celery embedding tasks.
|
||||
|
||||
Tasks are tested with CELERY_TASK_ALWAYS_EAGER=True for synchronous execution.
|
||||
"""
|
||||
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
from django.test import TestCase, override_settings
|
||||
|
||||
|
||||
@override_settings(CELERY_TASK_ALWAYS_EAGER=True)
|
||||
class EmbedItemTaskTests(TestCase):
|
||||
"""Tests for the embed_item task."""
|
||||
|
||||
@patch("library.tasks.EmbeddingPipeline")
|
||||
def test_embed_item_success(self, mock_pipeline_cls):
|
||||
from library.tasks import embed_item
|
||||
|
||||
mock_pipeline = MagicMock()
|
||||
mock_pipeline.process_item.return_value = {
|
||||
"chunks_created": 10,
|
||||
"images_stored": 2,
|
||||
"model_name": "test-model",
|
||||
}
|
||||
mock_pipeline_cls.return_value = mock_pipeline
|
||||
|
||||
result = embed_item("test-uid-123")
|
||||
|
||||
self.assertTrue(result["success"])
|
||||
self.assertEqual(result["item_uid"], "test-uid-123")
|
||||
mock_pipeline.process_item.assert_called_once()
|
||||
|
||||
@patch("library.tasks.EmbeddingPipeline")
|
||||
def test_embed_item_failure(self, mock_pipeline_cls):
|
||||
from library.tasks import embed_item
|
||||
|
||||
mock_pipeline = MagicMock()
|
||||
mock_pipeline.process_item.side_effect = ValueError("Item not found")
|
||||
mock_pipeline_cls.return_value = mock_pipeline
|
||||
|
||||
result = embed_item("nonexistent-uid")
|
||||
|
||||
self.assertFalse(result["success"])
|
||||
self.assertIn("error", result)
|
||||
|
||||
|
||||
@override_settings(CELERY_TASK_ALWAYS_EAGER=True)
|
||||
class ReembedItemTaskTests(TestCase):
|
||||
"""Tests for the reembed_item task."""
|
||||
|
||||
@patch("library.tasks.EmbeddingPipeline")
|
||||
def test_reembed_item_success(self, mock_pipeline_cls):
|
||||
from library.tasks import reembed_item
|
||||
|
||||
mock_pipeline = MagicMock()
|
||||
mock_pipeline.reprocess_item.return_value = {
|
||||
"chunks_created": 5,
|
||||
"images_stored": 1,
|
||||
"model_name": "test-model",
|
||||
}
|
||||
mock_pipeline_cls.return_value = mock_pipeline
|
||||
|
||||
result = reembed_item("test-uid-123")
|
||||
|
||||
self.assertTrue(result["success"])
|
||||
mock_pipeline.reprocess_item.assert_called_once()
|
||||
|
||||
|
||||
class ResolveUserTests(TestCase):
|
||||
"""Tests for the _resolve_user helper."""
|
||||
|
||||
def test_none_user_id(self):
|
||||
from library.tasks import _resolve_user
|
||||
|
||||
self.assertIsNone(_resolve_user(None))
|
||||
|
||||
def test_zero_user_id(self):
|
||||
from library.tasks import _resolve_user
|
||||
|
||||
self.assertIsNone(_resolve_user(0))
|
||||
|
||||
def test_invalid_user_id(self):
|
||||
from library.tasks import _resolve_user
|
||||
|
||||
self.assertIsNone(_resolve_user(999999))
|
||||
121
mnemosyne/library/tests/test_text_utils.py
Normal file
121
mnemosyne/library/tests/test_text_utils.py
Normal file
@@ -0,0 +1,121 @@
|
||||
"""
|
||||
Tests for text sanitization utilities.
|
||||
"""
|
||||
|
||||
from django.test import TestCase
|
||||
|
||||
from library.services.text_utils import (
|
||||
clean_pdf_artifacts,
|
||||
remove_excessive_whitespace,
|
||||
sanitize_text,
|
||||
truncate_text,
|
||||
)
|
||||
|
||||
|
||||
class SanitizeTextTests(TestCase):
|
||||
"""Tests for the sanitize_text function."""
|
||||
|
||||
def test_empty_string(self):
|
||||
self.assertEqual(sanitize_text("", log_changes=False), "")
|
||||
|
||||
def test_none_input(self):
|
||||
self.assertIsNone(sanitize_text(None, log_changes=False))
|
||||
|
||||
def test_clean_text_unchanged(self):
|
||||
text = "Hello, this is clean text."
|
||||
self.assertEqual(sanitize_text(text, log_changes=False), text)
|
||||
|
||||
def test_removes_null_bytes(self):
|
||||
text = "Hello\x00World"
|
||||
result = sanitize_text(text, log_changes=False)
|
||||
self.assertNotIn("\x00", result)
|
||||
self.assertEqual(result, "HelloWorld")
|
||||
|
||||
def test_removes_control_characters(self):
|
||||
text = "Hello\x07World\x0eTest"
|
||||
result = sanitize_text(text, log_changes=False)
|
||||
self.assertNotIn("\x07", result)
|
||||
self.assertNotIn("\x0e", result)
|
||||
|
||||
def test_preserves_newlines_and_tabs(self):
|
||||
text = "Hello\nWorld\tTest\r\n"
|
||||
result = sanitize_text(text, log_changes=False)
|
||||
self.assertIn("\n", result)
|
||||
self.assertIn("\t", result)
|
||||
|
||||
def test_removes_zero_width_characters(self):
|
||||
text = "Hello\u200bWorld"
|
||||
result = sanitize_text(text, log_changes=False)
|
||||
self.assertNotIn("\u200b", result)
|
||||
|
||||
def test_normalizes_unicode(self):
|
||||
# é as combining characters vs. precomposed
|
||||
combining = "e\u0301" # e + combining acute
|
||||
result = sanitize_text(combining, log_changes=False)
|
||||
self.assertEqual(result, "\u00e9") # precomposed é
|
||||
|
||||
def test_cleans_pdf_ligatures(self):
|
||||
text = "finding the flow of effort"
|
||||
result = sanitize_text(text, log_changes=False)
|
||||
self.assertIn("fi", result)
|
||||
self.assertIn("fl", result)
|
||||
self.assertIn("ff", result)
|
||||
|
||||
|
||||
class CleanPdfArtifactsTests(TestCase):
|
||||
"""Tests for clean_pdf_artifacts."""
|
||||
|
||||
def test_replaces_smart_quotes(self):
|
||||
text = "\u201cHello\u201d \u2018World\u2019"
|
||||
result = clean_pdf_artifacts(text)
|
||||
self.assertEqual(result, '"Hello" \'World\'')
|
||||
|
||||
def test_replaces_dashes(self):
|
||||
text = "word\u2013word\u2014end"
|
||||
result = clean_pdf_artifacts(text)
|
||||
self.assertEqual(result, "word-word-end")
|
||||
|
||||
def test_replaces_ellipsis(self):
|
||||
text = "wait\u2026"
|
||||
result = clean_pdf_artifacts(text)
|
||||
self.assertEqual(result, "wait...")
|
||||
|
||||
def test_replaces_nbsp(self):
|
||||
text = "non\u00a0breaking"
|
||||
result = clean_pdf_artifacts(text)
|
||||
self.assertEqual(result, "non breaking")
|
||||
|
||||
|
||||
class RemoveExcessiveWhitespaceTests(TestCase):
|
||||
"""Tests for remove_excessive_whitespace."""
|
||||
|
||||
def test_collapses_spaces(self):
|
||||
self.assertEqual(remove_excessive_whitespace("a b"), "a b")
|
||||
|
||||
def test_collapses_newlines(self):
|
||||
self.assertEqual(
|
||||
remove_excessive_whitespace("a\n\n\n\nb"), "a\n\nb"
|
||||
)
|
||||
|
||||
def test_strips_line_whitespace(self):
|
||||
self.assertEqual(
|
||||
remove_excessive_whitespace(" hello \n world "),
|
||||
"hello\nworld",
|
||||
)
|
||||
|
||||
|
||||
class TruncateTextTests(TestCase):
|
||||
"""Tests for truncate_text."""
|
||||
|
||||
def test_short_text_unchanged(self):
|
||||
self.assertEqual(truncate_text("hello", 100), "hello")
|
||||
|
||||
def test_truncates_at_word_boundary(self):
|
||||
text = "hello beautiful world"
|
||||
result = truncate_text(text, 15)
|
||||
self.assertTrue(result.endswith("..."))
|
||||
self.assertLessEqual(len(result), 15)
|
||||
|
||||
def test_custom_suffix(self):
|
||||
result = truncate_text("hello beautiful world", 15, suffix="…")
|
||||
self.assertTrue(result.endswith("…"))
|
||||
56
mnemosyne/library/urls.py
Normal file
56
mnemosyne/library/urls.py
Normal file
@@ -0,0 +1,56 @@
|
||||
"""
|
||||
URL patterns for the library app.
|
||||
|
||||
Provides both custom admin views (HTML CRUD) and DRF API endpoints.
|
||||
"""
|
||||
|
||||
from django.urls import include, path
|
||||
|
||||
from . import views
|
||||
|
||||
app_name = "library"
|
||||
|
||||
urlpatterns = [
|
||||
# Embedding Pipeline Dashboard
|
||||
path("embedding/", views.embedding_dashboard, name="embedding-dashboard"),
|
||||
path("embedding/embed-all/", views.embed_all_pending, name="embed-all-pending"),
|
||||
# Library CRUD
|
||||
path("", views.library_list, name="library-list"),
|
||||
path("create/", views.library_create, name="library-create"),
|
||||
path("<str:uid>/", views.library_detail, name="library-detail"),
|
||||
path("<str:uid>/edit/", views.library_edit, name="library-edit"),
|
||||
path("<str:uid>/delete/", views.library_delete, name="library-delete"),
|
||||
# Collection CRUD
|
||||
path(
|
||||
"<str:library_uid>/collections/create/",
|
||||
views.collection_create,
|
||||
name="collection-create",
|
||||
),
|
||||
path(
|
||||
"collections/<str:uid>/",
|
||||
views.collection_detail,
|
||||
name="collection-detail",
|
||||
),
|
||||
path(
|
||||
"collections/<str:uid>/edit/",
|
||||
views.collection_edit,
|
||||
name="collection-edit",
|
||||
),
|
||||
path(
|
||||
"collections/<str:uid>/delete/",
|
||||
views.collection_delete,
|
||||
name="collection-delete",
|
||||
),
|
||||
# Item CRUD
|
||||
path(
|
||||
"collections/<str:collection_uid>/items/create/",
|
||||
views.item_create,
|
||||
name="item-create",
|
||||
),
|
||||
path("items/<str:uid>/", views.item_detail, name="item-detail"),
|
||||
path("items/<str:uid>/edit/", views.item_edit, name="item-edit"),
|
||||
path("items/<str:uid>/reembed/", views.item_reembed, name="item-reembed"),
|
||||
path("items/<str:uid>/delete/", views.item_delete, name="item-delete"),
|
||||
# DRF API
|
||||
path("api/", include("library.api.urls")),
|
||||
]
|
||||
23
mnemosyne/library/utils.py
Normal file
23
mnemosyne/library/utils.py
Normal file
@@ -0,0 +1,23 @@
|
||||
"""
|
||||
Utility helpers for the library app.
|
||||
"""
|
||||
|
||||
import logging
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def neo4j_available():
|
||||
"""
|
||||
Check whether Neo4j is reachable.
|
||||
|
||||
Returns True if a simple Cypher query succeeds, False otherwise.
|
||||
Used to guard views/tests that require Neo4j.
|
||||
"""
|
||||
try:
|
||||
from neomodel import db
|
||||
|
||||
db.cypher_query("RETURN 1")
|
||||
return True
|
||||
except Exception:
|
||||
return False
|
||||
587
mnemosyne/library/views.py
Normal file
587
mnemosyne/library/views.py
Normal file
@@ -0,0 +1,587 @@
|
||||
"""
|
||||
Custom admin views for Library, Collection, and Item CRUD.
|
||||
|
||||
Since neomodel StructuredNodes cannot use Django's standard ModelAdmin,
|
||||
these FBVs provide CRUD operations rendered within Themis's template structure.
|
||||
All views require login.
|
||||
"""
|
||||
|
||||
import hashlib
|
||||
import logging
|
||||
import os
|
||||
|
||||
from django.contrib import messages
|
||||
from django.contrib.auth.decorators import login_required
|
||||
from django.core.files.base import ContentFile
|
||||
from django.core.files.storage import default_storage
|
||||
from django.shortcuts import redirect, render
|
||||
|
||||
from .content_types import get_library_type_config
|
||||
from .forms import CollectionForm, ItemForm, LibraryForm
|
||||
from .utils import neo4j_available
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Library views
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@login_required
|
||||
def library_list(request):
|
||||
"""List all libraries."""
|
||||
libraries = []
|
||||
error = None
|
||||
if neo4j_available():
|
||||
try:
|
||||
from .models import Library
|
||||
|
||||
libraries = Library.nodes.order_by("name")
|
||||
except Exception as e:
|
||||
error = f"Could not connect to Neo4j: {e}"
|
||||
logger.error(error)
|
||||
else:
|
||||
error = "Neo4j is not available."
|
||||
return render(
|
||||
request,
|
||||
"library/library_list.html",
|
||||
{"libraries": libraries, "error": error},
|
||||
)
|
||||
|
||||
|
||||
@login_required
|
||||
def library_create(request):
|
||||
"""Create a new library."""
|
||||
if request.method == "POST":
|
||||
form = LibraryForm(request.POST)
|
||||
if form.is_valid():
|
||||
try:
|
||||
from .models import Library
|
||||
|
||||
# If content-type fields are empty, populate from defaults
|
||||
library_type = form.cleaned_data["library_type"]
|
||||
defaults = get_library_type_config(library_type)
|
||||
|
||||
lib = Library(
|
||||
name=form.cleaned_data["name"],
|
||||
library_type=library_type,
|
||||
description=form.cleaned_data.get("description", ""),
|
||||
chunking_config=defaults["chunking_config"],
|
||||
embedding_instruction=(
|
||||
form.cleaned_data.get("embedding_instruction")
|
||||
or defaults["embedding_instruction"]
|
||||
),
|
||||
reranker_instruction=(
|
||||
form.cleaned_data.get("reranker_instruction")
|
||||
or defaults["reranker_instruction"]
|
||||
),
|
||||
llm_context_prompt=(
|
||||
form.cleaned_data.get("llm_context_prompt")
|
||||
or defaults["llm_context_prompt"]
|
||||
),
|
||||
)
|
||||
lib.save()
|
||||
messages.success(request, f'Library "{lib.name}" created.')
|
||||
return redirect("library:library-detail", uid=lib.uid)
|
||||
except Exception as e:
|
||||
messages.error(request, f"Error creating library: {e}")
|
||||
else:
|
||||
form = LibraryForm()
|
||||
return render(request, "library/library_form.html", {"form": form, "editing": False})
|
||||
|
||||
|
||||
@login_required
|
||||
def library_detail(request, uid):
|
||||
"""View library details and its collections."""
|
||||
try:
|
||||
from .models import Library
|
||||
|
||||
lib = Library.nodes.get(uid=uid)
|
||||
collections = lib.collections.all()
|
||||
except Exception as e:
|
||||
messages.error(request, f"Library not found: {e}")
|
||||
return redirect("library:library-list")
|
||||
return render(
|
||||
request,
|
||||
"library/library_detail.html",
|
||||
{"library": lib, "collections": collections},
|
||||
)
|
||||
|
||||
|
||||
@login_required
|
||||
def library_edit(request, uid):
|
||||
"""Edit an existing library."""
|
||||
try:
|
||||
from .models import Library
|
||||
|
||||
lib = Library.nodes.get(uid=uid)
|
||||
except Exception as e:
|
||||
messages.error(request, f"Library not found: {e}")
|
||||
return redirect("library:library-list")
|
||||
|
||||
if request.method == "POST":
|
||||
form = LibraryForm(request.POST)
|
||||
if form.is_valid():
|
||||
try:
|
||||
lib.name = form.cleaned_data["name"]
|
||||
lib.library_type = form.cleaned_data["library_type"]
|
||||
lib.description = form.cleaned_data.get("description", "")
|
||||
lib.embedding_instruction = form.cleaned_data.get(
|
||||
"embedding_instruction", ""
|
||||
)
|
||||
lib.reranker_instruction = form.cleaned_data.get(
|
||||
"reranker_instruction", ""
|
||||
)
|
||||
lib.llm_context_prompt = form.cleaned_data.get(
|
||||
"llm_context_prompt", ""
|
||||
)
|
||||
lib.save()
|
||||
messages.success(request, f'Library "{lib.name}" updated.')
|
||||
return redirect("library:library-detail", uid=lib.uid)
|
||||
except Exception as e:
|
||||
messages.error(request, f"Error updating library: {e}")
|
||||
else:
|
||||
form = LibraryForm(
|
||||
initial={
|
||||
"name": lib.name,
|
||||
"library_type": lib.library_type,
|
||||
"description": lib.description,
|
||||
"embedding_instruction": lib.embedding_instruction,
|
||||
"reranker_instruction": lib.reranker_instruction,
|
||||
"llm_context_prompt": lib.llm_context_prompt,
|
||||
}
|
||||
)
|
||||
return render(
|
||||
request,
|
||||
"library/library_form.html",
|
||||
{"form": form, "editing": True, "library": lib},
|
||||
)
|
||||
|
||||
|
||||
@login_required
|
||||
def library_delete(request, uid):
|
||||
"""Delete a library (and confirm)."""
|
||||
try:
|
||||
from .models import Library
|
||||
|
||||
lib = Library.nodes.get(uid=uid)
|
||||
except Exception as e:
|
||||
messages.error(request, f"Library not found: {e}")
|
||||
return redirect("library:library-list")
|
||||
|
||||
if request.method == "POST":
|
||||
name = lib.name
|
||||
lib.delete()
|
||||
messages.success(request, f'Library "{name}" deleted.')
|
||||
return redirect("library:library-list")
|
||||
return render(request, "library/library_confirm_delete.html", {"library": lib})
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Collection views
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@login_required
|
||||
def collection_create(request, library_uid):
|
||||
"""Create a new collection within a library."""
|
||||
try:
|
||||
from .models import Library
|
||||
|
||||
lib = Library.nodes.get(uid=library_uid)
|
||||
except Exception as e:
|
||||
messages.error(request, f"Library not found: {e}")
|
||||
return redirect("library:library-list")
|
||||
|
||||
if request.method == "POST":
|
||||
form = CollectionForm(request.POST)
|
||||
if form.is_valid():
|
||||
try:
|
||||
from .models import Collection
|
||||
|
||||
col = Collection(
|
||||
name=form.cleaned_data["name"],
|
||||
description=form.cleaned_data.get("description", ""),
|
||||
)
|
||||
col.save()
|
||||
lib.collections.connect(col)
|
||||
col.library.connect(lib)
|
||||
messages.success(request, f'Collection "{col.name}" created.')
|
||||
return redirect("library:collection-detail", uid=col.uid)
|
||||
except Exception as e:
|
||||
messages.error(request, f"Error creating collection: {e}")
|
||||
else:
|
||||
form = CollectionForm()
|
||||
return render(
|
||||
request,
|
||||
"library/collection_form.html",
|
||||
{"form": form, "library": lib, "editing": False},
|
||||
)
|
||||
|
||||
|
||||
@login_required
|
||||
def collection_detail(request, uid):
|
||||
"""View collection details and its items."""
|
||||
try:
|
||||
from .models import Collection
|
||||
|
||||
col = Collection.nodes.get(uid=uid)
|
||||
items = col.items.all()
|
||||
libraries = col.library.all()
|
||||
library = libraries[0] if libraries else None
|
||||
except Exception as e:
|
||||
messages.error(request, f"Collection not found: {e}")
|
||||
return redirect("library:library-list")
|
||||
return render(
|
||||
request,
|
||||
"library/collection_detail.html",
|
||||
{"collection": col, "items": items, "library": library},
|
||||
)
|
||||
|
||||
|
||||
@login_required
|
||||
def collection_edit(request, uid):
|
||||
"""Edit an existing collection."""
|
||||
try:
|
||||
from .models import Collection
|
||||
|
||||
col = Collection.nodes.get(uid=uid)
|
||||
libraries = col.library.all()
|
||||
library = libraries[0] if libraries else None
|
||||
except Exception as e:
|
||||
messages.error(request, f"Collection not found: {e}")
|
||||
return redirect("library:library-list")
|
||||
|
||||
if request.method == "POST":
|
||||
form = CollectionForm(request.POST)
|
||||
if form.is_valid():
|
||||
try:
|
||||
col.name = form.cleaned_data["name"]
|
||||
col.description = form.cleaned_data.get("description", "")
|
||||
col.save()
|
||||
messages.success(request, f'Collection "{col.name}" updated.')
|
||||
return redirect("library:collection-detail", uid=col.uid)
|
||||
except Exception as e:
|
||||
messages.error(request, f"Error updating collection: {e}")
|
||||
else:
|
||||
form = CollectionForm(
|
||||
initial={
|
||||
"name": col.name,
|
||||
"description": col.description,
|
||||
}
|
||||
)
|
||||
return render(
|
||||
request,
|
||||
"library/collection_form.html",
|
||||
{"form": form, "collection": col, "library": library, "editing": True},
|
||||
)
|
||||
|
||||
|
||||
@login_required
|
||||
def collection_delete(request, uid):
|
||||
"""Delete a collection."""
|
||||
try:
|
||||
from .models import Collection
|
||||
|
||||
col = Collection.nodes.get(uid=uid)
|
||||
except Exception as e:
|
||||
messages.error(request, f"Collection not found: {e}")
|
||||
return redirect("library:library-list")
|
||||
|
||||
if request.method == "POST":
|
||||
name = col.name
|
||||
col.delete()
|
||||
messages.success(request, f'Collection "{name}" deleted.')
|
||||
return redirect("library:library-list")
|
||||
return render(
|
||||
request,
|
||||
"library/collection_confirm_delete.html",
|
||||
{"collection": col},
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Item views
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@login_required
|
||||
def item_create(request, collection_uid):
|
||||
"""Create a new item within a collection, with optional file upload."""
|
||||
try:
|
||||
from .models import Collection
|
||||
|
||||
col = Collection.nodes.get(uid=collection_uid)
|
||||
libraries = col.library.all()
|
||||
library = libraries[0] if libraries else None
|
||||
except Exception as e:
|
||||
messages.error(request, f"Collection not found: {e}")
|
||||
return redirect("library:library-list")
|
||||
|
||||
if request.method == "POST":
|
||||
form = ItemForm(request.POST, request.FILES)
|
||||
if form.is_valid():
|
||||
try:
|
||||
from .models import Item
|
||||
|
||||
uploaded_file = request.FILES.get("file")
|
||||
file_type = form.cleaned_data.get("file_type", "")
|
||||
|
||||
# Infer file_type from upload if not explicitly set
|
||||
if uploaded_file and not file_type:
|
||||
_, ext = os.path.splitext(uploaded_file.name)
|
||||
file_type = ext.lstrip(".").lower()
|
||||
|
||||
item = Item(
|
||||
title=form.cleaned_data["title"],
|
||||
item_type=form.cleaned_data.get("item_type", ""),
|
||||
file_type=file_type,
|
||||
embedding_status="pending",
|
||||
)
|
||||
|
||||
# Handle file upload
|
||||
if uploaded_file:
|
||||
file_data = uploaded_file.read()
|
||||
item.file_size = len(file_data)
|
||||
item.content_hash = hashlib.sha256(file_data).hexdigest()
|
||||
item.save()
|
||||
|
||||
# Store in S3
|
||||
s3_key = f"items/{item.uid}/original.{file_type}"
|
||||
default_storage.save(s3_key, ContentFile(file_data))
|
||||
item.s3_key = s3_key
|
||||
item.save()
|
||||
else:
|
||||
item.save()
|
||||
|
||||
col.items.connect(item)
|
||||
|
||||
# Auto-trigger embedding if file uploaded and checkbox set
|
||||
auto_embed = form.cleaned_data.get("auto_embed", True)
|
||||
if uploaded_file and auto_embed:
|
||||
try:
|
||||
from .tasks import embed_item
|
||||
|
||||
task = embed_item.delay(item.uid, request.user.id)
|
||||
messages.info(
|
||||
request,
|
||||
f"Embedding queued (task: {task.id})",
|
||||
)
|
||||
except Exception as exc:
|
||||
logger.warning("Failed to queue embedding: %s", exc)
|
||||
|
||||
messages.success(request, f'Item "{item.title}" created.')
|
||||
return redirect("library:item-detail", uid=item.uid)
|
||||
except Exception as e:
|
||||
messages.error(request, f"Error creating item: {e}")
|
||||
else:
|
||||
form = ItemForm(initial={"auto_embed": True})
|
||||
return render(
|
||||
request,
|
||||
"library/item_form.html",
|
||||
{"form": form, "collection": col, "library": library, "editing": False},
|
||||
)
|
||||
|
||||
|
||||
@login_required
|
||||
def item_detail(request, uid):
|
||||
"""View item details."""
|
||||
try:
|
||||
from .models import Item
|
||||
|
||||
item = Item.nodes.get(uid=uid)
|
||||
chunks = item.chunks.all()
|
||||
images = item.images.all()
|
||||
concepts = item.concepts.all()
|
||||
except Exception as e:
|
||||
messages.error(request, f"Item not found: {e}")
|
||||
return redirect("library:library-list")
|
||||
return render(
|
||||
request,
|
||||
"library/item_detail.html",
|
||||
{"item": item, "chunks": chunks, "images": images, "concepts": concepts},
|
||||
)
|
||||
|
||||
|
||||
@login_required
|
||||
def item_edit(request, uid):
|
||||
"""Edit an existing item."""
|
||||
try:
|
||||
from .models import Item
|
||||
|
||||
item = Item.nodes.get(uid=uid)
|
||||
except Exception as e:
|
||||
messages.error(request, f"Item not found: {e}")
|
||||
return redirect("library:library-list")
|
||||
|
||||
if request.method == "POST":
|
||||
form = ItemForm(request.POST)
|
||||
if form.is_valid():
|
||||
try:
|
||||
item.title = form.cleaned_data["title"]
|
||||
item.item_type = form.cleaned_data.get("item_type", "")
|
||||
item.file_type = form.cleaned_data.get("file_type", "")
|
||||
item.save()
|
||||
messages.success(request, f'Item "{item.title}" updated.')
|
||||
return redirect("library:item-detail", uid=item.uid)
|
||||
except Exception as e:
|
||||
messages.error(request, f"Error updating item: {e}")
|
||||
else:
|
||||
form = ItemForm(
|
||||
initial={
|
||||
"title": item.title,
|
||||
"item_type": item.item_type,
|
||||
"file_type": item.file_type,
|
||||
}
|
||||
)
|
||||
return render(
|
||||
request,
|
||||
"library/item_form.html",
|
||||
{"form": form, "item": item, "editing": True},
|
||||
)
|
||||
|
||||
|
||||
@login_required
|
||||
def item_reembed(request, uid):
|
||||
"""Trigger re-embedding for an item."""
|
||||
try:
|
||||
from .models import Item
|
||||
|
||||
item = Item.nodes.get(uid=uid)
|
||||
except Exception as e:
|
||||
messages.error(request, f"Item not found: {e}")
|
||||
return redirect("library:library-list")
|
||||
|
||||
if request.method == "POST":
|
||||
try:
|
||||
from .tasks import reembed_item
|
||||
|
||||
task = reembed_item.delay(uid, request.user.id)
|
||||
messages.info(request, f"Re-embedding queued for \"{item.title}\" (task: {task.id})")
|
||||
except Exception as exc:
|
||||
messages.error(request, f"Failed to queue re-embedding: {exc}")
|
||||
return redirect("library:item-detail", uid=uid)
|
||||
|
||||
return redirect("library:item-detail", uid=uid)
|
||||
|
||||
|
||||
@login_required
|
||||
def item_delete(request, uid):
|
||||
"""Delete an item."""
|
||||
try:
|
||||
from .models import Item
|
||||
|
||||
item = Item.nodes.get(uid=uid)
|
||||
except Exception as e:
|
||||
messages.error(request, f"Item not found: {e}")
|
||||
return redirect("library:library-list")
|
||||
|
||||
if request.method == "POST":
|
||||
title = item.title
|
||||
item.delete()
|
||||
messages.success(request, f'Item "{title}" deleted.')
|
||||
return redirect("library:library-list")
|
||||
return render(request, "library/item_confirm_delete.html", {"item": item})
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Embedding Pipeline Dashboard
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@login_required
|
||||
def embedding_dashboard(request):
|
||||
"""
|
||||
Embedding pipeline dashboard — system model status, item embedding
|
||||
progress, knowledge graph node counts, and batch actions.
|
||||
"""
|
||||
context = {
|
||||
"system_embedding_model": None,
|
||||
"system_chat_model": None,
|
||||
"system_reranker_model": None,
|
||||
"status_counts": {},
|
||||
"node_counts": {},
|
||||
"total_items": 0,
|
||||
"embedded_chunks": 0,
|
||||
"total_chunks": 0,
|
||||
"neo4j_available": False,
|
||||
}
|
||||
|
||||
# Get system models from LLM Manager
|
||||
try:
|
||||
from llm_manager.models import LLMModel
|
||||
|
||||
context["system_embedding_model"] = LLMModel.get_system_embedding_model()
|
||||
context["system_chat_model"] = LLMModel.get_system_chat_model()
|
||||
context["system_reranker_model"] = LLMModel.get_system_reranker_model()
|
||||
except Exception as exc:
|
||||
logger.warning("Could not load system models: %s", exc)
|
||||
|
||||
# Get item status counts and node counts from Neo4j
|
||||
if neo4j_available():
|
||||
context["neo4j_available"] = True
|
||||
try:
|
||||
from neomodel import db
|
||||
|
||||
for status in ["pending", "processing", "completed", "failed"]:
|
||||
results, _ = db.cypher_query(
|
||||
"MATCH (i:Item {embedding_status: $status}) RETURN count(i)",
|
||||
{"status": status},
|
||||
)
|
||||
context["status_counts"][status] = results[0][0] if results else 0
|
||||
|
||||
results, _ = db.cypher_query("MATCH (i:Item) RETURN count(i)")
|
||||
context["total_items"] = results[0][0] if results else 0
|
||||
|
||||
for label in ["Library", "Collection", "Item", "Chunk", "Concept", "Image", "ImageEmbedding"]:
|
||||
results, _ = db.cypher_query(f"MATCH (n:{label}) RETURN count(n)")
|
||||
context["node_counts"][label] = results[0][0] if results else 0
|
||||
|
||||
results, _ = db.cypher_query(
|
||||
"MATCH (c:Chunk) WHERE c.embedding IS NOT NULL RETURN count(c)"
|
||||
)
|
||||
context["embedded_chunks"] = results[0][0] if results else 0
|
||||
context["total_chunks"] = context["node_counts"].get("Chunk", 0)
|
||||
|
||||
except Exception as exc:
|
||||
logger.warning("Could not query Neo4j for dashboard: %s", exc)
|
||||
messages.warning(request, f"Neo4j query error: {exc}")
|
||||
|
||||
return render(request, "library/embedding_dashboard.html", context)
|
||||
|
||||
|
||||
@login_required
|
||||
def embed_all_pending(request):
|
||||
"""
|
||||
Trigger embedding for all pending items with uploaded files.
|
||||
POST-only action, redirects back to dashboard.
|
||||
"""
|
||||
if request.method != "POST":
|
||||
return redirect("library:embedding-dashboard")
|
||||
|
||||
try:
|
||||
from neomodel import db
|
||||
|
||||
results, _ = db.cypher_query(
|
||||
"MATCH (i:Item {embedding_status: 'pending'}) "
|
||||
"WHERE i.s3_key IS NOT NULL AND i.s3_key <> '' "
|
||||
"RETURN i.uid"
|
||||
)
|
||||
item_uids = [row[0] for row in results]
|
||||
|
||||
if not item_uids:
|
||||
messages.info(request, "No pending items with files to embed.")
|
||||
else:
|
||||
from .tasks import batch_embed_items
|
||||
|
||||
task = batch_embed_items.delay(item_uids, request.user.id)
|
||||
messages.success(
|
||||
request,
|
||||
f"Queued embedding for {len(item_uids)} items (task: {task.id})",
|
||||
)
|
||||
except Exception as exc:
|
||||
logger.error("Failed to trigger batch embedding: %s", exc, exc_info=True)
|
||||
messages.error(request, f"Failed to trigger embedding: {exc}")
|
||||
|
||||
return redirect("library:embedding-dashboard")
|
||||
Reference in New Issue
Block a user