from celery import shared_task
from langchain.text_splitter import RecursiveCharacterTextSplitter
from shared.models.opencopilot_db.pdf_data_sources import (
    insert_pdf_data_source,
    update_pdf_data_source_status,
)

from langchain.document_loaders import UnstructuredMarkdownLoader
from shared.utils.opencopilot_utils import (
    get_embeddings,
    StoreOptions,
    get_file_path,
)
from shared.utils.opencopilot_utils.init_vector_store import init_vector_store
from workers.utils.remove_escape_sequences import remove_escape_sequences


@shared_task
def process_markdown(file_name: str, bot_id: str):
    try:
        insert_pdf_data_source(chatbot_id=bot_id, file_name=file_name, status="PENDING")
        loader = UnstructuredMarkdownLoader(get_file_path(file_name))
        raw_docs = loader.load()
        for doc in raw_docs:
            doc.page_content = remove_escape_sequences(doc.page_content)

        text_splitter = RecursiveCharacterTextSplitter(
            chunk_size=1000, chunk_overlap=200, length_function=len
        )
        docs = text_splitter.split_documents(raw_docs)
        embeddings = get_embeddings()
        init_vector_store(
            docs,
            StoreOptions(namespace="knowledgebase", metadata={"bot_id": bot_id}),
        )

        update_pdf_data_source_status(
            chatbot_id=bot_id, file_name=file_name, status="COMPLETED"
        )
    except Exception as e:
        update_pdf_data_source_status(
            chatbot_id=bot_id, file_name=file_name, status="FAILED"
        )
        print(f"Error processing {file_name}:", e)


@shared_task
def retry_failed_markdown_crawl(chatbot_id: str, file_name: str):
    """Re-runs a failed PDF crawl.

    Args:
      chatbot_id: The ID of the chatbot.
      file_name: The name of the PDF file to crawl.
    """

    update_pdf_data_source_status(
        chatbot_id=chatbot_id, file_name=file_name, status="PENDING"
    )
    try:
        process_markdown(file_name=file_name, bot_id=chatbot_id)
    except Exception as e:
        update_pdf_data_source_status(
            chatbot_id=chatbot_id, file_name=file_name, status="FAILED"
        )
        print(f"Error reprocessing {file_name}:", e)

langchain.text_splitter.RecursiveCharacterTextSplitter
[((1830, 1925), 'shared.models.opencopilot_db.pdf_data_sources.update_pdf_data_source_status', 'update_pdf_data_source_status', ([], {'chatbot_id': 'chatbot_id', 'file_name': 'file_name', 'status': '"""PENDING"""'}), "(chatbot_id=chatbot_id, file_name=file_name,\n    status='PENDING')\n", (1859, 1925), False, 'from shared.models.opencopilot_db.pdf_data_sources import insert_pdf_data_source, update_pdf_data_source_status\n'), ((630, 715), 'shared.models.opencopilot_db.pdf_data_sources.insert_pdf_data_source', 'insert_pdf_data_source', ([], {'chatbot_id': 'bot_id', 'file_name': 'file_name', 'status': '"""PENDING"""'}), "(chatbot_id=bot_id, file_name=file_name, status='PENDING'\n    )\n", (652, 715), False, 'from shared.models.opencopilot_db.pdf_data_sources import insert_pdf_data_source, update_pdf_data_source_status\n'), ((941, 1032), 'langchain.text_splitter.RecursiveCharacterTextSplitter', 'RecursiveCharacterTextSplitter', ([], {'chunk_size': '(1000)', 'chunk_overlap': '(200)', 'length_function': 'len'}), '(chunk_size=1000, chunk_overlap=200,\n    length_function=len)\n', (971, 1032), False, 'from langchain.text_splitter import RecursiveCharacterTextSplitter\n'), ((1127, 1143), 'shared.utils.opencopilot_utils.get_embeddings', 'get_embeddings', ([], {}), '()\n', (1141, 1143), False, 'from shared.utils.opencopilot_utils import get_embeddings, StoreOptions, get_file_path\n'), ((1290, 1383), 'shared.models.opencopilot_db.pdf_data_sources.update_pdf_data_source_status', 'update_pdf_data_source_status', ([], {'chatbot_id': 'bot_id', 'file_name': 'file_name', 'status': '"""COMPLETED"""'}), "(chatbot_id=bot_id, file_name=file_name,\n    status='COMPLETED')\n", (1319, 1383), False, 'from shared.models.opencopilot_db.pdf_data_sources import insert_pdf_data_source, update_pdf_data_source_status\n'), ((755, 779), 'shared.utils.opencopilot_utils.get_file_path', 'get_file_path', (['file_name'], {}), '(file_name)\n', (768, 779), False, 'from shared.utils.opencopilot_utils import get_embeddings, StoreOptions, get_file_path\n'), ((874, 915), 'workers.utils.remove_escape_sequences.remove_escape_sequences', 'remove_escape_sequences', (['doc.page_content'], {}), '(doc.page_content)\n', (897, 915), False, 'from workers.utils.remove_escape_sequences import remove_escape_sequences\n'), ((1201, 1269), 'shared.utils.opencopilot_utils.StoreOptions', 'StoreOptions', ([], {'namespace': '"""knowledgebase"""', 'metadata': "{'bot_id': bot_id}"}), "(namespace='knowledgebase', metadata={'bot_id': bot_id})\n", (1213, 1269), False, 'from shared.utils.opencopilot_utils import get_embeddings, StoreOptions, get_file_path\n'), ((1437, 1527), 'shared.models.opencopilot_db.pdf_data_sources.update_pdf_data_source_status', 'update_pdf_data_source_status', ([], {'chatbot_id': 'bot_id', 'file_name': 'file_name', 'status': '"""FAILED"""'}), "(chatbot_id=bot_id, file_name=file_name,\n    status='FAILED')\n", (1466, 1527), False, 'from shared.models.opencopilot_db.pdf_data_sources import insert_pdf_data_source, update_pdf_data_source_status\n'), ((2045, 2139), 'shared.models.opencopilot_db.pdf_data_sources.update_pdf_data_source_status', 'update_pdf_data_source_status', ([], {'chatbot_id': 'chatbot_id', 'file_name': 'file_name', 'status': '"""FAILED"""'}), "(chatbot_id=chatbot_id, file_name=file_name,\n    status='FAILED')\n", (2074, 2139), False, 'from shared.models.opencopilot_db.pdf_data_sources import insert_pdf_data_source, update_pdf_data_source_status\n')]