Update app.py

app.py

@@ -15,67 +15,109 @@ logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(
 
 # Constants
 CONTEXT_SIZES = {
-    "4K": 4096,  # Corrected to power of 2
-    "8K": 8192,  # Corrected to power of 2
-    "32K": 32768, # Corrected to power of 2
-    "64K": 65536, # Added 64k option
-    "128K": 131072 # Corrected to power of 2
+   "4K": 4096,  
+   "8K": 8192,  
+   "32K": 32768,
+   "64K": 65536,
+   "128K": 131072
 }
 
 MODEL_CONTEXT_SIZES = {
-    "OpenAI ChatGPT": 4096,
-    "HuggingFace Inference": 4096,
-    "Groq API": {
-        "llama-3.1-70b-versatile": 32768,
-        "mixtral-8x7b-32768": 32768,
-        "llama-3.1-8b-instant": 8192
-    }
+   "Clipboard only": 4096,
+   "OpenAI ChatGPT": {
+       "gpt-3.5-turbo": 4096,
+       "gpt-4": 8192, 
+       "gpt-4-32k": 32768
+   },
+   "HuggingFace Inference": {
+       "microsoft/phi-3-mini-4k-instruct": 4096,
+       "HuggingFaceH4/zephyr-7b-beta": 8192,
+       "deepseek-ai/DeepSeek-Coder-V2-Instruct": 8192,
+       "meta-llama/Llama-3-8b-Instruct": 8192,
+       "mistralai/Mistral-7B-Instruct-v0.3": 32768,
+       "NousResearch/Nous-Hermes-2-Mixtral-8x7B-DPO": 32768
+   },
+   "Groq API": {
+       "gemma-7b-it": 8192,
+       "llama-3.1-70b": 32768,
+       "mixtral-8x7b-32768": 32768,
+       "llama-3.1-8b": 8192
+   }
 }
 
 class ModelRegistry:
-    def __init__(self):
-        self.hf_models = {
-            "Phi-3 Mini 128k": "microsoft/Phi-3-mini-128k-instruct",
-            "Custom Model": ""
-        }
-        self.groq_models = self._fetch_groq_models()
-
-    def _fetch_groq_models(self) -> Dict[str, str]:
-        """Fetch available Groq models with proper error handling"""
-        try:
-            groq_api_key = os.getenv('GROQ_API_KEY')
-            if not groq_api_key:
-                logging.warning("No GROQ_API_KEY found in environment")
-                return self._get_default_groq_models()
-
-            headers = {
-                "Authorization": f"Bearer {groq_api_key}",
-                "Content-Type": "application/json"
-            }
-            response = requests.get("https://api.groq.com/openai/v1/models", headers=headers)
-            
-            if response.status_code == 200:
-                models = response.json().get("data", [])
-                return {model["id"]: model["id"] for model in models}
-            else:
-                logging.error(f"Failed to fetch Groq models: {response.status_code}")
-                return self._get_default_groq_models()
-        except Exception as e:
-            logging.error(f"Error fetching Groq models: {e}")
-            return self._get_default_groq_models()
-
-    def _get_default_groq_models(self) -> Dict[str, str]:
-        """Return default Groq models when API is unavailable"""
-        return {
-            "llama-3.1-70b-versatile": "llama-3.1-70b-versatile",
-            "mixtral-8x7b-32768": "mixtral-8x7b-32768",
-            "llama-3.1-8b-instant": "llama-3.1-8b-instant"
-        }
-
-    def refresh_groq_models(self) -> Dict[str, str]:
-        """Refresh the list of available Groq models"""
-        self.groq_models = self._fetch_groq_models()
-        return self.groq_models
+   def __init__(self):
+       # HuggingFace Models
+       self.hf_models = {
+           "Phi-3 Mini 4K": "microsoft/phi-3-mini-4k-instruct",
+           "Phi-3 Mini 128k": "microsoft/Phi-3-mini-128k-instruct", 
+           "Zephyr 7B Beta": "HuggingFaceH4/zephyr-7b-beta",
+           "DeepSeek Coder V2": "deepseek-ai/DeepSeek-Coder-V2-Instruct",
+           "Meta Llama 3.1 8B": "meta-llama/Llama-3-8b-Instruct",
+           "Meta Llama 3.1 70B": "meta-llama/Meta-Llama-3.1-70B-Instruct",
+           "Mixtral 7B": "mistralai/Mistral-7B-Instruct-v0.3", 
+           "Nous-Hermes": "NousResearch/Nous-Hermes-2-Mixtral-8x7B-DPO",
+           "Cohere Command R+": "CohereForAI/c4ai-command-r-plus",
+           "Aya 23-35B": "CohereForAI/aya-23-35B",
+           "Custom Model": ""
+       }
+       
+       # Default Groq Models
+       self.default_groq_models = {
+           "gemma-7b-it": "gemma-7b-it",
+           "llama-3.1-70b-8192": "llama-3.1-70b-8192",
+           "llama-3.1-70b-versatile": "llama-3.1-70b-versatile",
+           "mixtral-8x7b-32768": "mixtral-8x7b-32768",
+           "llama-3.1-8b-instant": "llama-3.1-8b-instant",
+           "llama-3.1-70b-8192-tool-use-preview": "llama3-groq-70b-8192-tool-use-preview"
+       }
+       
+       self.groq_models = self._fetch_groq_models()
+
+   def _fetch_groq_models(self) -> Dict[str, str]:
+       """Fetch available Groq models with proper error handling"""
+       try:
+           groq_api_key = os.getenv('GROQ_API_KEY')
+           if not groq_api_key:
+               logging.warning("No GROQ_API_KEY found in environment")
+               return self.default_groq_models
+
+           headers = {
+               "Authorization": f"Bearer {groq_api_key}",
+               "Content-Type": "application/json"
+           }
+           
+           response = requests.get(
+               "https://api.groq.com/openai/v1/models", 
+               headers=headers,
+               timeout=10
+           )
+           
+           if response.status_code == 200:
+               models = response.json().get("data", [])
+               model_dict = {model["id"]: model["id"] for model in models}
+               
+               # Merge with defaults to ensure all models are available
+               return {**self.default_groq_models, **model_dict}
+           else:
+               logging.error(f"Failed to fetch Groq models: {response.status_code}")
+               return self.default_groq_models
+               
+       except requests.exceptions.Timeout:
+           logging.error("Timeout while fetching Groq models")
+           return self.default_groq_models
+       except Exception as e:
+           logging.error(f"Error fetching Groq models: {e}")
+           return self.default_groq_models
+
+   def _get_default_groq_models(self) -> Dict[str, str]:
+       """Return default Groq models"""
+       return self.default_groq_models
+
+   def refresh_groq_models(self) -> Dict[str, str]:
+       """Refresh the list of available Groq models"""
+       self.groq_models = self._fetch_groq_models()
+       return self.groq_models
 
 # Initialize model registry
 model_registry = ModelRegistry()
@@ -208,69 +250,58 @@ def send_to_model_impl(prompt, model_selection, hf_model_choice, hf_custom_model
         return error_msg, []
     
 def send_to_hf_inference(prompt: str, model_name: str, api_key: str) -> str:
-    """Send prompt to HuggingFace using Inference API"""
-    try:
-        client = InferenceClient(token=api_key)
-        response = client.text_generation(
-            prompt,
-            model=model_name,
-            max_new_tokens=500,
-            temperature=0.7,
-            details=True,  # Get full response details
-            stream=False   # Don't stream output
-        )
-        return response.generated_text  # Return just the generated text
-    except Exception as e:
-        logging.error(f"Error with HF inference: {e}")
-        return f"Error with HF inference: {e}"
+   try:
+       client = InferenceClient(token=api_key)
+       response = client.text_generation(
+           prompt,
+           model=model_name,
+           max_new_tokens=500,
+           temperature=0.7,
+           top_p=0.95,
+           repetition_penalty=1.1
+       )
+       return str(response)
+   except Exception as e:
+       logging.error(f"Error with HF inference: {e}")
+       return f"Error with HF inference: {e}"
 
 def send_to_groq(prompt: str, model_name: str, api_key: str) -> str:
-    """Send prompt to Groq API"""
-    try:
-        headers = {
-            "Authorization": f"Bearer {api_key}",
-            "Content-Type": "application/json"
-        }
-        data = {
-            "model": model_name,
-            "messages": [{"role": "user", "content": prompt}],
-            "temperature": 0.7,
-            "max_tokens": 500
-        }
-        response = requests.post(
-            "https://api.groq.com/openai/v1/chat/completions",
-            headers=headers,
-            json=data
-        )
-        if response.status_code != 200:
-            return f"Error: Groq API returned status {response.status_code}"
-            
-        response_json = response.json()
-        if "choices" not in response_json or not response_json["choices"]:
-            return "Error: No response from Groq API"
-            
-        return response_json["choices"][0]["message"]["content"]
-    except Exception as e:
-        logging.error(f"Error with Groq API: {e}")
-        return f"Error with Groq API: {e}"
-
-def send_to_openai(prompt: str, api_key: str) -> str:
-    """Send prompt to OpenAI API"""
-    try:
-        import openai
-        openai.api_key = api_key
-        
-        response = openai.ChatCompletion.create(
-            model="gpt-3.5-turbo",
-            messages=[{"role": "user", "content": prompt}],
-            temperature=0.7,
-            max_tokens=500
-        )
-        
-        return response.choices[0].message.content
-    except Exception as e:
-        logging.error(f"Error with OpenAI API: {e}")
-        return f"Error with OpenAI API: {e}"
+   try:
+       client = Groq(api_key=api_key)
+       response = client.chat.completions.create(
+           model=model_name,
+           messages=[{
+               "role": "user", 
+               "content": prompt
+           }],
+           temperature=0.7,
+           max_tokens=500,
+           top_p=0.95
+       )
+       return response.choices[0].message.content
+   except Exception as e:
+       logging.error(f"Error with Groq API: {e}")
+       return f"Error with Groq API: {e}"
+
+def send_to_openai(prompt: str, api_key: str, model: str = "gpt-3.5-turbo") -> str:
+   try:
+       import openai
+       openai.api_key = api_key
+       
+       response = openai.ChatCompletion.create(
+           model=model,
+           messages=[
+               {"role": "system", "content": "You are a helpful assistant that provides detailed responses with examples and references where appropriate."},
+               {"role": "user", "content": prompt}
+           ],
+           temperature=0.7,
+           max_tokens=500,
+           top_p=0.95
+       )
+       return response.choices[0].message.content
+   except Exception as e:
+       logging.error(f"Error with OpenAI API: {e}")
+       return f"Error with OpenAI API: {e}"
 
 def copy_text_js(element_id: str) -> str:
     return f"""function() {{
@@ -460,50 +491,51 @@ with gr.Blocks(css="""
 
         # Tab 3: Model Processing
         with gr.Tab("3️⃣ Model Processing"):
-            with gr.Row():
-                with gr.Column(scale=1):
-                    model_choice = gr.Radio(
-                    choices=["Clipboard only", "OpenAI ChatGPT", "HuggingFace Inference", "Groq API"],
-                    value="Clipboard only",
-                    label="🤖 Model Selection"
-                )
-                    
-                    with gr.Column(visible=False) as openai_options:
-                        openai_api_key = gr.Textbox(
-                            label="🔑 OpenAI API Key",
-                            type="password"
-                        )
-                    
-                    with gr.Column(visible=False) as hf_options:
-                        hf_model = gr.Dropdown(
-                            choices=list(model_registry.hf_models.keys()),
-                            label="🔧 HuggingFace Model",
-                            value="Phi-3 Mini 128k"
-                        )
-                        hf_custom_model = gr.Textbox(
-                            label="Custom Model ID",
-                            visible=False
-                        )
-                        hf_api_key = gr.Textbox(
-                            label="🔑 HuggingFace API Key",
-                            type="password"
-                        )
+         with gr.Row():
+           with gr.Column(scale=1):
+               model_choice = gr.Radio(
+                   choices=list(MODEL_CONTEXT_SIZES.keys()),
+                   value="Clipboard only",
+                   label="🤖 Provider Selection"
+               )
+               
+               with gr.Column(visible=False) as openai_options:
+                   openai_model = gr.Dropdown(
+                       choices=list(MODEL_CONTEXT_SIZES["OpenAI ChatGPT"].keys()),
+                       value="gpt-3.5-turbo",
+                       label="OpenAI Model"
+                   )
+                   openai_api_key = gr.Textbox(
+                       label="🔑 OpenAI API Key",
+                       type="password"
+                   )
+               
+               with gr.Column(visible=False) as hf_options:
+                   hf_model = gr.Dropdown(
+                       choices=list(MODEL_CONTEXT_SIZES["HuggingFace Inference"].keys()),
+                       value="microsoft/phi-3-mini-4k-instruct",
+                       label="HuggingFace Model"
+                   )
+                   hf_api_key = gr.Textbox(
+                       label="🔑 HuggingFace API Key",
+                       type="password"
+                   )
+               
+               with gr.Column(visible=False) as groq_options:
+                   groq_model = gr.Dropdown(
+                       choices=list(MODEL_CONTEXT_SIZES["Groq API"].keys()),
+                       value="mixtral-8x7b-32768",
+                       label="Groq Model"
+                   )
+                   groq_api_key = gr.Textbox(
+                       label="🔑 Groq API Key",
+                       type="password"
+                   )
                     
-                    with gr.Column(visible=False) as groq_options:
-                        groq_model = gr.Dropdown(
-                            choices=list(model_registry.groq_models.keys()),
-                            label="🔧 Groq Model"
-                        )
-                        groq_refresh_btn = gr.Button("🔄 Refresh Models")
-                        groq_api_key = gr.Textbox(
-                            label="🔑 Groq API Key",
-                            type="password"
-                        )
+               send_to_model_btn = gr.Button("🚀 Send to Model", variant="primary")
+               open_chatgpt_button = gr.Button("🌐 Open ChatGPT")
                     
-                    send_to_model_btn = gr.Button("🚀 Send to Model", variant="primary")
-                    open_chatgpt_button = gr.Button("🌐 Open ChatGPT")
-                    
-                with gr.Column(scale=1):
+               with gr.Column(scale=1):
                     summary_output = gr.Textbox(
                         label="📝 Summary",
                         lines=15,
@@ -569,26 +601,31 @@ with gr.Blocks(css="""
 
     def toggle_custom_model(model_name):
         return gr.update(visible=model_name == "Custom Model")
-    
-    def handle_model_change(choice):
-        """Handle model selection change"""
-        return (
-            gr.update(visible=choice == "HuggingFace Inference"),
-            gr.update(visible=choice == "Groq API"),
-            gr.update(visible=choice == "OpenAI ChatGPT"),
-            update_context_size(choice)
-        )
 
     def handle_groq_model_change(model_name):
         """Handle Groq model selection change"""
         return update_context_size("Groq API", model_name)
 
     def handle_model_selection(choice):
+        """Handle model selection and update UI"""
+        ctx_size = MODEL_CONTEXT_SIZES.get(choice, {})
+        if isinstance(ctx_size, dict):
+            first_model = list(ctx_size.keys())[0]
+            ctx_size = ctx_size[first_model]
+            
+            # Update model dropdown based on provider
+            if choice == "OpenAI ChatGPT":
+                openai_model.update(choices=list(MODEL_CONTEXT_SIZES["OpenAI ChatGPT"].keys()))
+            elif choice == "HuggingFace Inference":
+                hf_model.update(choices=list(MODEL_CONTEXT_SIZES["HuggingFace Inference"].keys()))
+            elif choice == "Groq API":
+                groq_model.update(choices=list(MODEL_CONTEXT_SIZES["Groq API"].keys()))
+        
         return [
             gr.update(visible=choice == "HuggingFace Inference"),
             gr.update(visible=choice == "Groq API"),
             gr.update(visible=choice == "OpenAI ChatGPT"),
-            gr.update(value=get_model_context_size(choice))
+            gr.update(value=ctx_size)
         ]
     
     # PDF Processing Handlers
@@ -766,23 +803,29 @@ with gr.Blocks(css="""
     )
 
     # Download handlers
-    for btn, content, prefix in [
-        (download_full_text, pdf_content, "full_text"),
-        (download_snippet, generated_prompt, "snippet"),
-        (download_prompt, generated_prompt, "prompt"),
-        (download_summary, summary_output, "summary")
+    for btn, elem_id in [
+        (copy_prompt_button, "generated_prompt"),
+        (copy_summary_button, "summary_output")
     ]:
         btn.click(
-            lambda x, p=prefix: download_file(x, p) if x else [],
-            inputs=[content],
-            outputs=[download_files]
+            fn=None,
+            _js=f"""
+            () => {{
+                const el = document.getElementById('{elem_id}');
+                if (!el) return 'Element not found';
+                navigator.clipboard.writeText(el.value);
+                return 'Copied to clipboard!';
+            }}
+            """,
+            outputs=progress_status
         )
 
     def download_file(content: str, prefix: str) -> List[str]:
         if not content:
             return []
         try:
-            with tempfile.NamedTemporaryFile(delete=False, mode='w', suffix='.txt', prefix=prefix) as f:
+            filename = f"{prefix}_{int(time.time())}.txt"  # Add timestamp
+            with tempfile.NamedTemporaryFile(delete=False, mode='w', suffix='.txt', prefix=filename) as f:
                 f.write(content)
                 return [f.name]
         except Exception as e: