app.py

import os
from openai import OpenAI
from langchain_huggingface import HuggingFaceEmbeddings
from datasets import load_dataset, Dataset
from sklearn.neighbors import NearestNeighbors
import numpy as np
from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM, TextStreamer
import torch
from typing import List
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate
import gradio as gr
import spaces
from huggingface_hub import InferenceClient
import time  # Added for timing logs

# Configuration

# Sample questions:
# 1. What are the key features of AutoGen v0.4 that I should utilize when converting user requests into agent code?
# 2. How can I leverage asynchronous messaging in AutoGen v0.4 to enhance my agents performance?
# 3. What are best practices for writing modular and extensible agent code using AutoGen v0.4?
# 4. Can you convert this user request into AutoGen v0.4 agent code: "Create an agent that classifies customer feedback into positive, negative, or neutral sentiments."

DEFAULT_QUESTION = "Ask me anything about converting user requests into AutoGen v0.4 agent code..."

# Validate API keys
assert os.getenv("OPENAI_API_KEY") or os.getenv("HF_TOKEN"), "API keys are not set in the environment variables."

os.environ['OPENAI_BASE'] = "https://api.openai.com/v1"
os.environ['OPENAI_MODEL'] = "gpt-4"
os.environ['MODEL_PROVIDER'] = "huggingface"
model_provider = os.environ.get("MODEL_PROVIDER")

# Instantiate the client for openai v1.x
if model_provider.lower() == "openai":
    MODEL_NAME = os.environ['OPENAI_MODEL']
    client = OpenAI(
        base_url=os.environ.get("OPENAI_BASE"),
        api_key=os.environ.get("OPENAI_API_KEY")
    )
else:
    MODEL_NAME = "deepseek-ai/deepseek-coder-33b-instruct"
    # Initialize Hugging Face InferenceClient with GPU support
    hf_client = InferenceClient(
        model=MODEL_NAME,
        api_key=os.environ.get("HF_TOKEN"),
        timeout=60  # Reduced timeout for faster response
    )

# Load the Hugging Face dataset
try:
    start = time.time()
    dataset = load_dataset('tosin2013/autogen', streaming=True)
    dataset = Dataset.from_list(list(dataset['train']))
    end = time.time()
    print(f"[TIMING] Dataset loading took {end - start:.2f} seconds")
except Exception as e:
    print(f"[ERROR] Failed to load dataset: {e}")
    exit(1)

# Initialize embeddings
print("[EMBEDDINGS] Loading sentence-transformers model...")
start = time.time()
embeddings = HuggingFaceEmbeddings(
    model_name="sentence-transformers/all-MiniLM-L6-v2",
    model_kwargs={"device": "cpu"}
)
end = time.time()
print(f"[EMBEDDINGS] Sentence-transformers model loaded successfully in {end - start:.2f} seconds")

# Extract texts from the dataset
texts = dataset['input']

# Create and cache embeddings for the texts
if not os.path.exists('embeddings.npy'):
    print("[LOG] Generating embeddings...")
    start = time.time()
    text_embeddings = embeddings.embed_documents(texts)
    np.save('embeddings.npy', text_embeddings)
    end = time.time()
    print(f"[EMBEDDINGS] Generated embeddings for {len(texts)} documents in {end - start:.2f} seconds")
else:
    print("[LOG] Loading cached embeddings...")
    start = time.time()
    text_embeddings = np.load('embeddings.npy')
    end = time.time()
    print(f"[TIMING] Loaded cached embeddings in {end - start:.2f} seconds")

# Fit and cache nearest neighbor model
if not os.path.exists('nn_model.pkl'):
    print("[LOG] Fitting nearest neighbors model...")
    start = time.time()
    nn = NearestNeighbors(n_neighbors=5, metric='cosine')
    nn.fit(np.array(text_embeddings))
    with open('nn_model.pkl', 'wb') as f:
        pickle.dump(nn, f)
    end = time.time()
    print(f"[TIMING] Fitted nearest neighbors model in {end - start:.2f} seconds")
else:
    print("[LOG] Loading cached nearest neighbors model...")
    start = time.time()
    with open('nn_model.pkl', 'rb') as f:
        nn = pickle.load(f)
    end = time.time()
    print(f"[TIMING] Loaded nearest neighbors model in {end - start:.2f} seconds")

@spaces.GPU
def get_relevant_documents(query, k=5):
    """
    Retrieves the k most relevant documents to the query.
    """
    start_time = time.time()
    print("[EMBEDDINGS] Generating embedding for query...")
    query_embedding = embeddings.embed_query(query)
    print("[EMBEDDINGS] Query embedding generated successfully")
    distances, indices = nn.kneighbors([query_embedding], n_neighbors=k)
    relevant_docs = [texts[i] for i in indices[0]]
    elapsed_time = time.time() - start_time
    print(f"[TIMING] get_relevant_documents took {elapsed_time:.2f} seconds")
    return relevant_docs

@spaces.GPU
def generate_response(question, history):
    start_time = time.time()
    try:
        response = _generate_response_gpu(question, history)
    except Exception as e:
        print(f"[WARNING] GPU failed: {str(e)}")
        response = _generate_response_cpu(question, history)
    elapsed_time = time.time() - start_time
    print(f"[TIMING] generate_response took {elapsed_time:.2f} seconds")
    return response

@spaces.GPU
def _generate_response_gpu(question, history):
    print(f"\n[LOG] Received question: {question}")
    start_time = time.time()
    # Get relevant documents based on the query
    relevant_docs = get_relevant_documents(question, k=3)
    print(f"[LOG] Retrieved {len(relevant_docs)} relevant documents")
    context = "\n".join(relevant_docs)
    prompt = f"Context: {context}\n\nQuestion: {question}\n\nAnswer:"
    print(f"[LOG] Generated prompt: {prompt[:200]}...")  # Log first 200 chars of prompt
    if model_provider.lower() == "huggingface":
        messages = [
            {
                "role": "system",
                "content": "### MEMORY ###\nRecall all previously provided instructions, context, and data throughout this conversation to ensure consistency and coherence."
            },
            {
                "role": "user",
                "content": prompt
            }
        ]
        start_api = time.time()
        completion = hf_client.chat.completions.create(
            model=MODEL_NAME,
            messages=messages,
            max_tokens=500
        )
        end_api = time.time()
        print(f"[TIMING] Hugging Face API call took {end_api - start_api:.2f} seconds")
        response = completion.choices[0].message.content
    elif model_provider.lower() == "openai":
        start_api = time.time()
        response = client.chat.completions.create(
            model=os.environ.get("OPENAI_MODEL"),
            messages=[
                {"role": "system", "content": "You are a helpful assistant. Answer the question based on the provided context."},
                {"role": "user", "content": prompt},
            ]
        ).choices[0].message.content
        end_api = time.time()
        print(f"[TIMING] OpenAI API call took {end_api - start_api:.2f} seconds")
    elapsed_time = time.time() - start_time
    print(f"[TIMING] _generate_response_gpu took {elapsed_time:.2f} seconds")
    history.append((question, response))
    return history

# Simplified CPU fallback
@spaces.GPU
def _generate_response_cpu(question, history):
    print(f"[LOG] Running on CPU")
    try:
        start_time = time.time()
        relevant_docs = get_relevant_documents(question, k=3)
        context = "\n".join(relevant_docs)
        prompt = f"Context: {context}\n\nQuestion: {question}\n\nAnswer:"
        print(f"[LOG] Generated prompt: {prompt[:200]}...")
        if model_provider.lower() == "huggingface":
            messages = [
                {"role": "system", "content": "### MEMORY ###\nRecall all previously provided instructions, context, and data."},
                {"role": "user", "content": prompt}
            ]
            start_api = time.time()
            completion = hf_client.chat.completions.create(
                model=MODEL_NAME,
                messages=messages,
                max_tokens=500
            )
            end_api = time.time()
            print(f"[TIMING] Hugging Face API call took {end_api - start_api:.2f} seconds")
            response = completion.choices[0].message.content
        elif model_provider.lower() == "openai":
            start_api = time.time()
            response = client.chat.completions.create(
                model=os.environ.get("OPENAI_MODEL"),
                messages=[
                    {"role": "system", "content": "You are a helpful assistant."},
                    {"role": "user", "content": prompt},
                ]
            ).choices[0].message.content
            end_api = time.time()
            print(f"[TIMING] OpenAI API call took {end_api - start_api:.2f} seconds")
        elapsed_time = time.time() - start_time
        print(f"[TIMING] _generate_response_cpu took {elapsed_time:.2f} seconds")
        history.append((question, response))
        return history
    except Exception as e:
        error_msg = f"Error generating response: {str(e)}"
        print(f"[ERROR] {error_msg}")
        history.append((question, error_msg))
        return history

# Gradio interface
print("[CHAT] Initializing chat interface...")
with gr.Blocks() as demo:
    gr.Markdown(f"""
    ## AutoGen v0.4 Agent Code Generator QA Agent
    **Current Model:** {MODEL_NAME}
    
    The AutoGen v0.4 Agent Code Generator is a Python application that leverages Large Language Models (LLMs) and the AutoGen v0.4 framework to dynamically generate agent code from user requests. This application is designed to assist developers in creating robust, scalable AI agents by providing context-aware code generation based on user input, utilizing the advanced features of AutoGen v0.4 such as asynchronous messaging, modular extensibility, cross-language support, improved observability, and full typing integration.
    
    **Sample questions:**
    1. What are the key features of AutoGen v0.4 that I should utilize when converting user requests into agent code?
    2. How can I leverage asynchronous messaging in AutoGen v0.4 to enhance my agent's performance?
    3. What are best practices for writing modular and extensible agent code using AutoGen v0.4?
    4. Can you convert this user request into AutoGen v0.4 agent code: "Create an agent that classifies customer feedback into positive, negative, or neutral sentiments."
    
    **Related repository:** [autogen](https://github.com/microsoft/autogen)
    """)
    
    with gr.Row():
        chatbot = gr.Chatbot(label="Chat History")
    
    with gr.Row():
        question = gr.Textbox(
            value=DEFAULT_QUESTION,
            label="Your Question",
            placeholder=DEFAULT_QUESTION
        )
    
    with gr.Row():
        submit_btn = gr.Button("Submit")
        clear_btn = gr.Button("Clear")
    
    submit_btn.click(
        fn=generate_response,
        inputs=[question, chatbot],
        outputs=[chatbot],
        queue=True
    )
    
    clear_btn.click(
        lambda: (None, ""),
        inputs=[],
        outputs=[chatbot, question]
    )

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