Spaces:
Sleeping
Sleeping
| """ | |
| LLM Client for Overgrowth Pipeline | |
| Supports multiple providers: OpenAI, Anthropic, Blaxel, SambaNova, Nebius, Hugging Face, Modal | |
| """ | |
| import os | |
| import json | |
| import logging | |
| import uuid | |
| import time | |
| from typing import Dict, List, Optional, Iterator, Any | |
| from dataclasses import dataclass | |
| logger = logging.getLogger(__name__) | |
| # Import API monitor for tracking | |
| try: | |
| from agent.api_monitor import monitor | |
| except ImportError: | |
| logger.warning("API monitor not available - tracking disabled") | |
| monitor = None | |
| class LLMMessage: | |
| role: str # system, user, assistant | |
| content: str | |
| class LLMClient: | |
| """ | |
| Unified LLM client supporting multiple providers | |
| Falls back gracefully if API keys not available | |
| """ | |
| def __init__(self): | |
| def _get_env(names): | |
| for n in names: | |
| v = os.getenv(n) | |
| if v: | |
| return v.strip() | |
| return None | |
| # Support both standard and MCP hackathon naming conventions | |
| self.openai_key = _get_env(["OPENAI_API_KEY", "OPENAI_MCP_1ST_BDAY"]) | |
| self.anthropic_key = _get_env(["ANTHROPIC_API_KEY", "ANTHROPIC_MCP_1ST_BDAY"]) | |
| self.blaxel_key = _get_env(["BLAXEL_MCP_1ST_BDAY"]) | |
| self.sambanova_key = _get_env(["SAMBA_NOVA_MCP_1ST_BDAY"]) | |
| self.nebius_key = _get_env(["NEBIUS_MCP_1ST_BDAY"]) | |
| self.huggingface_key = _get_env(["HUGGING_FACE_MCP_1ST_BDAY"]) | |
| self.modal_key = _get_env(["MODAL_API_KEY", "MODAL_TOKEN"]) | |
| self.modal_base_url = os.getenv("MODAL_BASE_URL") | |
| self.modal_model = os.getenv("MODAL_MODEL", "gpt-4o-mini") | |
| # Determine which provider to use | |
| self.provider = self._detect_provider() | |
| if self.provider: | |
| logger.info(f"LLM client initialized with provider: {self.provider}") | |
| else: | |
| logger.warning("No LLM API keys found - using mock responses") | |
| def env_status() -> Dict[str, str]: | |
| """ | |
| Report which keys are present (without exposing values) and chosen provider. | |
| Helpful for UI/debug when Secrets are misconfigured. | |
| """ | |
| forced = os.getenv("OG_LLM_PROVIDER", "").strip().lower() | |
| def _present(names): | |
| for n in names: | |
| v = os.getenv(n) | |
| if v and v.strip(): | |
| return True | |
| return False | |
| status = { | |
| "openai_key": "present" if _present(["OPENAI_API_KEY", "OPENAI_MCP_1ST_BDAY"]) else "missing", | |
| "anthropic_key": "present" if _present(["ANTHROPIC_API_KEY", "ANTHROPIC_MCP_1ST_BDAY"]) else "missing", | |
| "blaxel_key": "present" if _present(["BLAXEL_MCP_1ST_BDAY"]) else "missing", | |
| "sambanova_key": "present" if _present(["SAMBA_NOVA_MCP_1ST_BDAY"]) else "missing", | |
| "nebius_key": "present" if _present(["NEBIUS_MCP_1ST_BDAY"]) else "missing", | |
| "huggingface_key": "present" if _present(["HUGGING_FACE_MCP_1ST_BDAY"]) else "missing", | |
| "modal_key": "present" if _present(["MODAL_API_KEY", "MODAL_TOKEN"]) else "missing", | |
| "provider": "unknown", | |
| "forced_provider": forced or "none", | |
| } | |
| if forced: | |
| status["provider"] = forced | |
| return status | |
| if status["anthropic_key"] == "present": | |
| status["provider"] = "anthropic" | |
| elif status["openai_key"] == "present": | |
| status["provider"] = "openai" | |
| elif status["blaxel_key"] == "present": | |
| status["provider"] = "blaxel" | |
| elif status["sambanova_key"] == "present": | |
| status["provider"] = "sambanova" | |
| elif status["nebius_key"] == "present": | |
| status["provider"] = "nebius" | |
| elif status["modal_key"] == "present": | |
| status["provider"] = "modal" | |
| elif status["huggingface_key"] == "present": | |
| status["provider"] = "huggingface" | |
| return status | |
| def _detect_provider(self) -> Optional[str]: | |
| """Detect which LLM provider is available""" | |
| forced = os.getenv("OG_LLM_PROVIDER", "").strip().lower() | |
| if forced: | |
| return forced | |
| # Prefer reliable Anthropic/OpenAI first; Blaxel optional | |
| if self.anthropic_key: | |
| return "anthropic" | |
| elif self.openai_key: | |
| return "openai" | |
| elif self.blaxel_key: | |
| return "blaxel" | |
| elif self.sambanova_key: | |
| return "sambanova" | |
| elif self.nebius_key: | |
| return "nebius" | |
| elif self.modal_key: | |
| return "modal" | |
| elif self.huggingface_key: | |
| return "huggingface" | |
| return None | |
| def chat( | |
| self, | |
| messages: List[LLMMessage], | |
| temperature: float = 0.7, | |
| max_tokens: int = 4000, | |
| stream: bool = False | |
| ) -> str: | |
| """ | |
| Send chat completion request | |
| Returns response text or yields chunks if streaming | |
| """ | |
| if not self.provider: | |
| raise RuntimeError("No LLM provider configured. Set ANTHROPIC_MCP_1ST_BDAY or OPENAI_MCP_1ST_BDAY.") | |
| if self.provider == "openai": | |
| return self._call_openai(messages, temperature, max_tokens, stream) | |
| elif self.provider == "anthropic": | |
| return self._call_anthropic(messages, temperature, max_tokens, stream) | |
| elif self.provider == "blaxel": | |
| return self._call_blaxel(messages, temperature, max_tokens) | |
| elif self.provider == "sambanova": | |
| return self._call_sambanova(messages, temperature, max_tokens) | |
| elif self.provider == "nebius": | |
| return self._call_nebius(messages, temperature, max_tokens) | |
| elif self.provider == "modal": | |
| return self._call_modal(messages, temperature, max_tokens) | |
| elif self.provider == "huggingface": | |
| return self._call_huggingface(messages, temperature, max_tokens) | |
| def chat_stream( | |
| self, | |
| messages: List[LLMMessage], | |
| temperature: float = 0.7, | |
| max_tokens: int = 4000 | |
| ) -> Iterator[str]: | |
| """Stream chat completion response""" | |
| if not self.provider: | |
| raise RuntimeError("No LLM provider configured. Set ANTHROPIC_MCP_1ST_BDAY or OPENAI_MCP_1ST_BDAY.") | |
| if self.provider == "openai": | |
| yield from self._stream_openai(messages, temperature, max_tokens) | |
| elif self.provider == "anthropic": | |
| yield from self._stream_anthropic(messages, temperature, max_tokens) | |
| else: | |
| # Other providers: no streaming support; fall back to single response | |
| yield self.chat(messages, temperature=temperature, max_tokens=max_tokens) | |
| def _call_openai(self, messages, temperature, max_tokens, stream): | |
| """Call OpenAI API""" | |
| call_id = str(uuid.uuid4()) | |
| model = "gpt-4o" | |
| if monitor: | |
| monitor.start_call(call_id, "llm", "openai", model, temperature=temperature) | |
| try: | |
| from openai import OpenAI | |
| client = OpenAI(api_key=self.openai_key) | |
| response = client.chat.completions.create( | |
| model=model, | |
| messages=[{"role": m.role, "content": m.content} for m in messages], | |
| temperature=temperature, | |
| max_tokens=max_tokens | |
| ) | |
| content = response.choices[0].message.content | |
| if monitor and response.usage: | |
| monitor.complete_call( | |
| call_id, | |
| success=True, | |
| input_tokens=response.usage.prompt_tokens, | |
| output_tokens=response.usage.completion_tokens | |
| ) | |
| return content | |
| except Exception as e: | |
| logger.error(f"OpenAI API error: {e}") | |
| if monitor: | |
| monitor.complete_call(call_id, success=False, error_message=str(e)) | |
| raise | |
| def _stream_openai(self, messages, temperature, max_tokens): | |
| """Stream from OpenAI""" | |
| call_id = str(uuid.uuid4()) | |
| model = "gpt-4o" | |
| if monitor: | |
| monitor.start_call(call_id, "llm", "openai", model, temperature=temperature) | |
| total_tokens_est = 0 | |
| try: | |
| from openai import OpenAI | |
| client = OpenAI(api_key=self.openai_key) | |
| stream = client.chat.completions.create( | |
| model=model, | |
| messages=[{"role": m.role, "content": m.content} for m in messages], | |
| temperature=temperature, | |
| max_tokens=max_tokens, | |
| stream=True | |
| ) | |
| for chunk in stream: | |
| if chunk.choices[0].delta.content: | |
| content = chunk.choices[0].delta.content | |
| total_tokens_est += len(content) // 4 | |
| yield content | |
| # Complete call after streaming | |
| if monitor: | |
| monitor.complete_call(call_id, success=True, | |
| input_tokens=total_tokens_est // 2, | |
| output_tokens=total_tokens_est // 2) | |
| except Exception as e: | |
| logger.error(f"OpenAI streaming error: {e}") | |
| if monitor: | |
| monitor.complete_call(call_id, success=False, error_message=str(e)) | |
| raise | |
| def _call_anthropic(self, messages, temperature, max_tokens, stream): | |
| """Call Anthropic API""" | |
| call_id = str(uuid.uuid4()) | |
| model = "claude-3-haiku-20240307" | |
| if monitor: | |
| monitor.start_call(call_id, "llm", "anthropic", model, temperature=temperature) | |
| try: | |
| import anthropic | |
| client = anthropic.Anthropic(api_key=self.anthropic_key) | |
| # Convert messages format | |
| system_msg = None | |
| user_messages = [] | |
| for m in messages: | |
| if m.role == "system": | |
| system_msg = m.content | |
| else: | |
| user_messages.append({"role": m.role, "content": m.content}) | |
| response = client.messages.create( | |
| model=model, | |
| max_tokens=max_tokens, | |
| temperature=temperature, | |
| system=system_msg if system_msg else "You are a helpful network automation assistant.", | |
| messages=user_messages | |
| ) | |
| content = response.content[0].text | |
| if monitor and response.usage: | |
| monitor.complete_call( | |
| call_id, | |
| success=True, | |
| input_tokens=response.usage.input_tokens, | |
| output_tokens=response.usage.output_tokens | |
| ) | |
| return content | |
| except Exception as e: | |
| err_msg = f"{e} | cause: {repr(getattr(e, '__cause__', ''))}" | |
| logger.error(f"Anthropic API error: {err_msg}") | |
| if monitor: | |
| monitor.complete_call(call_id, success=False, error_message=err_msg) | |
| # Fallback to OpenAI if available | |
| if self.openai_key: | |
| logger.info("Falling back to OpenAI due to Anthropic error") | |
| return self._call_openai(messages, temperature, max_tokens, stream=False) | |
| raise RuntimeError(f"Anthropic call failed: {err_msg}") | |
| def _stream_anthropic(self, messages, temperature, max_tokens): | |
| """Stream from Anthropic""" | |
| call_id = str(uuid.uuid4()) | |
| model = "claude-3-haiku-20240307" | |
| if monitor: | |
| monitor.start_call(call_id, "llm", "anthropic", model, temperature=temperature) | |
| total_input_tokens = 0 | |
| total_output_tokens = 0 | |
| try: | |
| import anthropic | |
| client = anthropic.Anthropic(api_key=self.anthropic_key) | |
| # Convert messages format | |
| system_msg = None | |
| user_messages = [] | |
| for m in messages: | |
| if m.role == "system": | |
| system_msg = m.content | |
| else: | |
| user_messages.append({"role": m.role, "content": m.content}) | |
| with client.messages.stream( | |
| model=model, | |
| max_tokens=max_tokens, | |
| temperature=temperature, | |
| system=system_msg if system_msg else "You are a helpful network automation assistant.", | |
| messages=user_messages | |
| ) as stream: | |
| for text in stream.text_stream: | |
| yield text | |
| # Get final usage stats | |
| final_message = stream.get_final_message() | |
| if final_message and final_message.usage and monitor: | |
| monitor.complete_call( | |
| call_id, | |
| success=True, | |
| input_tokens=final_message.usage.input_tokens, | |
| output_tokens=final_message.usage.output_tokens | |
| ) | |
| except Exception as e: | |
| err_msg = f"{e} | cause: {repr(getattr(e, '__cause__', ''))}" | |
| logger.error(f"Anthropic streaming error: {err_msg}") | |
| if monitor: | |
| monitor.complete_call(call_id, success=False, error_message=err_msg) | |
| if self.openai_key: | |
| logger.info("Falling back to OpenAI streaming due to Anthropic error") | |
| yield from self._stream_openai(messages, temperature, max_tokens) | |
| else: | |
| raise RuntimeError(f"Anthropic streaming failed: {err_msg}") | |
| # ----- Additional Providers (non-streaming) ----- | |
| def _call_blaxel(self, messages, temperature, max_tokens): | |
| """Call Blaxel sandbox API (OpenAI-style).""" | |
| import requests | |
| call_id = str(uuid.uuid4()) | |
| model = os.getenv("BLAXEL_MODEL", "blaxel/claude-3-haiku") | |
| base_url = os.getenv("BLAXEL_BASE_URL", "https://api.blaxel.ai/v0").rstrip("/") | |
| # Try a few endpoint variants in case the API version changes (v0 vs v1) | |
| endpoints = [ | |
| f"{base_url}/agents/query", # primary documented endpoint | |
| f"{base_url}/chat/completions", # OpenAI-style fallback | |
| ] | |
| if base_url.endswith("/v0"): | |
| endpoints.append(f"{base_url[:-3]}/v1/agents/query") | |
| endpoints.append(f"{base_url[:-3]}/v1/chat/completions") | |
| elif base_url.endswith("/v1"): | |
| endpoints.append(f"{base_url[:-3]}/v0/agents/query") | |
| endpoints.append(f"{base_url[:-3]}/v0/chat/completions") | |
| if monitor: | |
| monitor.start_call(call_id, "llm", "blaxel", model, temperature=temperature) | |
| thread_id = os.getenv("BLAXEL_THREAD_ID", call_id) # keep conversations grouped if provided | |
| payload = { | |
| "model": model, | |
| "messages": [{"role": m.role, "content": m.content} for m in messages], | |
| "temperature": temperature, | |
| "max_tokens": max_tokens, | |
| # Some Blaxel endpoints expect `inputs` instead of `messages`; include both for compatibility | |
| "inputs": " ".join(m.content for m in messages if m.content), | |
| } | |
| headers = { | |
| "X-Blaxel-Authorization": f"Bearer {self.blaxel_key}", | |
| "X-Blaxel-Thread-Id": thread_id, | |
| "Content-Type": "application/json", | |
| } | |
| def _fallback_to_anthropic(err): | |
| # If Blaxel is down but Anthropic is configured, fall back transparently | |
| if self.anthropic_key: | |
| logger.warning(f"Blaxel call failed ({err}); falling back to Anthropic") | |
| return self._call_anthropic(messages, temperature, max_tokens, stream=False) | |
| if self.openai_key: | |
| logger.warning(f"Blaxel call failed ({err}); falling back to OpenAI") | |
| return self._call_openai(messages, temperature, max_tokens, stream=False) | |
| raise err | |
| try: | |
| last_error = None | |
| for endpoint in endpoints: | |
| try: | |
| resp = requests.post(endpoint, json=payload, headers=headers, timeout=30) | |
| resp.raise_for_status() | |
| data = resp.json() | |
| content = data.get("choices", [{}])[0].get("message", {}).get("content", "") | |
| usage = data.get("usage", {}) | |
| if monitor: | |
| monitor.complete_call( | |
| call_id, | |
| success=True, | |
| input_tokens=usage.get("prompt_tokens"), | |
| output_tokens=usage.get("completion_tokens") | |
| ) | |
| return content | |
| except requests.HTTPError as e: | |
| last_error = e | |
| # Retry on 404 to handle versioned paths; otherwise break fast | |
| if resp.status_code != 404: | |
| raise | |
| continue | |
| # If all endpoints failed, raise the last error | |
| if last_error: | |
| return _fallback_to_anthropic(last_error) | |
| raise RuntimeError("Blaxel call failed: no endpoint attempted") | |
| except Exception as e: | |
| if monitor: | |
| monitor.complete_call(call_id, success=False, error_message=str(e)) | |
| return _fallback_to_anthropic(e) | |
| def _call_sambanova(self, messages, temperature, max_tokens): | |
| """Call SambaNova API (OpenAI-compatible).""" | |
| import requests | |
| call_id = str(uuid.uuid4()) | |
| model = os.getenv("SAMBA_NOVA_MODEL", "Meta-Llama-3-8B-Instruct") | |
| base_url = os.getenv("SAMBA_NOVA_BASE_URL", "https://api.sambanova.ai/v1") | |
| endpoint = f"{base_url.rstrip('/')}/chat/completions" | |
| if monitor: | |
| monitor.start_call(call_id, "llm", "sambanova", model, temperature=temperature) | |
| payload = { | |
| "model": model, | |
| "messages": [{"role": m.role, "content": m.content} for m in messages], | |
| "temperature": temperature, | |
| "max_tokens": max_tokens, | |
| } | |
| headers = { | |
| "Authorization": f"Bearer {self.sambanova_key}", | |
| "Content-Type": "application/json" | |
| } | |
| try: | |
| resp = requests.post(endpoint, json=payload, headers=headers, timeout=30) | |
| resp.raise_for_status() | |
| data = resp.json() | |
| content = data.get("choices", [{}])[0].get("message", {}).get("content", "") | |
| usage = data.get("usage", {}) | |
| if monitor: | |
| monitor.complete_call( | |
| call_id, | |
| success=True, | |
| input_tokens=usage.get("prompt_tokens"), | |
| output_tokens=usage.get("completion_tokens") | |
| ) | |
| return content | |
| except Exception as e: | |
| if monitor: | |
| monitor.complete_call(call_id, success=False, error_message=str(e)) | |
| raise | |
| def _call_nebius(self, messages, temperature, max_tokens): | |
| """Call Nebius Token Factory API (OpenAI-compatible).""" | |
| import requests | |
| call_id = str(uuid.uuid4()) | |
| model = os.getenv("NEBIUS_MODEL", "gpt-3.5-turbo") | |
| base_url = os.getenv("NEBIUS_BASE_URL", "https://api.studio.nebius.ai/v1") | |
| endpoint = f"{base_url.rstrip('/')}/chat/completions" | |
| if monitor: | |
| monitor.start_call(call_id, "llm", "nebius", model, temperature=temperature) | |
| payload = { | |
| "model": model, | |
| "messages": [{"role": m.role, "content": m.content} for m in messages], | |
| "temperature": temperature, | |
| "max_tokens": max_tokens, | |
| } | |
| headers = { | |
| "Authorization": f"Bearer {self.nebius_key}", | |
| "Content-Type": "application/json" | |
| } | |
| try: | |
| resp = requests.post(endpoint, json=payload, headers=headers, timeout=30) | |
| resp.raise_for_status() | |
| data = resp.json() | |
| content = data.get("choices", [{}])[0].get("message", {}).get("content", "") | |
| usage = data.get("usage", {}) | |
| if monitor: | |
| monitor.complete_call( | |
| call_id, | |
| success=True, | |
| input_tokens=usage.get("prompt_tokens"), | |
| output_tokens=usage.get("completion_tokens") | |
| ) | |
| return content | |
| except Exception as e: | |
| if monitor: | |
| monitor.complete_call(call_id, success=False, error_message=str(e)) | |
| raise | |
| def _call_modal(self, messages, temperature, max_tokens): | |
| """Call Modal's OpenAI-compatible endpoint (optional sponsor integration).""" | |
| import requests | |
| call_id = str(uuid.uuid4()) | |
| model = self.modal_model or "gpt-4o-mini" | |
| base_url = (self.modal_base_url or "https://api.modal.com/v1").rstrip("/") | |
| endpoint = f"{base_url}/chat/completions" | |
| if monitor: | |
| monitor.start_call(call_id, "llm", "modal", model, temperature=temperature) | |
| if not self.modal_key: | |
| raise RuntimeError("Modal provider selected but MODAL_API_KEY/MODAL_TOKEN is missing") | |
| payload = { | |
| "model": model, | |
| "messages": [{"role": m.role, "content": m.content} for m in messages], | |
| "temperature": temperature, | |
| "max_tokens": max_tokens, | |
| } | |
| headers = { | |
| "Authorization": f"Bearer {self.modal_key}", | |
| "Content-Type": "application/json", | |
| } | |
| try: | |
| resp = requests.post(endpoint, json=payload, headers=headers, timeout=45) | |
| resp.raise_for_status() | |
| data = resp.json() | |
| content = data.get("choices", [{}])[0].get("message", {}).get("content", "") | |
| usage = data.get("usage", {}) | |
| if monitor: | |
| monitor.complete_call( | |
| call_id, | |
| success=True, | |
| input_tokens=usage.get("prompt_tokens"), | |
| output_tokens=usage.get("completion_tokens"), | |
| ) | |
| return content | |
| except Exception as e: | |
| if monitor: | |
| monitor.complete_call(call_id, success=False, error_message=str(e)) | |
| raise | |
| def _call_huggingface(self, messages, temperature, max_tokens): | |
| """Call Hugging Face Inference API (text generation).""" | |
| import requests | |
| call_id = str(uuid.uuid4()) | |
| model = os.getenv("HUGGINGFACE_MODEL", "tiiuae/falcon-7b-instruct") | |
| endpoint = f"https://api-inference.huggingface.co/models/{model}" | |
| if monitor: | |
| monitor.start_call(call_id, "llm", "huggingface", model, temperature=temperature) | |
| # Simple prompt concatenation | |
| prompt = "\n".join(f"{m.role.upper()}: {m.content}" for m in messages) | |
| payload = { | |
| "inputs": prompt, | |
| "parameters": { | |
| "max_new_tokens": max_tokens, | |
| "temperature": temperature, | |
| "return_full_text": False | |
| } | |
| } | |
| headers = { | |
| "Authorization": f"Bearer {self.huggingface_key}", | |
| "Content-Type": "application/json" | |
| } | |
| try: | |
| resp = requests.post(endpoint, json=payload, headers=headers, timeout=60) | |
| resp.raise_for_status() | |
| data = resp.json() | |
| # Response can be list or dict | |
| text = "" | |
| if isinstance(data, list) and data: | |
| if isinstance(data[0], dict): | |
| text = data[0].get("generated_text", "") or data[0].get("generated_texts", "") | |
| else: | |
| text = str(data[0]) | |
| elif isinstance(data, dict): | |
| text = data.get("generated_text", "") or data.get("generated_texts", "") or "" | |
| if monitor: | |
| monitor.complete_call(call_id, success=True) | |
| return text | |
| except Exception as e: | |
| if monitor: | |
| monitor.complete_call(call_id, success=False, error_message=str(e)) | |
| raise | |
| def _mock_response(self, messages: List[LLMMessage]) -> str: | |
| """Deprecated: mocks disabled to avoid hiding real failures.""" | |
| raise RuntimeError("LLM mock responses are disabled. Provide a valid API key.") | |