app.py

import os
import sys
import logging
import gradio as gr
from huggingface_hub import InferenceClient
import re
import numpy as np
from sklearn.feature_extraction.text import CountVectorizer
from sklearn.naive_bayes import MultinomialNB
import asyncio
from crewai import Agent as CrewAgent, Task, Crew
import autogen
import openai

# Set up logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

# Check for OpenAI API key
if 'OPENAI_API_KEY' not in os.environ:
    logger.error("OPENAI_API_KEY environment variable is not set.")
    logger.info("Please set the OPENAI_API_KEY environment variable before running this script.")
    sys.exit(1)

# Set OpenAI API key
openai.api_key = os.environ['OPENAI_API_KEY']

# Initialize the client with the Mistral-7B-Instruct-v0.2 model
try:
    client = InferenceClient("mistralai/Mistral-7B-Instruct-v0.2")
except Exception as e:
    logger.error(f"Failed to initialize InferenceClient: {e}")
    sys.exit(1)

# Guardrail functions
def sanitize_input(input_text):
    return re.sub(r'[<>&\']', '', input_text)

approved_topics = ['account opening', 'trading', 'fees', 'platforms', 'funds', 'regulations', 'support']
vectorizer = CountVectorizer()
classifier = MultinomialNB()

X = vectorizer.fit_transform(approved_topics)
y = np.arange(len(approved_topics))
classifier.fit(X, y)

def is_relevant_topic(query):
    query_vector = vectorizer.transform([query])
    prediction = classifier.predict(query_vector)
    return prediction[0] in range(len(approved_topics))

def redact_sensitive_info(text):
    text = re.sub(r'\b\d{10,12}\b', '[REDACTED]', text)
    text = re.sub(r'[A-Z]{5}[0-9]{4}[A-Z]', '[REDACTED]', text)
    return text

def check_response_content(response):
    unauthorized_patterns = [
        r'\b(guarantee|assured|certain)\b.*\b(returns|profit)\b',
        r'\b(buy|sell)\b.*\b(specific stocks?|shares?)\b'
    ]
    return not any(re.search(pattern, response, re.IGNORECASE) for pattern in unauthorized_patterns)

def check_confidence(response):
    uncertain_phrases = ["I'm not sure", "It's possible", "I don't have enough information"]
    return not any(phrase.lower() in response.lower() for phrase in uncertain_phrases)

async def generate_response(prompt):
    try:
        return await client.text_generation(prompt, max_new_tokens=500, temperature=0.7)
    except Exception as e:
        logger.error(f"Error generating response: {e}")
        return "I apologize, but I'm having trouble generating a response at the moment. Please try again later."

def post_process_response(response):
    response = re.sub(r'\b(stupid|dumb|idiotic|foolish)\b', 'mistaken', response, flags=re.IGNORECASE)
    
    if not re.search(r'(Thank you|Is there anything else|Hope this helps|Let me know if you need more information)\s*$', response, re.IGNORECASE):
        response += "\n\nIs there anything else I can help you with regarding Zerodha's services?"
    
    if re.search(r'\b(invest|trade|buy|sell|market)\b', response, re.IGNORECASE):
        response += "\n\nPlease note that this information is for educational purposes only and should not be considered as financial advice. Always do your own research and consider consulting with a qualified financial advisor before making investment decisions."
    
    return response

# Gradio interface setup
demo = gr.ChatInterface(
    zerodha_support,
    chatbot=gr.Chatbot(height=600),
    textbox=gr.Textbox(placeholder="Ask your question about Zerodha here...", container=False, scale=7),
    title="Zerodha Support Assistant",
    description="Ask questions about Zerodha's services, trading, account management, and more. Our multi-agent system ensures respectful and empathetic responses.",
    theme="soft",
    examples=[
        "How do I open a Zerodha account?",
        "I'm frustrated with the recent changes to the Kite platform. Can you help?",
        "What are the risks involved in F&O trading?",
        "I think there's an error in my account statement. What should I do?",
        "Can you explain Zerodha's policy on intraday trading margins?",
        "I'm new to investing. What resources does Zerodha offer for beginners?",
        "How does Zerodha ensure the security of my investments and personal data?"
    ],
)

if __name__ == "__main__":
    demo.launch()