Your First API Call¶

#!/usr/bin/env python3
"""Your first Gaia API call — list datasets and ask a question."""

import asyncio
import os
import sys

from dotenv import load_dotenv

load_dotenv()


async def main():
    # Check for the SDK, fall back to raw httpx
    try:
        from gaia_sdk import GaiaClient
        use_sdk = True
    except ImportError:
        import httpx
        use_sdk = False
        print("gaia_sdk not found — using httpx directly\n")

    if use_sdk:
        await run_with_sdk()
    else:
        await run_with_httpx()


async def run_with_sdk():
    from gaia_sdk import GaiaClient

    async with GaiaClient.from_env() as gaia:
        # List datasets
        print("=== Available Datasets ===\n")
        datasets = await gaia.list_datasets()

        if not datasets:
            print("No datasets found. Create one in the Cohesity UI first.")
            sys.exit(1)

        for ds in datasets:
            print(f"  • {ds.name}  (status: {ds.status})")

        # Use the first ready dataset
        ready = [ds for ds in datasets if ds.status == "ready"]
        if not ready:
            print("\nNo datasets are ready for querying yet.")
            sys.exit(1)

        target = ready[0].name
        query = "Summarize the key information in this dataset."

        print(f"\n=== Asking: '{query}' ===")
        print(f"    Dataset: {target}\n")

        response = await gaia.ask(
            dataset_names=[target],
            query=query,
        )

        print(f"Answer:\n{response.response_string}\n")
        print(f"Query UID:       {response.query_uid}")
        print(f"Conversation ID: {response.conversation_id}")

        if response.documents:
            print(f"\nSources ({len(response.documents)} documents):")
            for doc in response.documents[:5]:
                score = f"{doc.score:.2f}" if doc.score else "N/A"
                print(f"  [{score}] {doc.filename or doc.filepath}")


async def run_with_httpx():
    import httpx

    base_url = os.environ.get(
        "GAIA_BASE_URL", "https://helios.cohesity.com/v2/mcm/gaia"
    )
    api_key = os.environ.get("GAIA_API_KEY")

    if not api_key:
        print("Error: GAIA_API_KEY environment variable is not set.")
        sys.exit(1)

    async with httpx.AsyncClient(
        base_url=base_url,
        headers={"apiKey": api_key, "Content-Type": "application/json"},
        timeout=60,
    ) as client:
        # List datasets
        print("=== Available Datasets ===\n")
        resp = await client.get("/datasets")
        resp.raise_for_status()
        data = resp.json()
        dataset_list = data.get("datasets", [])

        if not dataset_list:
            print("No datasets found.")
            sys.exit(1)

        for ds in dataset_list:
            print(f"  • {ds['name']}  (status: {ds.get('status', 'unknown')})")

        ready = [ds for ds in dataset_list if ds.get("status") == "ready"]
        if not ready:
            print("\nNo datasets are ready for querying yet.")
            sys.exit(1)

        target = ready[0]["name"]
        query = "Summarize the key information in this dataset."

        print(f"\n=== Asking: '{query}' ===")
        print(f"    Dataset: {target}\n")

        resp = await client.post("/ask", json={
            "datasetNames": [target],
            "queryString": query,
        })
        resp.raise_for_status()
        result = resp.json()

        print(f"Answer:\n{result.get('responseString', 'No response')}\n")
        print(f"Query UID:       {result.get('queryUid', 'N/A')}")
        print(f"Conversation ID: {result.get('conversationId', 'N/A')}")

        docs = result.get("documents", [])
        if docs:
            print(f"\nSources ({len(docs)} documents):")
            for doc in docs[:5]:
                score = f"{doc.get('score', 0):.2f}"
                name = doc.get("filename") or doc.get("filepath", "unknown")
                print(f"  [{score}] {name}")


if __name__ == "__main__":
    asyncio.run(main())