app/memory/db.py

from __future__ import annotations

from typing import Any

from app.memory.neon_db import (
    save_conversation as _save_conversation,
    get_conversations as _get_conversations,
    get_conversation as _get_conversation,
    delete_conversation as _delete_conversation,
    update_conversation_title as _update_conversation_title,
    save_message as _save_message,
    get_conversation_messages as _get_conversation_messages,
    save_document as _save_document,
    get_documents as _get_documents,
    delete_document as _delete_document,
    save_document_chunk as _save_document_chunk,
    get_document_chunks as _get_document_chunks,
)
from app.utils.logger import get_logger

logger = get_logger(__name__)


def save_conversation(
    conversation_id: str,
    user_id: str,
    title: str,
    created_at: str,
    updated_at: str,
) -> None:
    _save_conversation(conversation_id, user_id, title, created_at, updated_at)


def get_conversations(user_id: str, limit: int = 50) -> list[dict[str, Any]]:
    return _get_conversations(user_id, limit)


def get_conversation(conversation_id: str, user_id: str) -> dict[str, Any] | None:
    return _get_conversation(conversation_id, user_id)


def delete_conversation(conversation_id: str, user_id: str) -> None:
    _delete_conversation(conversation_id, user_id)


def update_conversation_title(conversation_id: str, user_id: str, title: str) -> None:
    _update_conversation_title(conversation_id, user_id, title)


def save_message(
    message_id: str,
    conversation_id: str,
    user_id: str,
    role: str,
    content: str,
    created_at: str,
    image_url: str | None = None,
) -> None:
    _save_message(message_id, conversation_id, user_id, role, content, created_at, image_url=image_url)


def get_conversation_messages(
    conversation_id: str, user_id: str
) -> list[dict[str, Any]]:
    return _get_conversation_messages(conversation_id, user_id)


def save_document(
    document_id: str,
    conversation_id: str,
    user_id: str,
    filename: str,
    content_type: str | None,
    size_bytes: int | None,
) -> None:
    _save_document(document_id, conversation_id, user_id, filename, content_type, size_bytes)


def get_documents(conversation_id: str, user_id: str) -> list[dict[str, Any]]:
    return _get_documents(conversation_id, user_id)


def delete_document(document_id: str, user_id: str) -> None:
    _delete_document(document_id, user_id)


def save_document_chunk(
    chunk_id: str,
    document_id: str,
    chunk_index: int,
    content: str,
    chunk_size: int,
) -> None:
    _save_document_chunk(chunk_id, document_id, chunk_index, content, chunk_size)


def get_document_chunks(document_id: str) -> list[dict[str, Any]]:
    return _get_document_chunks(document_id)