search_api.py

# OpenAI Search API Client

# A Python module for making authenticated requests to /openai/search with ztoken authentication.

# ## Usage

# ```bash
# # Basic search
# python search_api.py "How do I retire applications?"

# # Custom namespace
# python search_api.py "Deploy to production" --namespace deployment

# # Verbose output (shows full response structure)
# python search_api.py "Security guidelines" --verbose
# ```

# ## Setup

# 1. Install dependencies: `requests`
# 2. Set API host: `export SEARCH_API_HOST="your-api-host.com"`
# 3. Ensure `ztoken` is available in PATH

# ## Response Format

# ```json
# {
#   "output": "The AI-generated response",
#   "source_nodes": [],
#   "metadata": {},
#   "sources": []
# }
# ```

# ## Requirements

# - `requests` library
# - `ztoken` command for authentication



import logging
import subprocess
import json
import argparse
import sys
from typing import Dict, Any

import requests

import config

logger = logging.getLogger(__name__)

session = requests.Session()


def format_search_response(raw_response: Dict[str, Any]) -> Dict[str, Any]:
    """
    Format the raw API response into the expected structure.

    Args:
        raw_response (dict): The raw response from the API

    Returns:
        dict: Formatted response with consistent structure
    """
    # If the response already has the expected format, return as-is
    if all(key in raw_response for key in ["output", "source_nodes", "metadata", "sources"]):
        return raw_response

    # Otherwise, format the response into the expected structure
    return {
        "output": raw_response.get("output", raw_response.get("response", str(raw_response))),
        "source_nodes": raw_response.get("source_nodes", raw_response.get("sources", [])),
        "metadata": raw_response.get("metadata", {}),
        "sources": raw_response.get("sources", raw_response.get("references", []))
    }


def get_ztoken() -> str:
    """
    Get the current ztoken for authentication.

    Returns:
        str: The ztoken from running the `ztoken` command

    Raises:
        RuntimeError: If unable to get ztoken
    """
    try:
        result = subprocess.run(['ztoken'], capture_output=True, text=True, check=True)
        token = result.stdout.strip()
        if not token:
            raise RuntimeError("ztoken command returned empty token")
        return token
    except subprocess.CalledProcessError as e:
        raise RuntimeError(f"Failed to get ztoken: {e}") from e
    except FileNotFoundError:
        raise RuntimeError("ztoken command not found. Make sure it's installed and in PATH") from None


def search(text: str, namespace: str = "build") -> Dict[str, Any]:
    """
    Make a search request to the OpenAI search API.

    Args:
        text (str): The search query text
        namespace (str): The namespace to search in (default: "build")

    Returns:
        dict: Formatted response with structure:
            {
                "output": "The AI-generated response to your query",
                "source_nodes": [],  # Source documents used
                "metadata": {},      # Additional metadata
                "sources": []        # Formatted source references
            }

    Raises:
        requests.RequestException: If the HTTP request fails
        RuntimeError: If unable to get authentication token
    """
    url = f"https://{config.SEARCH_API_HOST}/openai/search"

    # Get authentication token
    try:
        token = get_ztoken()
    except RuntimeError as e:
        logger.error("Failed to get authentication token: %s", e)
        raise

    # Prepare request data
    payload = {
        "text": text,
        "namespace": namespace
    }

    headers = {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {token}"
    }

    logger.info("Making search request to %s with payload: %s", url, payload)

    try:
        response = session.post(url, json=payload, headers=headers)
        response.raise_for_status()

        response_data = response.json()

        # Log the full response structure for inspection
        logger.info("Search API response status: %d", response.status_code)
        logger.info("Search API response headers: %s", dict(response.headers))
        logger.info("Search API raw response data: %s",
                   json.dumps(response_data, indent=2, default=str))

        # Format the response into the expected structure
        formatted_response = format_search_response(response_data)

        logger.info("Formatted search response: %s",
                   json.dumps(formatted_response, indent=2, default=str))

        return formatted_response

    except requests.exceptions.RequestException as e:
        logger.error("HTTP request failed: %s", e)
        if hasattr(e, 'response') and e.response is not None:
            logger.error("Response status: %d", e.response.status_code)
            logger.error("Response text: %s", e.response.text)
        raise
    except json.JSONDecodeError as e:
        logger.error("Failed to decode JSON response: %s", e)
        logger.error("Response text: %s", response.text)
        raise


if __name__ == "__main__":
    # Command line interface for testing the search API
    parser = argparse.ArgumentParser(
        description="Search the OpenAI API",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  python -m support_bot.search_api "How do I retire applications?"
  python -m support_bot.search_api "How do I deploy to production?" --namespace deployment
  python -m support_bot.search_api "Database migration steps" --namespace db --verbose
        """
    )

    parser.add_argument(
        "query",
        help="The search query text"
    )

    parser.add_argument(
        "--namespace", "-n",
        default="build",
        help="The namespace to search in (default: build)"
    )

    parser.add_argument(
        "--verbose", "-v",
        action="store_true",
        help="Enable verbose logging to see full request/response details"
    )

    args = parser.parse_args()

    # Set up logging level based on verbose flag
    if args.verbose:
        logging.basicConfig(level=logging.INFO, format='%(levelname)s - %(message)s')
    else:
        # Only show warnings and errors by default
        logging.basicConfig(level=logging.WARNING, format='%(levelname)s - %(message)s')

    try:
        print(f"Searching for: '{args.query}' in namespace '{args.namespace}'")
        print("-" * 60)

        result = search(args.query, args.namespace)

        print("Search completed successfully!")
        print()

        print("Response:")
        print(result.get("output", "No output available"))
        print()

        if args.verbose:
            print("🔍 Full Response Structure:")
            print(json.dumps(result, indent=2, default=str))

    except KeyboardInterrupt:
        print("\n Search cancelled by user")
        sys.exit(1)
    except Exception as e:
        print(f" Error: {e}")
        sys.exit(1)