Create model_logic.py

model_logic.py

@@ -0,0 +1,402 @@
+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}"