mirror of
https://github.com/katanemo/plano.git
synced 2026-04-28 02:23:56 +02:00
109 lines
3.2 KiB
Python
109 lines
3.2 KiB
Python
from openai import OpenAI
|
|
from fastapi import FastAPI, HTTPException
|
|
from pydantic import BaseModel, Field
|
|
from typing import List, Optional
|
|
from common import create_gradio_app
|
|
import gradio as gr
|
|
import os
|
|
|
|
|
|
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,
|
|
)
|
|
|
|
gr.mount_gradio_app(
|
|
app, create_gradio_app(demo_description, client), path="/agent/chat"
|
|
)
|
|
|
|
if __name__ == "__main__":
|
|
app.run(debug=True)
|