motia-iii/services/xai_service.py

"""xAI Files & Collections Service"""
import os
import aiohttp
import logging
from typing import Optional, List

logger = logging.getLogger(__name__)

XAI_FILES_URL = "https://api.x.ai"
XAI_MANAGEMENT_URL = "https://management-api.x.ai"


class XAIService:
    """
    Client für xAI Files API und Collections Management API.

    Benötigte Umgebungsvariablen:
    - XAI_API_KEY         – regulärer API-Key für File-Uploads (api.x.ai)
    - XAI_MANAGEMENT_KEY  – Management-API-Key für Collection-Operationen (management-api.x.ai)
    """

    def __init__(self, ctx=None):
        self.api_key = os.getenv('XAI_API_KEY', '')
        self.management_key = os.getenv('XAI_MANAGEMENT_KEY', '')
        self.ctx = ctx
        self._session: Optional[aiohttp.ClientSession] = None

        if not self.api_key:
            raise ValueError("XAI_API_KEY not configured in environment")
        if not self.management_key:
            raise ValueError("XAI_MANAGEMENT_KEY not configured in environment")

    def _log(self, msg: str, level: str = 'info') -> None:
        if self.ctx:
            getattr(self.ctx.logger, level, self.ctx.logger.info)(msg)
        else:
            getattr(logger, level, logger.info)(msg)

    async def _get_session(self) -> aiohttp.ClientSession:
        if self._session is None or self._session.closed:
            self._session = aiohttp.ClientSession(
                timeout=aiohttp.ClientTimeout(total=120)
            )
        return self._session

    async def close(self) -> None:
        if self._session and not self._session.closed:
            await self._session.close()

    async def upload_file(
        self,
        file_content: bytes,
        filename: str,
        mime_type: str = 'application/octet-stream'
    ) -> str:
        """
        Lädt eine Datei zur xAI Files API hoch (multipart/form-data).

        POST https://api.x.ai/v1/files

        Returns:
            xAI file_id (str)

        Raises:
            RuntimeError: bei HTTP-Fehler oder fehlendem file_id in der Antwort
        """
        self._log(f"📤 Uploading {len(file_content)} bytes to xAI: {filename}")

        session = await self._get_session()
        url = f"{XAI_FILES_URL}/v1/files"
        headers = {"Authorization": f"Bearer {self.api_key}"}

        form = aiohttp.FormData()
        form.add_field('file', file_content, filename=filename, content_type=mime_type)

        async with session.post(url, data=form, headers=headers) as response:
            try:
                data = await response.json()
            except Exception:
                raw = await response.text()
                data = {"_raw": raw}

            if response.status not in (200, 201):
                raise RuntimeError(
                    f"xAI file upload failed ({response.status}): {data}"
                )

            file_id = data.get('id') or data.get('file_id')
            if not file_id:
                raise RuntimeError(
                    f"No file_id in xAI upload response: {data}"
                )

        self._log(f"✅ xAI file uploaded: {file_id}")
        return file_id

    async def add_to_collection(self, collection_id: str, file_id: str) -> None:
        """
        Fügt eine Datei einer xAI-Collection hinzu.

        POST https://management-api.x.ai/v1/collections/{collection_id}/documents/{file_id}

        Raises:
            RuntimeError: bei HTTP-Fehler
        """
        self._log(f"📚 Adding file {file_id} to collection {collection_id}")

        session = await self._get_session()
        url = f"{XAI_MANAGEMENT_URL}/v1/collections/{collection_id}/documents/{file_id}"
        headers = {
            "Authorization": f"Bearer {self.management_key}",
            "Content-Type": "application/json",
        }

        async with session.post(url, headers=headers) as response:
            if response.status not in (200, 201):
                raw = await response.text()
                raise RuntimeError(
                    f"Failed to add file to collection {collection_id} ({response.status}): {raw}"
                )

        self._log(f"✅ File {file_id} added to collection {collection_id}")

    async def remove_from_collection(self, collection_id: str, file_id: str) -> None:
        """
        Entfernt eine Datei aus einer xAI-Collection.
        Die Datei selbst wird NICHT gelöscht – sie kann in anderen Collections sein.

        DELETE https://management-api.x.ai/v1/collections/{collection_id}/documents/{file_id}

        Raises:
            RuntimeError: bei HTTP-Fehler
        """
        self._log(f"🗑️  Removing file {file_id} from collection {collection_id}")

        session = await self._get_session()
        url = f"{XAI_MANAGEMENT_URL}/v1/collections/{collection_id}/documents/{file_id}"
        headers = {"Authorization": f"Bearer {self.management_key}"}

        async with session.delete(url, headers=headers) as response:
            if response.status not in (200, 204):
                raw = await response.text()
                raise RuntimeError(
                    f"Failed to remove file from collection {collection_id} ({response.status}): {raw}"
                )

        self._log(f"✅ File {file_id} removed from collection {collection_id}")

    async def add_to_collections(self, collection_ids: List[str], file_id: str) -> List[str]:
        """
        Fügt eine Datei zu mehreren Collections hinzu.

        Returns:
            Liste der erfolgreich hinzugefügten Collection-IDs
        """
        added = []
        for collection_id in collection_ids:
            try:
                await self.add_to_collection(collection_id, file_id)
                added.append(collection_id)
            except Exception as e:
                self._log(
                    f"⚠️  Fehler beim Hinzufügen zu Collection {collection_id}: {e}",
                    level='warn'
                )
        return added

    async def remove_from_collections(self, collection_ids: List[str], file_id: str) -> None:
        """Entfernt eine Datei aus mehreren Collections (ignoriert Fehler pro Collection)."""
        for collection_id in collection_ids:
            try:
                await self.remove_from_collection(collection_id, file_id)
            except Exception as e:
                self._log(
                    f"⚠️  Fehler beim Entfernen aus Collection {collection_id}: {e}",
                    level='warn'
                )