agent/llm_client.py

"""
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
@dataclass
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")

    @staticmethod
    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.")