File size: 4,114 Bytes
a2aaac9
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
import os
import gradio as gr
from groq import Groq
from fpdf import FPDF

from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.document_loaders import PyPDFLoader, TextLoader
from langchain_community.vectorstores import FAISS
from langchain_community.embeddings import HuggingFaceEmbeddings

# -------------------------
# GROQ CLIENT
# -------------------------
client = Groq(api_key=os.environ.get("RAG_API_KEY"))

# -------------------------
# GLOBAL STORAGE
# -------------------------
vectorstore = None
chat_history = []

# -------------------------
# LOAD DOCUMENTS
# -------------------------
def load_documents(files):
    documents = []
    for file in files:
        if file.name.endswith(".pdf"):
            loader = PyPDFLoader(file.name)
        else:
            loader = TextLoader(file.name)
        documents.extend(loader.load())
    return documents

# -------------------------
# BUILD VECTOR STORE
# -------------------------
def build_vectorstore(files):
    global vectorstore
    docs = load_documents(files)

    splitter = RecursiveCharacterTextSplitter(
        chunk_size=500,
        chunk_overlap=100
    )
    chunks = splitter.split_documents(docs)

    embeddings = HuggingFaceEmbeddings(
        model_name="sentence-transformers/all-MiniLM-L6-v2"
    )

    vectorstore = FAISS.from_documents(chunks, embeddings)
    return "βœ… Documents processed successfully."

# -------------------------
# ASK QUESTION
# -------------------------
def ask_question(question):
    global chat_history, vectorstore

    if vectorstore is None:
        return "❌ Please upload documents first.", ""

    docs = vectorstore.similarity_search(question, k=4)
    context = "\n\n".join([d.page_content for d in docs])

    prompt = f"""
You are a helpful AI assistant.

First, use the document context below to answer.
If the document does not fully answer the question,
you may add relevant general knowledge.

Document Context:
{context}

Question:
{question}
"""

    response = client.chat.completions.create(
        model="llama-3.3-70b-versatile",
        messages=[{"role": "user", "content": prompt}],
    )

    answer = response.choices[0].message.content

    chat_history.append((question, answer))

    sources = "\n\n".join(
        [f"Source {i+1}:\n{d.page_content[:300]}..." for i, d in enumerate(docs)]
    )

    return answer, sources

# -------------------------
# EXPORT CHAT TO PDF
# -------------------------
def export_chat_pdf():
    pdf = FPDF()
    pdf.add_page()
    pdf.set_font("Arial", size=12)

    pdf.cell(0, 10, "AneesLLM - Chat Export", ln=True)
    pdf.ln(5)

    for q, a in chat_history:
        pdf.multi_cell(0, 8, f"Q: {q}")
        pdf.multi_cell(0, 8, f"A: {a}")
        pdf.ln(4)

    path = "/tmp/AneesLLM_Chat.pdf"
    pdf.output(path)
    return path

# -------------------------
# GRADIO UI
# -------------------------
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown(
        """
        # πŸ€– AneesLLM – RAG Based AI Assistant  
        Upload documents and ask intelligent questions with sources.
        """
    )

    with gr.Row():
        file_input = gr.File(
            file_types=[".pdf", ".txt"],
            file_count="multiple",
            label="πŸ“„ Upload Documents"
        )
        process_btn = gr.Button("πŸ“š Process Documents")

    status = gr.Textbox(label="Status")

    process_btn.click(
        build_vectorstore,
        inputs=file_input,
        outputs=status
    )

    gr.Markdown("## πŸ’¬ Ask a Question")

    question = gr.Textbox(
        placeholder="Ask something about your documents...",
        label="Your Question"
    )

    ask_btn = gr.Button("Ask")

    answer = gr.Textbox(label="Answer", lines=6)
    sources = gr.Textbox(label="Sources Used", lines=6)

    ask_btn.click(
        ask_question,
        inputs=question,
        outputs=[answer, sources]
    )

    export_btn = gr.Button("πŸ“„ Export Chat as PDF")
    pdf_file = gr.File(label="Download PDF")

    export_btn.click(export_chat_pdf, outputs=pdf_file)

demo.launch()