chat_backend/llm_be/chat_backend/admin.py

from django.contrib import admin
from django.db.models import Sum
from .models import (
    CustomUser,
    Announcement,
    AgentRun,
    AgentStep,
    Company,
    LLMModels,
    Conversation,
    Prompt,
    Feedback,
    PromptFeedback,
    PromptMetric,
    DocumentWorkspace,
    Document,
    DriveConnection,
    UserAuthEvent,
    OutboundEmail,
    OAuthIdentity,
)


class AnnouncmentAdmin(admin.ModelAdmin):
    model = Announcement


class CompanyAdmin(admin.ModelAdmin):
    model = Company


class UserAuthEventInline(admin.TabularInline):
    model = UserAuthEvent
    extra = 0
    can_delete = False
    fields = ("created", "event_type", "detail", "ip_address")
    readonly_fields = ("created", "event_type", "detail", "ip_address")
    ordering = ("-created",)
    show_change_link = False

    def has_add_permission(self, request, obj=None):
        return False


class OutboundEmailInline(admin.TabularInline):
    model = OutboundEmail
    extra = 0
    can_delete = False
    fields = ("created", "kind", "status", "subject", "to_email", "sent_at")
    readonly_fields = ("created", "kind", "status", "subject", "to_email", "sent_at")
    ordering = ("-created",)
    show_change_link = True

    def has_add_permission(self, request, obj=None):
        return False


class CustomUserAdmin(admin.ModelAdmin):
    model = CustomUser
    list_display = (
        "email",
        "username",
        "is_company_manager",
        "first_name",
        "last_name",
        "is_active",
        "is_staff",
        "has_usable_password",
        "deleted",
        "has_signed_tos",
        "use_conversation_context",
        "last_login",
        "slug",
        "get_set_password_url",
    )
    search_fields = ("email", "username", "first_name", "last_name", "slug")
    readonly_fields = ("last_login", "date_joined", "get_set_password_url", "slug")
    inlines = (UserAuthEventInline, OutboundEmailInline)


class UserAuthEventAdmin(admin.ModelAdmin):
    model = UserAuthEvent
    list_display = ("created", "user", "event_type", "detail", "ip_address")
    list_filter = ("event_type",)
    search_fields = ("user__email", "user__username", "detail", "ip_address")
    readonly_fields = ("user", "event_type", "created", "detail", "ip_address")
    ordering = ("-created",)

    def has_add_permission(self, request):
        return False

    def has_change_permission(self, request, obj=None):
        return False


class OutboundEmailAdmin(admin.ModelAdmin):
    model = OutboundEmail
    list_display = (
        "created",
        "kind",
        "status",
        "subject",
        "to_email",
        "user",
        "sent_at",
    )
    list_filter = ("kind", "status")
    search_fields = ("to_email", "subject", "user__email", "id")
    readonly_fields = (
        "id",
        "kind",
        "status",
        "to_email",
        "from_email",
        "subject",
        "html_template",
        "text_template",
        "context",
        "error_message",
        "user",
        "created",
        "updated",
        "sent_at",
    )
    ordering = ("-created",)
    date_hierarchy = "created"

    def has_add_permission(self, request):
        return False

    def has_change_permission(self, request, obj=None):
        return False


class FeedbackAdmin(admin.ModelAdmin):
    model = Feedback
    search_fields = ("status", "text", "get_user_email")
    list_display = ("status", "get_user_email", "title", "category")


class PromptFeedbackAdmin(admin.ModelAdmin):
    model = PromptFeedback
    list_display = ("id", "prompt", "user", "rating", "reason", "created")
    list_filter = ("rating", "reason")
    search_fields = ("user__email", "comment", "prompt__message")
    raw_id_fields = ("prompt", "user")


class LLMModelsAdmin(admin.ModelAdmin):
    model = LLMModels
    list_display = ("name", "port", "description")
    search_fields = ("name", "port", "description")


class PromptInline(admin.TabularInline):
    model = Prompt


class ConversationAdmin(admin.ModelAdmin):
    model = Conversation
    list_display = (
        "title",
        "get_user_email",
        "deleted",
        "tokens_in",
        "tokens_out",
        "tokens_total",
    )
    search_fields = ("title",)
    inlines = [
        PromptInline,
    ]

    def _token_sum(self, conversation, field):
        total = PromptMetric.objects.filter(conversation_id=conversation.id).aggregate(
            total=Sum(field)
        )["total"]
        return total or 0

    @admin.display(description="Tokens in")
    def tokens_in(self, conversation):
        return self._token_sum(conversation, "tokens_in")

    @admin.display(description="Tokens out")
    def tokens_out(self, conversation):
        return self._token_sum(conversation, "tokens_out")

    @admin.display(description="Tokens total")
    def tokens_total(self, conversation):
        return self.tokens_in(conversation) + self.tokens_out(conversation)


class PromptAdmin(admin.ModelAdmin):
    model = Prompt
    list_display = ("id", "message", "user_created", "get_conversation_title", "created")
    search_fields = ("message",)


class PromptMetricAdmin(admin.ModelAdmin):
    model = PromptMetric
    list_display = (
        "id",
        "event",
        "model_name",
        "prompt_length",
        "reponse_length",
        "tokens_in",
        "tokens_out",
        "has_file",
        "file_type",
        "get_duration",
        "created",
    )
    list_filter = ("event", "model_name", "has_file")


class DocumentWorkspaceAdmin(admin.ModelAdmin):
    model = DocumentWorkspace
    list_display = (
        "name",
        "company",
        "user",
    )


class DocumentAdmin(admin.ModelAdmin):
    model = Document
    list_display = (
        "file",
        "source",
        "active",
        "created",
        "processed",
        "drive_connection",
    )
    list_filter = ("source", "active", "processed")
    raw_id_fields = ("drive_connection",)


class DriveConnectionAdmin(admin.ModelAdmin):
    model = DriveConnection
    list_display = (
        "provider",
        "kind",
        "company",
        "user",
        "external_account_email",
        "is_active",
        "last_sync_status",
        "last_sync_at",
    )
    list_filter = ("provider", "kind", "is_active", "last_sync_status")
    search_fields = ("external_account_email", "company__name", "user__email")
    raw_id_fields = ("company", "user")
    readonly_fields = (
        "created",
        "last_modified",
        "access_token",
        "refresh_token",
    )


admin.site.register(Announcement, AnnouncmentAdmin)
admin.site.register(Company, CompanyAdmin)
admin.site.register(CustomUser, CustomUserAdmin)
admin.site.register(UserAuthEvent, UserAuthEventAdmin)
admin.site.register(OutboundEmail, OutboundEmailAdmin)

admin.site.register(LLMModels, LLMModelsAdmin)
admin.site.register(Conversation, ConversationAdmin)
admin.site.register(Prompt, PromptAdmin)
admin.site.register(PromptMetric, PromptMetricAdmin)
admin.site.register(Feedback, FeedbackAdmin)
admin.site.register(PromptFeedback, PromptFeedbackAdmin)

admin.site.register(DocumentWorkspace, DocumentWorkspaceAdmin)
admin.site.register(Document, DocumentAdmin)
admin.site.register(DriveConnection, DriveConnectionAdmin)


class OAuthIdentityAdmin(admin.ModelAdmin):
    model = OAuthIdentity
    list_display = (
        "provider",
        "email",
        "subject",
        "user",
        "token_expires_at",
        "created",
        "last_modified",
    )
    list_filter = ("provider",)
    search_fields = ("email", "subject", "user__email")
    readonly_fields = ("created", "last_modified", "raw_profile")


admin.site.register(OAuthIdentity, OAuthIdentityAdmin)


class AgentStepInline(admin.TabularInline):
    model = AgentStep
    fk_name = "run"
    extra = 0
    can_delete = False
    fields = ("index", "title", "status", "tool_name", "is_subagent", "started_at", "completed_at")
    readonly_fields = fields
    ordering = ("index",)

    def has_add_permission(self, request, obj=None):
        return False


class AgentRunAdmin(admin.ModelAdmin):
    model = AgentRun
    list_display = (
        "id",
        "user",
        "status",
        "title",
        "tool_call_count",
        "iteration_count",
        "cancel_requested",
        "created",
        "completed_at",
    )
    list_filter = ("status", "cancel_requested")
    search_fields = ("goal", "title", "user__email")
    raw_id_fields = ("user", "company", "conversation", "prompt")
    readonly_fields = ("created", "last_modified")
    inlines = (AgentStepInline,)


class AgentStepAdmin(admin.ModelAdmin):
    model = AgentStep
    list_display = ("id", "run", "index", "title", "status", "tool_name", "is_subagent")
    list_filter = ("status", "is_subagent")
    search_fields = ("title", "tool_name", "run__id")
    raw_id_fields = ("run", "parent_step")


admin.site.register(AgentRun, AgentRunAdmin)
admin.site.register(AgentStep, AgentStepAdmin)