stores/llm/providers/OpenRouterProvider.py

from stores.llm.LLMInterface import LLMInterface
import logging
import requests
import re
import os


class OpenRouterProvider(LLMInterface):
    def __init__(self, url: str = None, model: str = None,
                 default_input_max_characters: int = 1000,
                 default_generation_max_output_tokens: int = 1000,
                 default_generation_temperature: float = 0.1, api_key: str = None):
        self.url = url or "https://openrouter.ai/api/v1"
        self.api_key = api_key or os.getenv("OPENROUTER_API_KEY")
        self.model = model
        self.generation_model_id = None

        self.embedding_model = None
        self.embedding_model_id = None
        self.embedding_size = None

        self.default_input_max_characters = default_input_max_characters
        self.default_generation_max_output_tokens = default_generation_max_output_tokens
        self.default_generation_temperature = default_generation_temperature
        self.logger = logging.getLogger(__name__)

    def set_generation_model(self, model_id: str):
        if model_id:
            self.model = model_id

    def set_embedding_model(self, model_id: str, embedding_size: int):
        if model_id:
            self.embedding_model = model_id
            self.embedding_size = embedding_size
            self.embedding_model_id = model_id

    def process_text(self, text: str):
        if not text:
            return ""
        return str(text).strip()

    def generate_text(self, prompt: str, chat_history: list = None,
                      max_output_tokens: int = None, temperature: float = None):
        try:
            chat_history = chat_history or []
            clean_prompt = self.process_text(prompt)

            messages = []
            for entry in chat_history:
                messages.append({
                    "role": entry.get("role", "user"),
                    "content": entry.get("content", "")
                })
            messages.append({"role": "user", "content": clean_prompt})

            payload = {
                "model": self.model,
                "messages": messages,
                "max_tokens": int(max_output_tokens or self.default_generation_max_output_tokens),
                "temperature": float(temperature or self.default_generation_temperature),
            }

            url = self.url.rstrip("/") + "/chat/completions"
            headers = {
                "Authorization": f"Bearer {self.api_key}",
                "Content-Type": "application/json",
                # Recommended by OpenRouter for usage tracking
                "HTTP-Referer": os.getenv("OPENROUTER_SITE_URL", "http://localhost"),
                "X-Title": os.getenv("OPENROUTER_APP_NAME", "LLMApp"),
            }

            resp = requests.post(url, json=payload, headers=headers, timeout=6000)
            if resp.status_code != 200:
                self.logger.error("OpenRouter generate failed: %s %s", resp.status_code, resp.text)
                return None

            data = resp.json()

            try:
                generated_text = data["choices"][0]["message"]["content"].strip()
            except (KeyError, IndexError, TypeError):
                self.logger.error("Unexpected OpenRouter response structure: %s", data)
                return None

            if not generated_text:
                return None

            usage = data.get("usage", {})
            return {
                "model": data.get("model"),
                "response": generated_text,
                "tokens_generated": usage.get("completion_tokens"),
                "total_duration_ms": None,
                "prompt_eval_tokens": usage.get("prompt_tokens"),
            }

        except Exception as e:
            self.logger.exception("Error in OpenRouterProvider.generate_text: %s", e)
            return None

    def embed_text(self, text: str, document_type: str = None):
        """OpenRouter does not support embeddings natively — returns None."""
        self.logger.warning("OpenRouterProvider does not support embeddings.")
        return None

    def construct_prompt(self, prompt: str, role: str):
        return {
            "role": role,
            "content": self.process_text(prompt)
        }

    def embed_text_batch(self, texts: list[str], batch_size: int = 32):
        """OpenRouter does not support embeddings natively — returns None."""
        self.logger.warning("OpenRouterProvider does not support embeddings.")
        return None

    def clean_content(self, text: str) -> str:
        text = re.sub(r'\[.*?\]\(.*?\)', '', text)
        text = re.sub(r'\[[^\]]*\]', '', text)
        text = re.sub(r'\n+', '\n', text).strip()
        return text

    def web_search(self, query: str):
        """
        OpenRouter supports online models (e.g. perplexity/sonar-online) that have
        built-in web search. Route the query through one of those models if available,
        otherwise fall back to a not-supported notice.
        """
        try:
            online_model = os.getenv("OPENROUTER_SEARCH_MODEL", "perplexity/sonar-online")

            payload = {
                "model": online_model,
                "messages": [{"role": "user", "content": query}],
                "max_tokens": int(self.default_generation_max_output_tokens),
                "temperature": float(self.default_generation_temperature),
            }

            url = self.url.rstrip("/") + "/chat/completions"
            headers = {
                "Authorization": f"Bearer {self.api_key}",
                "Content-Type": "application/json",
                "HTTP-Referer": os.getenv("OPENROUTER_SITE_URL", "http://localhost"),
                "X-Title": os.getenv("OPENROUTER_APP_NAME", "LLMApp"),
            }

            resp = requests.post(url, json=payload, headers=headers, timeout=6000)
            if not resp or resp.status_code != 200:
                return {
                    "text": "No relevant external results found.",
                    "references": []
                }

            data = resp.json()

            combined_text = []
            references = set()

            try:
                text_content = data["choices"][0]["message"]["content"]
                combined_text.append(self.clean_content(text_content))
            except (KeyError, IndexError, TypeError):
                pass

            # Extract any URLs from the response text
            for found_url in re.findall(r"https?://[^\s)]+", "\n".join(combined_text)):
                references.add(found_url)

            return {
                "text": "\n\n".join(combined_text[:3]),
                "references": list(references)
            }

        except Exception as e:
            self.logger.error("OpenRouter web search failed: %s", e)
            return {
                "text": f"OpenRouter search error: {str(e)}",
                "references": []
            }