Spaces:
Sleeping
Sleeping
Update app.py
Browse files
app.py
CHANGED
@@ -1,125 +1,189 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
import os
|
|
|
|
|
|
|
2 |
import gradio as gr
|
|
|
|
|
|
|
3 |
from llama_index.core import SimpleDirectoryReader, VectorStoreIndex
|
4 |
from llama_index.embeddings.mixedbreadai import MixedbreadAIEmbedding
|
5 |
from llama_index.llms.groq import Groq
|
6 |
from llama_parse import LlamaParse
|
7 |
|
8 |
-
#
|
9 |
-
# 1.
|
10 |
-
#
|
11 |
-
|
12 |
-
|
13 |
-
|
14 |
|
15 |
-
if not (
|
16 |
raise EnvironmentError(
|
17 |
-
"LLAMA_CLOUD_API_KEY, GROQ_API_KEY and MXBAI_API_KEY must be set."
|
18 |
)
|
19 |
|
20 |
-
#
|
21 |
-
# 2.
|
22 |
-
#
|
23 |
-
LLM_MODEL = "llama-3.1-70b-versatile"
|
24 |
-
EMBED_MODEL = "mixedbread-ai/mxbai-embed-large-v1"
|
|
|
|
|
25 |
|
26 |
-
|
27 |
-
file_extractor = {ext: parser for ext in (
|
28 |
".pdf", ".docx", ".doc", ".txt", ".csv", ".xlsx",
|
29 |
".pptx", ".html", ".jpg", ".jpeg", ".png", ".webp", ".svg",
|
30 |
-
)
|
31 |
-
|
32 |
-
|
33 |
-
|
34 |
-
|
35 |
-
|
36 |
-
|
37 |
-
|
38 |
-
|
39 |
-
|
40 |
-
|
41 |
-
|
42 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
43 |
global vector_index
|
44 |
-
if
|
45 |
return "β οΈ No file selected."
|
46 |
|
47 |
-
if
|
48 |
-
|
49 |
-
|
50 |
|
51 |
docs = SimpleDirectoryReader(
|
52 |
-
input_files=[
|
|
|
53 |
).load_data()
|
54 |
|
55 |
-
|
56 |
-
|
|
|
|
|
|
|
|
|
57 |
|
58 |
|
59 |
-
def respond(message: str, history:
|
60 |
-
"""
|
61 |
if vector_index is None:
|
62 |
-
|
|
|
63 |
|
64 |
-
query_engine
|
65 |
-
|
66 |
|
67 |
partial = ""
|
68 |
-
for
|
69 |
-
partial +=
|
70 |
-
yield partial
|
71 |
|
72 |
|
73 |
-
def
|
74 |
-
"""Reset everything."""
|
75 |
global vector_index
|
76 |
vector_index = None
|
77 |
-
return
|
78 |
|
79 |
|
80 |
-
#
|
81 |
-
# 4. Gradio UI
|
82 |
-
#
|
83 |
with gr.Blocks(
|
84 |
theme=gr.themes.Default(
|
85 |
primary_hue="green",
|
86 |
secondary_hue="blue",
|
87 |
-
font=[gr.themes.GoogleFont("Poppins")]
|
88 |
),
|
89 |
-
css="footer {visibility:
|
90 |
) as demo:
|
91 |
|
92 |
-
gr.Markdown("<h1 style='text-align:center'>DataCamp Doc Q&A
|
|
|
93 |
with gr.Row():
|
94 |
with gr.Column(scale=1):
|
95 |
-
file_input = gr.File(
|
96 |
-
|
97 |
-
|
|
|
|
|
|
|
|
|
98 |
with gr.Row():
|
99 |
-
|
100 |
-
clear_btn = gr.Button("Clear")
|
101 |
-
status_box = gr.Markdown()
|
102 |
|
103 |
with gr.Column(scale=3):
|
104 |
-
|
105 |
-
|
106 |
-
|
107 |
-
|
108 |
-
|
109 |
-
placeholder="Ask a question about the uploaded documentβ¦",
|
110 |
-
container=False,
|
111 |
-
),
|
112 |
)
|
|
|
113 |
|
114 |
-
|
115 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
116 |
|
117 |
-
#
|
118 |
demo.queue(api_open=False)
|
119 |
|
120 |
-
#
|
121 |
# 5. Launch
|
122 |
-
#
|
123 |
if __name__ == "__main__":
|
124 |
-
# β¦and make a public share link so the container doesnβt choke on localhost
|
125 |
demo.launch(share=True, server_name="0.0.0.0", server_port=7860)
|
|
|
1 |
+
"""
|
2 |
+
Doc-Q&A app (Gradio 5.x + Llama-Index 0.12.x, June 2025)
|
3 |
+
|
4 |
+
Key upgrades
|
5 |
+
------------
|
6 |
+
βͺ Gradio 5.34βnew event system (`upload`, `clear` etc.)
|
7 |
+
βͺ Llama-Index 0.12.42β`VectorStoreIndex.from_documents` signature unchanged
|
8 |
+
βͺ MixedbreadAIEmbedding 0.3.0 β supports `batch_size`, `timeout`
|
9 |
+
βͺ Tenacity for exponential-back-off when MXBAI returns 5xx / rate limits
|
10 |
+
"""
|
11 |
+
|
12 |
+
from __future__ import annotations
|
13 |
+
|
14 |
import os
|
15 |
+
from pathlib import Path
|
16 |
+
from typing import List
|
17 |
+
|
18 |
import gradio as gr
|
19 |
+
from tenacity import retry, wait_exponential, stop_after_attempt
|
20 |
+
from mixedbread_ai.core.api_error import ApiError
|
21 |
+
|
22 |
from llama_index.core import SimpleDirectoryReader, VectorStoreIndex
|
23 |
from llama_index.embeddings.mixedbreadai import MixedbreadAIEmbedding
|
24 |
from llama_index.llms.groq import Groq
|
25 |
from llama_parse import LlamaParse
|
26 |
|
27 |
+
# ββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
28 |
+
# 1. Environment variables (fail-fast if missing)
|
29 |
+
# ββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
30 |
+
LLAMA_CLOUD_API_KEY = os.getenv("LLAMA_CLOUD_API_KEY")
|
31 |
+
GROQ_API_KEY = os.getenv("GROQ_API_KEY")
|
32 |
+
MXBAI_API_KEY = os.getenv("MXBAI_API_KEY")
|
33 |
|
34 |
+
if not all([LLAMA_CLOUD_API_KEY, GROQ_API_KEY, MXBAI_API_KEY]):
|
35 |
raise EnvironmentError(
|
36 |
+
"LLAMA_CLOUD_API_KEY, GROQ_API_KEY and MXBAI_API_KEY must be set in the env."
|
37 |
)
|
38 |
|
39 |
+
# ββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
40 |
+
# 2. Models & parsers (latest defaults - June 2025)
|
41 |
+
# ββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
42 |
+
LLM_MODEL = "llama-3.1-70b-versatile" # Groqβs best for Q&A
|
43 |
+
EMBED_MODEL = "mixedbread-ai/mxbai-embed-large-v1" # 1024-dim
|
44 |
+
|
45 |
+
parser = LlamaParse(api_key=LLAMA_CLOUD_API_KEY, result_type="markdown")
|
46 |
|
47 |
+
SUPPORTED_EXTS = (
|
|
|
48 |
".pdf", ".docx", ".doc", ".txt", ".csv", ".xlsx",
|
49 |
".pptx", ".html", ".jpg", ".jpeg", ".png", ".webp", ".svg",
|
50 |
+
)
|
51 |
+
file_extractor = {ext: parser for ext in SUPPORTED_EXTS}
|
52 |
+
|
53 |
+
embed_model = MixedbreadAIEmbedding(
|
54 |
+
api_key = MXBAI_API_KEY,
|
55 |
+
model_name = EMBED_MODEL,
|
56 |
+
batch_size = 8, # keep requests < 100 KB
|
57 |
+
timeout = 60, # generous server-side processing window
|
58 |
+
)
|
59 |
+
|
60 |
+
llm = Groq(model=LLM_MODEL, api_key=GROQ_API_KEY)
|
61 |
+
|
62 |
+
# A simple global cache (could be swapped for Redis, etc.)
|
63 |
+
vector_index: VectorStoreIndex | None = None
|
64 |
+
|
65 |
+
|
66 |
+
# ββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
67 |
+
# 3. Helper wrappers
|
68 |
+
# ββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
69 |
+
@retry(
|
70 |
+
wait=wait_exponential(multiplier=2, min=4, max=32),
|
71 |
+
stop=stop_after_attempt(4),
|
72 |
+
retry_error_callback=lambda retry_state: None, # bubble up as None
|
73 |
+
reraise=False,
|
74 |
+
)
|
75 |
+
def _safe_build_index(docs) -> VectorStoreIndex | None:
|
76 |
+
"""Retry MXBAI 503 / 429 transparently."""
|
77 |
+
try:
|
78 |
+
return VectorStoreIndex.from_documents(docs, embed_model=embed_model)
|
79 |
+
except ApiError as e:
|
80 |
+
# Tenacity will catch and retry unless non-5xx / non-429
|
81 |
+
if e.status_code not in (429, 500, 502, 503, 504):
|
82 |
+
raise
|
83 |
+
raise # trigger retry
|
84 |
+
|
85 |
+
|
86 |
+
def load_files(file: Path | None) -> str:
|
87 |
+
"""Parse uploaded file and build vector index (with retries)."""
|
88 |
global vector_index
|
89 |
+
if file is None:
|
90 |
return "β οΈ No file selected."
|
91 |
|
92 |
+
if file.suffix.lower() not in SUPPORTED_EXTS:
|
93 |
+
allow = ", ".join(SUPPORTED_EXTS)
|
94 |
+
return f"β οΈ Unsupported file type. Allowed: {allow}"
|
95 |
|
96 |
docs = SimpleDirectoryReader(
|
97 |
+
input_files=[str(file)],
|
98 |
+
file_extractor=file_extractor,
|
99 |
).load_data()
|
100 |
|
101 |
+
idx = _safe_build_index(docs)
|
102 |
+
if idx is None:
|
103 |
+
return "π§ Embedding service busy. Please retry in ~1 minute."
|
104 |
+
|
105 |
+
vector_index = idx
|
106 |
+
return f"β
Parsed **{file.name}** β you can start chatting!"
|
107 |
|
108 |
|
109 |
+
def respond(message: str, history: List[List[str]]):
|
110 |
+
"""Stream answer chunks to the Chatbot."""
|
111 |
if vector_index is None:
|
112 |
+
yield "β‘οΈ Please upload a document first."
|
113 |
+
return
|
114 |
|
115 |
+
query_engine = vector_index.as_query_engine(streaming=True, llm=llm)
|
116 |
+
response = query_engine.query(message)
|
117 |
|
118 |
partial = ""
|
119 |
+
for token in response.response_gen:
|
120 |
+
partial += token
|
121 |
+
yield partial
|
122 |
|
123 |
|
124 |
+
def clear():
|
125 |
+
"""Reset everything (file widget + status + index)."""
|
126 |
global vector_index
|
127 |
vector_index = None
|
128 |
+
return None, "", None # file_input, status_md, chatbot history
|
129 |
|
130 |
|
131 |
+
# ββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
132 |
+
# 4. Gradio UI (5.x syntax)
|
133 |
+
# ββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
134 |
with gr.Blocks(
|
135 |
theme=gr.themes.Default(
|
136 |
primary_hue="green",
|
137 |
secondary_hue="blue",
|
138 |
+
font=[gr.themes.GoogleFont("Poppins")],
|
139 |
),
|
140 |
+
css="footer {visibility:hidden}",
|
141 |
) as demo:
|
142 |
|
143 |
+
gr.Markdown("<h1 style='text-align:center'>DataCamp Doc Q&A π€π</h1>")
|
144 |
+
|
145 |
with gr.Row():
|
146 |
with gr.Column(scale=1):
|
147 |
+
file_input = gr.File(
|
148 |
+
label="Upload document",
|
149 |
+
file_count="single",
|
150 |
+
type="filepath",
|
151 |
+
show_label=True,
|
152 |
+
)
|
153 |
+
status_md = gr.Markdown()
|
154 |
with gr.Row():
|
155 |
+
clear_btn = gr.Button("Reset π", variant="secondary")
|
|
|
|
|
156 |
|
157 |
with gr.Column(scale=3):
|
158 |
+
chatbot = gr.Chatbot(height=340)
|
159 |
+
txt_box = gr.Textbox(
|
160 |
+
placeholder="Ask something about the uploaded documentβ¦",
|
161 |
+
container=False,
|
162 |
+
scale=7,
|
|
|
|
|
|
|
163 |
)
|
164 |
+
send_btn = gr.Button("Send", variant="primary")
|
165 |
|
166 |
+
# events (v5 style)
|
167 |
+
file_input.upload(
|
168 |
+
fn=load_files,
|
169 |
+
inputs=file_input,
|
170 |
+
outputs=status_md,
|
171 |
+
)
|
172 |
+
send_btn.click(
|
173 |
+
fn=respond,
|
174 |
+
inputs=[txt_box, chatbot],
|
175 |
+
outputs=chatbot,
|
176 |
+
)
|
177 |
+
clear_btn.click(
|
178 |
+
fn=clear,
|
179 |
+
outputs=[file_input, status_md, chatbot],
|
180 |
+
)
|
181 |
|
182 |
+
# optional: disable public OpenAPI schema (old crash guard)
|
183 |
demo.queue(api_open=False)
|
184 |
|
185 |
+
# ββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
186 |
# 5. Launch
|
187 |
+
# ββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
188 |
if __name__ == "__main__":
|
|
|
189 |
demo.launch(share=True, server_name="0.0.0.0", server_port=7860)
|