Spaces:
Paused
Paused
| """ | |
| OpenRouter endpoint implementation. | |
| This module provides integration with OpenRouter's API for LLM inference. | |
| """ | |
| import requests | |
| from .ai_endpoint import BaseAIEndpoint, AIEndpointRequestError | |
| DEFAULT_MODEL = "openai/gpt-4o-mini" | |
| class OpenRouterEndpoint(BaseAIEndpoint): | |
| """OpenRouter endpoint for cloud-based LLM inference.""" | |
| # Models that support structured output | |
| STRUCTURED_OUTPUT_MODELS = { | |
| "openai/gpt-4o", | |
| "openai/gpt-4o-mini", | |
| "openai/gpt-4-turbo", | |
| "anthropic/claude-3-5-sonnet", | |
| "deepseek/deepseek-r1:free" | |
| } | |
| def _initialize_client(self) -> None: | |
| """Initialize the OpenAI client.""" | |
| api_key = self.ai_config.get("api_key", "") | |
| if not api_key: | |
| raise AIEndpointRequestError("OpenRouter API key is required") | |
| def _get_default_model(self) -> str: | |
| """Get the default OpenAI model.""" | |
| return DEFAULT_MODEL | |
| def supports_structured_output(self) -> bool: | |
| """Check if the current model supports structured output.""" | |
| model = self.model or DEFAULT_MODEL | |
| return any(model.startswith(prefix.split('/')[0]) or model in self.STRUCTURED_OUTPUT_MODELS | |
| for prefix in self.STRUCTURED_OUTPUT_MODELS) | |
| def query(self, prompt: str, output_format: dict) -> str: | |
| """ | |
| Send a query to OpenRouter and return the response. | |
| Args: | |
| prompt: The prompt to send to the model (as messages list or string) | |
| output_format: Pydantic model for structured output | |
| Returns: | |
| The model's response as a string | |
| Raises: | |
| AIEndpointRequestError: If the request fails | |
| """ | |
| try: | |
| url = "https://openrouter.ai/api/v1/chat/completions" | |
| headers = { | |
| "Authorization": f"Bearer {self.ai_config.get('api_key')}", | |
| "Content-Type": "application/json" | |
| } | |
| messages = [{"role": "user", "content": prompt}] | |
| schema = output_format.model_json_schema() | |
| body = { | |
| "model": self.model or DEFAULT_MODEL, | |
| "max_tokens": self.max_tokens, | |
| "temperature": self.temperature, | |
| } | |
| # Handle structured output based on model support | |
| if self.supports_structured_output(): | |
| body["messages"] = messages | |
| body["response_format"] = { | |
| "type": "json_schema", | |
| "json_schema": { | |
| "name": "response", | |
| "schema": schema, | |
| "strict": True | |
| } | |
| } | |
| else: | |
| # If model does not support structured format, just send raw prompt | |
| body["messages"] = messages | |
| r = requests.post(url, headers=headers, json=body) | |
| if r.status_code >= 400: | |
| raise AIEndpointRequestError(f"OpenRouter error {r.status_code}: {r.text}") | |
| data = r.json() | |
| if self.supports_structured_output(): | |
| return self.parseStringToJson(data["choices"][0]["message"]["content"]) | |
| return data["choices"][0]["message"]["content"] | |
| except Exception as e: | |
| raise AIEndpointRequestError(f"OpenRouter request failed: {e}") |