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

# 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=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=30  # Reduced timeout for faster response
    )

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

# Initialize embeddings
embeddings = HuggingFaceEmbeddings(
    model_name="sentence-transformers/all-MiniLM-L6-v2",
    model_kwargs={"device": "cpu"}
)

# 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...")
    text_embeddings = embeddings.embed_documents(texts)
    np.save('embeddings.npy', text_embeddings)
else:
    print("[LOG] Loading cached embeddings...")
    text_embeddings = np.load('embeddings.npy')

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

def get_relevant_documents(query, k=5):
    """
    Retrieves the k most relevant documents to the query.
    """
    import time
    start_time = time.time()
    
    query_embedding = embeddings.embed_query(query)
    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):
    import time
    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 response

def _generate_response_gpu(question, history):
    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

def _generate_response_cpu(question, history):
    print(f"[LOG] Running on CPU")
    try:
        # 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":
            # Use CPU version of the model
            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 (CPU): {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(
        fn=generate_response,
        inputs=[question, chatbot],
        outputs=[chatbot],
        queue=True
    )
    
    clear_btn.click(
        lambda: (None, ""),
        inputs=[],
        outputs=[chatbot, question]
    )

import socket

def find_available_port(start_port=7860, end_port=7900):
    for port in range(start_port, end_port + 1):
        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
            try:
                s.bind(('', port))
                return port
            except OSError:
                continue
    raise OSError(f"No available ports between {start_port} and {end_port}")

if __name__ == "__main__":
    try:
        port = find_available_port()
        print(f"[LOG] Launching application on port {port}")
        demo.launch()
        # Verify server is actually running
        import time
        time.sleep(2)  # Give server time to start
        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
            if s.connect_ex(('localhost', port)) == 0:
                print(f"[SUCCESS] Server is running on port {port}")
            else:
                print(f"[ERROR] Failed to bind to port {port}")
    except Exception as e:
        print(f"[ERROR] Failed to start application: {str(e)}")