Update pipeline.py

pipeline.py

@@ -600,163 +600,164 @@ def run_with_chain(query: str) -> str:
     3) If not refused, check if query is aggression/ethical => route to chain
     4) Otherwise classify => brand/wellness/out-of-scope => RAG => tailor
     """
-    start_time = time.time()
-    try:
-        # 1) Validate
-        if not query or query.strip() == "":
-            return create_error_response("validation", "Empty query.")
-        if len(query.strip()) < 2:
-            return create_error_response("validation", "Too short.")
-        words_in_text = re.findall(r'\b\w+\b', query.lower())
-        if not any(w in english_words for w in words_in_text):
-            return create_error_response("validation", "Unclear words.")
-        if len(query) > 500:
-            return create_error_response("validation", "Too long (>500).")
-        if not handle_rate_limiting(pipeline_state):
-            return create_error_response("rate_limit")
-        # New: Check if the query is a greeting
-        if is_greeting(query):
-            greeting_response = "Hello there!! Welcome to Healthy AI Expert,  How may I assist you today?"
-            manage_cache(pipeline_state, query, greeting_response)
-            pipeline_state.update_metrics(start_time)
-            return greeting_response
-    
-        if not handle_rate_limiting(pipeline_state):
-            return create_error_response("rate_limit")
-
-        # Cache check
-        cached = manage_cache(pipeline_state, query)
-        if cached:
-            pipeline_state.update_metrics(start_time, is_cache_hit=True)
-            return cached
-
-        # 2) Mistral moderation
+    with tracer.new_trace(name="wellness_pipeline_run") as run:
+        start_time = time.time()
         try:
-            mod_res = moderate_text(query)
-            severity = compute_moderation_severity(mod_res)
-
-            # If self-harm => supportive
-            if mod_res.categories.get("selfharm", False):
-                logger.info("Self-harm flagged => providing supportive chain response.")
-                selfharm_resp = pipeline_state.self_harm_chain.run({"query": query})
-                final_tailored = pipeline_state.tailor_chain.run({"response": selfharm_resp}).strip()
+            # 1) Validate
+            if not query or query.strip() == "":
+                return create_error_response("validation", "Empty query.")
+            if len(query.strip()) < 2:
+                return create_error_response("validation", "Too short.")
+            words_in_text = re.findall(r'\b\w+\b', query.lower())
+            if not any(w in english_words for w in words_in_text):
+                return create_error_response("validation", "Unclear words.")
+            if len(query) > 500:
+                return create_error_response("validation", "Too long (>500).")
+            if not handle_rate_limiting(pipeline_state):
+                return create_error_response("rate_limit")
+            # New: Check if the query is a greeting
+            if is_greeting(query):
+                greeting_response = "Hello there!! Welcome to Healthy AI Expert,  How may I assist you today?"
+                manage_cache(pipeline_state, query, greeting_response)
+                pipeline_state.update_metrics(start_time)
+                return greeting_response
+        
+            if not handle_rate_limiting(pipeline_state):
+                return create_error_response("rate_limit")
+    
+            # Cache check
+            cached = manage_cache(pipeline_state, query)
+            if cached:
+                pipeline_state.update_metrics(start_time, is_cache_hit=True)
+                return cached
+    
+            # 2) Mistral moderation
+            try:
+                mod_res = moderate_text(query)
+                severity = compute_moderation_severity(mod_res)
+    
+                # If self-harm => supportive
+                if mod_res.categories.get("selfharm", False):
+                    logger.info("Self-harm flagged => providing supportive chain response.")
+                    selfharm_resp = pipeline_state.self_harm_chain.run({"query": query})
+                    final_tailored = pipeline_state.tailor_chain.run({"response": selfharm_resp}).strip()
+                    manage_cache(pipeline_state, query, final_tailored)
+                    pipeline_state.update_metrics(start_time)
+                    return final_tailored
+    
+                # If hate => refuse
+                if mod_res.categories.get("hate", False):
+                    logger.info("Hate content => refusal.")
+                    refusal_resp = pipeline_state.refusal_chain.run({"topic": "moderation_flagged"})
+                    manage_cache(pipeline_state, query, refusal_resp)
+                    pipeline_state.update_metrics(start_time)
+                    return refusal_resp
+    
+                # If "dangerous" or "violence" is flagged, we might still want to 
+                # provide a "non-violent advice" approach (like revenge queries).
+                # So we won't automatically refuse. We'll rely on the 
+                # is_ethical_conflict() check below.
+    
+            except Exception as e:
+                logger.error(f"Moderation error: {e}")
+                severity = 0.0
+    
+            # 3) Check for aggression or ethical conflict
+            if is_aggressive_or_harsh(query):
+                logger.info("Detected harsh/aggressive language => frustration_chain.")
+                frustration_resp = pipeline_state.frustration_chain.run({"query": query})
+                final_tailored = pipeline_state.tailor_chain.run({"response": frustration_resp}).strip()
                 manage_cache(pipeline_state, query, final_tailored)
                 pipeline_state.update_metrics(start_time)
                 return final_tailored
-
-            # If hate => refuse
-            if mod_res.categories.get("hate", False):
-                logger.info("Hate content => refusal.")
-                refusal_resp = pipeline_state.refusal_chain.run({"topic": "moderation_flagged"})
-                manage_cache(pipeline_state, query, refusal_resp)
+    
+            if is_ethical_conflict(query):
+                logger.info("Detected ethical dilemma => ethical_conflict_chain.")
+                ethical_resp = pipeline_state.ethical_conflict_chain.run({"query": query})
+                final_tailored = pipeline_state.tailor_chain.run({"response": ethical_resp}).strip()
+                manage_cache(pipeline_state, query, final_tailored)
                 pipeline_state.update_metrics(start_time)
-                return refusal_resp
-
-            # If "dangerous" or "violence" is flagged, we might still want to 
-            # provide a "non-violent advice" approach (like revenge queries).
-            # So we won't automatically refuse. We'll rely on the 
-            # is_ethical_conflict() check below.
-
-        except Exception as e:
-            logger.error(f"Moderation error: {e}")
-            severity = 0.0
-
-        # 3) Check for aggression or ethical conflict
-        if is_aggressive_or_harsh(query):
-            logger.info("Detected harsh/aggressive language => frustration_chain.")
-            frustration_resp = pipeline_state.frustration_chain.run({"query": query})
-            final_tailored = pipeline_state.tailor_chain.run({"response": frustration_resp}).strip()
-            manage_cache(pipeline_state, query, final_tailored)
-            pipeline_state.update_metrics(start_time)
-            return final_tailored
-
-        if is_ethical_conflict(query):
-            logger.info("Detected ethical dilemma => ethical_conflict_chain.")
-            ethical_resp = pipeline_state.ethical_conflict_chain.run({"query": query})
-            final_tailored = pipeline_state.tailor_chain.run({"response": ethical_resp}).strip()
-            manage_cache(pipeline_state, query, final_tailored)
-            pipeline_state.update_metrics(start_time)
-            return final_tailored
-
-        # 4) Standard path: classification => brand/wellness/out-of-scope
-        try:
-            class_out = pipeline_state.classification_chain.run({"query": query})
-            classification = class_out.strip().lower()
-        except Exception as e:
-            logger.error(f"Classification error: {e}")
-            if not pipeline_state.handle_error(e):
-                return create_error_response("processing", "Classification error.")
-            return create_error_response("processing")
-
-        if classification in ["outofscope", "out_of_scope"]:
+                return final_tailored
+    
+            # 4) Standard path: classification => brand/wellness/out-of-scope
             try:
-                # Politely refuse if truly out-of-scope
-                refusal_text = pipeline_state.refusal_chain.run({"topic": query})
-                tailored_refusal = pipeline_state.tailor_chain.run({"response": refusal_text}).strip()
-                manage_cache(pipeline_state, query, tailored_refusal)
-                pipeline_state.update_metrics(start_time)
-                return tailored_refusal
+                class_out = pipeline_state.classification_chain.run({"query": query})
+                classification = class_out.strip().lower()
             except Exception as e:
-                logger.error(f"Refusal chain error: {e}")
+                logger.error(f"Classification error: {e}")
                 if not pipeline_state.handle_error(e):
-                    return create_error_response("processing", "Refusal error.")
+                    return create_error_response("processing", "Classification error.")
                 return create_error_response("processing")
-
-        # brand vs wellness
-        if classification == "brand":
-            rag_chain_main = pipeline_state.brand_rag_chain
-            # rag_chain_fallback = pipeline_state.brand_rag_chain_fallback
-        else:
-            rag_chain_main = pipeline_state.wellness_rag_chain
-            # rag_chain_fallback = pipeline_state.wellness_rag_chain_fallback
-
-        # RAG with fallback
-        try:
+    
+            if classification in ["outofscope", "out_of_scope"]:
+                try:
+                    # Politely refuse if truly out-of-scope
+                    refusal_text = pipeline_state.refusal_chain.run({"topic": query})
+                    tailored_refusal = pipeline_state.tailor_chain.run({"response": refusal_text}).strip()
+                    manage_cache(pipeline_state, query, tailored_refusal)
+                    pipeline_state.update_metrics(start_time)
+                    return tailored_refusal
+                except Exception as e:
+                    logger.error(f"Refusal chain error: {e}")
+                    if not pipeline_state.handle_error(e):
+                        return create_error_response("processing", "Refusal error.")
+                    return create_error_response("processing")
+    
+            # brand vs wellness
+            if classification == "brand":
+                rag_chain_main = pipeline_state.brand_rag_chain
+                # rag_chain_fallback = pipeline_state.brand_rag_chain_fallback
+            else:
+                rag_chain_main = pipeline_state.wellness_rag_chain
+                # rag_chain_fallback = pipeline_state.wellness_rag_chain_fallback
+    
+            # RAG with fallback
             try:
-                rag_output = rag_chain_main({"query": query})
-            except Exception as e_main:
-                if "resource exhausted" in str(e_main).lower():
-                    logger.warning("Gemini resource exhausted. Falling back to Groq.")
-                    # rag_output = rag_chain_fallback({"query": query})
+                try:
+                    rag_output = rag_chain_main({"query": query})
+                except Exception as e_main:
+                    if "resource exhausted" in str(e_main).lower():
+                        logger.warning("Gemini resource exhausted. Falling back to Groq.")
+                        # rag_output = rag_chain_fallback({"query": query})
+                    else:
+                        raise
+    
+                if isinstance(rag_output, dict) and "result" in rag_output:
+                    csv_ans = rag_output["result"].strip()
                 else:
-                    raise
-
-            if isinstance(rag_output, dict) and "result" in rag_output:
-                csv_ans = rag_output["result"].strip()
-            else:
-                csv_ans = str(rag_output).strip()
-
-            # If not enough => web
-            if "not enough context" in csv_ans.lower() or len(csv_ans) < 40:
-                logger.info("Insufficient RAG => web search.")
-                web_info = do_web_search(query)
-                if web_info:
-                    csv_ans += f"\n\nAdditional info:\n{web_info}"
-        except Exception as e:
-            logger.error(f"RAG error: {e}")
-            if not pipeline_state.handle_error(e):
-                return create_error_response("processing", "RAG error.")
-            return create_error_response("processing")
-
-        # Tailor final
-        try:
-            final_tailored = pipeline_state.tailor_chainWellnessBrand.run({"response": csv_ans}).strip()
-            if severity > 0.5:
-                final_tailored += "\n\n(Please note: This may involve sensitive content.)"
-
-            manage_cache(pipeline_state, query, final_tailored)
-            pipeline_state.update_metrics(start_time)
-            return final_tailored
+                    csv_ans = str(rag_output).strip()
+    
+                # If not enough => web
+                if "not enough context" in csv_ans.lower() or len(csv_ans) < 40:
+                    logger.info("Insufficient RAG => web search.")
+                    web_info = do_web_search(query)
+                    if web_info:
+                        csv_ans += f"\n\nAdditional info:\n{web_info}"
+            except Exception as e:
+                logger.error(f"RAG error: {e}")
+                if not pipeline_state.handle_error(e):
+                    return create_error_response("processing", "RAG error.")
+                return create_error_response("processing")
+    
+            # Tailor final
+            try:
+                final_tailored = pipeline_state.tailor_chainWellnessBrand.run({"response": csv_ans}).strip()
+                if severity > 0.5:
+                    final_tailored += "\n\n(Please note: This may involve sensitive content.)"
+    
+                manage_cache(pipeline_state, query, final_tailored)
+                pipeline_state.update_metrics(start_time)
+                return final_tailored
+            except Exception as e:
+                logger.error(f"Tailor chain error: {e}")
+                if not pipeline_state.handle_error(e):
+                    return create_error_response("processing", "Tailoring error.")
+                return create_error_response("processing")
+    
         except Exception as e:
-            logger.error(f"Tailor chain error: {e}")
-            if not pipeline_state.handle_error(e):
-                return create_error_response("processing", "Tailoring error.")
-            return create_error_response("processing")
-
-    except Exception as e:
-        logger.error(f"Critical error in run_with_chain: {e}")
-        pipeline_state.metrics.errors += 1
-        return create_error_response("general")
+            logger.error(f"Critical error in run_with_chain: {e}")
+            pipeline_state.metrics.errors += 1
+            return create_error_response("general")
 
 # -------------------------------------------------------
 # Health & Utility