app.py

from typing import TypedDict, Annotated, List, Optional, Dict, Any
import os
from dotenv import load_dotenv
import chainlit as cl
import json
import sys

# Debug mode flag with verbose option
DEBUG = True
VERBOSE = True  # Set to True for even more detailed output

def debug_print(*args, **kwargs):
    if DEBUG:
        print("\033[94m[DEBUG]\033[0m", *args, **kwargs)
        if VERBOSE and len(args) > 0 and isinstance(args[0], str):
            if "response" in args[0].lower() or "result" in args[0].lower():
                print("\033[93m[CONTENT]\033[0m", args[1] if len(args) > 1 else "No content")

from langchain_core.messages import HumanMessage, AIMessage
from langchain_core.tools import tool
from langchain_openai import ChatOpenAI
from langchain_qdrant import Qdrant as QdrantVectorStore
from langchain_community.tools.tavily_search import TavilySearchResults
from langgraph.graph import StateGraph, END
from langgraph.graph.message import add_messages
from langgraph.prebuilt import ToolNode
from langgraph.errors import GraphRecursionError
from qdrant_client import QdrantClient
from langchain_huggingface import HuggingFaceEmbeddings

from pregnancy_kb.config import QDRANT_PATH, COLLECTION_NAME, LLM_MODEL
from pregnancy_kb.prompts import (
    PREGNANCY_ADVISOR_PROMPT, 
    WELCOME_MESSAGE, 
    TAVILY_SEARCH_PROMPT,
    NO_INFO_MESSAGE,
    FOLLOW_UP_PROMPT
)

# Load environment variables
load_dotenv()

# Initialize Qdrant client with local storage
try:
    client = QdrantClient(path=str(QDRANT_PATH))
    debug_print("Qdrant client initialized successfully")
    
    # Check if collection exists
    collections = client.get_collections().collections
    if not any(c.name == COLLECTION_NAME for c in collections):
        debug_print(f"Warning: Collection '{COLLECTION_NAME}' not found")
except Exception as e:
    debug_print(f"Error connecting to Qdrant: {e}")
    client = QdrantClient(":memory:")

# Initialize vector store
try:
    vector_store = QdrantVectorStore(
        client=client,
        collection_name=COLLECTION_NAME,
        embeddings=HuggingFaceEmbeddings(model_name="AkshaySandbox/pregnancy-mpnet-embeddings")
    )
    debug_print("Vector store initialized successfully")
except Exception as e:
    debug_print(f"Error initializing vector store: {e}")
    from langchain_community.vectorstores import FAISS
    vector_store = FAISS(
        embeddings=HuggingFaceEmbeddings(model_name="AkshaySandbox/pregnancy-mpnet-embeddings"), 
        index=None, 
        docstore={}, 
        index_to_docstore_id={}
    )

# Initialize tools
@tool
def canada_pregnancy_search(query: str) -> str:
    """
    Search the internet for pregnancy, childbirth, and parenting information specific to Canada.
    Focuses on official Canadian sources like government websites, health authorities, and Canadian
    medical associations.
    """
    debug_print(f"Canada pregnancy search called with query: {query}")
    try:
        canadian_query = f"{query} Canada official pregnancy childbirth parenting information"
        
        tavily_tool = TavilySearchResults(
            max_results=5,
            k=5,
            search_depth="advanced",
            include_domains=[
                "canada.ca", 
                "healthycanadians.gc.ca",
                "pregnancyinfo.ca",
                "caringforkids.cps.ca",
                "sogc.org",
                "cmaj.ca",
                "phac-aspc.gc.ca"
            ]
        )
        
        debug_print(f"Invoking Tavily search with query: {canadian_query}")
        results = tavily_tool.invoke(canadian_query)
        debug_print(f"Received {len(results)} results from Tavily")
        
        if not results:
            debug_print("No results found from Tavily search")
            return "I couldn't find specific Canadian information on this topic. Please try a different search or consult with your healthcare provider."
        
        # Format the results with metadata
        sources = []
        content_parts = []
        
        for result in results:
            try:
                # Extract content first - if no content, skip this result
                content = result.get('content', '').strip()
                if not content:
                    continue
                
                # Create source metadata with safe defaults
                source_metadata = {
                    "title": result.get('title') or "Canadian Health Resource",  # Safe default if title is missing
                    "section": "Web Search",
                    "category": "Canadian Resources",
                    "url": result.get('url', '')  # Empty string if URL is missing
                }
                
                # Only add sources that have either a title or URL
                if source_metadata["title"] or source_metadata["url"]:
                    sources.append(source_metadata)
                    content_parts.append(content)
                    
            except Exception as e:
                debug_print(f"Error processing individual search result: {e}")
                continue  # Skip this result and continue with others
        
        if not content_parts:
            return """I found some Canadian resources but couldn't extract meaningful information. 
            
Here are some reliable Canadian pregnancy resources you can check directly:
- Health Canada: www.canada.ca/en/public-health/services/pregnancy.html
- Canadian Paediatric Society: www.caringforkids.cps.ca
- Society of Obstetricians and Gynaecologists of Canada: www.pregnancyinfo.ca"""
        
        # Combine and format content
        combined_content = "\n\n".join(content_parts)
        
        # Create a summary prompt for better formatting
        summary_prompt = f"""Summarize the following information about {query} in a clear, organized way:

{combined_content}

Format the response with clear sections and bullet points where appropriate."""
        
        # Get a well-formatted summary
        llm = ChatOpenAI(model=LLM_MODEL, temperature=0)
        debug_print("Invoking LLM with pregnancy advisor prompt")
        summary_response = llm.invoke([HumanMessage(content=summary_prompt)])
        
        formatted_response = format_response_with_metadata(
            content=summary_response.content,
            sources=sources,
            query=query
        )
        
        debug_print(f"Formatted response: \n\n{formatted_response[:200]}...\n\n")
        return formatted_response
        
    except Exception as e:
        debug_print(f"Error in canada_pregnancy_search: {str(e)}")
        return """I apologize, but I'm having trouble accessing Canadian pregnancy information at the moment. 

Here are some reliable Canadian resources you can check directly:
- Health Canada: www.canada.ca/en/public-health/services/pregnancy.html
- Canadian Paediatric Society: www.caringforkids.cps.ca
- Society of Obstetricians and Gynaecologists of Canada: www.pregnancyinfo.ca

You can also consult with your healthcare provider for specific information."""

def format_response_with_metadata(content: str, sources: List[Dict[str, str]], query: str) -> str:
    """Format the response with metadata, sources, and key points."""
    try:
        # Extract key points using bullet points or numbered lists
        key_points = []
        for line in content.split('\n'):
            if line.strip().startswith(('•', '-', '*', '1.', '2.', '3.')):
                key_points.append(line.strip().lstrip('•-* ').strip())
        
        # If no bullet points found, try to extract sentences that look like key points
        if not key_points:
            sentences = [s.strip() for s in content.split('.') if len(s.strip()) > 20]
            key_points = [s for s in sentences if any(kw in s.lower() for kw in ['important', 'should', 'recommend', 'key', 'essential', 'crucial'])][:3]
        
        # Format the response
        formatted_response = "### Response\n\n"
        formatted_response += content
        
        # Add key takeaways section
        formatted_response += "\n\n### Key Takeaways\n"
        if key_points:
            for i, point in enumerate(key_points[:5], 1):  # Limit to top 5 key points
                formatted_response += f"{i}. {point}\n"
        else:
            formatted_response += "• " + content.split('.')[0] + "\n"  # Use first sentence if no key points found
        
        # Add sources section with better formatting
        formatted_response += "\n### Sources\n"
        unique_sources = {}
        for source in sources:
            title = source.get('title', 'Untitled')
            if title not in unique_sources:
                unique_sources[title] = source
        
        for title, source in unique_sources.items():
            formatted_response += f"- **{title}**"
            
            # Add section information if available
            if source.get('section'):
                formatted_response += f"\n  Section: {source['section']}"
            
            # Add category information if available
            if source.get('category'):
                formatted_response += f"\n  Category: {source['category']}"
            
            # Add URL for web sources
            if source.get('url'):
                formatted_response += f"\n  Link: {source['url']}"
            
            # Add document information for knowledge base sources
            if source.get('document_info'):
                formatted_response += f"\n  Document: {source['document_info']}"
            
            formatted_response += "\n\n"  # Add extra line break between sources
        
        return formatted_response
    except Exception as e:
        debug_print(f"Error in format_response_with_metadata: {str(e)}")
        # Return the original content if formatting fails
        return content

@tool
def pregnancy_knowledge_base(query: str, category: Optional[str] = None) -> str:
    """
    Search the pregnancy and early parenthood knowledge base for relevant information.
    
    Args:
        query: The user's question about pregnancy or early parenthood
        category: Optional category to filter results
    """
    debug_print(f"Pregnancy knowledge base search called with query: {query}, category: {category}")
    try:
        # Prepare filter if category is provided
        filter_condition = None
        if category:
            filter_condition = {
                "must": [
                    {
                        "key": "category",
                        "match": {"value": category}
                    }
                ]
            }
        
        # Search with optional filter
        debug_print(f"Searching vector store with query: {query}")
        docs = vector_store.similarity_search(
            query, 
            k=4,
            filter=filter_condition
        )
        debug_print(f"Retrieved {len(docs)} documents from vector store")
        
        if not docs:
            debug_print("No documents found in knowledge base")
            return NO_INFO_MESSAGE
        
        # Prepare context and collect source information
        context_parts = []
        sources = []
        for doc in docs:
            # Create a detailed document info string
            doc_info = []
            if doc.metadata.get("title"):
                doc_info.append(doc.metadata["title"])
            if doc.metadata.get("date"):
                doc_info.append(f"Updated: {doc.metadata['date']}")
            if doc.metadata.get("version"):
                doc_info.append(f"Version: {doc.metadata['version']}")
            
            metadata = {
                'title': doc.metadata.get("title", "Untitled"),
                'section': doc.metadata.get("section", ""),
                'category': doc.metadata.get("category", "general").replace("_", " ").title(),
                'document_info': " | ".join(doc_info) if doc_info else None
            }
            sources.append(metadata)
            context_parts.append(f"[From: {metadata['title']} | Section: {metadata['section']} | Category: {metadata['category']}]\n{doc.page_content}")
        
        context = "\n\n---\n\n".join(context_parts)
        debug_print(f"Prepared context with {len(context_parts)} parts")
        
        # Use the enhanced prompt
        llm = ChatOpenAI(model=LLM_MODEL, temperature=0)
        debug_print("Invoking LLM with pregnancy advisor prompt")
        response = llm.invoke(
            [HumanMessage(content=PREGNANCY_ADVISOR_PROMPT.format(query=query, context=context))]
        )
        
        # Format the response with metadata
        formatted_response = format_response_with_metadata(
            content=response.content,
            sources=sources,
            query=query
        )
        
        debug_print(f"Formatted response: \n\n{formatted_response[:200]}...\n\n")
        return formatted_response
    except Exception as e:
        debug_print(f"Error in pregnancy_knowledge_base: {e}")
        import traceback
        debug_print(f"Detailed error: {traceback.format_exc()}")
        
        # Return a more helpful error message
        return f"""I'm having trouble accessing the pregnancy knowledge base at the moment. This might be due to a technical issue.

Here are some reliable Canadian resources you can check in the meantime:

- Health Canada Pregnancy Resources: www.canada.ca/en/public-health/services/pregnancy.html
- Canadian Paediatric Society: www.caringforkids.cps.ca
- Society of Obstetricians and Gynaecologists of Canada: www.pregnancyinfo.ca

For information about birth plans in Canada specifically, you might want to check:
- The Society of Obstetricians and Gynaecologists of Canada (SOGC): www.pregnancyinfo.ca
- Health Canada's Healthy Pregnancy Guide: www.canada.ca/en/public-health/services/pregnancy/healthy-pregnancy-guide.html

You can also try asking me a different question or rephrasing your current one."""

# Setup tools
tool_belt = [
    canada_pregnancy_search,
    pregnancy_knowledge_base
]

# Initialize model with tools
model = ChatOpenAI(model=LLM_MODEL, temperature=0)
model = model.bind_tools(tool_belt)

class AgentState(TypedDict):
    messages: Annotated[list, add_messages]
    conversation_history: List[Dict[str, Any]]

# Create graph nodes
tool_node = ToolNode(tool_belt)

def call_model(state):
    messages = state["messages"]
    conversation_history = state.get("conversation_history", [])
    
    debug_print(f"Call model received state with {len(messages)} messages")
    debug_print("Messages content:", [m.content[:100] + "..." for m in messages])
    
    try:
        if conversation_history:
            history_context = "\n\n".join([
                f"User: {item['user']}\nAssistant: {item['assistant'][:150]}..." 
                for item in conversation_history[-3:]
            ])
            
            if messages and isinstance(messages[0], HumanMessage):
                enhanced_message = HumanMessage(
                    content=FOLLOW_UP_PROMPT.format(
                        conversation_history=history_context,
                        user_query=messages[0].content
                    )
                )
                messages = [enhanced_message]
        
        # Add system message to guide the model
        system_message = """You are a knowledgeable pregnancy advisor focusing on Canadian healthcare and parenting information. 
        Use the available tools when needed:
        - canada_pregnancy_search: For finding official Canadian pregnancy and parenting information
        - pregnancy_knowledge_base: For searching the knowledge base about pregnancy and early parenthood
        
        Always provide helpful, accurate information and use the tools when you need to find specific information."""
        
        messages = [HumanMessage(content=system_message)] + messages
        
        debug_print("Invoking model with messages:", [m.content[:100] + "..." for m in messages])
        response = model.invoke(messages)
        debug_print("Raw model response:", response)
        
        # Check for tool calls first
        if response.additional_kwargs.get("tool_calls"):
            debug_print("Tool calls detected in response")
            return {
                "messages": [response],
                "conversation_history": conversation_history
            }
            
        # Handle regular response
        if not response.content:
            debug_print("WARNING: Empty content received from model")
            return {
                "messages": [AIMessage(content="I apologize, but I received an empty response. Please try asking your question again.")],
                "conversation_history": conversation_history
            }
        
        content = response.content
        debug_print("Response content:", content[:200] + "..." if content else "No content")
        
        # Process regular response
        if len(content) > 300 and "##" not in content and "**" not in content:
            lines = content.split("\n")
            formatted_lines = []
            for i, line in enumerate(lines):
                if i > 0 and line.strip() and len(line) < 80 and line.strip()[-1] not in ".,:;?!":
                    formatted_lines.append(f"\n### {line}")
                else:
                    formatted_lines.append(line)
            content = "\n".join(formatted_lines)
        
        if "follow-up" not in content.lower() and "next steps" not in content.lower():
            content += "\n\n### Follow-up Information\n"
            content += "If you have more questions about this topic, feel free to ask! "
            content += "I can provide additional details on specific aspects or related topics that might be helpful for your situation."
        
        return {
            "messages": [AIMessage(content=content)],
            "conversation_history": conversation_history
        }
    except Exception as e:
        debug_print(f"Error in call_model: {str(e)}")
        import traceback
        debug_print("Traceback:", traceback.format_exc())
        return {
            "messages": [AIMessage(content=f"I encountered an error: {str(e)}. Please try again.")],
            "conversation_history": conversation_history
        }

def should_continue(state):
    """Determine if we should continue processing or end the conversation."""
    try:
        last_message = state["messages"][-1]
        debug_print(f"Checking if should continue. Last message: {last_message}")
        
        # Check if the message has tool calls
        if hasattr(last_message, "additional_kwargs") and \
           "tool_calls" in last_message.additional_kwargs and \
           last_message.additional_kwargs["tool_calls"]:
            
            # Get the tool calls
            tool_calls = last_message.additional_kwargs["tool_calls"]
            debug_print(f"Tool calls detected: {len(tool_calls)}")
            
            # Check if any tool calls are valid
            valid_tools = [t.name for t in tool_belt]
            has_valid_tool = False
            
            for tool_call in tool_calls:
                if "function" in tool_call and "name" in tool_call["function"]:
                    tool_name = tool_call["function"]["name"]
                    if tool_name in valid_tools:
                        has_valid_tool = True
                        debug_print(f"Valid tool call detected: {tool_name}")
                        break
            
            if has_valid_tool:
                debug_print("Valid tool calls detected, continuing to action node")
                return "action"
            else:
                debug_print("No valid tool calls detected, ending conversation turn")
                return "end"
        
        debug_print("No tool calls detected, ending conversation turn")
        return "end"
    except Exception as e:
        debug_print(f"Error in should_continue: {e}")
        import traceback
        debug_print(f"Traceback: {traceback.format_exc()}")
        # If there's any error in processing, end the conversation
        return "end"

# Build graph
graph = StateGraph(AgentState)
graph.add_node("agent", call_model)
graph.add_node("action", tool_node)
graph.set_entry_point("agent")
graph.add_conditional_edges(
    "agent",
    should_continue,
    {
        "action": "action",
        "end": END
    }
)
graph.add_edge("action", "agent")

# Compile graph without recursion limit
compiled_graph = graph.compile()

@cl.on_chat_start
async def start():
    """Initialize the chat session."""
    try:
        # Initialize with fresh graph and empty conversation history
        # Instead of clearing the session, just set the variables directly
        cl.user_session.set("graph", compiled_graph)
        cl.user_session.set("conversation_history", [])
        
        debug_print("Chat session initialized")
        await cl.Message(content=WELCOME_MESSAGE).send()
    except Exception as e:
        debug_print(f"Chat start error: {e}")
        await cl.Message(content=f"Startup error: {str(e)}. Please refresh.").send()

@cl.on_message
async def handle(message: cl.Message):
    """Process user messages."""
    debug_print(f"Received message: {message.content}")
    try:
        graph = cl.user_session.get("graph")
        if not graph:
            debug_print("WARNING: Graph not found in session, reinitializing")
            cl.user_session.set("graph", compiled_graph)
            graph = compiled_graph
            
        conversation_history = cl.user_session.get("conversation_history", [])
        debug_print(f"Current conversation history has {len(conversation_history)} entries")
        
        state = {
            "messages": [HumanMessage(content=message.content)],
            "conversation_history": conversation_history
        }
        
        msg = cl.Message(content="")
        await msg.send()
        
        full_response = ""
        has_tool_calls = False
        
        try:
            debug_print("Starting graph stream")
            async for chunk in graph.astream(state, {"recursion_limit": 50}):
                debug_print(f"Received chunk: {chunk}")
                for node, values in chunk.items():
                    debug_print(f"Processing node: {node}")
                    if node == "agent" and values.get("messages"):
                        message = values["messages"][-1]
                        debug_print("Message content:", message.content[:200] + "..." if message.content else "Empty content")
                        debug_print("Message kwargs:", message.additional_kwargs)
                        
                        # Check for tool calls
                        if message.additional_kwargs.get("tool_calls"):
                            has_tool_calls = True
                            debug_print("Tool call detected, waiting for results...")
                            continue
                            
                        response = message.content
                        if response:
                            debug_print("Streaming response chunk:", response[:100] + "...")
                            await msg.stream_token(response)
                            full_response = response
                        else:
                            debug_print("WARNING: Empty response chunk received")
                    elif node == "action":
                        debug_print("Processing action node result")
                        if values.get("messages"):
                            action_response = values["messages"][-1].content
                            if action_response:
                                debug_print("Action response:", action_response[:100] + "...")
                                await msg.stream_token(action_response)
                                full_response = action_response
                            else:
                                debug_print("WARNING: Empty action response")
        
        except GraphRecursionError as e:
            debug_print(f"Graph recursion error: {e}")
            error_message = """
I apologize, but I'm having trouble processing your request due to a technical limitation. 

Here's what I can tell you about birth plans in Canada:

1. Birth plans are personal documents that outline your preferences for labor and delivery.
2. They typically include preferences for pain management, delivery positions, and who you want present.
3. In Canada, birth plans are respected by healthcare providers but may need to be flexible based on medical needs.
4. It's recommended to discuss your birth plan with your healthcare provider well before your due date.

For more detailed information, please consider:
- Discussing with your healthcare provider
- Visiting Health Canada's pregnancy resources: www.canada.ca/en/public-health/services/pregnancy.html
- Checking the Society of Obstetricians and Gynaecologists of Canada: www.pregnancyinfo.ca

You can also try asking a more specific question about birth plans.
"""
            await msg.stream_token(error_message)
            full_response = error_message
            
        except Exception as e:
            debug_print(f"Error in graph streaming: {e}")
            import traceback
            debug_print("Traceback:", traceback.format_exc())
            error_message = f"\n\nI encountered an error while processing your request: {str(e)}. Please try again with a different question."
            await msg.stream_token(error_message)
            full_response = error_message
        
        if not full_response and not has_tool_calls:
            debug_print("WARNING: No response generated and no tool calls detected")
            full_response = "I apologize, but I wasn't able to generate a response. Please try asking your question again."
            await msg.stream_token(full_response)
        
        # Finalize the message - use send() instead of update() which might not be supported
        await msg.send()
        
        if full_response:  # Only update history if we have a response
            conversation_history.append({
                "user": message.content,
                "assistant": full_response
            })
            
            if len(conversation_history) > 10:
                conversation_history = conversation_history[-10:]
                
            # Update session with new conversation history
            cl.user_session.set("conversation_history", conversation_history)
            
            # Ensure graph is reset for next question by reinitializing it
            cl.user_session.set("graph", compiled_graph)
            
            # Only generate follow-ups for non-error responses
            if not full_response.startswith("I apologize") and len(conversation_history) >= 2:
                await generate_follow_up_suggestions(conversation_history)
                
    except Exception as e:
        debug_print(f"Error in message handler: {e}")
        await cl.Message(content=f"An error occurred: {str(e)}. Please try again.").send()

async def generate_follow_up_suggestions(conversation_history):
    """Generate and display follow-up suggestions."""
    try:
        follow_up_llm = ChatOpenAI(model="gpt-3.5-turbo", temperature=0.7)
        
        history_text = "\n".join([
            f"User: {item['user']}\nAssistant: {item['assistant'][:150]}..." 
            for item in conversation_history[-3:]
        ])
        
        follow_up_prompt = f"""Based on this conversation about pregnancy and parenting, suggest 3 natural follow-up questions the user might want to ask next. Make them specific to the conversation context and helpful for a new parent or expecting parent in Canada.

Conversation:
{history_text}

Provide exactly 3 follow-up questions, each on a new line starting with a bullet point (•).
"""
        
        debug_print("Generating follow-up suggestions")
        follow_up_response = follow_up_llm.invoke([HumanMessage(content=follow_up_prompt)])
        debug_print(f"Follow-up response: {follow_up_response.content}")
        
        # Extract questions from response
        questions = []
        for line in follow_up_response.content.split("\n"):
            line = line.strip()
            if line.startswith("•") or line.startswith("-") or line.startswith("*"):
                question = line[1:].strip()
                if question and len(question) > 10:  # Ensure it's a valid question
                    questions.append(question)
        
        # Limit to 3 questions
        questions = questions[:3]
        
        if questions:
            debug_print(f"Extracted {len(questions)} follow-up questions")
            
            # Send each question as a separate message with an action button
            await cl.Message(content="You might also want to ask:").send()
            
            for i, question in enumerate(questions):
                action_name = f"ask_{i}"
                await cl.Message(
                    content=question,
                    actions=[
                        cl.Action(
                            name=action_name,
                            value=question,
                            label="Ask",
                            description="Ask this follow-up question"
                        )
                    ]
                ).send()
        else:
            debug_print("No valid follow-up questions extracted")
    except Exception as e:
        debug_print(f"Error generating follow-up suggestions: {e}")
        # Don't send an error message to the user, just log it
        # This ensures the conversation can continue even if follow-ups fail

@cl.action_callback("ask")
async def on_action(action):
    """Handle action callbacks for follow-up questions."""
    try:
        debug_print(f"Action received: {action.name} with value: {action.value}")
        await cl.Message(content=action.value, author="User").send()
        await handle(cl.Message(content=action.value))
    except Exception as e:
        debug_print(f"Error in action callback: {e}")
        await cl.Message(content=f"Error processing follow-up question: {str(e)}").send()

@cl.action_callback("ask_0")
async def on_action_0(action):
    try:
        debug_print(f"Action 0 received with value: {action.value}")
        await cl.Message(content=action.value, author="User").send()
        await handle(cl.Message(content=action.value))
    except Exception as e:
        debug_print(f"Error in action_0 callback: {e}")
        await cl.Message(content=f"Error processing follow-up question: {str(e)}").send()

@cl.action_callback("ask_1")
async def on_action_1(action):
    try:
        debug_print(f"Action 1 received with value: {action.value}")
        await cl.Message(content=action.value, author="User").send()
        await handle(cl.Message(content=action.value))
    except Exception as e:
        debug_print(f"Error in action_1 callback: {e}")
        await cl.Message(content=f"Error processing follow-up question: {str(e)}").send()

@cl.action_callback("ask_2")
async def on_action_2(action):
    try:
        debug_print(f"Action 2 received with value: {action.value}")
        await cl.Message(content=action.value, author="User").send()
        await handle(cl.Message(content=action.value))
    except Exception as e:
        debug_print(f"Error in action_2 callback: {e}")
        await cl.Message(content=f"Error processing follow-up question: {str(e)}").send()