meshai/meshai/backends/openai_backend.py

"""OpenAI-compatible LLM backend with rolling summary memory."""

import asyncio
import logging
from typing import Optional

from openai import AsyncOpenAI

from ..config import LLMConfig
from ..memory import RollingSummaryMemory
from .base import LLMBackend

logger = logging.getLogger(__name__)

_SUMMARIZE_PROMPT = """Summarize this conversation in 2-3 concise sentences. Focus on:
- Main topics discussed
- Important context or user preferences
- Key information to remember

Conversation:
{conversation}

Summary (2-3 sentences):"""


class OpenAIBackend(LLMBackend):
    """OpenAI-compatible backend (works with OpenAI, LiteLLM, local models)."""

    def __init__(
        self,
        config: LLMConfig,
        api_key: str,
        window_size: int = 4,
        summarize_threshold: int = 8,
    ):
        """Initialize OpenAI backend.

        Args:
            config: LLM configuration
            api_key: API key to use
            window_size: Recent message pairs to keep in full
            summarize_threshold: Messages before re-summarizing
        """
        self.config = config
        self._client = AsyncOpenAI(
            api_key=api_key,
            base_url=config.base_url,
        )

        # Initialize rolling summary memory with OpenAI summarize function
        self._memory = RollingSummaryMemory(
            summarize_fn=self._summarize_messages,
            window_size=window_size,
            summarize_threshold=summarize_threshold,
        )

    async def _summarize_messages(self, messages: list[dict]) -> str:
        """Summarize messages using OpenAI API."""
        if not messages:
            return "No previous conversation."

        conversation = "\n".join(
            [f"{msg['role'].upper()}: {msg['content']}" for msg in messages]
        )
        prompt = _SUMMARIZE_PROMPT.format(conversation=conversation)

        try:
            response = await self._client.chat.completions.create(
                model=self.config.model,
                messages=[{"role": "user", "content": prompt}],
                max_tokens=150,
                temperature=0.3,
            )
            content = response.choices[0].message.content
            return content.strip() if content else f"Previous conversation: {len(messages)} messages."
        except Exception as e:
            logger.warning(f"Failed to generate summary: {e}")
            return f"Previous conversation: {len(messages)} messages about various topics."

    async def generate(
        self,
        messages: list[dict],
        system_prompt: str,
        max_tokens: int = 300,
        user_id: Optional[str] = None,
    ) -> str:
        """Generate a response using OpenAI-compatible API.

        Args:
            messages: Conversation history
            system_prompt: System prompt
            max_tokens: Maximum tokens to generate
            user_id: User identifier (enables memory optimization)

        Returns:
            Generated response
        """
        # Use memory manager to optimize context if user_id provided
        if user_id and len(messages) > self._memory._window_size * 2:
            summary, recent_messages = await self._memory.get_context_messages(
                user_id=user_id,
                full_history=messages,
            )

            if summary:
                # Long conversation: system + summary + recent
                enhanced_system = f"{system_prompt}\n\nPrevious conversation summary: {summary}"
                full_messages = [{"role": "system", "content": enhanced_system}]
                full_messages.extend(recent_messages)

                logger.debug(
                    f"Using summary + {len(recent_messages)} recent messages "
                    f"(total history: {len(messages)})"
                )
            else:
                # Short conversation: system + all messages
                full_messages = [{"role": "system", "content": system_prompt}]
                full_messages.extend(messages)
        else:
            # No user_id or short conversation - use full history
            full_messages = [{"role": "system", "content": system_prompt}]
            full_messages.extend(messages)

        try:
            # Build request kwargs
            request_kwargs = {
                "model": self.config.model,
                "messages": full_messages,
                "max_tokens": max_tokens,
                "temperature": 0.7,
            }

            # Enable web search if configured (Open WebUI feature)
            # Uses features.web_search parameter
            if getattr(self.config, 'web_search', False):
                request_kwargs["extra_body"] = {"features": {"web_search": True}}

            response = await asyncio.wait_for(
                self._client.chat.completions.create(**request_kwargs),
                timeout=self.config.timeout,
            )

            content = response.choices[0].message.content
            return content.strip() if content else ""

        except asyncio.TimeoutError:
            logger.error(f"OpenAI API timed out after {self.config.timeout}s")
            raise
        except Exception as e:
            logger.error(f"OpenAI API error: {e}")
            raise

    def get_memory(self) -> RollingSummaryMemory:
        """Get the memory manager instance."""
        return self._memory

    async def close(self) -> None:
        """Close the client."""
        await self._client.close()