app.py

import os
from openai import OpenAI
from langchain_community.embeddings 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
from huggingface_hub import InferenceClient

# 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..."

# Set API keys (make sure these are set in your environment)
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=api_key
    )
else:
    MODEL_NAME = "meta-llama/Llama-3.3-70B-Instruct"
    # Initialize Hugging Face InferenceClient with GPU support
    hf_client = InferenceClient(
        model=MODEL_NAME,
        api_key=os.environ.get("HF_TOKEN"),
        timeout=120  # Increased timeout for GPU inference
    )

# Load the Hugging Face dataset
dataset = load_dataset('tosin2013/autogen', streaming=True)
dataset = Dataset.from_list(list(dataset['train']))

# Initialize embeddings with GPU support if available
device = "cuda" if torch.cuda.is_available() else "cpu"
embeddings = HuggingFaceEmbeddings(
    model_name="sentence-transformers/all-MiniLM-L6-v2",
    model_kwargs={"device": device}
)

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

# Create embeddings for the texts
text_embeddings = embeddings.embed_documents(texts)

# Fit a nearest neighbor model
nn = NearestNeighbors(n_neighbors=5, metric='cosine')
nn.fit(np.array(text_embeddings))

def get_relevant_documents(query, k=5):
    """
    Retrieves the k most relevant documents to the query.
    """
    query_embedding = embeddings.embed_query(query)
    distances, indices = nn.kneighbors([query_embedding], n_neighbors=k)
    relevant_docs = [texts[i] for i in indices[0]]
    return relevant_docs

def generate_response(question, history):
    try:
        print(f"\n[LOG] Received question: {question}")
        
        # Get relevant documents based on the query
        relevant_docs = get_relevant_documents(question, k=3)
        print(f"[LOG] Retrieved {len(relevant_docs)} relevant documents")
        
        # Create the prompt for the LLM
        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 ###
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 ###
'''
                },
                {
                    "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.lower() == "openai":
            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},
                ]
            )
            response = response.choices[0].message.content
            print(f"[LOG] Using OpenAI model: {os.environ.get('OPENAI_MODEL')}")
            print(f"[LOG] OpenAI response: {response[:200]}...")  # Log first 200 chars of response
        
        # Update chat history with new message pair
        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

# Create Gradio 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")
    
    # Event handlers
    submit_btn.click(
        generate_response,
        inputs=[question, chatbot],
        outputs=[chatbot]
    )
    
    clear_btn.click(
        lambda: (None, ""),
        inputs=[],
        outputs=[chatbot, question]
    )

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