Update model_logic.py

model_logic.py

@@ -1,7 +1,12 @@
+# model_handler.py
 import os
 import requests
 import json
 import logging
+from dotenv import load_dotenv
+
+# Load environment variables from .env file
+load_dotenv()
 
 logging.basicConfig(
     level=logging.INFO,
@@ -9,15 +14,16 @@ logging.basicConfig(
 )
 logger = logging.getLogger(__name__)
 
-API_KEYS = {
-  "HUGGINGFACE": 'HF_TOKEN',
+# Maps provider name (uppercase) to environment variable name for API key
+API_KEYS_ENV_VARS = {
+  "HUGGINGFACE": 'HF_TOKEN', # Note: HF_TOKEN is often used for general HF auth
   "GROQ": 'GROQ_API_KEY',
   "OPENROUTER": 'OPENROUTER_API_KEY',
   "TOGETHERAI": 'TOGETHERAI_API_KEY',
   "COHERE": 'COHERE_API_KEY',
   "XAI": 'XAI_API_KEY',
   "OPENAI": 'OPENAI_API_KEY',
-  "GOOGLE": 'GOOGLE_API_KEY',
+  "GOOGLE": 'GOOGLE_API_KEY', # Or GOOGLE_GEMINI_API_KEY etc.
 }
 
 API_URLS = {
@@ -25,12 +31,13 @@ API_URLS = {
   "GROQ": 'https://api.groq.com/openai/v1/chat/completions',
   "OPENROUTER": 'https://openrouter.ai/api/v1/chat/completions',
   "TOGETHERAI": 'https://api.together.ai/v1/chat/completions',
-  "COHERE": 'https://api.cohere.ai/v1/chat',
+  "COHERE": 'https://api.cohere.ai/v1/chat', # v1 is common for chat, was v2 in ai-learn
   "XAI": 'https://api.x.ai/v1/chat/completions',
   "OPENAI": 'https://api.openai.com/v1/chat/completions',
   "GOOGLE": 'https://generativelanguage.googleapis.com/v1beta/models/',
 }
 
+# Structure: provider_key: { "default": "model_id", "models": {"Display Name": "model_id", ...} }
 MODELS_BY_PROVIDER = {
     "groq": {
         "default": "llama3-8b-8192",
@@ -42,15 +49,20 @@ MODELS_BY_PROVIDER = {
         }
     },
     "openrouter": {
-         "default": "nousresearch/llama-3-8b-instruct",
+         "default": "nousresearch/llama-3-8b-instruct", # Updated default
          "models": {
             "Nous Llama-3 8B Instruct (OpenRouter)": "nousresearch/llama-3-8b-instruct",
-            "Mistral 7B Instruct v0.2 (OpenRouter)": "mistralai/mistral-7b-instruct:free",
-            "Gemma 7B Instruct (OpenRouter)": "google/gemma-7b-it:free",
-            "Mixtral 8x7B Instruct v0.1 (OpenRouter)": "mistralai/mixtral-8x7b-instruct",
-            "Llama 2 70B Chat (OpenRouter)": "meta-llama/llama-2-70b-chat",
-            "Neural Chat 7B v3.1 (OpenRouter)": "intel/neural-chat-7b-v3-1",
-            "Goliath 120B (OpenRouter)": "twob/goliath-v2-120b",
+            "Mistral 7B Instruct v0.3 (OpenRouter)": "mistralai/mistral-7b-instruct-v0.3", # v0.3 is newer
+            "Mistral 7B Instruct (Free/OpenRouter)": "mistralai/mistral-7b-instruct:free", # Keep free tier if distinct
+            "Gemma 2 9B Instruct (OpenRouter)": "google/gemma-2-9b-it", # Gemma 2
+            "Gemma 7B Instruct (Free/OpenRouter)": "google/gemma-7b-it:free",
+            "Llama 3.1 8B Instruct (OpenRouter)": "meta-llama/llama-3.1-8b-instruct", # Llama 3.1
+            "Llama 3.1 70B Instruct (OpenRouter)": "meta-llama/llama-3.1-70b-instruct",
+            "OpenAI GPT-4o mini (OpenRouter)": "openai/gpt-4o-mini",
+            "OpenAI GPT-4o (OpenRouter)": "openai/gpt-4o",
+            "Claude 3.5 Sonnet (OpenRouter)": "anthropic/claude-3.5-sonnet",
+            "Mixtral 8x7B Instruct v0.1 (OpenRouter)": "mistralai/mixtral-8x7b-instruct", # Older Mixtral
+            "Qwen 2 72B Instruct (OpenRouter)": "qwen/qwen-2-72b-instruct",
          }
     },
      "togetherai": {
@@ -60,7 +72,7 @@ MODELS_BY_PROVIDER = {
              "Llama 3 70B Chat (TogetherAI)": "meta-llama/Llama-3-70b-chat-hf",
              "Mixtral 8x7B Instruct (TogetherAI)": "mistralai/Mixtral-8x7B-Instruct-v0.1",
              "Gemma 7B Instruct (TogetherAI)": "google/gemma-7b-it",
-             "RedPajama INCITE Chat 3B (TogetherAI)": "togethercomputer/RedPajama-INCITE-Chat-3B-v1",
+             "Qwen1.5-72B-Chat (TogetherAI)": "qwen/Qwen1.5-72B-Chat",
          }
      },
      "google": {
@@ -68,335 +80,428 @@ MODELS_BY_PROVIDER = {
           "models": {
               "Gemini 1.5 Flash (Latest)": "gemini-1.5-flash-latest",
               "Gemini 1.5 Pro (Latest)": "gemini-1.5-pro-latest",
+              # "Gemini 1.0 Pro": "gemini-pro" # Older model example
           }
      },
      "cohere": {
-         "default": "command-light",
+         "default": "command-r", # command-r is generally better than light
          "models": {
              "Command R (Cohere)": "command-r",
              "Command R+ (Cohere)": "command-r-plus",
              "Command Light (Cohere)": "command-light",
-             "Command (Cohere)": "command",
          }
      },
-     "huggingface": {
-         "default": "HuggingFaceH4/zephyr-7b-beta",
+     "huggingface": { # Direct HF Inference API is tricky for chat, often better via OpenRouter/TogetherAI
+         "default": "mistralai/Mistral-7B-Instruct-v0.2", # A common TGI compatible model
          "models": {
-             "Zephyr 7B Beta (H4/HF Inf.)": "HuggingFaceH4/zephyr-7b-beta",
              "Mistral 7B Instruct v0.2 (HF Inf.)": "mistralai/Mistral-7B-Instruct-v0.2",
-             "Llama 2 13B Chat (Meta/HF Inf.)": "meta-llama/Llama-2-13b-chat-hf",
-             "OpenAssistant/oasst-sft-4-pythia-12b (HF Inf.)": "OpenAssistant/oasst-sft-4-pythia-12b",
+             "Llama 3 8B Instruct (HF Inf.)": "meta-llama/Meta-Llama-3-8B-Instruct", # Ensure this specific ID is for TGI
+             # "Zephyr 7B Beta (H4/HF Inf.)": "HuggingFaceH4/zephyr-7b-beta", # Older model
          }
      },
      "openai": {
-          "default": "gpt-3.5-turbo",
+          "default": "gpt-4o-mini", # New default
           "models": {
               "GPT-4o (OpenAI)": "gpt-4o",
               "GPT-4o mini (OpenAI)": "gpt-4o-mini",
-              "GPT-4 Turbo (OpenAI)": "gpt-4-turbo",
-              "GPT-3.5 Turbo (OpenAI)": "gpt-3.5-turbo",
+              "GPT-4 Turbo (OpenAI)": "gpt-4-turbo", # Refers to latest gpt-4-turbo variant
+              "GPT-3.5 Turbo (OpenAI)": "gpt-3.5-turbo", # Refers to latest gpt-3.5-turbo variant
           }
      },
-     "xai": {
-          "default": "grok-1",
+     "xai": { # Assuming xAI might expand model list
+          "default": "grok-1.5-flash", # Assuming Grok 1.5 flash is available
           "models": {
-              "Grok-1 (xAI)": "grok-1",
+              "Grok 1.5 Flash (xAI)": "grok-1.5-flash",
+              # "Grok-1 (xAI)": "grok-1", # Older model
           }
      }
 }
 
-def _get_api_key(provider: str, ui_api_key_override: str = None) -> str:
-    if ui_api_key_override:
+def _get_api_key(provider: str, ui_api_key_override: str = None) -> str | None:
+    """
+    Retrieves API key for a given provider.
+    Priority: UI Override > Environment Variable from API_KEYS_ENV_VARS > Specific (e.g. HF_TOKEN for HuggingFace).
+    """
+    provider_upper = provider.upper()
+    if ui_api_key_override and ui_api_key_override.strip():
+        logger.debug(f"Using UI-provided API key for {provider_upper}.")
         return ui_api_key_override.strip()
 
-    env_var_name = API_KEYS.get(provider.upper())
+    env_var_name = API_KEYS_ENV_VARS.get(provider_upper)
     if env_var_name:
         env_key = os.getenv(env_var_name)
-        if env_key:
+        if env_key and env_key.strip():
+            logger.debug(f"Using API key from env var '{env_var_name}' for {provider_upper}.")
             return env_key.strip()
 
-    if provider.lower() == 'huggingface':
-         hf_token = os.getenv("HF_TOKEN")
-         if hf_token: return hf_token.strip()
+    # Specific fallback for HuggingFace if HF_TOKEN is set and API_KEYS_ENV_VARS['HUGGINGFACE'] wasn't specific enough
+    if provider_upper == 'HUGGINGFACE':
+         hf_token_fallback = os.getenv("HF_TOKEN")
+         if hf_token_fallback and hf_token_fallback.strip():
+             logger.debug("Using HF_TOKEN as fallback for HuggingFace provider.")
+             return hf_token_fallback.strip()
 
-    logger.warning(f"API Key not found for provider '{provider}'. Checked UI override and environment variable '{env_var_name or 'N/A'}'.")
+    logger.warning(f"API Key not found for provider '{provider_upper}'. Checked UI override and environment variable '{env_var_name or 'N/A'}'.")
     return None
 
 def get_available_providers() -> list[str]:
+    """Returns a sorted list of available provider names (e.g., 'groq', 'openai')."""
     return sorted(list(MODELS_BY_PROVIDER.keys()))
 
-def get_models_for_provider(provider: str) -> list[str]:
+def get_model_display_names_for_provider(provider: str) -> list[str]:
+    """Returns a sorted list of model display names for a given provider."""
     return sorted(list(MODELS_BY_PROVIDER.get(provider.lower(), {}).get("models", {}).keys()))
 
-def get_default_model_for_provider(provider: str) -> str | None:
-    models_dict = MODELS_BY_PROVIDER.get(provider.lower(), {}).get("models", {})
-    default_model_id = MODELS_BY_PROVIDER.get(provider.lower(), {}).get("default")
-    if default_model_id:
-        for display_name, model_id in models_dict.items():
-            if model_id == default_model_id:
+def get_default_model_display_name_for_provider(provider: str) -> str | None:
+    """Gets the default model's display name for a provider."""
+    provider_data = MODELS_BY_PROVIDER.get(provider.lower(), {})
+    models_dict = provider_data.get("models", {})
+    default_model_id = provider_data.get("default")
+
+    if default_model_id and models_dict:
+        for display_name, model_id_val in models_dict.items():
+            if model_id_val == default_model_id:
                 return display_name
+    
+    # Fallback to the first model in the sorted list if default not found or not set
     if models_dict:
-        return sorted(list(models_dict.keys()))[0]
+        sorted_display_names = sorted(list(models_dict.keys()))
+        if sorted_display_names:
+            return sorted_display_names[0]
     return None
 
 def get_model_id_from_display_name(provider: str, display_name: str) -> str | None:
+    """Gets the actual model ID from its display name for a given provider."""
     models = MODELS_BY_PROVIDER.get(provider.lower(), {}).get("models", {})
     return models.get(display_name)
 
-def generate_stream(provider: str, model_display_name: str, api_key_override: str, messages: list[dict]) -> iter:
+
+def call_model_stream(provider: str, model_display_name: str, messages: list[dict], api_key_override: str = None, temperature: float = 0.7, max_tokens: int = None) -> iter:
+    """
+    Calls the specified model via its provider and streams the response.
+    Handles provider-specific request formatting and error handling.
+    Yields chunks of the response text or an error string.
+    """
     provider_lower = provider.lower()
     api_key = _get_api_key(provider_lower, api_key_override)
-
     base_url = API_URLS.get(provider.upper())
     model_id = get_model_id_from_display_name(provider_lower, model_display_name)
 
     if not api_key:
-        env_var_name = API_KEYS.get(provider.upper(), 'N/A')
-        yield f"Error: API Key not found for {provider}. Please set it in the UI override or environment variable '{env_var_name}'."
+        env_var_name = API_KEYS_ENV_VARS.get(provider.upper(), 'N/A')
+        yield f"Error: API Key not found for {provider}. Please set it in the UI or env var '{env_var_name}'."
         return
     if not base_url:
         yield f"Error: Unknown provider '{provider}' or missing API URL configuration."
         return
     if not model_id:
-         yield f"Error: Unknown model '{model_display_name}' for provider '{provider}'. Please select a valid model."
+         yield f"Error: Model ID not found for '{model_display_name}' under provider '{provider}'. Check configuration."
          return
 
     headers = {}
     payload = {}
     request_url = base_url
 
-    logger.info(f"Calling {provider}/{model_display_name} (ID: {model_id}) stream...")
+    logger.info(f"Streaming from {provider}/{model_display_name} (ID: {model_id})...")
+    
+    # --- Standard OpenAI-compatible providers ---
+    if provider_lower in ["groq", "openrouter", "togetherai", "openai", "xai"]:
+        headers = {"Authorization": f"Bearer {api_key}", "Content-Type": "application/json"}
+        payload = {"model": model_id, "messages": messages, "stream": True, "temperature": temperature}
+        if max_tokens: payload["max_tokens"] = max_tokens
 
-    try:
-        if provider_lower in ["groq", "openrouter", "togetherai", "openai", "xai"]:
-            headers = {"Authorization": f"Bearer {api_key}", "Content-Type": "application/json"}
-            payload = {
-                "model": model_id,
-                "messages": messages,
-                "stream": True
-            }
-            if provider_lower == "openrouter":
-                 headers["HTTP-Referer"] = os.getenv("SPACE_HOST") or "https://github.com/your_username/ai-space-builder"
-                 headers["X-Title"] = "AI Space Builder"
+        if provider_lower == "openrouter":
+             headers["HTTP-Referer"] = os.getenv("OPENROUTER_REFERRER") or "http://localhost/gradio" # Example Referer
+             headers["X-Title"] = os.getenv("OPENROUTER_X_TITLE") or "Gradio AI Researcher"      # Example Title
 
+        try:
             response = requests.post(request_url, headers=headers, json=payload, stream=True, timeout=180)
             response.raise_for_status()
 
-            byte_buffer = b""
-            for chunk in response.iter_content(chunk_size=8192):
-                byte_buffer += chunk
-                while b'\n' in byte_buffer:
-                    line, byte_buffer = byte_buffer.split(b'\n', 1)
-                    decoded_line = line.decode('utf-8', errors='ignore')
-                    if decoded_line.startswith('data: '):
-                        data = decoded_line[6:]
-                        if data == '[DONE]':
-                            byte_buffer = b''
-                            break
-                        try:
-                            event_data = json.loads(data)
-                            if event_data.get("choices") and len(event_data["choices"]) > 0:
-                                delta = event_data["choices"][0].get("delta")
-                                if delta and delta.get("content"):
-                                    yield delta["content"]
-                        except json.JSONDecodeError:
-                            logger.warning(f"Failed to decode JSON from stream line: {decoded_line}")
-                        except Exception as e:
-                             logger.error(f"Error processing stream data: {e}, Data: {decoded_line}")
-            if byte_buffer:
-                 remaining_line = byte_buffer.decode('utf-8', errors='ignore')
-                 if remaining_line.startswith('data: '):
-                     data = remaining_line[6:]
-                     if data != '[DONE]':
-                         try:
-                              event_data = json.loads(data)
-                              if event_data.get("choices") and len(event_data["choices"]) > 0:
-                                  delta = event_data["choices"][0].get("delta")
-                                  if delta and delta.get("content"):
-                                      yield delta["content"]
-                         except json.JSONDecodeError:
-                               logger.warning(f"Failed to decode final stream buffer JSON: {remaining_line}")
-                         except Exception as e:
-                                logger.error(f"Error processing final stream buffer data: {e}, Data: {remaining_line}")
-
-
-        elif provider_lower == "google":
-            system_instruction = None
-            filtered_messages = []
-            for msg in messages:
-                if msg["role"] == "system":
-                    system_instruction = msg["content"]
-                else:
-                    role = "model" if msg["role"] == "assistant" else msg["role"]
-                    filtered_messages.append({"role": role, "parts": [{"text": msg["content"]}]})
-
-            payload = {
-                 "contents": filtered_messages,
-                 "safetySettings": [
-                     {"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT", "threshold": "BLOCK_NONE"},
-                     {"category": "HARM_CATEGORY_HATE_SPEECH", "threshold": "BLOCK_NONE"},
-                     {"category": "HARM_CATEGORY_HARASSMENT", "threshold": "BLOCK_NONE"},
-                     {"category": "HARM_CATEGORY_DANGEROUS_CONTENT", "threshold": "BLOCK_NONE"},
-                 ],
-                 "generationConfig": {
-                    "temperature": 0.7,
-                 }
-            }
-            if system_instruction:
-                payload["system_instruction"] = {"parts": [{"text": system_instruction}]}
-
-            request_url = f"{base_url}{model_id}:streamGenerateContent"
-            headers = {"Content-Type": "application/json"}
-            request_url = f"{request_url}?key={api_key}"
+            # More robust SSE parsing
+            buffer = ""
+            for chunk in response.iter_content(chunk_size=None): # Process raw bytes
+                buffer += chunk.decode('utf-8', errors='replace')
+                while '\n\n' in buffer:
+                    event_str, buffer = buffer.split('\n\n', 1)
+                    if not event_str.strip(): continue
+
+                    content_chunk = ""
+                    for line in event_str.splitlines():
+                        if line.startswith('data: '):
+                            data_json = line[len('data: '):].strip()
+                            if data_json == '[DONE]':
+                                return # Stream finished
+                            try:
+                                data = json.loads(data_json)
+                                if data.get("choices") and len(data["choices"]) > 0:
+                                    delta = data["choices"][0].get("delta", {})
+                                    if delta and delta.get("content"):
+                                        content_chunk += delta["content"]
+                            except json.JSONDecodeError:
+                                logger.warning(f"Failed to decode JSON from stream line: {data_json}")
+                    if content_chunk:
+                        yield content_chunk
+            # Process any remaining buffer content (less common with '\n\n' delimiter)
+            if buffer.strip():
+                logger.debug(f"Remaining buffer after OpenAI-like stream: {buffer}")
+
+
+        except requests.exceptions.HTTPError as e:
+            err_msg = f"API HTTP Error ({e.response.status_code}): {e.response.text[:500]}"
+            logger.error(f"{err_msg} for {provider}/{model_id}", exc_info=False)
+            yield f"Error: {err_msg}"
+        except requests.exceptions.RequestException as e:
+            logger.error(f"API Request Error for {provider}/{model_id}: {e}", exc_info=False)
+            yield f"Error: Could not connect to {provider} ({e})"
+        except Exception as e:
+            logger.exception(f"Unexpected error during {provider} stream:")
+            yield f"Error: An unexpected error occurred: {e}"
+        return
+
+    # --- Google Gemini ---
+    elif provider_lower == "google":
+        system_instruction = None
+        filtered_messages = []
+        for msg in messages:
+            if msg["role"] == "system": system_instruction = {"parts": [{"text": msg["content"]}]}
+            else:
+                role = "model" if msg["role"] == "assistant" else msg["role"]
+                filtered_messages.append({"role": role, "parts": [{"text": msg["content"]}]})
+
+        payload = {
+             "contents": filtered_messages,
+             "safetySettings": [ # Example: more permissive settings
+                 {"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT", "threshold": "BLOCK_NONE"},
+                 {"category": "HARM_CATEGORY_HATE_SPEECH", "threshold": "BLOCK_NONE"},
+                 {"category": "HARM_CATEGORY_HARASSMENT", "threshold": "BLOCK_NONE"},
+                 {"category": "HARM_CATEGORY_DANGEROUS_CONTENT", "threshold": "BLOCK_NONE"},
+             ],
+             "generationConfig": {"temperature": temperature}
+        }
+        if max_tokens: payload["generationConfig"]["maxOutputTokens"] = max_tokens
+        if system_instruction: payload["system_instruction"] = system_instruction
+        
+        request_url = f"{base_url}{model_id}:streamGenerateContent?key={api_key}" # API key in query param
+        headers = {"Content-Type": "application/json"}
 
+        try:
             response = requests.post(request_url, headers=headers, json=payload, stream=True, timeout=180)
             response.raise_for_status()
-
-            byte_buffer = b""
-            for chunk in response.iter_content(chunk_size=8192):
-                 byte_buffer += chunk
-                 while b'\n' in byte_buffer:
-                    line, byte_buffer = byte_buffer.split(b'\n', 1)
-                    decoded_line = line.decode('utf-8', errors='ignore')
-
-                    if decoded_line.startswith('data: '):
-                        decoded_line = decoded_line[6:].strip()
-
-                    if not decoded_line: continue
+            
+            # Google's stream is a bit different, often newline-delimited JSON arrays/objects
+            buffer = ""
+            for chunk in response.iter_content(chunk_size=None):
+                buffer += chunk.decode('utf-8', errors='replace')
+                # Google might send chunks that are not complete JSON objects, or multiple objects
+                # A common pattern is [ {obj1} , {obj2} ] where chunks split mid-array or mid-object.
+                # This parsing needs to be robust. A simple split by '\n' might not always work if JSON is pretty-printed.
+                # The previous code's `json.loads(f"[{decoded_line}]")` was an attempt to handle this.
+                # For now, let's assume newline delimited for simplicity, but this is a known tricky part.
+                
+                while '\n' in buffer:
+                    line, buffer = buffer.split('\n', 1)
+                    line = line.strip()
+                    if not line: continue
+                    if line.startswith(','): line = line[1:] # Handle leading commas if splitting an array
 
                     try:
-                         event_data_list = json.loads(f"[{decoded_line}]")
-                         if not isinstance(event_data_list, list): event_data_list = [event_data_list]
+                        # Remove "data: " prefix if present (less common for Gemini direct API but good practice)
+                        if line.startswith('data: '): line = line[len('data: '):]
+                        
+                        # Gemini often streams an array of objects, or just one object.
+                        # Try to parse as a single object first. If fails, try as array.
+                        parsed_data = None
+                        try:
+                            parsed_data = json.loads(line)
+                        except json.JSONDecodeError:
+                            # If it's part of an array, it might be missing brackets.
+                            # This heuristic is fragile. A proper SSE parser or stateful JSON parser is better.
+                            if line.startswith('{') and line.endswith('}'): # Looks like a complete object
+                                pass # already tried json.loads
+                            # Try to wrap with [] if it seems like a list content without brackets
+                            elif line.startswith('{') or line.endswith('}'):
+                                try:
+                                    temp_parsed_list = json.loads(f"[{line}]")
+                                    if temp_parsed_list and isinstance(temp_parsed_list, list):
+                                        parsed_data = temp_parsed_list[0] # take first if it becomes a list
+                                except json.JSONDecodeError:
+                                    logger.warning(f"Google: Still can't parse line even with array wrap: {line}")
+
+                        if parsed_data:
+                            data_to_process = [parsed_data] if isinstance(parsed_data, dict) else parsed_data # Ensure list
+                            for event_data in data_to_process:
+                                if not isinstance(event_data, dict): continue
+                                if event_data.get("candidates"):
+                                    for candidate in event_data["candidates"]:
+                                        if candidate.get("content", {}).get("parts"):
+                                            for part in candidate["content"]["parts"]:
+                                                if part.get("text"):
+                                                    yield part["text"]
+                    except json.JSONDecodeError:
+                        logger.warning(f"Google: JSONDecodeError for line: {line}")
+                    except Exception as e_google_proc:
+                        logger.error(f"Google: Error processing stream data: {e_google_proc}, Line: {line}")
+
+        except requests.exceptions.HTTPError as e:
+            err_msg = f"Google API HTTP Error ({e.response.status_code}): {e.response.text[:500]}"
+            logger.error(err_msg, exc_info=False)
+            yield f"Error: {err_msg}"
+        except Exception as e:
+            logger.exception(f"Unexpected error during Google stream:")
+            yield f"Error: An unexpected error occurred with Google API: {e}"
+        return
 
-                         for event_data in event_data_list:
-                             if not isinstance(event_data, dict): continue
+    # --- Cohere ---
+    elif provider_lower == "cohere":
+        headers = {"Authorization": f"Bearer {api_key}", "Content-Type": "application/json", "Accept": "application/json"}
+        
+        # Cohere message format
+        chat_history_cohere = []
+        preamble_cohere = None
+        user_message_cohere = ""
+
+        temp_messages = list(messages) # Work with a copy
+        if temp_messages and temp_messages[0]["role"] == "system":
+            preamble_cohere = temp_messages.pop(0)["content"]
+        
+        if temp_messages:
+            user_message_cohere = temp_messages.pop()["content"] # Last message is the current user query
+            for msg in temp_messages: # Remaining are history
+                role = "USER" if msg["role"] == "user" else "CHATBOT"
+                chat_history_cohere.append({"role": role, "message": msg["content"]})
+        
+        if not user_message_cohere:
+            yield "Error: User message is empty for Cohere."
+            return
 
-                             if event_data.get("candidates") and len(event_data["candidates"]) > 0:
-                                 candidate = event_data["candidates"][0]
-                                 if candidate.get("content") and candidate["content"].get("parts"):
-                                     full_text_chunk = "".join(part.get("text", "") for part in candidate["content"]["parts"])
-                                     if full_text_chunk:
-                                         yield full_text_chunk
+        payload = {
+            "model": model_id, 
+            "message": user_message_cohere, 
+            "stream": True, 
+            "temperature": temperature
+        }
+        if max_tokens: payload["max_tokens"] = max_tokens # Cohere uses max_tokens
+        if chat_history_cohere: payload["chat_history"] = chat_history_cohere
+        if preamble_cohere: payload["preamble"] = preamble_cohere
+        
+        try:
+            response = requests.post(base_url, headers=headers, json=payload, stream=True, timeout=180)
+            response.raise_for_status()
+            
+            # Cohere SSE format is event: type\ndata: {json}\n\n
+            buffer = ""
+            for chunk_bytes in response.iter_content(chunk_size=None):
+                buffer += chunk_bytes.decode('utf-8', errors='replace')
+                while '\n\n' in buffer:
+                    event_str, buffer = buffer.split('\n\n', 1)
+                    if not event_str.strip(): continue
+                    
+                    event_type = None
+                    data_json_str = None
+                    for line in event_str.splitlines():
+                        if line.startswith("event:"): event_type = line[len("event:"):].strip()
+                        elif line.startswith("data:"): data_json_str = line[len("data:"):].strip()
+                    
+                    if data_json_str:
+                        try:
+                            data = json.loads(data_json_str)
+                            if event_type == "text-generation" and "text" in data:
+                                yield data["text"]
+                            elif event_type == "stream-end":
+                                logger.debug(f"Cohere stream ended. Finish reason: {data.get('finish_reason')}")
+                                return 
+                        except json.JSONDecodeError:
+                            logger.warning(f"Cohere: Failed to decode JSON: {data_json_str}")
+            if buffer.strip():
+                 logger.debug(f"Cohere: Remaining buffer: {buffer.strip()}")
+
+
+        except requests.exceptions.HTTPError as e:
+            err_msg = f"Cohere API HTTP Error ({e.response.status_code}): {e.response.text[:500]}"
+            logger.error(err_msg, exc_info=False)
+            yield f"Error: {err_msg}"
+        except Exception as e:
+            logger.exception(f"Unexpected error during Cohere stream:")
+            yield f"Error: An unexpected error occurred with Cohere API: {e}"
+        return
 
-                    except json.JSONDecodeError:
-                        logger.warning(f"Failed to decode JSON from Google stream chunk: {decoded_line}. Accumulating buffer.")
-                        pass
-
-                    except Exception as e:
-                        logger.error(f"Error processing Google stream data: {e}, Data: {decoded_line}")
-
-            if byte_buffer:
-                 remaining_line = byte_buffer.decode('utf-8', errors='ignore').strip()
-                 if remaining_line:
-                      try:
-                           event_data_list = json.loads(f"[{remaining_line}]")
-                           if not isinstance(event_data_list, list): event_data_list = [event_data_list]
-                           for event_data in event_data_list:
-                               if not isinstance(event_data, dict): continue
-                               if event_data.get("candidates") and len(event_data["candidates"]) > 0:
-                                   candidate = event_data["candidates"][0]
-                                   if candidate.get("content") and candidate["content"].get("parts"):
-                                       full_text_chunk = "".join(part.get("text", "") for part in candidate["content"]["parts"])
-                                       if full_text_chunk:
-                                           yield full_text_chunk
-                      except json.JSONDecodeError:
-                           logger.warning(f"Failed to decode final Google stream buffer JSON: {remaining_line}")
-                      except Exception as e:
-                           logger.error(f"Error processing final Google stream buffer data: {e}, Data: {remaining_line}")
-
-
-        elif provider_lower == "cohere":
-            headers = {"Authorization": f"Bearer {api_key}", "Content-Type": "application/json"}
-            request_url = f"{base_url}"
-
-            chat_history_for_cohere = []
-            system_prompt_for_cohere = None
-            current_message_for_cohere = ""
-
-            temp_history = []
-            for msg in messages:
-                 if msg["role"] == "system":
-                      system_prompt_for_cohere = msg["content"]
-                 elif msg["role"] == "user" or msg["role"] == "assistant":
-                     temp_history.append(msg)
-
-            if temp_history:
-                 current_message_for_cohere = temp_history[-1]["content"]
-                 chat_history_for_cohere = [{"role": ("chatbot" if m["role"] == "assistant" else m["role"]), "message": m["content"]} for m in temp_history[:-1]]
-
-            if not current_message_for_cohere:
-                 yield "Error: User message not found for Cohere API call."
-                 return
-
-            payload = {
-                "model": model_id,
-                "message": current_message_for_cohere,
-                "stream": True,
-                "temperature": 0.7
-            }
-            if chat_history_for_cohere:
-                 payload["chat_history"] = chat_history_for_cohere
-            if system_prompt_for_cohere:
-                 payload["preamble"] = system_prompt_for_cohere
+    # --- HuggingFace Inference API (Basic TGI support) ---
+    # This is very basic and might not work for all models or complex scenarios.
+    # Assumes model is deployed with Text Generation Inference (TGI) and supports streaming.
+    elif provider_lower == "huggingface":
+        headers = {"Authorization": f"Bearer {api_key}", "Content-Type": "application/json"}
+        # Construct prompt string for TGI (often needs specific formatting)
+        # This is a generic attempt, specific models might need <|user|>, <|assistant|> etc.
+        prompt_parts = []
+        for msg in messages:
+            role_prefix = ""
+            if msg['role'] == 'system': role_prefix = "System: " # Or might be ignored/handled differently
+            elif msg['role'] == 'user': role_prefix = "User: "
+            elif msg['role'] == 'assistant': role_prefix = "Assistant: "
+            prompt_parts.append(f"{role_prefix}{msg['content']}")
+        
+        # TGI typically expects a final "Assistant: " to start generating from
+        tgi_prompt = "\n".join(prompt_parts) + "\nAssistant: "
+
+        payload = {
+            "inputs": tgi_prompt,
+            "parameters": {
+                "temperature": temperature if temperature > 0 else 0.01, # TGI needs temp > 0 for sampling
+                "max_new_tokens": max_tokens or 1024, # Default TGI max_new_tokens
+                "return_full_text": False, # We only want generated part
+                "do_sample": True if temperature > 0 else False,
+            },
+            "stream": True
+        }
+        request_url = f"{base_url}{model_id}" # Model ID is part of URL path for HF
 
+        try:
             response = requests.post(request_url, headers=headers, json=payload, stream=True, timeout=180)
             response.raise_for_status()
 
-            byte_buffer = b""
-            for chunk in response.iter_content(chunk_size=8192):
-                byte_buffer += chunk
-                while b'\n\n' in byte_buffer:
-                    event_chunk, byte_buffer = byte_buffer.split(b'\n\n', 1)
-                    lines = event_chunk.strip().split(b'\n')
-                    event_type = None
-                    event_data = None
-
-                    for l in lines:
-                         if l.startswith(b"event: "): event_type = l[7:].strip().decode('utf-8', errors='ignore')
-                         elif l.startswith(b"data: "):
-                              try: event_data = json.loads(l[6:].strip().decode('utf-8', errors='ignore'))
-                              except json.JSONDecodeError: logger.warning(f"Cohere: Failed to decode event data JSON: {l[6:].strip()}")
-
-                    if event_type == "text-generation" and event_data and "text" in event_data:
-                        yield event_data["text"]
-                    elif event_type == "stream-end":
-                        byte_buffer = b''
-                        break
-
-            if byte_buffer:
-                 event_chunk = byte_buffer.strip()
-                 if event_chunk:
-                      lines = event_chunk.split(b'\n')
-                      event_type = None
-                      event_data = None
-                      for l in lines:
-                           if l.startswith(b"event: "): event_type = l[7:].strip().decode('utf-8', errors='ignore')
-                           elif l.startswith(b"data: "):
-                                try: event_data = json.loads(l[6:].strip().decode('utf-8', errors='ignore'))
-                                except json.JSONDecodeError: logger.warning(f"Cohere: Failed to decode final event data JSON: {l[6:].strip()}")
-
-                      if event_type == "text-generation" and event_data and "text" in event_data:
-                          yield event_data["text"]
-                      elif event_type == "stream-end":
-                           pass
-
-
-        elif provider_lower == "huggingface":
-             yield f"Error: Direct Hugging Face Inference API streaming for chat models is experimental and model-dependent. Consider using OpenRouter or TogetherAI for HF models with standardized streaming."
-             return
-
-        else:
-            yield f"Error: Unsupported provider '{provider}' for streaming chat."
-            return
+            # TGI SSE stream: data: {"token": {"id": ..., "text": "...", "logprob": ..., "special": ...}}
+            # Or sometimes just data: "text_chunk" for simpler models/configs
+            buffer = ""
+            for chunk_bytes in response.iter_content(chunk_size=None):
+                buffer += chunk_bytes.decode('utf-8', errors='replace')
+                while '\n' in buffer: # TGI often uses single newline
+                    line, buffer = buffer.split('\n', 1)
+                    line = line.strip()
+                    if not line: continue
+
+                    if line.startswith('data:'):
+                        data_json_str = line[len('data:'):].strip()
+                        try:
+                            data = json.loads(data_json_str)
+                            if "token" in data and "text" in data["token"]:
+                                yield data["token"]["text"]
+                            elif "generated_text" in data and data.get("details") is None: # Sometimes a final non-streaming like object might appear
+                                 # This case is tricky, if it's the *only* thing then it's not really streaming
+                                 pass # For now, ignore if it's not a token object
+                            # Some TGI might send raw text if not fully SSE compliant for stream
+                            # elif isinstance(data, str): yield data 
+
+                        except json.JSONDecodeError:
+                            # If it's not JSON, it might be a raw string (less common for TGI stream=True)
+                            # For safety, only yield if it's a clear text string
+                            if not data_json_str.startswith('{') and not data_json_str.startswith('['):
+                                yield data_json_str
+                            else:
+                                logger.warning(f"HF: Failed to decode JSON and not raw string: {data_json_str}")
+            if buffer.strip():
+                 logger.debug(f"HF: Remaining buffer: {buffer.strip()}")
+
+
+        except requests.exceptions.HTTPError as e:
+            err_msg = f"HF API HTTP Error ({e.response.status_code}): {e.response.text[:500]}"
+            logger.error(err_msg, exc_info=False)
+            yield f"Error: {err_msg}"
+        except Exception as e:
+            logger.exception(f"Unexpected error during HF stream:")
+            yield f"Error: An unexpected error occurred with HF API: {e}"
+        return
 
-    except requests.exceptions.HTTPError as e:
-        status_code = e.response.status_code if e.response is not None else 'N/A'
-        error_text = e.response.text if e.response is not None else 'No response text'
-        logger.error(f"HTTP error during streaming for {provider}/{model_id}: {e}")
-        yield f"API HTTP Error ({status_code}): {error_text}\nDetails: {e}"
-    except requests.exceptions.RequestException as e:
-        logger.error(f"Request error during streaming for {provider}/{model_id}: {e}")
-        yield f"API Request Error: Could not connect or receive response from {provider} ({e})"
-    except Exception as e:
-        logger.exception(f"Unexpected error during streaming for {provider}/{model_id}:")
-        yield f"An unexpected error occurred: {e}"
+    else:
+        yield f"Error: Provider '{provider}' is not configured for streaming in this handler."
+        return