File size: 29,360 Bytes
0fc4dff
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
145f8c2
51c3a33
145f8c2
 
 
51c3a33
 
145f8c2
 
 
 
 
 
 
 
 
0fc4dff
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
145f8c2
 
 
0fc4dff
145f8c2
 
 
 
 
 
 
 
 
 
0fc4dff
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
import os
import logging
import re
import time
import gc
from datetime import datetime
from typing import Optional, List, Dict, Any
from collections import OrderedDict

import pandas as pd
from pydantic import BaseModel, Field, ValidationError, validator

# NLTK for input validation
import nltk
from nltk.corpus import words
try:
    english_words = set(words.words())
except LookupError:
    nltk.download('words')
    english_words = set(words.words())

# LangChain / Groq / LLM imports
from langchain_groq import ChatGroq
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain_community.vectorstores import FAISS
from langchain.chains import RetrievalQA, LLMChain
from langchain.prompts import PromptTemplate
from langchain.docstore.document import Document

# Custom chain imports
from classification_chain import get_classification_chain
from refusal_chain import get_refusal_chain
from tailor_chain import get_tailor_chain
from cleaner_chain import get_cleaner_chain
from tailor_chain_wellnessBrand import get_tailor_chain_wellnessBrand

# Mistral moderation
from mistralai import Mistral

# Google Gemini LLM
from langchain_google_genai import ChatGoogleGenerativeAI

# Web search
# from smolagents import DuckDuckGoSearchTool, ManagedAgent, HfApiModel, CodeAgent
# from openinference.instrumentation.smolagents import SmolagentsInstrumentor
from phoenix.otel import register


# register()
# SmolagentsInstrumentor().instrument(skip_dep_check=True)


from smolagents import (
    CodeAgent,
    DuckDuckGoSearchTool,
    HfApiModel,
    ToolCallingAgent,
    VisitWebpageTool,
)

# Import new prompts
from prompts import (
    selfharm_prompt, frustration_prompt, ethical_conflict_prompt,
    classification_prompt, refusal_prompt, tailor_prompt, cleaner_prompt
)

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# -------------------------------------------------------
# Basic Models
# -------------------------------------------------------
class QueryInput(BaseModel):
    query: str = Field(..., min_length=1)

    @validator('query')
    def check_query_is_string(cls, v):
        if not isinstance(v, str):
            raise ValueError("Query must be a valid string")
        if not v.strip():
            raise ValueError("Query cannot be empty or whitespace")
        return v.strip()

class ProcessingMetrics(BaseModel):
    total_requests: int = 0
    cache_hits: int = 0
    errors: int = 0
    average_response_time: float = 0.0
    last_reset: Optional[datetime] = None
    
    def update_metrics(self, processing_time: float, is_cache_hit: bool = False):
        self.total_requests += 1
        if is_cache_hit:
            self.cache_hits += 1
        self.average_response_time = (
            (self.average_response_time * (self.total_requests - 1) + processing_time)
            / self.total_requests
        )

# -------------------------------------------------------
# Mistral Moderation
# -------------------------------------------------------
class ModerationResult(BaseModel):
    is_safe: bool
    categories: Dict[str, bool]
    original_text: str

mistral_api_key = os.environ.get("MISTRAL_API_KEY")
client = Mistral(api_key=mistral_api_key)

def moderate_text(query: str) -> ModerationResult:
    """
    Uses Mistral's moderation to detect unsafe content.
    """
    try:
        query_input = QueryInput(query=query)
        response = client.classifiers.moderate_chat(
            model="mistral-moderation-latest",
            inputs=[{"role": "user", "content": query_input.query}]
        )
        
        is_safe = True
        categories = {}
        
        if hasattr(response, 'results') and response.results:
            cats = response.results[0].categories
            categories = {
                "violence": cats.get("violence_and_threats", False),
                "hate": cats.get("hate_and_discrimination", False),
                "dangerous": cats.get("dangerous_and_criminal_content", False),
                "selfharm": cats.get("selfharm", False)
            }
            is_safe = not any(categories.values())
        
        return ModerationResult(
            is_safe=is_safe,
            categories=categories,
            original_text=query_input.query
        )
    except ValidationError as ve:
        raise ValueError(f"Moderation input validation failed: {ve}")
    except Exception as e:
        raise RuntimeError(f"Moderation failed: {e}")

def compute_moderation_severity(mresult: ModerationResult) -> float:
    severity = 0.0
    for flag in mresult.categories.values():
        if flag:
            severity += 0.3
    return min(severity, 1.0)

# -------------------------------------------------------
# Models
# -------------------------------------------------------
GROQ_MODELS = {
    "default":        "llama3-70b-8192",
    "classification": "mixtral-8x7b-32768",
    "moderation":     "mistral-moderation-latest",
    "combination":    "llama-3.3-70b-versatile"
}

MAX_RETRIES = 3
RATE_LIMIT_REQUESTS = 60  
CACHE_SIZE_LIMIT = 1000

# Google Gemini (primary)
GEMINI_API_KEY = os.environ.get("GEMINI_API_KEY")
gemini_llm = ChatGoogleGenerativeAI(
    model="gemini-1.0-pro",
    temperature=0.5,
    max_retries=2,
    google_api_key=GEMINI_API_KEY
)

# Fallback
fallback_groq_api_key = os.environ.get("GROQ_API_KEY_FALLBACK", "YOUR_GROQ_API_KEY")
groq_fallback_llm = ChatGroq(
    model=GROQ_MODELS["default"],
    temperature=0.7,
    groq_api_key=fallback_groq_api_key,
    max_tokens=2048
)

# -------------------------------------------------------
# Rate-limit & Cache
# -------------------------------------------------------
def handle_rate_limiting(state: "PipelineState") -> bool:
    current_time = time.time()
    one_min_ago = current_time - 60
    state.request_timestamps = [t for t in state.request_timestamps if t > one_min_ago]
    if len(state.request_timestamps) >= RATE_LIMIT_REQUESTS:
        return False
    state.request_timestamps.append(current_time)
    return True

def manage_cache(state: "PipelineState", query: str, response: str = None) -> Optional[str]:
    cache_key = query.strip().lower()
    if response is None:
        return state.cache.get(cache_key)
    if cache_key in state.cache:
        state.cache.move_to_end(cache_key)
    state.cache[cache_key] = response
    if len(state.cache) > CACHE_SIZE_LIMIT:
        state.cache.popitem(last=False)
    return None

def create_error_response(error_type: str, details: str = "") -> str:
    templates = {
        "validation": "I couldn't process your query: {details}",
        "processing": "I encountered an error while processing: {details}",
        "rate_limit": "Too many requests. Please try again soon.",
        "general":    "Apologies, but something went wrong."
    }
    return templates.get(error_type, templates["general"]).format(details=details)

# -------------------------------------------------------
# Web Search
# -------------------------------------------------------
web_search_cache: Dict[str, str] = {}

def store_websearch_result(query: str, result: str):
    web_search_cache[query.strip().lower()] = result

def retrieve_websearch_result(query: str) -> Optional[str]:
    return web_search_cache.get(query.strip().lower())

def do_web_search(query: str) -> str:
    try:
        cached = retrieve_websearch_result(query)
        if cached:
            logger.info("Using cached web search result.")
            return cached

        logger.info("Performing a new web search for: '%s'", query)
        # model = HfApiModel()
        # search_tool = DuckDuckGoSearchTool()
        # web_agent = CodeAgent(tools=[search_tool], model=model)
        
        # managed_web_agent = ManagedAgent(
        #     agent=web_agent,
        #     name="web_search",
        #     description="Runs a web search. Provide your query."
        # )
        search_agent = ToolCallingAgent(
        tools=[DuckDuckGoSearchTool(), VisitWebpageTool()],
        model=HfApiModel(),
        name="search_agent",
        description="This is an agent that can do web search.",
        )
        
        manager_agent = CodeAgent(
            tools=[],
            model=model,
            managed_agents=[managed_web_agent]
        )
        
        new_search_result = manager_agent.run(f"Search for information about: {query}")
        store_websearch_result(query, new_search_result)
        return str(new_search_result).strip()
    except Exception as e:
        logger.error(f"Web search failed: {e}")
        return ""

def is_greeting(query: str) -> bool:
    """
    Returns True if the query is a greeting. This check is designed to be
    lenient enough to catch common greetings even with minor spelling mistakes
    or punctuation.
    """
    # Define a set of common greeting words (you can add variants or use fuzzy matching if needed)
    greetings = {"hello", "hi", "hey", "hii", "hola", "greetings"}
    
    # Remove punctuation and extra whitespace, and lower the case.
    cleaned = re.sub(r'[^\w\s]', '', query).strip().lower()
    
    # Split the cleaned text into words.
    words_in_query = set(cleaned.split())
    
    # Return True if any of the greeting words are in the query.
    return not words_in_query.isdisjoint(greetings)


# -------------------------------------------------------
# Vector Stores & RAG
# -------------------------------------------------------
def build_or_load_vectorstore(csv_path: str, store_dir: str) -> FAISS:
    if os.path.exists(store_dir):
        logger.info(f"Loading existing FAISS store from {store_dir}")
        embeddings = HuggingFaceEmbeddings(
            model_name="sentence-transformers/multi-qa-mpnet-base-dot-v1"
        )
        return FAISS.load_local(store_dir, embeddings)
    else:
        logger.info(f"Building new FAISS store from {csv_path}")
        df = pd.read_csv(csv_path)
        df = df.loc[:, ~df.columns.str.contains('^Unnamed')]
        df.columns = df.columns.str.strip()

        if "Answer" in df.columns:
            df.rename(columns={"Answer": "Answers"}, inplace=True)
        if "Question " in df.columns and "Question" not in df.columns:
            df.rename(columns={"Question ": "Question"}, inplace=True)
        if "Question" not in df.columns or "Answers" not in df.columns:
            raise ValueError("CSV must have 'Question' and 'Answers' columns.")

        docs = []
        for _, row in df.iterrows():
            question_text = str(row["Question"]).strip()
            ans = str(row["Answers"]).strip()
            doc = Document(page_content=ans, metadata={"question": question_text})
            docs.append(doc)

        embeddings = HuggingFaceEmbeddings(
            model_name="sentence-transformers/multi-qa-mpnet-base-dot-v1"
        )
        vectorstore = FAISS.from_documents(docs, embedding=embeddings)
        vectorstore.save_local(store_dir)
        return vectorstore
#rag chain is for wellness
def build_rag_chain(vectorstore: FAISS, llm) -> RetrievalQA:
    prompt = PromptTemplate(
        template="""
        [INST] You are a helpful AI specialized in Wellness & Well-being topics.
            Please use the following context to provide a detailed, helpful answer.
            If the context doesn't fully address the question, acknowledge this and provide the best possible information.
            
            Context: {context}
            
            Question: {question}
            
            Guidelines for responses:
            1. Start with a clear introduction establishing the wellness topic
            2. Present information using numbered lists for actionable steps
            3. Include evidence-based examples and practical applications
            4. Provide specific, implementable suggestions
            5. End with clear takeaways or next steps
            
            Additional considerations:
            - All recommendations should be grounded in current wellness research
            - Focus on sustainable, long-term lifestyle modifications
            - Acknowledge individual differences in wellness journeys
            - Emphasize holistic approaches to health and well-being
            - Include relevant studies or research when applicable
            [/INST]

        """,
        input_variables=["context", "question"]
    )
    retriever = vectorstore.as_retriever(search_type="similarity", search_kwargs={"k": 3})

    chain = RetrievalQA.from_chain_type(
        llm=llm,
        chain_type="stuff",
        retriever=retriever,
        return_source_documents=True,
        chain_type_kwargs={
            "prompt": prompt,
            "verbose": False,
            "document_variable_name": "context"
        }
    )
    return chain
#rag chain to is for brand
def build_rag_chain2(vectorstore: FAISS, llm) -> RetrievalQA:
    prompt = PromptTemplate(
        template="""
         [INST] You are the Brand Strategy Specialist for Daily Wellness AI.
            Please provide detailed, strategic guidance specific to Daily Wellness AI's brand development and market positioning.
            If additional context is needed, acknowledge this while maintaining focus on our company's objectives.
            
            Context: {context}
            
            Question: {question}
            
            Guidelines for Daily Wellness AI specific responses:
            1. Begin with addressing specific Daily Wellness AI brand challenges or opportunities
            2. Align recommendations with our core mission of democratizing personalized wellness
            3. Include competitive analysis within the AI wellness space
            4. Provide actionable steps that reflect our technological capabilities
            5. Conclude with KPIs aligned with our growth objectives
            
            Brand Pillars to Address:
            - AI-Driven Personalization
            - Scientific Credibility
            - User-Centric Design
            - Innovation Leadership
            - Community Building
            [/INST]


        """,
        input_variables=["context", "question"]
    )
    retriever = vectorstore.as_retriever(search_type="similarity", search_kwargs={"k": 3})

    chain = RetrievalQA.from_chain_type(
        llm=llm,
        chain_type="stuff",
        retriever=retriever,
        return_source_documents=True,
        chain_type_kwargs={
            "prompt": prompt,
            "verbose": False,
            "document_variable_name": "context"
        }
    )
    return chain

# -------------------------------------------------------
# PipelineState
# -------------------------------------------------------
class PipelineState:
    _instance = None
    
    def __new__(cls):
        if cls._instance is None:
            cls._instance = super(PipelineState, cls).__new__(cls)
            cls._instance._initialized = False
        return cls._instance
    
    def __init__(self):
        if self._initialized:
            return
        self._initialize()

    def _initialize(self):
        try:
            self.metrics = ProcessingMetrics()
            self.error_count = 0
            self.request_timestamps = []
            self.cache = OrderedDict()

            self._setup_chains()

            self._initialized = True
            self.metrics.last_reset = datetime.now()
            logger.info("Pipeline state initialized successfully.")
        except Exception as e:
            logger.error(f"Failed to initialize pipeline: {e}")
            raise RuntimeError("Pipeline initialization failed.") from e

    def _setup_chains(self):
        # Existing custom chains
        self.tailor_chainWellnessBrand = get_tailor_chain_wellnessBrand()
        self.classification_chain = get_classification_chain()
        self.refusal_chain        = get_refusal_chain()
        self.tailor_chain         = get_tailor_chain()
        self.cleaner_chain        = get_cleaner_chain()

        # Specialized chain for self-harm
        from prompts import selfharm_prompt
        self.self_harm_chain = LLMChain(llm=gemini_llm, prompt=selfharm_prompt, verbose=False)
        
        # NEW: chain for frustration/harsh queries
        from prompts import frustration_prompt
        self.frustration_chain = LLMChain(llm=gemini_llm, prompt=frustration_prompt, verbose=False)
        
        # NEW: chain for ethical conflict queries
        from prompts import ethical_conflict_prompt
        self.ethical_conflict_chain = LLMChain(llm=gemini_llm, prompt=ethical_conflict_prompt, verbose=False)

        # Build brand & wellness vectorstores
        brand_csv    = "BrandAI.csv"
        brand_store  = "faiss_brand_store"
        wellness_csv = "AIChatbot.csv"
        wellness_store = "faiss_wellness_store"

        brand_vs    = build_or_load_vectorstore(brand_csv,    brand_store)
        wellness_vs = build_or_load_vectorstore(wellness_csv, wellness_store)

        # Default LLM & fallback
        self.gemini_llm = gemini_llm
        self.groq_fallback_llm = groq_fallback_llm

        self.brand_rag_chain    = build_rag_chain2(brand_vs,    self.gemini_llm)
        self.wellness_rag_chain = build_rag_chain(wellness_vs, self.gemini_llm)

        self.brand_rag_chain_fallback    = build_rag_chain2(brand_vs,    self.groq_fallback_llm)
        self.wellness_rag_chain_fallback = build_rag_chain(wellness_vs, self.groq_fallback_llm)

    def handle_error(self, error: Exception) -> bool:
        self.error_count += 1
        self.metrics.errors += 1
        if self.error_count >= MAX_RETRIES:
            logger.warning("Max error reached, resetting pipeline.")
            self.reset()
            return False
        return True

    def reset(self):
        try:
            logger.info("Resetting pipeline state.")
            old_metrics = self.metrics
            self._initialized = False
            self.__init__()
            self.metrics = old_metrics
            self.metrics.last_reset = datetime.now()
            self.error_count = 0
            gc.collect()
            logger.info("Pipeline state reset done.")
        except Exception as e:
            logger.error(f"Reset pipeline failed: {e}")
            raise RuntimeError("Failed to reset pipeline.")

    def get_metrics(self) -> Dict[str, Any]:
        uptime = (datetime.now() - self.metrics.last_reset).total_seconds() / 3600
        return {
            "total_requests": self.metrics.total_requests,
            "cache_hits": self.metrics.cache_hits,
            "error_rate": self.metrics.errors / max(self.metrics.total_requests, 1),
            "average_response_time": self.metrics.average_response_time,
            "uptime_hours": uptime
        }

    def update_metrics(self, start_time: float, is_cache_hit: bool = False):
        duration = time.time() - start_time
        self.metrics.update_metrics(duration, is_cache_hit)

pipeline_state = PipelineState()

# -------------------------------------------------------
# Helper checks: detect aggression or ethical conflict
# -------------------------------------------------------

def is_aggressive_or_harsh(query: str) -> bool:
    """
    Very naive check: If user is insulting AI, complaining about worthless answers, etc.
    You can refine with better logic or a small LLM classifier.
    """
    triggers = ["useless", "worthless", "you cannot do anything", "so bad at answering"]
    for t in triggers:
        if t in query.lower():
            return True
    return False

def is_ethical_conflict(query: str) -> bool:
    """
    Check if user is asking about lying, revenge, or other moral dilemmas.
    You can expand or refine as needed.
    """
    ethics_keywords = ["should i lie", "should i cheat", "revenge", "get back at", "hurt them back"]
    q_lower = query.lower()
    return any(k in q_lower for k in ethics_keywords)


# -------------------------------------------------------
# Main Pipeline
# -------------------------------------------------------
def run_with_chain(query: str) -> str:
    """
    Overall flow:
    1) Validate & rate-limit
    2) Mistral moderation => 
       - If self-harm => self_harm_chain
       - If hate => refusal
       - If violence/dangerous => we STILL produce a guided response (ethics) unless it's extreme
    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 DailyWellness,  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 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"]:
            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:
            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:
                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"Critical error in run_with_chain: {e}")
        pipeline_state.metrics.errors += 1
        return create_error_response("general")

# -------------------------------------------------------
# Health & Utility
# -------------------------------------------------------
# def reset_pipeline():
#     try:
#         pipeline_state.reset()
#         return {"status": "success", "message": "Pipeline reset successful"}
#     except Exception as e:
#         logger.error(f"Reset pipeline error: {e}")
#         return {"status": "error", "message": str(e)}

# def get_pipeline_health() -> Dict[str, Any]:
#     try:
#         stats = pipeline_state.get_metrics()
#         healthy = stats["error_rate"] < 0.1
#         return {
#             **stats,
#             "is_healthy": healthy,
#             "status": "healthy" if healthy else "degraded"
#         }
#     except Exception as e:
#         logger.error(f"Health check error: {e}")
#         return {"is_healthy": False, "status": "error", "error": str(e)}

# def health_check() -> Dict[str, Any]:
#     try:
#         _ = run_with_chain("Test query for pipeline health check.")
#         return {
#             "status": "ok",
#             "timestamp": datetime.now().isoformat(),
#             "metrics": get_pipeline_health()
#         }
#     except Exception as e:
#         return {
#             "status": "error",
#             "timestamp": datetime.now().isoformat(),
#             "error": str(e)
#         }

logger.info("Pipeline initialization complete!")