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 huggingface_hub import InferenceClient
import gradio as gr
import spaces
import pickle
import time

# Configuration
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 (either OPENAI_API_KEY or HF_TOKEN)."

# Model Provider Configuration
MODEL_PROVIDER = os.getenv("MODEL_PROVIDER", "huggingface").lower()  # Default to Hugging Face
if MODEL_PROVIDER == "openai":
    OPENAI_BASE = os.getenv("OPENAI_BASE", "https://api.openai.com/v1")
    OPENAI_MODEL = os.getenv("OPENAI_MODEL", "gpt-4")
    client = OpenAI(base_url=OPENAI_BASE, api_key=os.getenv("OPENAI_API_KEY"))
    MODEL_NAME = OPENAI_MODEL
elif MODEL_PROVIDER == "huggingface":
    HF_MODEL_NAME = os.getenv("HF_MODEL_NAME", "deepseek-ai/deepseek-coder-33b-instruct")
    HF_TOKEN = os.getenv("HF_TOKEN")
    hf_client = InferenceClient(model=HF_MODEL_NAME, token=HF_TOKEN, timeout=120)
    MODEL_NAME = HF_MODEL_NAME
else:
    raise ValueError(f"Unsupported MODEL_PROVIDER: {MODEL_PROVIDER}. Choose 'openai' or 'huggingface'.")

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

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

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

# Create and load embeddings and Nearest Neighbors model (in-memory caching)
EMBEDDINGS_FILE = 'embeddings.npy'
NN_MODEL_FILE = 'nn_model.pkl'
text_embeddings = None
nn = None

if os.path.exists(EMBEDDINGS_FILE):
    print("[LOG] Loading cached embeddings...")
    text_embeddings = np.load(EMBEDDINGS_FILE)
else:
    print("[LOG] Generating embeddings...")
    text_embeddings = embeddings.embed_documents(texts)
    print(f"[EMBEDDINGS] Generated embeddings for {len(texts)} documents")
    np.save(EMBEDDINGS_FILE, text_embeddings)

if os.path.exists(NN_MODEL_FILE):
    print("[LOG] Loading cached nearest neighbors model...")
    with open(NN_MODEL_FILE, 'rb') as f:
        nn = pickle.load(f)
else:
    print("[LOG] Fitting nearest neighbors model...")
    nn = NearestNeighbors(n_neighbors=5, metric='cosine')
    nn.fit(np.array(text_embeddings))
    with open(NN_MODEL_FILE, 'wb') as f:
        pickle.dump(nn, f)

@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"[PERF] get_relevant_documents took {elapsed_time:.2f} seconds")
    return relevant_docs

def generate_response(question, history):
    """Generates a response to the user's question, handling GPU/CPU fallback."""
    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"[PERF] generate_response took {elapsed_time:.2f} seconds")
    return history, history  # Return updated history twice for Gradio

@spaces.GPU
def _generate_response_gpu(question, history):
    """Generates a response using the GPU."""
    print(f"\n[LOG] Received question: {question}")
    relevant_docs = get_relevant_documents(question, k=3)
    print(f"[LOG] Retrieved {len(relevant_docs)} relevant documents")
    context = "\n".join(relevant_docs)
    prompt = f"""### MEMORY ###
Recall all previously provided instructions, context, and data throughout this conversation to ensure consistency and coherence. Use the details from the last interaction to guide your response.
### VISIONARY GUIDANCE ###
This prompt is designed to empower users to seamlessly convert their requests into AutoGen v0.4 agent code. By harnessing the advanced features of AutoGen v0.4, we aim to provide a scalable and flexible solution that is both user-friendly and technically robust. The collaborative effort of the personas ensures a comprehensive, innovative, and user-centric approach to meet the user's objectives.
### CONTEXT ###
AutoGen v0.4 is a comprehensive rewrite aimed at building robust, scalable, and cross-language AI agents. Key features include asynchronous messaging, scalable distributed agents support, modular extensibility, cross-language capabilities, improved observability, and full typing integration.
### OBJECTIVE ###
Translate user requests into AutoGen v0.4 agent code that leverages the framework's new features. Ensure the code is syntactically correct, scalable, and aligns with best practices.
### STYLE ###
Professional, clear, and focused on code quality.
### TONE ###
Informative, helpful, and user-centric.
### AUDIENCE ###
Users seeking to implement their requests using AutoGen v0.4 agents.
### RESPONSE FORMAT ###
Provide the AutoGen v0.4 agent code that fulfills the user's request. Utilize features like asynchronous messaging and modular design where appropriate. Include comments to explain key components and enhance understandability.
### TEAM PERSONAS’ CONTRIBUTIONS ###
- **Analyst:** Ensured the prompt provides clear, structured instructions to accurately convert user requests into code, emphasizing full typing integration for precision.
- **Creative:** Suggested incorporating comments and explanations within the code to foster innovative usage and enhance user engagement with AutoGen v0.4 features.
- **Strategist:** Focused on aligning the prompt with long-term scalability by encouraging the use of modular and extensible design principles inherent in AutoGen v0.4.
- **Empathizer:** Enhanced the prompt to be user-centric, ensuring it addresses user needs effectively and makes the code accessible and easy to understand.
- **Researcher:** Integrated the latest information about AutoGen v0.4, ensuring the prompt and generated code reflect current capabilities and best practices.
### SYSTEM GUARDRAILS ###
- If unsure about the user's request, ask clarifying questions rather than making assumptions.
- Do not fabricate data or features not supported by AutoGen v0.4.
- Ensure the code is scalable, modular, and adheres to best practices.
### START ###
Context: {context}\n\nQuestion: {question}\n\nAnswer:"""
    print(f"[LOG] Generated prompt: {prompt[:200]}...")

    if MODEL_PROVIDER == "huggingface":
        messages = [{"role": "user", "content": prompt}]
        completion = hf_client.chat.completions.create(model=MODEL_NAME, messages=messages, max_tokens=500)
        response = completion.choices[0].message.content
        print(f"[LOG] Using Hugging Face model (serverless): {MODEL_NAME}")
        print(f"[LOG] Hugging Face response: {response[:200]}...")
    elif MODEL_PROVIDER == "openai":
        response = client.chat.completions.create(
            model=OPENAI_MODEL,
            messages=[{"role": "user", "content": prompt}]
        ).choices[0].message.content
        print(f"[LOG] Using OpenAI model: {OPENAI_MODEL}")
        print(f"[LOG] OpenAI response: {response[:200]}...")

    history.append((question, response))
    return history

def _generate_response_cpu(question, history):
    """Generates a response using the CPU (fallback)."""
    print(f"[LOG] Running on CPU")
    try:
        relevant_docs = get_relevant_documents(question, k=3)
        context = "\n".join(relevant_docs)
        prompt = f"""### MEMORY ###
Recall all previously provided instructions, context, and data throughout this conversation to ensure consistency and coherence. Use the details from the last interaction to guide your response.
### SYSTEM GUARDRAILS ###
If unsure about the user's request, ask clarifying questions rather than making assumptions.
Do not fabricate data or features not supported by AutoGen v0.4.
Ensure the code is scalable, modular, and adheres to best practices.
Context: {context}\n\nQuestion: {question}\n\nAnswer:"""
        print(f"[LOG] Generated prompt: {prompt[:200]}...")

        if MODEL_PROVIDER == "huggingface":
            messages = [{"role": "user", "content": prompt}]
            completion = hf_client.chat.completions.create(model=MODEL_NAME, messages=messages, max_tokens=500)
            response = completion.choices[0].message.content
        elif MODEL_PROVIDER == "openai":
            response = client.chat.completions.create(
                model=OPENAI_MODEL,
                messages=[{"role": "user", "content": prompt}]
            ).choices[0].message.content

        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}
    **Model Provider:** {MODEL_PROVIDER.upper()}

    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)
    """)

    chatbot = gr.Chatbot(label="Chat History")
    question_textbox = gr.Textbox(value=DEFAULT_QUESTION, label="Your Question", placeholder=DEFAULT_QUESTION)
    submit_button = gr.Button("Submit")
    clear_button = gr.Button("Clear")

    submit_button.click(
        fn=generate_response,
        inputs=[question_textbox, chatbot],
        outputs=[chatbot, chatbot],  # Output the updated history to the chatbot
        queue=True
    )

    clear_button.click(
        lambda: (None, ""),
        inputs=[],
        outputs=[chatbot, question_textbox]
    )

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