Update app.py

app.py

@@ -15,20 +15,24 @@ from langchain_huggingface import HuggingFaceEmbeddings
 from langchain_community.vectorstores import Chroma
 
 # ==========================================
-# 2. CONNECT TO AI APIS (OpenAI-Compatible)
+# 2. CONNECT TO AI APIS (Replaces Local Models)
 # ==========================================
 print("⏳ Connecting to Hugging Face APIs...")
 
+# Get token from environment variable (Set this in HF Spaces Secrets)
 HF_TOKEN = os.environ.get("HF_TOKEN")
 
 if not HF_TOKEN:
     print("⚠️ WARNING: HF_TOKEN not found! The AI will not be able to generate responses.")
 
-# Initialize a single, generic client 
-# We do NOT bind the model name here to prevent the "text-generation" tagging error
-hf_client = InferenceClient(api_key=HF_TOKEN)
+# --- A. Mistral-7B API (The Writer/Scientist) ---
+text_client = InferenceClient("mistralai/Mistral-7B-Instruct-v0.3", token=HF_TOKEN)
 
-# Local Embeddings & Whisper
+# --- B. Qwen2-VL API (The Eye) ---
+# We use the 7B version since the cloud API handles the compute!
+vision_client = InferenceClient("Qwen/Qwen2-VL-7B-Instruct", token=HF_TOKEN)
+
+# --- C. Local Embeddings & Whisper (Runs fine on CPU) ---
 embedding_model = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
 whisper_model = whisper.load_model("base")
 
@@ -55,6 +59,7 @@ def save_to_json(user_msg, combined_ans, mode):
     with open(chat_history_file, "w") as f: json.dump(history, f, indent=4)
 
 def process_pdf_to_markdown(pdf_path, extract_images=True):
+    """Converts PDF to Markdown. Optionally extracts images."""
     global main_extracted_images
     output_image_dir = "extracted_images"
 
@@ -75,6 +80,7 @@ def process_pdf_to_markdown(pdf_path, extract_images=True):
     except Exception as e:
         return ""
 
+# --- UPLOAD HANDLER 1: MAIN PAPER ---
 def process_main_paper(file_obj):
     global main_paper_retriever
     main_paper_retriever = None
@@ -91,6 +97,7 @@ def process_main_paper(file_obj):
     except Exception as e:
         return f"❌ Error: {str(e)}"
 
+# --- UPLOAD HANDLER 2: REFERENCE SET ---
 def process_brainstorm_papers(file_list):
     global brainstorm_retriever
     brainstorm_retriever = None
@@ -119,23 +126,27 @@ def transcribe_audio(audio_path):
     return whisper_model.transcribe(audio_path)["text"].strip()
 
 # ==========================================
-# 4. INTELLIGENCE LAYERS (STRICT CHAT ROUTING)
+# 4. INTELLIGENCE LAYERS (API WRAPPERS)
 # ==========================================
 
+# Helper function to call Mistral API
+# Helper function to call Mistral API (Updated for Conversational Task)
 def ask_mistral(prompt):
     try:
+        # Format the prompt as a chat message
         messages = [{"role": "user", "content": prompt}]
-        # This explicitly hits the /v1/chat/completions route
-        response = hf_client.chat.completions.create(
-            model="mistralai/Mistral-7B-Instruct-v0.3",
-            messages=messages,
-            max_tokens=1000,
+        
+        # Use chat_completion instead of text_generation
+        response = text_client.chat_completion(
+            messages=messages, 
+            max_tokens=1000, 
             temperature=0.3
         )
         return response.choices[0].message.content
     except Exception as e:
         return f"⚠️ API Error (Mistral): {str(e)}"
 
+# Helper function to call Qwen API
 def ask_qwen(prompt, image_paths):
     try:
         messages = [{"role": "user", "content": []}]
@@ -148,26 +159,23 @@ def ask_qwen(prompt, image_paths):
             })
         messages[0]["content"].append({"type": "text", "text": prompt})
         
-        # Enforce chat completions route for Vision model too
-        response = hf_client.chat.completions.create(
-            model="Qwen/Qwen2-VL-7B-Instruct",
-            messages=messages,
-            max_tokens=150
-        )
+        response = vision_client.chat_completion(messages=messages, max_tokens=150)
         return response.choices[0].message.content
     except Exception as e:
-        return f"⚠️ API Error (Qwen - Server busy): {str(e)}"
+        return f"⚠️ API Error (Qwen - Server might be busy): {str(e)}"
 
 # MODE 1: CHAT WITH MAIN PAPER
 def get_main_paper_response(question):
     global main_paper_retriever, main_extracted_images
     vision_context = ""
 
+    # Vision Pass
     if main_extracted_images:
         images_to_process = main_extracted_images[:3]
         vision_prompt = f"Relate these images to: {question}"
         vision_context = ask_qwen(vision_prompt, images_to_process)
 
+    # Text Pass
     if main_paper_retriever:
         docs = main_paper_retriever.invoke(question)
         text_context = "\n\n".join(d.page_content for d in docs)
@@ -220,6 +228,7 @@ Context: {context}
 Task: Generate a draft for a paper about: {question} [/INST]"""
     return ask_mistral(prompt)
 
+
 # ==========================================
 # 5. GRADIO UI
 # ==========================================
@@ -230,6 +239,7 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
     gr.Markdown("Pipeline: Analyze -> Find Novelty -> Design Setup -> Write Draft")
 
     with gr.Row():
+        # --- LEFT: CONTROLS ---
         with gr.Column(scale=1):
 
             mode_radio = gr.Radio(
@@ -254,6 +264,7 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
 
             clear_btn = gr.Button("🗑️ Clear Workspace")
 
+        # --- RIGHT: WORKSPACE ---
         with gr.Column(scale=2):
             chatbot = gr.Chatbot(label="Lab Assistant", height=700)
             audio_input = gr.Audio(sources=["microphone"], type="filepath", label="🎤 Dictate Idea")
@@ -262,15 +273,18 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
                 msg_input = gr.Textbox(placeholder="Enter your query or research topic...", scale=4)
                 send_btn = gr.Button("🚀 Execute", variant="primary", scale=1)
 
+    # --- HANDLERS ---
     file_main.change(fn=process_main_paper, inputs=file_main, outputs=status_main)
     file_refs.change(fn=process_brainstorm_papers, inputs=file_refs, outputs=status_refs)
     audio_input.stop_recording(fn=transcribe_audio, inputs=audio_input, outputs=msg_input)
     clear_btn.click(fn=reset_chat, outputs=chatbot)
 
+    # --- MAIN ROUTER ---
     def respond(message, history, mode):
         if not message.strip(): return "", history
         if history is None: history = []
 
+        # Route based on selected Stage
         if mode == "1. Chat with Paper":
             response = get_main_paper_response(message)
         elif mode == "2. Brainstorm Novelty":
@@ -282,6 +296,7 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
         else:
             response = "Error: Unknown Mode"
 
+        # Log & Update
         final_ans = f"**[{mode}]**\n{response}"
         save_to_json(message, final_ans, mode)
         history.append({"role": "user", "content": message})
@@ -293,4 +308,5 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
     send_btn.click(respond, [msg_input, chatbot, mode_radio], [msg_input, chatbot])
 
 print("🚀 Launching Production Research Scientist Lab...")
+# In HF Spaces, we don't need share=True
 demo.launch()