sdk/rag_client.py

"""
RAG Template Python SDK Client.

Cliente Python para interagir com a API REST do RAG Template.
"""

from typing import List, Dict, Any, Optional
import requests
from pathlib import Path


class RAGClient:
    """Cliente para API REST do RAG Template."""

    def __init__(
        self,
        base_url: str = "http://localhost:8000",
        api_key: Optional[str] = None
    ):
        """
        Inicializa cliente.

        Args:
            base_url: URL base da API
            api_key: API key para autenticacao
        """
        self.base_url = base_url.rstrip("/")
        self.api_key = api_key
        self.session = requests.Session()

        if api_key:
            self.session.headers.update({"X-API-Key": api_key})

    def _make_request(
        self,
        method: str,
        endpoint: str,
        **kwargs
    ) -> Dict[str, Any]:
        """
        Faz requisicao HTTP.

        Args:
            method: Metodo HTTP (GET, POST, DELETE, etc)
            endpoint: Endpoint da API
            **kwargs: Parametros adicionais para requests

        Returns:
            Response JSON

        Raises:
            requests.HTTPError: Se requisicao falhar
        """
        url = f"{self.base_url}{endpoint}"
        response = self.session.request(method, url, **kwargs)
        response.raise_for_status()
        return response.json()

    def health_check(self) -> Dict[str, Any]:
        """
        Verifica health do sistema.

        Returns:
            Status do sistema
        """
        return self._make_request("GET", "/api/v1/health")

    def ingest_text(
        self,
        text: str,
        title: str,
        chunk_size: int = 1000,
        chunk_overlap: int = 200,
        strategy: str = "recursive",
        metadata: Optional[Dict[str, Any]] = None
    ) -> Dict[str, Any]:
        """
        Ingere texto no sistema.

        Args:
            text: Texto do documento
            title: Titulo do documento
            chunk_size: Tamanho dos chunks
            chunk_overlap: Overlap entre chunks
            strategy: Estrategia de chunking (fixed, sentence, semantic, recursive)
            metadata: Metadados opcionais

        Returns:
            Resposta da ingestao com document_id
        """
        data = {
            "text": text,
            "title": title,
            "chunk_size": chunk_size,
            "chunk_overlap": chunk_overlap,
            "strategy": strategy
        }

        if metadata:
            data["metadata"] = metadata

        return self._make_request("POST", "/api/v1/ingest", json=data)

    def upload_file(
        self,
        file_path: str,
        chunk_size: int = 1000,
        chunk_overlap: int = 200,
        strategy: str = "recursive"
    ) -> Dict[str, Any]:
        """
        Upload e ingere arquivo.

        Args:
            file_path: Caminho do arquivo (PDF ou TXT)
            chunk_size: Tamanho dos chunks
            chunk_overlap: Overlap entre chunks
            strategy: Estrategia de chunking

        Returns:
            Resposta da ingestao com document_id
        """
        file_path = Path(file_path)

        if not file_path.exists():
            raise FileNotFoundError(f"File not found: {file_path}")

        with open(file_path, "rb") as f:
            files = {"file": (file_path.name, f, "application/octet-stream")}
            params = {
                "chunk_size": chunk_size,
                "chunk_overlap": chunk_overlap,
                "strategy": strategy
            }

            return self._make_request(
                "POST",
                "/api/v1/upload",
                files=files,
                params=params
            )

    def query(
        self,
        query: str,
        top_k: int = 5,
        temperature: float = 0.3,
        max_tokens: int = 512,
        model: Optional[str] = None,
        filters: Optional[Dict[str, Any]] = None
    ) -> Dict[str, Any]:
        """
        Executa query RAG.

        Args:
            query: Pergunta do usuario
            top_k: Numero de contextos a recuperar
            temperature: Temperature para geracao
            max_tokens: Tokens maximos na resposta
            model: Modelo LLM (opcional)
            filters: Filtros de metadata (opcional)

        Returns:
            Resposta com query, response, contexts e metadata
        """
        data = {
            "query": query,
            "top_k": top_k,
            "temperature": temperature,
            "max_tokens": max_tokens
        }

        if model:
            data["model"] = model
        if filters:
            data["filters"] = filters

        return self._make_request("POST", "/api/v1/query", json=data)

    def list_documents(
        self,
        limit: int = 100,
        offset: int = 0,
        session_id: Optional[str] = None
    ) -> List[Dict[str, Any]]:
        """
        Lista documentos no sistema.

        Args:
            limit: Numero maximo de documentos
            offset: Offset para paginacao
            session_id: Filtrar por session_id (opcional)

        Returns:
            Lista de documentos
        """
        params = {"limit": limit, "offset": offset}
        if session_id:
            params["session_id"] = session_id

        return self._make_request("GET", "/api/v1/documents", params=params)

    def delete_document(self, document_id: int) -> Dict[str, Any]:
        """
        Deleta documento.

        Args:
            document_id: ID do documento

        Returns:
            Mensagem de confirmacao
        """
        return self._make_request("DELETE", f"/api/v1/documents/{document_id}")

    def get_stats(self) -> Dict[str, Any]:
        """
        Obtem estatisticas do sistema.

        Returns:
            Estatisticas do banco e metadados
        """
        return self._make_request("GET", "/api/v1/stats")

    def __repr__(self) -> str:
        return f"RAGClient(base_url='{self.base_url}')"


# Funcoes de conveniencia

def create_client(base_url: str = "http://localhost:8000", api_key: Optional[str] = None) -> RAGClient:
    """
    Cria cliente RAG.

    Args:
        base_url: URL base da API
        api_key: API key para autenticacao

    Returns:
        Cliente RAG configurado
    """
    return RAGClient(base_url=base_url, api_key=api_key)