model_logic.py

import os
import requests
import json
import logging

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
)
logger = logging.getLogger(__name__)

API_KEYS = {
  "HUGGINGFACE": 'HF_TOKEN',
  "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',
}

API_URLS = {
  "HUGGINGFACE": 'https://api-inference.huggingface.co/models/',
  "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',
  "XAI": 'https://api.x.ai/v1/chat/completions',
  "OPENAI": 'https://api.openai.com/v1/chat/completions',
  "GOOGLE": 'https://generativelanguage.googleapis.com/v1beta/models/',
}

MODELS_BY_PROVIDER = {
    "groq": {
        "default": "llama3-8b-8192",
        "models": {
            "Llama 3 8B (Groq)": "llama3-8b-8192",
            "Llama 3 70B (Groq)": "llama3-70b-8192",
            "Mixtral 8x7B (Groq)": "mixtral-8x7b-32768",
            "Gemma 7B (Groq)": "gemma-7b-it",
        }
    },
    "openrouter": {
         "default": "nousresearch/llama-3-8b-instruct",
         "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",
         }
    },
     "togetherai": {
         "default": "meta-llama/Llama-3-8b-chat-hf",
         "models": {
             "Llama 3 8B Chat (TogetherAI)": "meta-llama/Llama-3-8b-chat-hf",
             "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",
         }
     },
     "google": {
          "default": "gemini-1.5-flash-latest",
          "models": {
              "Gemini 1.5 Flash (Latest)": "gemini-1.5-flash-latest",
              "Gemini 1.5 Pro (Latest)": "gemini-1.5-pro-latest",
          }
     },
     "cohere": {
         "default": "command-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",
         "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",
         }
     },
     "openai": {
          "default": "gpt-3.5-turbo",
          "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",
          }
     },
     "xai": {
          "default": "grok-1",
          "models": {
              "Grok-1 (xAI)": "grok-1",
          }
     }
}

def _get_api_key(provider: str, ui_api_key_override: str = None) -> str:
    if ui_api_key_override:
        return ui_api_key_override.strip()

    env_var_name = API_KEYS.get(provider.upper())
    if env_var_name:
        env_key = os.getenv(env_var_name)
        if env_key:
            return env_key.strip()

    if provider.lower() == 'huggingface':
         hf_token = os.getenv("HF_TOKEN")
         if hf_token: return hf_token.strip()

    logger.warning(f"API Key not found for provider '{provider}'. Checked UI override and environment variable '{env_var_name or 'N/A'}'.")
    return None

def get_available_providers() -> list[str]:
    return sorted(list(MODELS_BY_PROVIDER.keys()))

def get_models_for_provider(provider: str) -> list[str]:
    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:
                return display_name
    if models_dict:
        return sorted(list(models_dict.keys()))[0]
    return None

def get_model_id_from_display_name(provider: str, display_name: str) -> str | None:
    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:
    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}'."
        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."
         return

    headers = {}
    payload = {}
    request_url = base_url

    logger.info(f"Calling {provider}/{model_display_name} (ID: {model_id}) stream...")

    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"

            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}"

            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

                    try:
                         event_data_list = json.loads(f"[{decoded_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 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

            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

    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}"