pipeline.py

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!")