Update app.py

app.py

@@ -114,6 +114,16 @@ class Model:
         export_to_ply(images[0], ply_path.name)
         return self.to_glb(ply_path.name)
 
+# -----------------------------------------------------------------------------
+# Helper function for 3D generation using the Model class
+# -----------------------------------------------------------------------------
+
+def generate_3d_fn(prompt: str, seed: int, guidance_scale: float, num_steps: int, randomize_seed: bool):
+    seed = randomize_seed_fn(seed, randomize_seed)
+    model_3d = Model()
+    glb_path = model_3d.run_text(prompt, seed=seed, guidance_scale=guidance_scale, num_steps=num_steps)
+    return glb_path, seed
+
 # -----------------------------------------------------------------------------
 # Gradio UI configuration
 # -----------------------------------------------------------------------------
@@ -367,6 +377,9 @@ def generate(
       - "@3d": triggers 3D model generation using the ShapE pipeline.
       - "@web": triggers a web command. Use "visit" to visit a URL (e.g., "@web visit https://example.com")
                 or "search" to perform a DuckDuckGo search (e.g., "@web search AI news").
+    
+    Additionally, for every default (plain text) query (i.e. no special command), after the answer
+    the bot will append reference links for further reading.
     """
     text = input_dict["text"]
     files = input_dict.get("files", [])
@@ -389,7 +402,6 @@ def generate(
         new_filename = f"mesh_{uuid.uuid4()}.glb"
         new_filepath = os.path.join(static_folder, new_filename)
         shutil.copy(glb_path, new_filepath)
-        
         yield gr.File(new_filepath)
         return
 
@@ -477,6 +489,7 @@ def generate(
             time.sleep(0.01)
             yield buffer
     else:
+        # --- Default plain text branch ---
         input_ids = tokenizer.apply_chat_template(conversation, add_generation_prompt=True, return_tensors="pt")
         if input_ids.shape[1] > MAX_INPUT_TOKEN_LENGTH:
             input_ids = input_ids[:, -MAX_INPUT_TOKEN_LENGTH:]
@@ -498,6 +511,7 @@ def generate(
         t.start()
 
         outputs = []
+        yield "Thinking..."
         for new_text in streamer:
             outputs.append(new_text)
             yield "".join(outputs)
@@ -505,6 +519,16 @@ def generate(
         final_response = "".join(outputs)
         yield final_response
 
+        # --- Append Reference Links after the answer ---
+        try:
+            # Use the original query as the search term to fetch reference links (limit to 3 results)
+            search_tool = DuckDuckGoSearchTool(max_results=3)
+            reference_links = search_tool.forward(input_dict["text"])
+            reference_message = "\n\nFor more info, visit:\n" + reference_links
+            yield reference_message
+        except Exception as e:
+            yield f"\n\n[Error retrieving reference links: {str(e)}]"
+
         if is_tts and voice:
             output_file = asyncio.run(text_to_speech(final_response, voice))
             yield gr.Audio(output_file, autoplay=True)