plano/demos/network_agent/main.py

import os
from typing import List, Optional

from fastapi import FastAPI, HTTPException
from openai import OpenAI
from pydantic import BaseModel, Field

app = FastAPI()
DEMO_DESCRIPTION = """This demo illustrates how **Arch** can be used to perform function calling
 with network-related tasks. In this demo, you act as a **network assistant** that provides factual
 information, without offering advice on manufacturers or purchasing decisions."""


# Define the request model
class DeviceSummaryRequest(BaseModel):
    device_ids: List[int]
    time_range: Optional[int] = Field(
        default=7, description="Time range in days, defaults to 7"
    )


# Define the response model
class DeviceStatistics(BaseModel):
    device_id: int
    time_range: str
    data: str


class DeviceSummaryResponse(BaseModel):
    statistics: List[DeviceStatistics]

    # Request model for device reboot


class DeviceRebootRequest(BaseModel):
    device_ids: List[int]


# Response model for the device reboot
class CoverageResponse(BaseModel):
    status: str
    summary: dict


@app.post("/agent/device_reboot", response_model=CoverageResponse)
def reboot_network_device(request_data: DeviceRebootRequest):
    """
    Endpoint to reboot network devices based on device IDs and an optional time range.
    """

    # Access data from the Pydantic model
    device_ids = request_data.device_ids

    # Validate 'device_ids'
    # (This is already validated by Pydantic, but additional logic can be added if needed)
    if not device_ids:
        raise HTTPException(
            status_code=400, detail="'device_ids' parameter is required"
        )

    # Simulate reboot operation and return the response
    statistics = []
    for device_id in device_ids:
        # Placeholder for actual data retrieval or device reboot logic
        stats = {"data": f"Device {device_id} has been successfully rebooted."}
        statistics.append(stats)

    # Return the response with a summary
    return CoverageResponse(status="success", summary={"device_ids": device_ids})


# Post method for device summary
@app.post("/agent/device_summary", response_model=DeviceSummaryResponse)
def get_device_summary(request: DeviceSummaryRequest):
    """
    Endpoint to retrieve device statistics based on device IDs and an optional time range.
    """

    # Extract 'device_ids' and 'time_range' from the request
    device_ids = request.device_ids
    time_range = request.time_range

    # Simulate retrieving statistics for the given device IDs and time range
    statistics = []
    minutes = 1
    for device_id in device_ids:
        stats = {
            "device_id": device_id,
            "time_range": f"Last {time_range} days",
            "data": f"""Device {device_id} over the last {time_range} days experienced {minutes}
             minutes of downtime.""",
        }
        minutes += 1
        statistics.append(DeviceStatistics(**stats))

    return DeviceSummaryResponse(statistics=statistics)


CHAT_COMPLETION_ENDPOINT = os.getenv("CHAT_COMPLETION_ENDPOINT")
client = OpenAI(
    api_key="--",
    base_url=CHAT_COMPLETION_ENDPOINT,
)