app.py

import gradio as gr
import os
from huggingface_hub import InferenceClient
from langchain.document_loaders import TextLoader
from langchain.text_splitter import CharacterTextSplitter
from langchain.embeddings import HuggingFaceEmbeddings
from langchain.vectorstores import Chroma
import tempfile

# Initialize InferenceClient with your fine-tuned model
client = InferenceClient(
    model="francis-amazark/chatbotmodel-v.1.0",  # Your adapter repo ID
    token=os.environ.get("HF_TOKEN", ""),
    timeout=120
)

# Configuration
GUIDE_FILE = "staff_guide.txt"

# Initialize vector DB for RAG
def initialize_rag():
    if not os.path.exists(GUIDE_FILE):
        print("No staff_guide.txt found - using empty knowledge base")
        return None
    
    try:
        # Load and process document
        loader = TextLoader(GUIDE_FILE)
        documents = loader.load()
        
        text_splitter = CharacterTextSplitter(
            chunk_size=500,
            chunk_overlap=100
        )
        texts = text_splitter.split_documents(documents)
        
        # Use lightweight embeddings
        embeddings = HuggingFaceEmbeddings(
            model_name="sentence-transformers/all-MiniLM-L6-v2"
        )
        
        # Create in-memory vector store
        vector_db = Chroma.from_documents(
            documents=texts,
            embedding=embeddings,
            persist_directory=None  # No disk persistence
        )
        
        return vector_db
        
    except Exception as e:
        print(f"Error initializing RAG: {e}")
        return None

# Retrieve relevant context from staff guide
def get_relevant_context(query, vector_db, k=3):
    if vector_db is None:
        return "No staff guide information available."
    
    try:
        results = vector_db.similarity_search(query, k=k)
        context = "\n\n".join([doc.page_content for doc in results])
        return context
    except Exception as e:
        return f"Error retrieving context: {str(e)}"

def chat_with_rag(message, history, vector_db):
    """
    Function to handle chat interactions with RAG
    """
    # Get relevant context from staff guide
    context = get_relevant_context(message, vector_db)
    
    # Format conversation history
    formatted_messages = []
    
    # System prompt with RAG context
    system_prompt = f"""You are a helpful staff assistant. Use the following information from the staff guide to answer the question.

RELEVANT INFORMATION FROM STAFF GUIDE:
{context}

USER QUESTION: {message}

Please provide a helpful and accurate response based on the staff guide. If the information isn't in the guide, say so politely and offer general help."""
    
    formatted_messages.append({"role": "system", "content": system_prompt})
    
    # Add conversation history
    for user_msg, assistant_msg in history:
        formatted_messages.append({"role": "user", "content": user_msg})
        formatted_messages.append({"role": "assistant", "content": assistant_msg})
    
    # Add current message
    formatted_messages.append({"role": "user", "content": message})
    
    try:
        # Generate response using InferenceClient
        response = client.chat_completion(
            messages=formatted_messages,
            max_tokens=300,
            temperature=0.7,
            top_p=0.9,
            stream=False
        )
        
        # Extract the assistant's response
        assistant_response = response.choices[0].message.content
        return assistant_response
        
    except Exception as e:
        return f"❌ Error: {str(e)}. Please try again."

def create_chat_interface():
    """
    Create the Gradio chat interface with RAG
    """
    # Initialize RAG system
    vector_db = initialize_rag()
    
    with gr.Blocks(
        title="Amazark Automation",
        theme=gr.themes.Soft(),
        css="""
        .gradio-container {max-width: 800px; margin: auto;}
        .chatbot {min-height: 400px;}
        """
    ) as demo:
        
        gr.Markdown("""
        # 🏢 Amazark Automation
        """)
        
        # Chatbot interface
        chatbot = gr.Chatbot(
            label="Conversation",
            height=400,
            show_copy_button=True
        )
        
        with gr.Row():
            msg = gr.Textbox(
                label="Your question",
                placeholder="e.g., How do I order a commission?",
                scale=4,
                max_lines=2
            )
            submit_btn = gr.Button("Send", variant="primary", scale=1)
        
        clear_btn = gr.Button("Clear Chat", variant="secondary")
        
        # Examples
        gr.Examples(
            examples=[
                "What is Amazark?",
                "How to order a commission?",
                "How to make a partnership?",
                "Can I host a giveaway?",
                "What are the company policies?"
            ],
            inputs=msg,
            label="Try these questions:"
        )
        
        # Status indicator
        status = gr.Textbox(
            label="Status",
            value="✅ Ready with staff guide" if vector_db else "⚠️ Limited mode (no staff guide found)",
            interactive=False
        )
        
        # Chat function with RAG
        def respond(message, chat_history):
            if not message.strip():
                return "", chat_history
                
            # Update status to show processing
            status.value = "⏳ Searching staff guide and generating response..."
            
            # Get response from model with RAG
            response = chat_with_rag(message, chat_history, vector_db)
            
            # Update chat history
            chat_history.append((message, response))
            
            # Reset status
            status.value = "✅ Ready with staff guide" if vector_db else "⚠️ Limited mode"
            
            return "", chat_history
        
        # Event handlers
        msg.submit(respond, [msg, chatbot], [msg, chatbot])
        submit_btn.click(respond, [msg, chatbot], [msg, chatbot])
        clear_btn.click(lambda: None, None, chatbot, queue=False)
    
    return demo

# Main function
def main():
    print("Starting Staff Guide Assistant with RAG...")
    
    # Create and launch the interface
    demo = create_chat_interface()
    
    # Get port from environment (for Spaces deployment)
    server_port = int(os.environ.get('PORT', 7860))
    
    demo.launch(
        server_name="0.0.0.0",
        server_port=server_port,
        share=False
    )

if __name__ == "__main__":
    main()