task/callbacks.py at main · Abdulrahman3119/task · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
import chainlit as cl
from typing import Dict, Any
from utils import load_pdf_and_create_vectorstore
from agent_setup import setup_agent

@cl.on_chat_start
async def on_chat_start():
    cl.user_session.set("chat_history", [])
    files = None
    while not files:
        files = await cl.AskFileMessage(
            content="Please upload a PDF file to begin!",
            accept=["application/pdf"],
            max_size_mb=20
        ).send()
    pdf_file = files[0]
    vectorstore = load_pdf_and_create_vectorstore(pdf_file.path, cl.user_session.get("id"))
    cl.user_session.set("vectorstore", vectorstore)
    agent = setup_agent(vectorstore)
    cl.user_session.set("agent", agent)
    cl.user_session.set("last_file", {"name": pdf_file.name, "path": pdf_file.path})
    await cl.Message(content=f"✅ `{pdf_file.name}` processed! Ask me anything about the document.").send()

@cl.on_chat_resume
async def on_chat_resume(thread: Dict[str, Any]):
    chat_history = []
    for message in thread.get("steps", []):
        role = "user" if message.get("type") == "user_message" else "assistant"
        chat_history.append({"role": role, "content": message.get("output", "")})
    cl.user_session.set("chat_history", chat_history)
    metadata = thread.get("metadata", {})
    last_file = metadata.get("last_file")
    if last_file:
        try:
            vectorstore = load_pdf_and_create_vectorstore(last_file["path"], cl.user_session.get("id"))
            cl.user_session.set("vectorstore", vectorstore)
            agent = setup_agent(vectorstore)
            cl.user_session.set("agent", agent)
            await cl.Message(content=f"✅ Session restored with document: `{last_file['name']}`").send()
        except Exception as e:
            await cl.Message(content=f"⚠️ Error restoring session: {str(e)}").send()
            await on_chat_start()
    else:
        await cl.Message(content="No previous document found. Please upload a new one.").send()
        await on_chat_start()

@cl.on_message
async def on_message(message: cl.Message):
    agent = cl.user_session.get("agent")
    if not agent:
        await cl.Message(content="Agent not initialized. Please upload a document.").send()
        return
    chat_history = cl.user_session.get("chat_history", [])
    chat_history.append({"role": "user", "content": message.content})
    try:
        with cl.Step(name="Agent", type="tool") as step:
            response = await agent.ainvoke({
                "input": message.content,
                "chat_history": chat_history
            })
            step.output = response
        if isinstance(response, dict):
            response = response.get("output", "")
        elif not isinstance(response, str):
            response = str(response)
        chat_history.append({"role": "assistant", "content": response})
        cl.user_session.set("chat_history", chat_history)
    except Exception as e:
        response = f"An error occurred: {str(e)}"
    await cl.Message(content=response).send()
    metadata = {
        "last_file": cl.user_session.get("last_file"),
        "chat_history": chat_history
    }
    cl.user_session.set("metadata", metadata)