app.py

import gradio as gr
import openai
from qdrant_client import QdrantClient
from qdrant_client.models import Distance, VectorParams, PointStruct
import numpy as np
from typing import List, Tuple, Dict
import os
from datetime import datetime
import json

# Initialize clients
try:
    openai_client = openai.OpenAI(api_key=os.getenv("OPENAI_API_KEY"))
except TypeError as e:
    # Fallback for compatibility issues
    openai.api_key = os.getenv("OPENAI_API_KEY")
    openai_client = openai.OpenAI()
except Exception as e:
    print(f"Warning: Could not initialize OpenAI client: {e}")
    openai_client = None

try:
    qdrant_client = QdrantClient(
        url=os.getenv("QDRANT_URL", "localhost:6333"),
        api_key=os.getenv("QDRANT_API_KEY", None)
    )
except Exception as e:
    print(f"Warning: Could not initialize Qdrant client: {e}")
    qdrant_client = None

# Collection name for vector store
COLLECTION_NAME = "chatbot_knowledge"
EMBEDDING_DIM = 1536  # OpenAI embedding dimension

# Initialize Qdrant collection if it doesn't exist
if qdrant_client:
    try:
        collections = qdrant_client.get_collections().collections
        if COLLECTION_NAME not in [c.name for c in collections]:
            qdrant_client.create_collection(
                collection_name=COLLECTION_NAME,
                vectors_config=VectorParams(size=EMBEDDING_DIM, distance=Distance.COSINE)
            )
    except Exception as e:
        print(f"Warning: Could not initialize Qdrant collection: {e}")

def get_embedding(text: str) -> List[float]:
    """Get embedding for text using OpenAI API"""
    if not openai_client:
        return None
    
    try:
        response = openai_client.embeddings.create(
            model="text-embedding-3-small",
            input=text
        )
        return response.data[0].embedding
    except Exception as e:
        print(f"Error getting embedding: {e}")
        return None

def search_vector_store(query: str, limit: int = 3) -> List[Dict]:
    """Search Qdrant vector store for relevant context"""
    if not qdrant_client:
        return []
    
    try:
        query_embedding = get_embedding(query)
        if not query_embedding:
            return []
        
        search_result = qdrant_client.search(
            collection_name=COLLECTION_NAME,
            query_vector=query_embedding,
            limit=limit
        )
        
        contexts = []
        for hit in search_result:
            if hasattr(hit, 'payload') and hit.payload:
                contexts.append({
                    'content': hit.payload.get('content', ''),
                    'metadata': hit.payload.get('metadata', {}),
                    'score': hit.score
                })
        
        return contexts
    except Exception as e:
        print(f"Error searching vector store: {e}")
        return []

def add_to_vector_store(content: str, metadata: Dict = None):
    """Add content to Qdrant vector store"""
    try:
        embedding = get_embedding(content)
        if not embedding:
            return False
        
        point_id = int(datetime.now().timestamp() * 1000)
        
        qdrant_client.upsert(
            collection_name=COLLECTION_NAME,
            points=[
                PointStruct(
                    id=point_id,
                    vector=embedding,
                    payload={
                        'content': content,
                        'metadata': metadata or {},
                        'timestamp': datetime.now().isoformat()
                    }
                )
            ]
        )
        return True
    except Exception as e:
        print(f"Error adding to vector store: {e}")
        return False

def format_context(contexts: List[Dict]) -> str:
    """Format vector store search results for inclusion in prompt"""
    if not contexts:
        return ""
    
    formatted = "Relevant context from knowledge base:\n"
    for i, ctx in enumerate(contexts):
        formatted += f"\n[Context {i+1}] (relevance: {ctx['score']:.2f}):\n{ctx['content']}\n"
    
    return formatted

def generate_response(message: str, history: List[Tuple[str, str]], system_prompt: str, temperature: float, max_tokens: int):
    """Generate response using OpenAI GPT-4o with context from vector store"""
    
    if not openai_client:
        yield "Error: OpenAI client not initialized. Please check your API key."
        return
    
    # Search vector store for relevant context
    contexts = search_vector_store(message)
    context_str = format_context(contexts)
    
    # Build conversation history
    messages = [{"role": "system", "content": system_prompt}]
    
    # Add context from vector store if available
    if context_str:
        messages.append({
            "role": "system", 
            "content": context_str
        })
    
    # Add conversation history
    for user_msg, assistant_msg in history:
        messages.append({"role": "user", "content": user_msg})
        messages.append({"role": "assistant", "content": assistant_msg})
    
    # Add current message
    messages.append({"role": "user", "content": message})
    
    try:
        # Call OpenAI API
        response = openai_client.chat.completions.create(
            model="gpt-4o",
            messages=messages,
            temperature=temperature,
            max_tokens=max_tokens,
            stream=True
        )
        
        # Stream response
        full_response = ""
        for chunk in response:
            if chunk.choices[0].delta.content:
                full_response += chunk.choices[0].delta.content
                yield full_response
    
    except Exception as e:
        yield f"Error: {str(e)}"

def chat_interface(message: str, history: List[Tuple[str, str]], system_prompt: str, temperature: float, max_tokens: int, add_to_kb: bool):
    """Main chat interface function"""
    
    # Optionally add user message to knowledge base
    if add_to_kb and message.strip():
        success = add_to_vector_store(
            content=message,
            metadata={"type": "user_input", "timestamp": datetime.now().isoformat()}
        )
        if success:
            print(f"Added to knowledge base: {message[:50]}...")
    
    # Generate response
    response_generator = generate_response(message, history, system_prompt, temperature, max_tokens)
    
    # Yield responses as they come
    for response in response_generator:
        yield response

def add_knowledge(text: str):
    """Add knowledge to vector store"""
    if not text.strip():
        return "Please enter some text to add to the knowledge base."
    
    success = add_to_vector_store(
        content=text,
        metadata={"type": "manual_entry", "timestamp": datetime.now().isoformat()}
    )
    
    if success:
        return f"Successfully added to knowledge base: {text[:100]}..."
    else:
        return "Error adding to knowledge base. Please check your Qdrant connection."

# Create Gradio interface
with gr.Blocks(title="GPT-4o Chatbot with Vector Store") as demo:
    gr.Markdown("""
    # GPT-4o Chatbot with Qdrant Vector Store
    
    This chatbot uses OpenAI's GPT-4o model with conversation memory and external knowledge from a Qdrant vector database.
    
    ## Features:
    - 💬 Conversation history maintained throughout the session
    - 🧠 External knowledge retrieval from Qdrant vector store
    - ➕ Add new knowledge to the vector store
    - 🎛️ Customizable parameters (temperature, max tokens, system prompt)
    """)
    
    with gr.Tab("Chat"):
        with gr.Row():
            with gr.Column(scale=3):
                chatbot = gr.Chatbot(
                    label="Conversation",
                    height=500,
                    show_copy_button=True
                )
                
                with gr.Row():
                    msg = gr.Textbox(
                        label="Your message",
                        placeholder="Type your message here...",
                        lines=2,
                        scale=4
                    )
                    submit_btn = gr.Button("Send", variant="primary", scale=1)
                
                with gr.Row():
                    clear_btn = gr.Button("Clear Chat")
                    add_to_kb_checkbox = gr.Checkbox(
                        label="Add messages to knowledge base",
                        value=False
                    )
            
            with gr.Column(scale=1):
                gr.Markdown("### Settings")
                
                system_prompt = gr.Textbox(
                    label="System Prompt",
                    value="You are a helpful AI assistant with access to a knowledge base. Use the provided context when relevant to answer questions accurately.",
                    lines=4
                )
                
                temperature = gr.Slider(
                    label="Temperature",
                    minimum=0,
                    maximum=2,
                    value=0.7,
                    step=0.1
                )
                
                max_tokens = gr.Slider(
                    label="Max Tokens",
                    minimum=50,
                    maximum=4000,
                    value=1000,
                    step=50
                )
    
    with gr.Tab("Knowledge Base"):
        gr.Markdown("""
        ### Add Knowledge to Vector Store
        
        Add information to the Qdrant vector store that will be used as context for the chatbot.
        """)
        
        knowledge_input = gr.Textbox(
            label="Knowledge Entry",
            placeholder="Enter information to add to the knowledge base...",
            lines=5
        )
        
        add_knowledge_btn = gr.Button("Add to Knowledge Base", variant="primary")
        knowledge_output = gr.Textbox(label="Status", interactive=False)
        
        gr.Markdown("""
        ### Tips:
        - Add relevant information, FAQs, documentation, or any context that would help the chatbot
        - The chatbot will search this knowledge base for relevant information when answering questions
        - Each entry is converted to embeddings and stored in Qdrant for semantic search
        """)
    
    with gr.Tab("Configuration"):
        gr.Markdown("""
        ### Environment Variables Required:
        
        1. **OPENAI_API_KEY**: Your OpenAI API key
        2. **QDRANT_URL**: Qdrant server URL (default: localhost:6333)
        3. **QDRANT_API_KEY**: Qdrant API key (optional for local instances)
        
        ### Setup Instructions:
        
        1. Set up a Qdrant instance (local or cloud)
        2. Set the required environment variables
        3. The app will automatically create a collection named 'chatbot_knowledge'
        
        ### Deployment on HuggingFace Spaces:
        
        1. Create a new Space on HuggingFace
        2. Add this code as `app.py`
        3. Add `requirements.txt` with the dependencies
        4. Set the environment variables in the Space settings
        """)
    
    # Event handlers
    def user_submit(message, history):
        return "", history + [[message, None]]
    
    def bot_response(history, system_prompt, temperature, max_tokens, add_to_kb):
        if history and history[-1][1] is None:
            user_message = history[-1][0]
            history[-1][1] = ""
            
            for response in chat_interface(user_message, history[:-1], system_prompt, temperature, max_tokens, add_to_kb):
                history[-1][1] = response
                yield history
    
    msg.submit(user_submit, [msg, chatbot], [msg, chatbot], queue=False).then(
        bot_response, [chatbot, system_prompt, temperature, max_tokens, add_to_kb_checkbox], chatbot
    )
    
    submit_btn.click(user_submit, [msg, chatbot], [msg, chatbot], queue=False).then(
        bot_response, [chatbot, system_prompt, temperature, max_tokens, add_to_kb_checkbox], chatbot
    )
    
    clear_btn.click(lambda: None, None, chatbot, queue=False)
    
    add_knowledge_btn.click(
        add_knowledge,
        inputs=[knowledge_input],
        outputs=[knowledge_output]
    )

# Launch the app
if __name__ == "__main__":
    demo.queue()
    demo.launch()