Spaces:
Sleeping
Sleeping
| """ | |
| 🏆 Premium Unified AI Client | |
| Supports: OpenAI, Google Gemini, Anthropic Claude, and OpenRouter (6 free models) | |
| """ | |
| import httpx | |
| from typing import Optional, Dict, Any, AsyncIterator | |
| import logging | |
| from src.config import Config, AIProvider, ModelConfig | |
| logger = logging.getLogger(__name__) | |
| class UnifiedAIClient: | |
| """ | |
| 🏆 Premium AI client supporting multiple providers | |
| - OpenAI GPT-5 (⭐ PRO) | |
| - Google Gemini 3 (⭐ PRO) | |
| - Anthropic Claude Sonnet 4.5 (⭐ PRO) | |
| - 6 Free OpenRouter Models (🆓 FREE) | |
| """ | |
| def __init__(self, model_name: str = "grok-4.1", api_key: Optional[str] = None): | |
| self.model_name = model_name | |
| self.model_config: Optional[ModelConfig] = Config.get_model_config(model_name) | |
| self.custom_api_key = api_key | |
| if not self.model_config: | |
| raise ValueError(f"Unknown model: {model_name}") | |
| # Initialize HTTP client | |
| self.http_client = httpx.AsyncClient(timeout=60.0) | |
| logger.info(f"Initialized AI client: {self.model_config.display_name}") | |
| async def generate( | |
| self, | |
| prompt: str, | |
| system_prompt: Optional[str] = None, | |
| max_tokens: int = 4096, | |
| temperature: float = 0.7, | |
| stream: bool = False | |
| ) -> str: | |
| """ | |
| Generate AI response using the configured model | |
| Args: | |
| prompt: User prompt | |
| system_prompt: Optional system prompt | |
| max_tokens: Maximum tokens to generate | |
| temperature: Sampling temperature | |
| stream: Whether to stream response | |
| Returns: | |
| Generated text response | |
| """ | |
| try: | |
| if self.model_config.provider == AIProvider.GOOGLE_GEMINI: | |
| return await self._generate_gemini(prompt, system_prompt, max_tokens, temperature) | |
| elif self.model_config.provider == AIProvider.OPENAI: | |
| return await self._generate_openai(prompt, system_prompt, max_tokens, temperature, stream) | |
| elif self.model_config.provider == AIProvider.ANTHROPIC: | |
| return await self._generate_anthropic(prompt, system_prompt, max_tokens, temperature) | |
| else: | |
| # All other models use OpenRouter | |
| return await self._generate_openrouter(prompt, system_prompt, max_tokens, temperature, stream) | |
| except Exception as e: | |
| logger.error(f"AI generation failed: {e}", exc_info=True) | |
| raise | |
| async def _generate_gemini( | |
| self, | |
| prompt: str, | |
| system_prompt: Optional[str], | |
| max_tokens: int, | |
| temperature: float | |
| ) -> str: | |
| """Generate using Google Gemini""" | |
| try: | |
| import google.generativeai as genai | |
| api_key = self.custom_api_key or Config.GOOGLE_API_KEY | |
| if not api_key: | |
| raise ValueError("Google API key required for Gemini") | |
| genai.configure(api_key=api_key) | |
| model = genai.GenerativeModel(self.model_config.model_id) | |
| full_prompt = f"{system_prompt}\n\n{prompt}" if system_prompt else prompt | |
| response = model.generate_content( | |
| full_prompt, | |
| generation_config=genai.GenerationConfig( | |
| max_output_tokens=max_tokens, | |
| temperature=temperature | |
| ) | |
| ) | |
| return response.text | |
| except Exception as e: | |
| logger.error(f"Gemini generation failed: {e}") | |
| raise | |
| async def _generate_openai( | |
| self, | |
| prompt: str, | |
| system_prompt: Optional[str], | |
| max_tokens: int, | |
| temperature: float, | |
| stream: bool | |
| ) -> str: | |
| """Generate using OpenAI""" | |
| api_key = self.custom_api_key or Config.OPENAI_API_KEY | |
| if not api_key: | |
| raise ValueError("OpenAI API key required") | |
| messages = [] | |
| if system_prompt: | |
| messages.append({"role": "system", "content": system_prompt}) | |
| messages.append({"role": "user", "content": prompt}) | |
| response = await self.http_client.post( | |
| "https://api.openai.com/v1/chat/completions", | |
| headers={ | |
| "Authorization": f"Bearer {api_key}", | |
| "Content-Type": "application/json" | |
| }, | |
| json={ | |
| "model": self.model_config.model_id, | |
| "messages": messages, | |
| "max_tokens": max_tokens, | |
| "temperature": temperature, | |
| "stream": False | |
| } | |
| ) | |
| response.raise_for_status() | |
| data = response.json() | |
| return data["choices"][0]["message"]["content"] | |
| async def _generate_anthropic( | |
| self, | |
| prompt: str, | |
| system_prompt: Optional[str], | |
| max_tokens: int, | |
| temperature: float | |
| ) -> str: | |
| """Generate using Anthropic Claude""" | |
| api_key = self.custom_api_key or Config.ANTHROPIC_API_KEY | |
| if not api_key: | |
| raise ValueError("Anthropic API key required") | |
| payload = { | |
| "model": self.model_config.model_id, | |
| "max_tokens": max_tokens, | |
| "temperature": temperature, | |
| "messages": [{"role": "user", "content": prompt}] | |
| } | |
| if system_prompt: | |
| payload["system"] = system_prompt | |
| response = await self.http_client.post( | |
| "https://api.anthropic.com/v1/messages", | |
| headers={ | |
| "x-api-key": api_key, | |
| "anthropic-version": "2023-06-01", | |
| "Content-Type": "application/json" | |
| }, | |
| json=payload | |
| ) | |
| response.raise_for_status() | |
| data = response.json() | |
| return data["content"][0]["text"] | |
| async def _generate_openrouter( | |
| self, | |
| prompt: str, | |
| system_prompt: Optional[str], | |
| max_tokens: int, | |
| temperature: float, | |
| stream: bool | |
| ) -> str: | |
| """ | |
| 🆓 Generate using OpenRouter (Free Models) | |
| Supports: Grok, KAT-Coder, Qwen3, LongCat, GPT-OSS, Kimi | |
| """ | |
| api_key = Config.OPENROUTER_API_KEY | |
| messages = [] | |
| if system_prompt: | |
| messages.append({"role": "system", "content": system_prompt}) | |
| messages.append({"role": "user", "content": prompt}) | |
| logger.info(f"🆓 Using OpenRouter model: {self.model_config.display_name}") | |
| response = await self.http_client.post( | |
| "https://openrouter.ai/api/v1/chat/completions", | |
| headers={ | |
| "Authorization": f"Bearer {api_key}", | |
| "Content-Type": "application/json", | |
| "HTTP-Referer": "https://codelint-mcp.com", | |
| "X-Title": "CodeLint MCP Server" | |
| }, | |
| json={ | |
| "model": self.model_config.model_id, | |
| "messages": messages, | |
| "max_tokens": max_tokens, | |
| "temperature": temperature, | |
| "stream": False | |
| } | |
| ) | |
| if response.status_code != 200: | |
| error_text = response.text | |
| logger.error(f"OpenRouter error: {error_text}") | |
| raise Exception(f"OpenRouter API error: {response.status_code}") | |
| data = response.json() | |
| if "choices" not in data or len(data["choices"]) == 0: | |
| raise Exception(f"No response from OpenRouter: {data}") | |
| return data["choices"][0]["message"]["content"] | |
| async def close(self): | |
| """Close HTTP client""" | |
| await self.http_client.aclose() | |
| async def generate_ai_response( | |
| prompt: str, | |
| model_name: str = "grok-4.1", | |
| system_prompt: Optional[str] = None, | |
| api_key: Optional[str] = None, | |
| max_tokens: int = 4096, | |
| temperature: float = 0.7 | |
| ) -> str: | |
| """ | |
| 🚀 Convenience function for generating AI responses | |
| Args: | |
| prompt: User prompt | |
| model_name: Model to use (default: grok-4.1 free) | |
| system_prompt: Optional system instructions | |
| api_key: Optional API key for premium models | |
| max_tokens: Maximum tokens | |
| temperature: Sampling temperature | |
| Returns: | |
| Generated text | |
| """ | |
| client = UnifiedAIClient(model_name=model_name, api_key=api_key) | |
| try: | |
| response = await client.generate( | |
| prompt=prompt, | |
| system_prompt=system_prompt, | |
| max_tokens=max_tokens, | |
| temperature=temperature | |
| ) | |
| return response | |
| finally: | |
| await client.close() | |