Update app.py

app.py

@@ -4,112 +4,81 @@ 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
+from huggingface_hub import InferenceClient
 import gradio as gr
 import spaces
-from huggingface_hub import InferenceClient
-import time  # Added for timing logs
+import pickle
+import time
 
 # 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")
-    )
+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:
-    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
-    )
+    raise ValueError(f"Unsupported MODEL_PROVIDER: {MODEL_PROVIDER}. Choose 'openai' or 'huggingface'.")
 
 # 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")
+print("[EMBEDDINGS] Sentence-transformers model loaded successfully")
 
 # Extract texts from the dataset
 texts = dataset['input']
 
-# Create and cache embeddings for the texts
-if not os.path.exists('embeddings.npy'):
+# 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...")
-    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")
+    print(f"[EMBEDDINGS] Generated embeddings for {len(texts)} documents")
+    np.save(EMBEDDINGS_FILE, text_embeddings)
 
-# Fit and cache nearest neighbor model
-if not os.path.exists('nn_model.pkl'):
+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...")
-    start = time.time()
     nn = NearestNeighbors(n_neighbors=5, metric='cosine')
     nn.fit(np.array(text_embeddings))
-    with open('nn_model.pkl', 'wb') as f:
+    with open(NN_MODEL_FILE, '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.
-    """
+    """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)
@@ -117,11 +86,11 @@ def get_relevant_documents(query, k=5):
     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")
+    print(f"[PERF] get_relevant_documents took {elapsed_time:.2f} seconds")
     return relevant_docs
 
-@spaces.GPU
 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)
@@ -129,92 +98,88 @@ def generate_response(question, history):
         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
+    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}")
-    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")
+    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
-    elif model_provider.lower() == "openai":
-        start_api = time.time()
+        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=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},
-            ]
+            model=OPENAI_MODEL,
+            messages=[{"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")
+        print(f"[LOG] Using OpenAI model: {OPENAI_MODEL}")
+        print(f"[LOG] OpenAI response: {response[:200]}...")
+
     history.append((question, response))
     return history
 
-# Simplified CPU fallback
-@spaces.GPU
 def _generate_response_cpu(question, history):
+    """Generates a response using the CPU (fallback)."""
     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:"
+        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.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")
+
+        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.lower() == "openai":
-            start_api = time.time()
+        elif MODEL_PROVIDER == "openai":
             response = client.chat.completions.create(
-                model=os.environ.get("OPENAI_MODEL"),
-                messages=[
-                    {"role": "system", "content": "You are a helpful assistant."},
-                    {"role": "user", "content": prompt},
-                ]
+                model=OPENAI_MODEL,
+                messages=[{"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:
@@ -223,50 +188,42 @@ def _generate_response_cpu(question, history):
         history.append((question, error_msg))
         return history
 
-# Gradio interface
+# 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)
     """)
-    
-    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(
+
+    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, chatbot],
-        outputs=[chatbot],
+        inputs=[question_textbox, chatbot],
+        outputs=[chatbot, chatbot],  # Output the updated history to the chatbot
         queue=True
     )
-    
-    clear_btn.click(
+
+    clear_button.click(
         lambda: (None, ""),
         inputs=[],
-        outputs=[chatbot, question]
+        outputs=[chatbot, question_textbox]
     )
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()