Update app.py

app.py

@@ -75,6 +75,7 @@ HF_API_KEY = os.environ.get("HF_API_KEY", "")
 TOGETHER_API_KEY = os.environ.get("TOGETHER_API_KEY", "")
 GOOGLEAI_API_KEY = os.environ.get("GOOGLEAI_API_KEY", "")
 ANTHROPIC_API_KEY = os.environ.get("ANTHROPIC_API_KEY", "")
+POE_API_KEY = os.environ.get("POE_API_KEY", "")
 
 # Print application startup message with timestamp
 current_time = time.strftime("%Y-%m-%d %H:%M:%S")
@@ -218,6 +219,48 @@ VISION_MODELS = {
     "GoogleAI": ["gemini-1.5-pro", "gemini-1.0-pro", "gemini-1.5-flash", "gemini-2.0-pro", "gemini-2.5-pro"]
 }
 
+# POE MODELS
+POE_MODELS = {
+    "claude_3_igloo": 4000,               # Claude-3.5-Sonnet
+    "claude_2_1_cedar": 4000,             # Claude-3-Opus
+    "claude_2_1_bamboo": 4000,            # Claude-3-Sonnet
+    "claude_3_haiku": 4000,               # Claude-3-Haiku
+    "claude_3_igloo_200k": 200000,        # Claude-3.5-Sonnet-200k
+    "claude_3_opus_200k": 200000,         # Claude-3-Opus-200k
+    "claude_3_sonnet_200k": 200000,       # Claude-3-Sonnet-200k
+    "claude_3_haiku_200k": 200000,        # Claude-3-Haiku-200k
+    "claude_2_short": 4000,               # Claude-2
+    "a2_2": 100000,                       # Claude-2-100k
+    "a2": 9000,                           # Claude-instant
+    "a2_100k": 100000,                    # Claude-instant-100k
+    "chinchilla": 4000,                   # GPT-3.5-Turbo
+    "gpt3_5": 2000,                       # GPT-3.5-Turbo-Raw
+    "chinchilla_instruct": 2000,          # GPT-3.5-Turbo-Instruct
+    "agouti": 16000,                      # ChatGPT-16k
+    "gpt4_classic": 2000,                 # GPT-4-Classic
+    "beaver": 4000,                       # GPT-4-Turbo
+    "vizcacha": 128000,                   # GPT-4-Turbo-128k
+    "gpt4_o": 4000,                       # GPT-4o
+    "gpt4_o_128k": 128000,                # GPT-4o-128k
+    "gpt4_o_mini": 4000,                  # GPT-4o-Mini
+    "gpt4_o_mini_128k": 128000,           # GPT-4o-Mini-128k
+    "acouchy": 8000,                      # Google-PaLM
+    "code_llama_13b_instruct": 4000,      # Code-Llama-13b
+    "code_llama_34b_instruct": 4000,      # Code-Llama-34b
+    "upstage_solar_0_70b_16bit": 2000,    # Solar-Mini
+    "gemini_pro_search": 4000,            # Gemini-1.5-Flash-Search
+    "gemini_1_5_pro_1m": 2000000,         # Gemini-1.5-Pro-2M
+}
+
+# Add vision-capable models to vision models list
+POE_VISION_MODELS = [
+    "claude_3_igloo", "claude_2_1_cedar", "claude_2_1_bamboo", "claude_3_haiku",
+    "claude_3_igloo_200k", "claude_3_opus_200k", "claude_3_sonnet_200k", "claude_3_haiku_200k",
+    "gpt4_o", "gpt4_o_128k", "gpt4_o_mini", "gpt4_o_mini_128k", "beaver", "vizcacha"
+]
+
+VISION_MODELS["Poe"] = POE_VISION_MODELS
+
 # OPENAI MODELS
 OPENAI_MODELS = {
     "gpt-3.5-turbo": 16385,
@@ -618,6 +661,9 @@ def get_model_info(provider, model_choice):
         for name, model_id, ctx_size in OPENROUTER_ALL_MODELS:
             if name == model_choice:
                 return model_id, ctx_size
+    elif provider == "Poe":
+        if model_choice in POE_MODELS:
+            return model_choice, POE_MODELS[model_choice]
     elif provider == "OpenAI":
         if model_choice in OPENAI_MODELS:
             return model_choice, OPENAI_MODELS[model_choice]
@@ -787,7 +833,73 @@ def call_anthropic_api(payload, api_key_override=None):
     except Exception as e:
         logger.error(f"Anthropic API error: {str(e)}")
         raise e
-    
+
+def call_poe_api(payload, api_key_override=None):
+    """Make a call to Poe API with error handling"""
+    try:
+        # Try to import fastapi_poe
+        try:
+            import fastapi_poe as fp
+        except ImportError:
+            raise ImportError("fastapi_poe package not installed. Install it with: pip install fastapi_poe")
+            
+        api_key = api_key_override if api_key_override else os.environ.get("POE_API_KEY", "")
+        if not api_key:
+            raise ValueError("Poe API key is required")
+            
+        # Extract parameters from payload
+        model = payload.get("model", "chinchilla")  # Default to GPT-3.5-Turbo
+        messages = payload.get("messages", [])
+        
+        # Convert messages to Poe format
+        poe_messages = []
+        for msg in messages:
+            role = msg["role"]
+            content = msg["content"]
+            
+            # Skip system messages as Poe doesn't support them directly
+            if role == "system":
+                continue
+                
+            # Convert content format
+            if isinstance(content, list):
+                # Handle multimodal content (images)
+                text_parts = []
+                for item in content:
+                    if item["type"] == "text":
+                        text_parts.append(item["text"])
+                
+                # For images, we'll need to extract and handle them separately
+                # This is a simplified approach - in reality, you'd need to handle images properly
+                content = "\n".join(text_parts)
+                
+            # Add message to Poe messages
+            poe_messages.append(fp.ProtocolMessage(role=role, content=content))
+        
+        # Make synchronous request to Poe
+        response_content = ""
+        for partial in fp.get_bot_response_sync(messages=poe_messages, bot_name=model, api_key=api_key):
+            if hasattr(partial, "text"):
+                response_content += partial.text
+        
+        # Create a response object with a structure similar to other APIs
+        response = {
+            "id": f"poe-{int(time.time())}",
+            "choices": [
+                {
+                    "message": {
+                        "role": "assistant",
+                        "content": response_content
+                    }
+                }
+            ]
+        }
+        
+        return response
+    except Exception as e:
+        logger.error(f"Poe API error: {str(e)}")
+        raise e
+
 def call_openrouter_api(payload, api_key_override=None):
     """Make a call to OpenRouter API with error handling"""
     try:
@@ -1030,7 +1142,12 @@ def extract_ai_response(result, provider):
                 return text_content
             else:
                 return "No response content from Cohere"
-                
+
+        elif provider == "Poe":
+            if isinstance(result, dict) and "choices" in result and len(result["choices"]) > 0:
+                return result["choices"][0]["message"]["content"]
+            return "No response content from Poe"
+
         elif provider == "Together":
             # Handle response from Together's native client
             if hasattr(result, "choices") and len(result.choices) > 0:
@@ -1580,6 +1697,43 @@ def ask_ai(message, history, provider, model_choice, temperature, max_tokens, to
                     {"role": "assistant", "content": error_message}
                 ]
             
+        elif provider == "Poe":
+            # Get model ID from registry
+            model_id, _ = get_model_info(provider, model_choice)
+            if not model_id:
+                error_message = f"Error: Model '{model_choice}' not found in Poe"
+                return history + [
+                    {"role": "user", "content": message},
+                    {"role": "assistant", "content": error_message}
+                ]
+                
+            # Build Poe payload
+            payload = {
+                "model": model_id,
+                "messages": messages
+                # Poe doesn't support most parameters directly
+            }
+            
+            # Call Poe API
+            logger.info(f"Sending request to Poe model: {model_id}")
+            
+            try:
+                response = call_poe_api(payload, api_key_override)
+                
+                # Extract response
+                ai_response = extract_ai_response(response, provider)
+                return history + [
+                    {"role": "user", "content": message},
+                    {"role": "assistant", "content": ai_response}
+                ]
+            except Exception as e:
+                error_message = f"Poe API Error: {str(e)}"
+                logger.error(error_message)
+                return history + [
+                    {"role": "user", "content": message},
+                    {"role": "assistant", "content": error_message}
+                ]
+            
         elif provider == "Anthropic":
             # Get model ID from registry
             model_id, _ = get_model_info(provider, model_choice)
@@ -2101,9 +2255,9 @@ def create_app():
                 with gr.Group(elem_classes="provider-selection"):
                     gr.Markdown("### Provider Selection")
                     
-                    # Provider selection - removed OVH and Cerebras, added Anthropic
+                    # Provider selection
                     provider_choice = gr.Radio(
-                        choices=["OpenRouter", "OpenAI", "HuggingFace", "Groq", "Cohere", "Together", "Anthropic", "GoogleAI"],
+                        choices=["OpenRouter", "OpenAI", "HuggingFace", "Groq", "Cohere", "Together", "Anthropic", "Poe", "GoogleAI"],
                         value="OpenRouter",
                         label="AI Provider"
                     )
@@ -2118,6 +2272,13 @@ def create_app():
                             type="password",
                             value=OPENROUTER_API_KEY if OPENROUTER_API_KEY else ""
                         )
+
+                        poe_api_key = gr.Textbox(
+                            placeholder="Enter Poe API key",
+                            label="Poe API Key",
+                            type="password",
+                            value=os.environ.get("POE_API_KEY", "")
+                        )
                         
                         openai_api_key = gr.Textbox(
                             placeholder="Enter OpenAI API key",
@@ -2187,6 +2348,15 @@ def create_app():
                         elem_id="openrouter-model-choice",
                         visible=True
                     )
+
+                    # Add Poe model dropdown
+                    poe_model = gr.Dropdown(
+                        choices=list(POE_MODELS.keys()),
+                        value="chinchilla" if "chinchilla" in POE_MODELS else None,
+                        label="Poe Model",
+                        elem_id="poe-model-choice",
+                        visible=False
+                    )
                     
                     openai_model = gr.Dropdown(
                         choices=list(OPENAI_MODELS.keys()),
@@ -2437,6 +2607,7 @@ def create_app():
                 cohere_model: gr.update(visible=(provider == "Cohere")),
                 together_model: gr.update(visible=(provider == "Together")),
                 anthropic_model: gr.update(visible=(provider == "Anthropic")),
+                poe_model: gr.update(visible=(provider == "Poe")),
                 googleai_model: gr.update(visible=(provider == "GoogleAI"))
             }
             
@@ -2612,7 +2783,7 @@ def create_app():
 
             
         def get_current_model(provider, openrouter_model, openai_model, hf_model, groq_model, cohere_model, 
-                            together_model, anthropic_model, googleai_model):
+                            together_model, anthropic_model, poe_model, googleai_model):
             """Get the currently selected model based on provider"""
             if provider == "OpenRouter":
                 return openrouter_model
@@ -2628,11 +2799,12 @@ def create_app():
                 return together_model
             elif provider == "Anthropic":
                 return anthropic_model
+            elif provider == "Poe":
+                return poe_model
             elif provider == "GoogleAI":
                 return googleai_model
             return None
         
-        
         # Process uploaded images
         image_upload_btn.upload(
             fn=lambda files: files,
@@ -2864,16 +3036,16 @@ def create_app():
         
         # Set up submission event
         def submit_message(message, history, provider, 
-                openrouter_model, openai_model, hf_model, groq_model, cohere_model, together_model, anthropic_model, googleai_model, 
+                openrouter_model, openai_model, hf_model, groq_model, cohere_model, together_model, anthropic_model, poe_model, googleai_model, 
                 temperature, max_tokens, top_p, frequency_penalty, presence_penalty, repetition_penalty, 
                 top_k, min_p, seed, top_a, stream_output, response_format,
                 images, documents, reasoning_effort, system_message, transforms, 
-                openrouter_api_key, openai_api_key, hf_api_key, groq_api_key, cohere_api_key, together_api_key, anthropic_api_key, googleai_api_key):
+                openrouter_api_key, openai_api_key, hf_api_key, groq_api_key, cohere_api_key, together_api_key, anthropic_api_key, poe_api_key, googleai_api_key):
 
             """Submit message to selected provider and model"""
             # Get the currently selected model
             model_choice = get_current_model(provider, openrouter_model, openai_model, hf_model, groq_model, cohere_model, 
-                                    together_model, anthropic_model, googleai_model)
+                                    together_model, anthropic_model, poe_model, googleai_model)
             
             # Check if model is selected
             if not model_choice:
@@ -2899,6 +3071,8 @@ def create_app():
                 api_key_override = together_api_key
             elif provider == "Anthropic" and anthropic_api_key:
                 api_key_override = anthropic_api_key
+            elif provider_choice == "Poe" and poe_api_key:
+                api_key_override = poe_api_key
             elif provider == "GoogleAI" and googleai_api_key:
                 api_key_override = googleai_api_key
             
@@ -2933,11 +3107,11 @@ def create_app():
             fn=submit_message,
             inputs=[
                 message, chatbot, provider_choice, 
-                openrouter_model, openai_model, hf_model, groq_model, cohere_model, together_model, anthropic_model, googleai_model,
+                openrouter_model, openai_model, hf_model, groq_model, cohere_model, together_model, anthropic_model, poe_model, googleai_model,
                 temperature, max_tokens, top_p, frequency_penalty, presence_penalty, repetition_penalty, 
                 top_k, min_p, seed, top_a, stream_output, response_format,
                 images, documents, reasoning_effort, system_message, transforms,
-                openrouter_api_key, openai_api_key, hf_api_key, groq_api_key, cohere_api_key, together_api_key, anthropic_api_key, googleai_api_key
+                openrouter_api_key, openai_api_key, hf_api_key, groq_api_key, cohere_api_key, together_api_key, anthropic_api_key, poe_api_key, googleai_api_key
             ],
             outputs=chatbot,
             show_progress="minimal",
@@ -2952,11 +3126,11 @@ def create_app():
             fn=submit_message,
             inputs=[
                 message, chatbot, provider_choice, 
-                openrouter_model, openai_model, hf_model, groq_model, cohere_model, together_model, anthropic_model, googleai_model,
+                openrouter_model, openai_model, hf_model, groq_model, cohere_model, together_model, anthropic_model, poe_model, googleai_model,
                 temperature, max_tokens, top_p, frequency_penalty, presence_penalty, repetition_penalty, 
                 top_k, min_p, seed, top_a, stream_output, response_format,
                 images, documents, reasoning_effort, system_message, transforms,
-                openrouter_api_key, openai_api_key, hf_api_key, groq_api_key, cohere_api_key, together_api_key, anthropic_api_key, googleai_api_key
+                openrouter_api_key, openai_api_key, hf_api_key, groq_api_key, cohere_api_key, together_api_key, anthropic_api_key, poe_api_key, googleai_api_key
             ],
             outputs=chatbot,
             show_progress="minimal",
@@ -2989,6 +3163,11 @@ if __name__ == "__main__":
         logger.warning("WARNING: OPENROUTER_API_KEY environment variable is not set")
         missing_keys.append("OpenRouter")
 
+    # Add Poe
+    if not POE_API_KEY:
+        logger.warning("WARNING: POE_API_KEY environment variable is not set")
+        missing_keys.append("Poe")
+
     if not ANTHROPIC_API_KEY:
         logger.warning("WARNING: ANTHROPIC_API_KEY environment variable is not set")
         missing_keys.append("Anthropic")