Code Browser¶

#!/usr/bin/env python3
"""Hello Gaia — your first Gaia application.

This script demonstrates the three fundamental Gaia operations:
1. List datasets
2. Ask a RAG question
3. Read the response and source documents
"""

import asyncio
import os
import sys

import httpx
from dotenv import load_dotenv

load_dotenv()

GAIA_API_KEY = os.getenv("GAIA_API_KEY")
GAIA_BASE_URL = os.getenv("GAIA_BASE_URL", "https://helios.cohesity.com/v2/mcm/gaia")
GAIA_VERIFY_SSL = os.getenv("GAIA_VERIFY_SSL", "true").lower() == "true"


def banner(text: str) -> None:
    width = max(len(text) + 4, 50)
    print(f"\n{'=' * width}")
    print(f"  {text}")
    print(f"{'=' * width}\n")


async def main() -> None:
    if not GAIA_API_KEY:
        print("Error: GAIA_API_KEY is not set.")
        print("Copy .env.example to .env and add your API key.")
        sys.exit(1)

    headers = {
        "apiKey": GAIA_API_KEY,
        "Content-Type": "application/json",
        "Accept": "application/json",
    }

    async with httpx.AsyncClient(
        base_url=GAIA_BASE_URL,
        headers=headers,
        timeout=60,
        verify=GAIA_VERIFY_SSL,
    ) as client:

        # ── Step 1: List Datasets ─────────────────────────────────────
        banner("Step 1: Listing Datasets")

        response = await client.get("/datasets")
        response.raise_for_status()
        data = response.json()

        datasets = data.get("datasets", data) if isinstance(data, dict) else data
        if not datasets:
            print("No datasets found. Create a dataset in the Cohesity UI first.")
            sys.exit(0)

        print(f"Found {len(datasets)} dataset(s):\n")
        for i, ds in enumerate(datasets):
            name = ds.get("name", "Unknown")
            status = ds.get("status", "unknown")
            print(f"  {i + 1}. {name} (status: {status})")

        # ── Step 2: Ask a Question ────────────────────────────────────
        first_dataset = datasets[0].get("name")
        banner(f"Step 2: Asking a Question (dataset: {first_dataset})")

        query = "Give me a brief summary of the data in this dataset."
        print(f"Query: {query}\n")

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

        # ── Step 3: Read the Response ─────────────────────────────────
        banner("Step 3: Response")

        answer = result.get("responseString", "No response received.")
        print(f"Answer:\n{answer}\n")

        documents = result.get("documents", [])
        if documents:
            print(f"\nSource Documents ({len(documents)}):")
            for i, doc in enumerate(documents[:5]):
                filename = doc.get("filename", "Unknown")
                score = doc.get("score", 0)
                print(f"  {i + 1}. {filename} (relevance: {score:.2f})")
        else:
            print("No source documents returned.")

        query_uid = result.get("queryUid")
        conversation_id = result.get("conversationId")
        if query_uid:
            print(f"\nQuery UID: {query_uid}")
        if conversation_id:
            print(f"Conversation ID: {conversation_id}")
            print("(Use this ID to continue the conversation with follow-up questions)")

    banner("Done! You've made your first Gaia API calls.")
    print("Next steps:")
    print("  - Try the chat-app example for a full web interface")
    print("  - Read the documentation at docs/")
    print("  - Build your own application!\n")


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

"""FastAPI dependencies for session-based auth."""

from typing import Optional

from fastapi import Header, HTTPException, status

from backend.services.session_service import get_api_key_for_session
from backend.settings import Settings, get_settings


def require_session(
    x_session_id: Optional[str] = Header(default=None, alias="X-SESSION-ID"),
) -> str:
    """Validate the session header and return the stored Gaia API key."""

    if not x_session_id:
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail="Missing X-SESSION-ID header.",
        )
    api_key = get_api_key_for_session(x_session_id)
    if not api_key:
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail="Invalid or expired session.",
        )
    return api_key

"""API routes for the Gaia Chat App."""

import logging
from typing import Optional

from fastapi import APIRouter, Depends, Header, HTTPException, status
from fastapi.responses import StreamingResponse

from backend.api.dependencies import require_session
from backend.models.api_models import (
    AskRequest,
    AskResponse,
    LoginRequest,
    LoginResponse,
)
from backend.services.session_service import create_session, delete_session
from backend.settings import get_settings
from gaia_sdk import GaiaClient
from gaia_sdk.exceptions import GaiaAuthError, GaiaError

logger = logging.getLogger(__name__)
router = APIRouter()


def _make_client(api_key: str) -> GaiaClient:
    s = get_settings()
    return GaiaClient(
        api_key=api_key,
        base_url=s.gaia_base_url,
        timeout=int(s.request_timeout_seconds),
        verify_ssl=s.gaia_verify_ssl,
    )


# ── Auth ─────────────────────────────────────────────────────────────

@router.post("/auth/login", response_model=LoginResponse, tags=["Auth"])
async def login(body: LoginRequest):
    """Validate the API key against Gaia and create a session."""
    async with _make_client(body.api_key) as gaia:
        try:
            await gaia.list_datasets()
        except GaiaAuthError as exc:
            raise HTTPException(
                status_code=status.HTTP_401_UNAUTHORIZED,
                detail=f"API key validation failed: {exc}",
            )

    settings = get_settings()
    session_id = create_session(body.api_key, ttl_minutes=settings.session_ttl_minutes)
    return LoginResponse(session_id=session_id)


@router.post("/auth/logout", status_code=status.HTTP_204_NO_CONTENT, tags=["Auth"])
async def logout(
    x_session_id: Optional[str] = Header(default=None, alias="X-SESSION-ID"),
):
    if x_session_id:
        delete_session(x_session_id)


# ── Datasets ─────────────────────────────────────────────────────────

@router.get("/datasets", tags=["Datasets"])
async def list_datasets(api_key: str = Depends(require_session)):
    async with _make_client(api_key) as gaia:
        datasets = await gaia.list_datasets()
    return {"datasets": [d.model_dump(by_alias=True) for d in datasets]}


# ── Ask ──────────────────────────────────────────────────────────────

@router.post("/ask", response_model=AskResponse, tags=["Ask"])
async def ask(body: AskRequest, api_key: str = Depends(require_session)):
    async with _make_client(api_key) as gaia:
        result = await gaia.ask(
            dataset_names=body.dataset_names,
            query=body.query,
            conversation_id=body.conversation_id,
        )
    return result.model_dump(by_alias=True)


@router.post("/ask/stream", tags=["Ask"])
async def ask_stream(body: AskRequest, api_key: str = Depends(require_session)):
    """Streaming RAG query — proxies SSE from Gaia."""

    async def event_generator():
        try:
            async with _make_client(api_key) as gaia:
                async for chunk in gaia.ask_stream_iter(
                    dataset_names=body.dataset_names,
                    query=body.query,
                    conversation_id=body.conversation_id,
                ):
                    if chunk.event and chunk.event != "message":
                        yield f"event: {chunk.event}\n"
                    yield f"data: {chunk.data}\n\n"
        except GaiaError:
            yield "data: [ERROR]\n\n"

    return StreamingResponse(
        event_generator(),
        media_type="text/event-stream",
        headers={
            "Cache-Control": "no-cache",
            "X-Accel-Buffering": "no",
        },
    )


# ── Conversations ────────────────────────────────────────────────────

@router.get("/conversations", tags=["Conversations"])
async def list_conversations(api_key: str = Depends(require_session)):
    async with _make_client(api_key) as gaia:
        return await gaia.list_conversations()


@router.get("/conversations/{conversation_id}/history", tags=["Conversations"])
async def get_chat_history(
    conversation_id: str, api_key: str = Depends(require_session)
):
    async with _make_client(api_key) as gaia:
        return await gaia.get_chat_history(conversation_id)

"""FastAPI application for the Gaia Chat App example."""

import contextlib

from fastapi import FastAPI, Request
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse

from backend.api.routes import router
from backend.services.session_service import cleanup_expired
from backend.settings import get_settings
from gaia_sdk.exceptions import GaiaAuthError, GaiaError, GaiaNotFoundError, GaiaRateLimitError


@contextlib.asynccontextmanager
async def lifespan(_app: FastAPI):
    cleanup_expired()
    yield


def create_app() -> FastAPI:
    settings = get_settings()

    app = FastAPI(
        title="Gaia Chat App",
        version="0.1.0",
        lifespan=lifespan,
    )

    app.add_middleware(
        CORSMiddleware,
        allow_origins=[settings.allow_cors_origin],
        allow_credentials=True,
        allow_methods=["*"],
        allow_headers=["*"],
    )

    @app.exception_handler(GaiaError)
    async def gaia_error_handler(request: Request, exc: GaiaError) -> JSONResponse:
        if isinstance(exc, GaiaAuthError):
            status_code = exc.status_code or 401
        elif isinstance(exc, GaiaNotFoundError):
            status_code = 404
        elif isinstance(exc, GaiaRateLimitError):
            status_code = 429
        else:
            status_code = 502
        return JSONResponse(status_code=status_code, content={"detail": str(exc)})

    app.include_router(router, prefix="/api/v1")

    return app


app = create_app()

"""Pydantic request/response models for the Chat App API."""

from typing import List, Optional

from pydantic import BaseModel, ConfigDict, Field


class LoginRequest(BaseModel):
    api_key: str = Field(..., alias="apiKey")
    model_config = ConfigDict(populate_by_name=True)


class LoginResponse(BaseModel):
    session_id: str = Field(..., alias="sessionId")
    model_config = ConfigDict(populate_by_name=True)


class AskRequest(BaseModel):
    query: str
    dataset_names: List[str] = Field(..., alias="datasetNames")
    conversation_id: Optional[str] = Field(default=None, alias="conversationId")
    model_config = ConfigDict(populate_by_name=True)


class AskResponse(BaseModel):
    """Thin pass-through of the Gaia /ask response."""

    response_string: str = Field(default="", alias="responseString")
    query_uid: str = Field(default="", alias="queryUid")
    conversation_id: Optional[str] = Field(default=None, alias="conversationId")
    documents: Optional[List[dict]] = None
    model_config = ConfigDict(populate_by_name=True)


class DatasetItem(BaseModel):
    name: str = ""
    dataset_id: str = Field(default="", alias="datasetId")
    description: Optional[str] = None
    model_config = ConfigDict(populate_by_name=True)


class DatasetListResponse(BaseModel):
    datasets: List[DatasetItem] = Field(default_factory=list)

"""In-memory session store keyed by session ID → API key."""

import uuid
from datetime import datetime, timedelta, timezone
from typing import Dict, Optional, Tuple

_sessions: Dict[str, Tuple[str, datetime]] = {}


def create_session(api_key: str, ttl_minutes: int = 60) -> str:
    session_id = uuid.uuid4().hex
    expires = datetime.now(timezone.utc) + timedelta(minutes=ttl_minutes)
    _sessions[session_id] = (api_key, expires)
    return session_id


def get_api_key_for_session(session_id: str) -> Optional[str]:
    entry = _sessions.get(session_id)
    if not entry:
        return None
    api_key, expires = entry
    if datetime.now(timezone.utc) > expires:
        _sessions.pop(session_id, None)
        return None
    return api_key


def delete_session(session_id: str) -> None:
    _sessions.pop(session_id, None)


def cleanup_expired() -> int:
    now = datetime.now(timezone.utc)
    expired = [sid for sid, (_, exp) in _sessions.items() if now > exp]
    for sid in expired:
        del _sessions[sid]
    return len(expired)

"""Configuration loaded from environment variables or .env file."""

from functools import lru_cache

from pydantic import Field
from pydantic_settings import BaseSettings, SettingsConfigDict


class Settings(BaseSettings):
    model_config = SettingsConfigDict(env_file=".env", case_sensitive=False)

    gaia_base_url: str = Field(
        default="https://helios.cohesity.com/v2/mcm/gaia",
        validation_alias="GAIA_BASE_URL",
    )
    gaia_verify_ssl: bool = Field(default=True, validation_alias="GAIA_VERIFY_SSL")
    request_timeout_seconds: float = Field(
        default=60.0, validation_alias="REQUEST_TIMEOUT_SECONDS"
    )
    allow_cors_origin: str = Field(
        default="http://localhost:5173", validation_alias="ALLOW_CORS_ORIGIN"
    )
    session_ttl_minutes: int = Field(default=60, validation_alias="SESSION_TTL_MINUTES")


@lru_cache()
def get_settings() -> Settings:
    return Settings()

"""Shared error helpers."""

from fastapi import HTTPException


def raise_http_error(status_code: int, detail: str) -> None:
    raise HTTPException(status_code=status_code, detail=detail)