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