app.py

import re
import gradio as gr
import requests
import os
import json
from typing import List, Dict, Any
from langchain_chroma import Chroma
from langchain_huggingface import HuggingFaceEmbeddings
from langchain.docstore.document import Document
import spaces
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig

# Import smolagents components
from smolagents import TransformersModel, Tool
from smolagents.agents import CodeAgent

# Configuration
MODEL_NAME = "Qwen/Qwen2.5-7B-Instruct" 
CHROMA_DB_PATH = "./character_monologues"
EMBEDDING_MODEL = "sentence-transformers/all-MiniLM-L6-v2"

# Initialize embeddings for RAG
embeddings = HuggingFaceEmbeddings(model_name=EMBEDDING_MODEL)


def load_monologues(dir_path: str, db_path: str = CHROMA_DB_PATH, embedding_model: str = EMBEDDING_MODEL) -> Chroma:
    """
    Load monologues from a directory and save them to a ChromaDB database.

    Args:
        dir_path (str): The path to the directory containing the monologue files.
    """
    embeddings = HuggingFaceEmbeddings(model_name=embedding_model)
    os.makedirs(db_path, exist_ok=True)

    all_docs = []
    for filename in os.listdir(dir_path):
        if filename.endswith(".txt"):
            file_path = os.path.join(dir_path, filename)
            print(f"Processing file: {file_path}")
            
            with open(file_path, 'r', encoding='utf-8') as file:
                content = file.read()
            
            # Split the content by "---" to get individual entries
            entries = content.split("---")
            
            for entry in entries:
                entry = entry.strip()
                if not entry:
                    continue
                
                # Extract the person's name and monologue using the exact format
                person_match = re.match(r"I am ([^\.]+)\.(.*)", entry, re.DOTALL)
                
                if person_match:
                    character = person_match.group(1).strip()
                    monologue = person_match.group(2).strip()
                    
                    if not monologue:  # Skip if no monologue content
                        print(f"Warning: No monologue content found for {character}")
                        continue
                        
                    # Create a document
                    document = Document(
                        page_content=monologue,
                        metadata={"character": character}
                    )
                    
                    all_docs.append(document)
                    print(f"Successfully processed monologue for {character}")
                else:
                    print(f"Warning: Could not extract character name from entry: {entry[:50]}...")

    if not all_docs:
        raise ValueError("No valid monologues were processed. Please check the input data format.")

    # Save the documents to a ChromaDB database
    vector_store = Chroma.from_documents(
        documents=all_docs,
        embedding=embeddings,
        persist_directory=db_path
    )
    print(f"Saved {len(all_docs)} monologues to {db_path}")
    return vector_store


# Load or create vector store
def get_vector_store(monologues_dir="data/monologues"):
    # Check if ChromaDB already exists
    if os.path.exists(CHROMA_DB_PATH):
        print(f"Loading existing ChromaDB from {CHROMA_DB_PATH}")
        return Chroma(persist_directory=CHROMA_DB_PATH, embedding_function=embeddings)
    else:
        print(f"ChromaDB not found at {CHROMA_DB_PATH}")
        # Check if monologues directory exists
        if os.path.exists(monologues_dir):
            print(f"Loading monologues from {monologues_dir}")
            # Create new ChromaDB from monologues
            return load_monologues(monologues_dir, CHROMA_DB_PATH, EMBEDDING_MODEL)
        else:
            print("Warning: Neither ChromaDB nor monologues directory found. RAG functionality will be limited.")
            return None
        

# Define a Wikipedia Tool for smolagents
class WikipediaTool(Tool):
    name = "wikipedia"
    description = "Fetches information about a philosophical topic from Wikipedia."
    inputs = {
        "topic": {
            "type": "string",
            "description": "The philosophical topic or concept to search for on Wikipedia."
        }
    }
    output_type = "string"

    def forward(self, topic: str) -> str:
        try:
            # Clean the topic for URL
            topic = topic.strip().replace(" ", "_")
            url = f"https://en.wikipedia.org/api/rest_v1/page/summary/{topic}"
            response = requests.get(url)
            
            if response.status_code == 200:
                data = response.json()
                return f"Wikipedia information about '{topic}':\n\n{data.get('extract', 'No detailed information found.')}"
            else:
                return f"Error retrieving information from Wikipedia. Status code: {response.status_code}"
        except Exception as e:
            return f"Error querying Wikipedia: {str(e)}"

# Define a Character Dialogue Tool for smolagents
class CharacterToneTool(Tool):
    name = "character_dialogue"
    description = "Retrieves a monologue from a specific character to define their tone and style."
    inputs = {
        "character": {
            "type": "string",
            "description": "The philosopher character whose tone and style should be emulated."
        }
    }
    output_type = "string"

    def __init__(self, vector_store, **kwargs):
        super().__init__(**kwargs)
        self.vector_store = vector_store

    def forward(self, character: str) -> str:
        if self.vector_store is None:
            return f"No character dialogue database found for {character}."
        
        # Construct a search query combining the character and philosophical query
        search_query = f"{character}"
        
        # Retrieve relevant documents
        docs = self.vector_store.similarity_search(search_query, k=3)
        
        if not docs:
            return f"No specific dialogue examples found for {character} discussing this topic."
        
        # Format the dialogue examples
        result = f"Monologue from {character}:\n\n"
        for i, doc in enumerate(docs):
            result += f"Example {i+1}:\n{doc.page_content}\n\n"
        
        return result

# List of characters (add or modify based on your database)
CHARACTERS = [
    "Socrates",
    "Confucius",
    "Albert Camus",
    "Frodo Baggins",
    "Harry Potter",
    "Sherlock Holmes",
    "Joker",
    "Martin Luther King Jr.",
    "Mahatma Gandhi",
    "Albert Einstein",
    "Jack Sparrow",
    "Don Quixote",
    "Leonardo da Vinci",
    "Princess Diana",
]

# Main function for Gradio interface
def main():
    # Load vector store for RAG
    vector_store = get_vector_store()
    
    # Initialize tools
    wikipedia_tool = WikipediaTool()
    character_dialogue_tool = CharacterToneTool(vector_store)

    # Initialize quantization configuration
    bnb_config = BitsAndBytesConfig(
        load_in_4bit=True,
        bnb_4bit_quant_type="nf4",
        bnb_4bit_compute_dtype="float16",
        bnb_4bit_use_double_quant=True,
    )

    model = TransformersModel(
        model_id=MODEL_NAME,
        max_new_tokens=4096,
        device_map="auto",
        # quantization_config=bnb_config  # doesn't work yet
    )

    # Define the chat function using smolagents
    @spaces.GPU(duration=60)
    def chat_function(message: str, character: str) -> str:
        # Create a CodeAgent for this specific query
        agent = CodeAgent(
            tools=[wikipedia_tool, character_dialogue_tool],
            model=model,
            max_steps=3,  # Limit the number of steps for faster responses
            verbosity_level=1,  # Set to 0 in production for less verbose outputs
        )
        
        prompt = f"""You are a fun philosophical AI assistant speaking in the tone and style of {character}.
        
The user has asked: "{message}"

Please follow these steps:
1. Use the wikipedia tool to get information about this philosophical topic.
2. Use the character_dialogue tool to see how {character} typically discusses this kind of topic.
3. Synthesize a response that answers the philosophical query in the authentic tone and style of {character}, drawing on both the factual information and the character's unique perspective.

Maintain {character}'s unique voice, philosophical outlook, and typical expressions throughout your response.
"""
        
        response = agent.run(prompt)

        return response.strip()
    
    # Create Gradio interface
    with gr.Blocks(title="Philosophical Character Chat") as demo:
        gr.Markdown("# Philosophical Character Chat")
        gr.Markdown("Ask philosophical questions and get answers in the tone of different historical philosophers.")
        
        with gr.Row():
            with gr.Column(scale=4):
                chatbot = gr.Chatbot(height=600)
                msg = gr.Textbox(
                    label="Your philosophical query",
                    placeholder="What is knowledge?",
                    lines=2
                )
                submit_btn = gr.Button("Ask")
            
            with gr.Column(scale=1):
                character = gr.Dropdown(
                    choices=CHARACTERS,
                    label="Choose a Philosopher",
                    value=CHARACTERS[0]
                )
                gr.Markdown("## About this app")
                gr.Markdown(
                    """This application answers philosophical queries in the style of different philosophers.
                    
                    It uses:
                    - Qwen2.5-7B-Instruct language model via smolagents
                    - RAG with ChromaDB for character dialogue examples
                    - Wikipedia API for philosophical information
                    """
                )
        
        def respond(message, chat_history, character):
            if not message.strip():
                return chat_history
            
            # Get model response using the agent
            bot_response = chat_function(message, character)
            
            # Update chat history
            chat_history.append((message, bot_response))
            return "", chat_history
        
        submit_btn.click(
            respond,
            inputs=[msg, chatbot, character],
            outputs=[msg, chatbot]
        )
        
        msg.submit(
            respond,
            inputs=[msg, chatbot, character],
            outputs=[msg, chatbot]
        )
    
    # Launch the demo
    demo.launch()

if __name__ == "__main__":
    main()