Update app.py

app.py

@@ -121,16 +121,22 @@ def load_paper_summaries() -> Dict[str, str]:
         return {}
 
 
-def hf_inference(model_name, prompt, max_tokens=2000, retries=5):
+def hf_inference(model_name, prompt, max_tokens=2000, retries=5, stream=False):  # Added stream parameter
     for attempt in range(retries):
         try:
             messages = [{"role": "user", "content": prompt}]
-            response = client.chat.completions.create(
+            response_generator = client.chat.completions.create(
                 model=model_name,
                 messages=messages,
-                max_tokens=max_tokens
+                max_tokens=max_tokens,
+                stream=stream  # Pass the stream parameter
             )
-            return {"generated_text": response.choices[0].message.content}
+            if stream:
+                return response_generator # Return the generator directly
+            else:
+                # If not streaming, get the full response
+                response = next(response_generator) # Consume the first chunk to get complete object
+                return {"generated_text": response.choices[0].message.content}
         except Exception as e:
             if attempt == retries - 1:
                 logger.error(f"Request failed after {retries} retries: {e}")
@@ -138,34 +144,51 @@ def hf_inference(model_name, prompt, max_tokens=2000, retries=5):
             time.sleep(RETRY_DELAY * (1 + attempt))
     return {"error": "Request failed after multiple retries."}
 
-def ensemble_inference(prompt, models=ENSEMBLE_MODELS, max_tokens=1500):
+def ensemble_inference(prompt, models=ENSEMBLE_MODELS, max_tokens=1500, stream=False): #Added stream
     results = []
-    with ThreadPoolExecutor(max_workers=len(models)) as executor:
-        future_to_model = {executor.submit(hf_inference, model, prompt, max_tokens): model for model in models}
-        for future in as_completed(future_to_model):
-            model = future_to_model[future]
-            try:
-                result = future.result()
-                if "generated_text" in result:
-                    results.append({"model": model, "text": result["generated_text"]})
-            except Exception as e:
-                logger.error(f"Error with model {model}: {e}")
-
-    if not results:
-        return {"error": "All models failed to generate responses"}
 
-    if len(results) == 1:
-        return {"generated_text": results[0]["text"]}
-
-    synthesis_prompt = "Synthesize these expert responses into a single coherent answer:\n\n"
-    for result in results:
-        synthesis_prompt += f"Expert {results.index(result) + 1} ({result['model'].split('/')[-1]}):\n{result['text']}\n\n"
-
-    synthesis = hf_inference(MAIN_LLM_MODEL, synthesis_prompt) # Use a consistent model for final synthesis
-    if "generated_text" in synthesis:
-        return synthesis
-    else:
-        return {"generated_text": max(results, key=lambda x: len(x["text"]))["text"]} # Fallback
+    if stream: # If streaming, return a generator that yields from each model
+        def generate_responses():
+            with ThreadPoolExecutor(max_workers=len(models)) as executor:
+                futures = {executor.submit(hf_inference, model, prompt, max_tokens, stream=True): model for model in models}
+
+                for future in as_completed(futures):
+                    model = future_to_model[future]
+                    try:
+                        for chunk in future.result():  # Iterate through chunks
+                            yield {"model": model, "text": chunk.choices[0].delta.content} #yield the content of the chunk
+                    except Exception as e:
+                        logger.error(f"Error with model {model}: {e}")
+                        yield {"model": model, "text": f"Error: {e}"}
+        return generate_responses() # return the generator
+
+    else: #Non-streaming behavior
+        with ThreadPoolExecutor(max_workers=len(models)) as executor:
+            future_to_model = {executor.submit(hf_inference, model, prompt, max_tokens, stream=False): model for model in models}
+            for future in as_completed(future_to_model):
+                model = future_to_model[future]
+                try:
+                    result = future.result()
+                    if "generated_text" in result:
+                        results.append({"model": model, "text": result["generated_text"]})
+                except Exception as e:
+                    logger.error(f"Error with model {model}: {e}")
+
+        if not results:
+            return {"error": "All models failed to generate responses"}
+
+        if len(results) == 1:
+            return {"generated_text": results[0]["text"]}
+
+        synthesis_prompt = "Synthesize these expert responses into a single coherent answer:\n\n"
+        for result in results:
+            synthesis_prompt += f"Expert {results.index(result) + 1} ({result['model'].split('/')[-1]}):\n{result['text']}\n\n"
+
+        synthesis = hf_inference(MAIN_LLM_MODEL, synthesis_prompt) # Use a consistent model for final synthesis
+        if "generated_text" in synthesis:
+            return synthesis
+        else:
+            return {"generated_text": max(results, key=lambda x: len(x["text"]))["text"]} # Fallback
 
 def tool_search_web(query: str, num_results: int = NUM_RESULTS, safesearch: str = "moderate",
                    time_filter: Optional[str] = None, region: str = "wt-wt", language: str = "en-us") -> list:
@@ -788,7 +811,7 @@ Output:
     return prompt
 
 def deep_research(prompt):
-    task_description = "You are an advanced research assistant. Use available tools iteratively, focus on different aspects, follow promising leads, critically evaluate your findings, and build up a comprehensive understanding. Utilize the FAISS index to avoid redundant searches and build a persistent knowledge base."
+    task_description = "You are an advanced research assistant, designed to be as comprehensive as possible.  Use available tools iteratively, focus on different aspects, explore promising leads thoroughly, critically evaluate your findings, and build up a comprehensive understanding of the research topic. Utilize the FAISS index to avoid redundant searches and to build a persistent knowledge base."
     research_data = load_research_data()
     paper_summaries = load_paper_summaries()  # Load paper summaries
 
@@ -818,6 +841,7 @@ def deep_research(prompt):
     if key_entities:
         context.append(f"Identified key entities: {key_entities}")
         intermediate_output += f"Identified key entities for focused research: {key_entities_with_descriptions}\n"
+        yield "Identifying key entities... (Completed)"
 
     # Initialize progress tracking for each entity.
     entity_progress = {entity: {'queries': [], 'insights': []} for entity in key_entities}
@@ -829,7 +853,9 @@ def deep_research(prompt):
 
     if not focus_areas:  # Corrected placement: outside the loop
         initial_focus_areas = tool_identify_focus_areas(prompt=prompt)
+        yield "Identifying initial focus areas...(Completed)"
         research_plan = tool_draft_research_plan(prompt=prompt, entities=key_entities, focus_areas=initial_focus_areas)
+        yield "Drafting initial research plan...(Completed)"
         context.append(f"Initial Research Plan: {research_plan[:200]}...") # Add plan to context
         intermediate_output += f"Initial Research Plan:\n{research_plan}\n\n"
         focus_areas = initial_focus_areas
@@ -859,6 +885,7 @@ def deep_research(prompt):
 
         if i == 0: #Initial broad search
             initial_query = tool_generate_search_query(prompt=prompt)
+            yield f"Generating initial search query... (Iteration {i+1})"
             if initial_query:
                 previous_queries.append(initial_query)
                 entity_progress['general']['queries'].append(initial_query)
@@ -875,12 +902,14 @@ def deep_research(prompt):
                     search_results = []
                     for future in as_completed(futures):
                         search_results.extend(future.result())
+                yield f"Performing initial searches... (Iteration {i+1})"
 
                 filtered_search_results = filter_results(search_results, prompt)
 
                 if filtered_search_results:
                     context.append(f"Initial Search Results: {len(filtered_search_results)} items found")
                     reasoning_output = tool_reason(prompt, filtered_search_results)
+                    yield f"Reasoning about initial search results... (Iteration {i+1})"
                     if reasoning_output:
                         all_insights.append(reasoning_output)
                         entity_progress['general']['insights'].append(reasoning_output)
@@ -897,6 +926,7 @@ def deep_research(prompt):
                 previous_queries=entity_progress[current_entity]['queries'],
                 focus_areas=focus_areas
             )
+            yield f"Generating search query for entity: {current_entity}... (Iteration {i+1})"
 
             if entity_query:
                 previous_queries.append(entity_query)
@@ -916,7 +946,7 @@ def deep_research(prompt):
                     for future in as_completed(futures):
                         search_results.extend(future.result())
 
-
+                yield f"Searching for information on entity: {current_entity}... (Iteration {i+1})"
                 filtered_search_results = filter_results(search_results,
                                                         f"{prompt} {current_entity}",
                                                         previous_snippets=seen_snippets) # Pass existing snippets
@@ -930,6 +960,7 @@ def deep_research(prompt):
                         reasoning_context=entity_progress[current_entity]['insights'], # Use entity-specific context
                         focus_areas=focus_areas
                     )
+                    yield f"Reasoning about entity: {current_entity}... (Iteration {i+1})"
 
                     if entity_reasoning:
                         all_insights.append(entity_reasoning)
@@ -946,23 +977,32 @@ def deep_research(prompt):
                     context.append(f"Entity query for {current_entity} yielded no relevant results")
 
         llm_prompt = create_prompt(task_description, prompt, tools, context)
-        llm_response = hf_inference(MAIN_LLM_MODEL, llm_prompt)
+        llm_response = hf_inference(MAIN_LLM_MODEL, llm_prompt, stream=True)  # Use streaming
 
         if isinstance(llm_response, dict) and "error" in llm_response:
             intermediate_output += f"LLM Error: {llm_response['error']}\n"
+            yield f"LLM Error (Iteration {i+1}): {llm_response['error']}"  # Display error in output
             continue
 
-        if not isinstance(llm_response, dict) or "generated_text" not in llm_response:
-            intermediate_output += "Error: Invalid LLM response.\n"
+        # Process streaming response
+        response_text = ""
+        try:
+            for chunk in llm_response:
+                if chunk.choices and chunk.choices[0].delta and chunk.choices[0].delta.content:
+                    response_text += chunk.choices[0].delta.content
+                    yield f"Iteration {i+1} - Thinking... {response_text}" # Real time output
+
+        except Exception as e:
+            intermediate_output += f"Streaming Error: {str(e)}\n"
+            yield f"Streaming Error (Iteration {i+1}): {str(e)}" #Error
             continue
 
         try:
-            response_text = llm_response["generated_text"].strip()
             response_json = json.loads(response_text)  # Parse the JSON response.
             intermediate_output += f"Iteration {i+1} - Focus: {current_entity} - Action: {response_text}\n"
         except json.JSONDecodeError:
-            intermediate_output += f"Iteration {i+1} - LLM Response (Invalid JSON): {llm_response['generated_text'][:100]}...\n"
-            context.append(f"Invalid JSON: {llm_response['generated_text'][:100]}...") # Add invalid JSON to context
+            intermediate_output += f"Iteration {i+1} - LLM Response (Invalid JSON): {response_text[:100]}...\n"
+            context.append(f"Invalid JSON: {response_text[:100]}...") # Add invalid JSON to context
             continue
 
         tool_name = response_json.get("tool")
@@ -988,6 +1028,7 @@ def deep_research(prompt):
                 parameters['failed_queries'] = failed_queries
                 parameters['focus_areas'] = focus_areas
                 result = tool["function"](**parameters)
+                yield f"Iteration {i+1} - Generated search query: {result}"
 
                 if current_entity != 'general':
                     entity_progress[current_entity]['queries'].append(result) # Add entity-specific
@@ -1027,6 +1068,7 @@ def deep_research(prompt):
                     parameters['focus_areas'] = focus_areas
 
                 result = tool["function"](**parameters)
+                yield f"Iteration {i+1} - Reasoning about information..."
 
                 if current_entity != 'general':
                     entity_progress[current_entity]['insights'].append(result)
@@ -1049,6 +1091,7 @@ def deep_research(prompt):
                         parameters['prompt'] = prompt
 
                     result = tool["function"](**parameters)
+                    yield f"Iteration {i+1} - Critiquing reasoning..."
                     previous_critiques.append(result)
                     context.append(f"Critique: {result[:200]}...")
                 else:
@@ -1056,6 +1099,7 @@ def deep_research(prompt):
 
             elif tool_name == "identify_contradictions":
                 result = tool["function"](**parameters)
+                yield f"Iteration {i+1} - Identifying contradictions..."
                 if result:
                     contradictions = result  # Keep track of contradictions.
                     context.append(f"Identified contradictions: {result}")
@@ -1064,6 +1108,7 @@ def deep_research(prompt):
                 if 'failed_areas' not in parameters:
                     parameters['failed_areas'] = failed_areas
                 result = tool["function"](**parameters)
+                yield f"Iteration {i+1} - Identifying focus areas..."
                 if result:
                     old_focus = set(focus_areas)
                     focus_areas = result  # Update focus areas
@@ -1072,6 +1117,7 @@ def deep_research(prompt):
 
             elif tool_name == "extract_article":
                 result = tool["function"](**parameters)
+                yield f"Iteration {i+1} - Extracting article content..."
                 if result:
                     context.append(f"Extracted article content from {parameters['url']}: {result[:200]}...")
                     # Reason specifically about the extracted article.
@@ -1082,6 +1128,7 @@ def deep_research(prompt):
             
             elif tool_name == "summarize_paper":
                 result = tool["function"](**parameters)
+                yield f"Iteration {i+1} - Summarizing paper..."
                 if result:
                     paper_summaries[parameters['paper_text'][:100]] = result  # Store by a snippet of the text
                     save_paper_summaries(paper_summaries)
@@ -1095,6 +1142,7 @@ def deep_research(prompt):
                 if 'prompt' not in parameters:
                     parameters['prompt'] = prompt
                 result = tool["function"](**parameters)
+                yield f"Iteration {i+1} - Performing meta-analysis..."
                 if result:
                     all_insights.append(result)  # Add meta-analysis to overall insights.
                     context.append(f"Meta-analysis across entities: {result[:200]}...")
@@ -1182,7 +1230,8 @@ def deep_research(prompt):
     full_output += f"Total iterations: {i+1}\n"
     full_output += f"Total insights generated: {len(all_insights)}\n"
 
-    return full_output
+    yield full_output  # Final output
+    
 
 custom_css = """
 .gradio-container {
@@ -1192,6 +1241,7 @@ custom_css = """
     font-family: 'Segoe UI', Tahoma, Geneva, Verdana, sans-serif;
     line-height: 1.5;
     font-size: 14px;
+    white-space: pre-wrap; /* Preserve newlines and spacing */
 }
 h3 {
     color: #2c3e50;
@@ -1210,11 +1260,11 @@ iface = gr.Interface(
     inputs=[
         gr.Textbox(lines=5, placeholder="Enter your research question...", label="Research Question")
     ],
-    outputs=gr.Textbox(lines=30, placeholder="Research results will appear here...", label="Research Results", elem_classes=["output-box"]),
+    outputs=gr.Markdown(label="Research Results",  elem_classes=["output-box"]), #Changed to markdown
     title="Advanced Multi-Stage Research Assistant",
     description="""This tool performs deep, multi-faceted research, leveraging multiple search engines, 
                    specialized academic databases, and advanced AI models. It incorporates a persistent knowledge 
-                   base using FAISS indexing to avoid redundant searches and build upon previous findings.""",
+                   base using FAISS indexing to avoid redundant searches and build upon previous findings.  Progress is shown in real-time.""",
     examples=[
         ["What are the key factors affecting urban tree survival and how do they vary between developing and developed countries?"],
         ["Compare and contrast the economic policies of China and the United States over the past two decades, analyzing their impacts on global trade."],
@@ -1226,6 +1276,7 @@ iface = gr.Interface(
     cache_examples=False,
     css=custom_css,
     allow_flagging="never",
+    live=True #for real time streaming
 )
 
 if __name__ == "__main__":