File size: 4,361 Bytes
71c916b
f1fd3e0
 
 
765ede8
 
ffc2ed9
765ede8
c1ca5a1
71c916b
 
765ede8
72ba547
f1fd3e0
765ede8
215277a
f1fd3e0
765ede8
 
 
ffc2ed9
f1fd3e0
 
 
 
 
 
 
 
ffc2ed9
f1fd3e0
ffc2ed9
95989dc
 
 
 
f1fd3e0
 
 
62d5359
 
fdd2048
95989dc
6863650
95989dc
 
6863650
 
95989dc
6863650
 
 
 
 
 
 
95989dc
fdd2048
 
95989dc
fdd2048
95989dc
fdd2048
 
 
 
ffc2ed9
fdd2048
 
 
95989dc
fdd2048
ffc2ed9
fdd2048
 
 
 
f1fd3e0
fdd2048
 
95989dc
 
fdd2048
 
 
 
c9eadbe
fdd2048
95989dc
 
 
 
fdd2048
 
ffc2ed9
fdd2048
 
 
ffc2ed9
fdd2048
 
ffc2ed9
fdd2048
 
 
95989dc
fdd2048
 
 
 
62d5359
fdd2048
 
 
 
 
 
 
ffc2ed9
fdd2048
 
 
 
ffc2ed9
fdd2048
 
 
 
 
 
 
 
 
 
 
 
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
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
import os
import shutil
import gradio as gr
import qdrant_client
from getpass import getpass


openai_api_key = os.getenv('OPENAI_API_KEY')

from llama_index.llms.openai import OpenAI
from llama_index.embeddings.openai import OpenAIEmbedding
from llama_index.core import Settings

Settings.llm = OpenAI(model="gpt-3.5-turbo", temperature=0.4)
Settings.embed_model = OpenAIEmbedding(model="text-embedding-ada-002")

from llama_index.core import SimpleDirectoryReader, VectorStoreIndex, StorageContext
from llama_index.vector_stores.qdrant import QdrantVectorStore
from llama_index.core.memory import ChatMemoryBuffer


chat_engine = None
index = None
query_engine = None
memory = None
client = None
vector_store = None
storage_context = None


def process_upload(files):

    upload_dir = "uploaded_files"
    if not os.path.exists(upload_dir):
        os.makedirs(upload_dir)
    
    for file_path in files:
        file_name = os.path.basename(file_path)
        dest = os.path.join(upload_dir, file_name)
        if not os.path.exists(dest):
            shutil.copy(file_path, dest)
    
    documents = SimpleDirectoryReader(upload_dir).load_data()
    
    global client, vector_store, storage_context, index, query_engine, memory, chat_engine
    client = qdrant_client.QdrantClient(location=":memory:")
    
    vector_store = QdrantVectorStore(
        collection_name="paper",
        client=client,
        enable_hybrid=True,
        batch_size=20,
    )
    
    storage_context = StorageContext.from_defaults(vector_store=vector_store)
    
    index = VectorStoreIndex.from_documents(documents, storage_context=storage_context)
    
    query_engine = index.as_query_engine(vector_store_query_mode="hybrid")
    
    memory = ChatMemoryBuffer.from_defaults(token_limit=3000)
    
    chat_engine = index.as_chat_engine(
        chat_mode="context",
        memory=memory,
        system_prompt=(
            "You are an AI assistant who answers the user questions"
        ),
    )
    
    return "Documents uploaded and index built successfully!"


def chat_with_ai(user_input, chat_history):
    global chat_engine
    if chat_engine is None:
        return chat_history, "Please upload documents first."
    
    response = chat_engine.chat(user_input)
    references = response.source_nodes
    ref, pages = [], []
    
    for node in references:
        file_name = node.metadata.get('file_name')
        if file_name and file_name not in ref:
            ref.append(file_name)
    
    complete_response = str(response) + "\n\n"
    if ref or pages:
        chat_history.append((user_input, complete_response))
    else:
        chat_history.append((user_input, str(response)))
    return chat_history, ""


def clear_history():
    return [], ""


def gradio_interface():
    with gr.Blocks() as demo:
        gr.Markdown("# AI Assistant")
        
        with gr.Tab("Upload Documents"):
            gr.Markdown("Upload PDF, Excel, CSV, DOC/DOCX, or TXT files below:")
            # The file upload widget: we specify allowed file types.
            file_upload = gr.File(
                label="Upload Files",
                file_count="multiple",
                file_types=[".pdf", ".csv", ".txt", ".xlsx", ".xls", ".doc", ".docx"],
                type="filepath"
            )
            upload_status = gr.Textbox(label="Upload Status", interactive=False)
            upload_button = gr.Button("Process Upload")
            
            upload_button.click(process_upload, inputs=file_upload, outputs=upload_status)
        
        with gr.Tab("Chat"):
            chatbot = gr.Chatbot(label="Chatbot Assistant")
            user_input = gr.Textbox(
                placeholder="Ask a question...", label="Enter your question"
            )
            submit_button = gr.Button("Send")
            btn_clear = gr.Button("Restart")
            
            # A State to hold the chat history.
            chat_history = gr.State([])
            
            submit_button.click(chat_with_ai, inputs=[user_input, chat_history], outputs=[chatbot, user_input])
            user_input.submit(chat_with_ai, inputs=[user_input, chat_history], outputs=[chatbot, user_input])
            btn_clear.click(clear_history, outputs=[chatbot, user_input])
    
    return demo

# Launch the Gradio app.
gradio_interface().launch(debug=True)