kingabzpro commited on
Commit
f3913f7
Β·
verified Β·
1 Parent(s): b34b849

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +136 -72
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. Check environment variables
10
- # ────────────────────────────────
11
- llama_cloud_key = os.getenv("LLAMA_CLOUD_API_KEY")
12
- groq_key = os.getenv("GROQ_API_KEY")
13
- mxbai_key = os.getenv("MXBAI_API_KEY")
14
 
15
- if not (llama_cloud_key and groq_key and mxbai_key):
16
  raise EnvironmentError(
17
- "LLAMA_CLOUD_API_KEY, GROQ_API_KEY and MXBAI_API_KEY must be set."
18
  )
19
 
20
- # ────────────────────────────────
21
- # 2. Model / parser setup
22
- # ────────────────────────────────
23
- LLM_MODEL = "llama-3.1-70b-versatile"
24
- EMBED_MODEL = "mixedbread-ai/mxbai-embed-large-v1"
 
 
25
 
26
- parser = LlamaParse(api_key=llama_cloud_key, result_type="markdown")
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
- embed_model = MixedbreadAIEmbedding(api_key=mxbai_key, model_name=EMBED_MODEL)
32
- llm = Groq(model=LLM_MODEL, api_key=groq_key)
33
-
34
- # Global cache for the current document
35
- vector_index = None
36
-
37
-
38
- # ────────────────────────────────
39
- # 3. Helper functions
40
- # ────────────────────────────────
41
- def load_files(file_path: str) -> str:
42
- """Parse the uploaded document and build a vector index."""
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
43
  global vector_index
44
- if not file_path:
45
  return "⚠️ No file selected."
46
 
47
- if not any(file_path.endswith(ext) for ext in file_extractor):
48
- return ("⚠️ Unsupported file type. "
49
- f"Allowed: {', '.join(file_extractor.keys())}")
50
 
51
  docs = SimpleDirectoryReader(
52
- input_files=[file_path], file_extractor=file_extractor
 
53
  ).load_data()
54
 
55
- vector_index = VectorStoreIndex.from_documents(docs, embed_model=embed_model)
56
- return f"βœ… Parsed **{os.path.basename(file_path)}**. Ask away!"
 
 
 
 
57
 
58
 
59
- def respond(message: str, history: list) -> str:
60
- """Chat handler. Streams partial tokens back to the UI."""
61
  if vector_index is None:
62
- return "➑️ Please upload a document first."
 
63
 
64
- query_engine = vector_index.as_query_engine(streaming=True, llm=llm)
65
- streaming_resp = query_engine.query(message)
66
 
67
  partial = ""
68
- for chunk in streaming_resp.response_gen:
69
- partial += chunk
70
- yield partial # <─ streaming to the frontend
71
 
72
 
73
- def clear_state():
74
- """Reset everything."""
75
  global vector_index
76
  vector_index = None
77
- return [None, ""]
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: hidden}"
90
  ) as demo:
91
 
92
- gr.Markdown("<h1 style='text-align:center'>DataCamp Doc Q&A πŸ€–πŸ“ƒ</h1>")
 
93
  with gr.Row():
94
  with gr.Column(scale=1):
95
- file_input = gr.File(file_count="single",
96
- type="filepath",
97
- label="Upload document")
 
 
 
 
98
  with gr.Row():
99
- submit_btn = gr.Button("Submit", variant="primary")
100
- clear_btn = gr.Button("Clear")
101
- status_box = gr.Markdown()
102
 
103
  with gr.Column(scale=3):
104
- chat = gr.ChatInterface(
105
- fn=respond,
106
- chatbot=gr.Chatbot(height=300),
107
- show_progress="full", # keep the nice progress bar
108
- textbox=gr.Textbox(
109
- placeholder="Ask a question about the uploaded document…",
110
- container=False,
111
- ),
112
  )
 
113
 
114
- submit_btn.click(load_files, inputs=file_input, outputs=status_box)
115
- clear_btn.click(clear_state, outputs=[file_input, status_box])
 
 
 
 
 
 
 
 
 
 
 
 
 
116
 
117
- # Disable OpenAPI generation (avoids the bool/β€˜const’ bug) …
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)