Fix openrouter and Azure endpoint integration

app.py

@@ -210,17 +210,39 @@ with st.sidebar:
         api_version: str = ''
     else:
         # The online LLMs
-        llm_provider_to_use = st.sidebar.selectbox(
+        selected_option = st.sidebar.selectbox(
             label='2: Select a suitable LLM to use:\n\n(Gemini and Mistral-Nemo are recommended)',
             options=[f'{k} ({v["description"]})' for k, v in GlobalConfig.VALID_MODELS.items()],
             index=GlobalConfig.DEFAULT_MODEL_INDEX,
             help=GlobalConfig.LLM_PROVIDER_HELP,
             on_change=reset_api_key
-        ).split(' ')[0]
+        )
+        
+        # Extract provider key more robustly using regex
+        provider_match = GlobalConfig.PROVIDER_REGEX.match(selected_option)
+        if provider_match:
+            llm_provider_to_use = selected_option  # Use full string for get_provider_model
+        else:
+            # Fallback: try to extract the key before the first space
+            llm_provider_to_use = selected_option.split(' ')[0]
+            logger.warning(f"Could not parse provider from selectbox option: {selected_option}")
 
         # --- Automatically fetch API key from .env if available ---
         provider_match = GlobalConfig.PROVIDER_REGEX.match(llm_provider_to_use)
-        selected_provider = provider_match.group(1) if provider_match else llm_provider_to_use
+        if provider_match:
+            selected_provider = provider_match.group(1)
+        else:
+            # If regex doesn't match, try to extract provider from the beginning
+            selected_provider = llm_provider_to_use.split(' ')[0] if ' ' in llm_provider_to_use else llm_provider_to_use
+            logger.warning(f"Provider regex did not match for: {llm_provider_to_use}, using: {selected_provider}")
+        
+        # Validate that the selected provider is valid
+        if selected_provider not in GlobalConfig.VALID_PROVIDERS:
+            logger.error(f"Invalid provider: {selected_provider}")
+            handle_error(f"Invalid provider selected: {selected_provider}", True)
+            st.error(f"Invalid provider selected: {selected_provider}")
+            st.stop()
+        
         env_key_name = GlobalConfig.PROVIDER_ENV_KEYS.get(selected_provider)
         default_api_key = os.getenv(env_key_name, "") if env_key_name else ""
 
@@ -372,6 +394,15 @@ def set_up_chat_ui():
             use_ollama=RUN_IN_OFFLINE_MODE
         )
 
+        # Validate that provider and model were parsed successfully
+        if not provider or not llm_name:
+            handle_error(
+                f'Failed to parse provider and model from: "{llm_provider_to_use}". '
+                f'Please select a valid LLM from the dropdown.',
+                True
+            )
+            return
+
         user_key = api_key_token.strip()
         az_deployment = azure_deployment.strip()
         az_endpoint = azure_endpoint.strip()

helpers/llm_helper.py

@@ -70,8 +70,20 @@ def get_provider_model(provider_model: str, use_ollama: bool) -> Tuple[str, str]
         if match:
             inside_brackets = match.group(1)
             outside_brackets = match.group(2)
+            
+            # Validate that the provider is in the valid providers list
+            if inside_brackets not in GlobalConfig.VALID_PROVIDERS:
+                logger.warning(f"Provider '{inside_brackets}' not in VALID_PROVIDERS: {GlobalConfig.VALID_PROVIDERS}")
+                return '', ''
+            
+            # Validate that the model name is not empty
+            if not outside_brackets.strip():
+                logger.warning(f"Empty model name for provider '{inside_brackets}'")
+                return '', ''
+            
             return inside_brackets, outside_brackets
 
+    logger.warning(f"Could not parse provider_model: '{provider_model}' (use_ollama={use_ollama})")
     return '', ''
 
 
@@ -183,7 +195,14 @@ def stream_litellm_completion(
         raise ImportError("LiteLLM is not installed. Please install it with: pip install litellm")
     
     # Convert to LiteLLM model name
-    litellm_model = get_litellm_model_name(provider, model)
+    if provider == GlobalConfig.PROVIDER_AZURE_OPENAI:
+        # For Azure OpenAI, use the deployment name as the model
+        # This is consistent with Azure OpenAI's requirement to use deployment names
+        if not azure_deployment_name:
+            raise ValueError("Azure deployment name is required for Azure OpenAI provider")
+        litellm_model = f"azure/{azure_deployment_name}"
+    else:
+        litellm_model = get_litellm_model_name(provider, model)
     
     # Prepare the request parameters
     request_params = {
@@ -196,13 +215,20 @@ def stream_litellm_completion(
     
     # Set API key and any provider-specific params
     if provider != GlobalConfig.PROVIDER_OLLAMA:
-        api_key_to_use = get_litellm_api_key(provider, api_key)
-        request_params["api_key"] = api_key_to_use
-
-        if provider == GlobalConfig.PROVIDER_AZURE_OPENAI:
-            request_params["azure_endpoint"] = azure_endpoint_url
-            request_params["azure_deployment"] = azure_deployment_name
-            request_params["api_version"] = azure_api_version
+        # For OpenRouter, set environment variable as per documentation
+        if provider == GlobalConfig.PROVIDER_OPENROUTER:
+            os.environ["OPENROUTER_API_KEY"] = api_key
+            # Optional: Set base URL if different from default
+            # os.environ["OPENROUTER_API_BASE"] = "https://openrouter.ai/api/v1"
+        elif provider == GlobalConfig.PROVIDER_AZURE_OPENAI:
+            # For Azure OpenAI, set environment variables as per documentation
+            os.environ["AZURE_API_KEY"] = api_key
+            os.environ["AZURE_API_BASE"] = azure_endpoint_url
+            os.environ["AZURE_API_VERSION"] = azure_api_version
+        else:
+            # For other providers, pass API key as parameter
+            api_key_to_use = get_litellm_api_key(provider, api_key)
+            request_params["api_key"] = api_key_to_use
     
     logger.debug('Streaming completion via LiteLLM: %s', litellm_model)