From e3628b57c506fa172d698a051370d7b3baeec73a Mon Sep 17 00:00:00 2001 From: kamalgovindgj Date: Tue, 7 Jan 2025 23:09:40 -0500 Subject: [PATCH] Add AISPM service and configuration updates --- javelin_cli/_internal/commands.py | 120 ++++++++++++++++++++ javelin_cli/cli.py | 97 +++++++++++++++++ javelin_sdk/client.py | 59 +++++++--- javelin_sdk/models.py | 130 ++++++++++++++++++++++ javelin_sdk/services/aispm_service.py | 151 ++++++++++++++++++++++++++ 5 files changed, 543 insertions(+), 14 deletions(-) create mode 100644 javelin_sdk/services/aispm_service.py diff --git a/javelin_cli/_internal/commands.py b/javelin_cli/_internal/commands.py index 049082f..01cd768 100644 --- a/javelin_cli/_internal/commands.py +++ b/javelin_cli/_internal/commands.py @@ -1,5 +1,7 @@ import json import os +import asyncio + from pathlib import Path from pydantic import ValidationError @@ -28,6 +30,7 @@ Secrets, Template, Templates, + Customer, AWSConfig, AzureConfig, UsageResponse, AlertResponse,Request,HttpMethod, ) @@ -87,6 +90,123 @@ def get_javelin_client(): return JavelinClient(config) +#aispm commands +def create_customer(args): + try: + client = get_javelin_client() + customer = Customer( + name=args.name, + description=args.description, + metrics_interval=args.metrics_interval, + security_interval=args.security_interval + ) + response = client._send_request_sync(Request( + method=HttpMethod.POST, + route="v1/admin/aispm/customer", + data=customer.dict() + )) + print(f"Customer '{args.name}' created successfully.") + except Exception as e: + print(f"Error creating customer: {e}") + +def get_customer(args): + try: + client = get_javelin_client() + route = "v1/admin/aispm/customer" + print(f"Making request to: {client.config.base_url}{route}") + print(f"Headers: {client._headers}") # Access _headers directly + + response = client._send_request_sync(Request( + method=HttpMethod.GET, + route=route + )) + print(f"Status: {response.status_code}") + print(f"Response: {response.content.decode('utf-8')}") + except Exception as e: + print(f"Error getting customer: {e}") + +def configure_aws(args): + try: + client = get_javelin_client() + config = json.loads(args.config) + configs = [AWSConfig(**config)] + result = client.aispm.configure_aws(configs) + print(f"AWS configuration created successfully.") + except Exception as e: + print(f"Error configuring AWS: {e}") + +def get_aws_config(args): + try: + client = get_javelin_client() + request = Request( + method=HttpMethod.GET, + route="v1/admin/aispm/config/aws" + ) + response = client._send_request_sync(request) + print(json.dumps(response.json(), indent=2)) + except Exception as e: + print(f"Error getting AWS config: {e}") + +def delete_aws_config(args): + try: + client = get_javelin_client() + request = Request( + method=HttpMethod.DELETE, + route=f"v1/admin/aispm/config/aws/{args.name}" + ) + response = client._send_request_sync(request) + print(f"AWS configuration '{args.name}' deleted successfully.") + except Exception as e: + print(f"Error deleting AWS config: {e}") + +def get_azure_config(args): + try: + client = get_javelin_client() + request = Request( + method=HttpMethod.GET, + route="v1/admin/aispm/config/azure" + ) + response = client._send_request_sync(request) + print(json.dumps(response.json(), indent=2)) + except Exception as e: + print(f"Error getting Azure config: {e}") + +def configure_azure(args): + try: + client = get_javelin_client() + config = json.loads(args.config) + configs = [AzureConfig(**config)] + result = client.aispm.configure_azure(configs) + print(f"Azure configuration created successfully.") + except Exception as e: + print(f"Error configuring Azure: {e}") + +def get_usage(args): + try: + client = get_javelin_client() + usage = client.aispm.get_usage( + provider=args.provider, + cloud_account=args.account, + model=args.model, + region=args.region + ) + print(json.dumps(usage.dict(), indent=2)) + except Exception as e: + print(f"Error getting usage: {e}") + +def get_alerts(args): + try: + client = get_javelin_client() + alerts = client.aispm.get_alerts( + provider=args.provider, + cloud_account=args.account, + model=args.model, + region=args.region + ) + print(json.dumps(alerts.dict(), indent=2)) + except Exception as e: + print(f"Error getting alerts: {e}") + def create_gateway(args): try: client = get_javelin_client() diff --git a/javelin_cli/cli.py b/javelin_cli/cli.py index 4481f76..dff023e 100644 --- a/javelin_cli/cli.py +++ b/javelin_cli/cli.py @@ -38,9 +38,30 @@ update_route, update_secret, update_template, + create_customer, + get_customer, + configure_aws, + configure_azure, + get_usage, + get_alerts, + get_aws_config, + get_azure_config, + delete_aws_config ) + +#def check_permissions(): +# """Check if user has permissions""" +# home_dir = Path.home() +# cache_file = home_dir / ".javelin" / "cache.json" + +# if not cache_file.exists(): +# print("❌ Not authenticated. Please run 'javelin auth' first.") +# sys.exit(1) + +# return True # Skip role check + def check_permissions(): """Check if user has superadmin permissions""" home_dir = Path.home() @@ -89,6 +110,80 @@ def main(): auth_parser = subparsers.add_parser("auth", help="Authenticate with Javelin.") auth_parser.add_argument("--force", action="store_true", help="Force re-authentication, overriding existing credentials") auth_parser.set_defaults(func=authenticate) + #aispm CRUD + # AISPM commands + aispm_parser = subparsers.add_parser("aispm", help="Manage AISPM functionality") + aispm_subparsers = aispm_parser.add_subparsers() + + # Customer commands + customer_parser = aispm_subparsers.add_parser("customer", help="Manage customers") + customer_subparsers = customer_parser.add_subparsers() + + customer_create = customer_subparsers.add_parser("create", help="Create customer") + customer_create.add_argument("--name", required=True, help="Customer name") + customer_create.add_argument("--description", help="Customer description") + customer_create.add_argument("--metrics-interval", default="5m", help="Metrics interval") + customer_create.add_argument("--security-interval", default="1m", help="Security interval") + customer_create.set_defaults(func=create_customer) + + customer_get = customer_subparsers.add_parser("get", help="Get customer details") + customer_get.set_defaults(func=get_customer) + + + # Cloud config commands + config_parser = aispm_subparsers.add_parser("config", help="Manage cloud configurations") + config_subparsers = config_parser.add_subparsers() + + aws_parser = config_subparsers.add_parser("aws", help="Configure AWS") + + azure_parser = config_subparsers.add_parser("azure", help="Configure Azure") + + #azure_parser.add_argument("--config", type=str, required=True, help="Azure config JSON") + #azure_parser.set_defaults(func=configure_azure) + + + aws_subparsers = aws_parser.add_subparsers() + + # GET AWS Config + aws_get_parser = aws_subparsers.add_parser("get", help="Get AWS configuration") + aws_get_parser.set_defaults(func=get_aws_config) + + # Existing AWS Config (for creating) + aws_config_parser = aws_subparsers.add_parser("create", help="Configure AWS") + aws_config_parser.add_argument("--config", type=str, required=True, help="AWS config JSON") + aws_config_parser.set_defaults(func=configure_aws) + + aws_delete_parser = aws_subparsers.add_parser("delete", help="Delete AWS configuration") + aws_delete_parser.add_argument("--name", type=str, required=True, help="Name of AWS configuration to delete") + aws_delete_parser.set_defaults(func=delete_aws_config) + + azure_subparsers = azure_parser.add_subparsers(dest='azure_command') + +# Get Azure Config + azure_get_parser = azure_subparsers.add_parser("get", help="Get Azure configuration") + azure_get_parser.set_defaults(func=get_azure_config) + + # Create Azure Config + azure_create_parser = azure_subparsers.add_parser("create", help="Configure Azure") + azure_create_parser.add_argument("--config", type=str, required=True, help="Azure config JSON") + azure_create_parser.set_defaults(func=configure_azure) + + + # Usage metrics + usage_parser = aispm_subparsers.add_parser("usage", help="Get usage metrics") + usage_parser.add_argument("--provider", help="Cloud provider") + usage_parser.add_argument("--account", help="Cloud account name") + usage_parser.add_argument("--model", help="Model ID") + usage_parser.add_argument("--region", help="Region") + usage_parser.set_defaults(func=get_usage) + + # Alerts + alerts_parser = aispm_subparsers.add_parser("alerts", help="Get alerts") + alerts_parser.add_argument("--provider", help="Cloud provider") + alerts_parser.add_argument("--account", help="Cloud account name") + alerts_parser.add_argument("--model", help="Model ID") + alerts_parser.add_argument("--region", help="Region") + alerts_parser.set_defaults(func=get_alerts) # Gateway CRUD gateway_parser = subparsers.add_parser( "gateway", @@ -373,6 +468,8 @@ def main(): ) template_delete.set_defaults(func=delete_template) + + args = parser.parse_args() if hasattr(args, "func"): diff --git a/javelin_sdk/client.py b/javelin_sdk/client.py index c5a6a1f..be019be 100644 --- a/javelin_sdk/client.py +++ b/javelin_sdk/client.py @@ -15,6 +15,8 @@ from javelin_sdk.services.secret_service import SecretService from javelin_sdk.services.template_service import TemplateService from javelin_sdk.services.trace_service import TraceService +from javelin_sdk.services.aispm_service import AISPMService + API_BASEURL = "https://api-dev.javelin.live" API_BASE_PATH = "/v1" @@ -45,6 +47,9 @@ def __init__(self, config: JavelinConfig) -> None: self.chat = Chat(self) self.completions = Completions(self) + self.aispm = AISPMService(self) + + @property def client(self): if self._client is None: @@ -84,24 +89,41 @@ def close(self): self._client.close() def _prepare_request(self, request: Request) -> tuple: - url = self._construct_url( - gateway_name=request.gateway, - provider_name=request.provider, - route_name=request.route, - secret_name=request.secret, - template_name=request.template, - trace=request.trace, - query=request.is_query, - archive=request.archive, - query_params=request.query_params, - is_transformation_rules=request.is_transformation_rules, - is_reload=request.is_reload, - ) + if request.route.startswith("v1/admin/aispm"): + url = f"{self.config.base_url.rstrip('/')}/{request.route}" + if request.query_params: + query_string = "&".join(f"{k}={v}" for k, v in request.query_params.items()) + url += f"?{query_string}" + + + else: + url = self._construct_url( + gateway_name=request.gateway, + provider_name=request.provider, + route_name=request.route, + secret_name=request.secret, + template_name=request.template, + trace=request.trace, + query=request.is_query, + archive=request.archive, + query_params=request.query_params, + is_transformation_rules=request.is_transformation_rules, + is_reload=request.is_reload, + ) + headers = {**self._headers, **(request.headers or {})} return url, headers def _send_request_sync(self, request: Request) -> httpx.Response: - return self._core_send_request(self.client, request) + url, headers = self._prepare_request(request) + print(f"Making request to: {url}") + print(f"With headers: {headers}") + response = self._core_send_request(self.client, request) + print(f"Response status: {response.status_code}") + print(f"Response body: {response.text}") + return response + + async def _send_request_async(self, request: Request) -> httpx.Response: return await self._core_send_request(self.aclient, request) @@ -199,6 +221,15 @@ def _construct_url( return url + def _construct_aispm_url(self, request: Request) -> str: + url = request.route + + if request.query_params: + query_string = "&".join(f"{k}={v}" for k, v in request.query_params.items()) + url += f"?{query_string}" + + return url + # Gateway methods create_gateway = lambda self, gateway: self.gateway_service.create_gateway(gateway) acreate_gateway = lambda self, gateway: self.gateway_service.acreate_gateway( diff --git a/javelin_sdk/models.py b/javelin_sdk/models.py index 80c8d7c..6a71f45 100644 --- a/javelin_sdk/models.py +++ b/javelin_sdk/models.py @@ -1,6 +1,12 @@ from enum import Enum, auto from typing import Any, Dict, List, Optional +from datetime import datetime + +from typing import Dict, List, Optional + + + from pydantic import BaseModel, Field, field_validator from javelin_sdk.exceptions import UnauthorizedError @@ -554,3 +560,127 @@ class EndpointType(str, Enum): INVOKE_STREAM = "invoke_stream" CONVERSE_STREAM = "converse_stream" ALL = "all" + + +#aispm models + +class TimeRange(BaseModel): + start_time: str # Change from datetime to str + end_time: str # Change from datetime to str + +class BaseResponse(BaseModel): + message: Optional[str] = None + +# Customer Models +class Customer(BaseModel): + name: str + description: Optional[str] + metrics_interval: str = "5m" + security_interval: str = "1m" + initial_scan: str = "24h" + +class CustomerResponse(Customer): + status: str + created_at: datetime + modified_at: datetime + +# Cloud Config Models +class BaseCloudConfig(BaseModel): + cloud_account_name: str + team: str + +class AWSConfig(BaseCloudConfig): + role_arn: str + region: Optional[str] = None # Make region optional + +class AzureConfig(BaseCloudConfig): + subscription_id: str + tenant_id: str + client_id: str + client_secret: str + location: str + +class GCPConfig(BaseCloudConfig): + project_id: str + service_account_key: str + +class CloudConfigResponse(BaseModel): + name: Optional[str] = Field(None, alias='cloud_account_name') + provider: str + status: str + created_at: datetime + modified_at: datetime + +# Usage Models +class ModelMetrics(BaseModel): + latency_avg_ms: float + cost_per_request: float + tokens_per_request: float + attempt_count: int + failure_count: int + success_count: int + success_rate_pct: float + cost_total: float + request_count: int + token_count: int + +class CloudAccountUsage(BaseModel): + region_count: int + regions: List[str] + model_count: int + models: List[str] + model_metrics: ModelMetrics + +class UsageResponse(BaseModel): + cloud_provider: Dict[str, Any] # Change to allow any structure + time_range: Optional[TimeRange] = None + +# Alert Models +class AlertSeverity(str, Enum): + CRITICAL = "CRITICAL" + HIGH = "HIGH" + MEDIUM = "MEDIUM" + LOW = "LOW" + +class AlertState(str, Enum): + ALARM = "ALARM" + OK = "OK" + INSUFFICIENT_DATA = "INSUFFICIENT_DATA" + +class AlertScope(str, Enum): + GLOBAL = "GLOBAL" + MODEL = "MODEL" + REGION = "REGION" + +class AlertMetrics(BaseModel): + total_alerts: int + active_alerts: int + resolved_alerts: int + critical_alerts: int + high_alerts: int + medium_alerts: int + low_alerts: int + +class Alert(BaseModel): + title: str + state: AlertState + state_reason: str + severity: AlertSeverity + scope: AlertScope + region: Optional[str] + model_id: Optional[str] + detected_at: datetime + +class CloudProviderAlerts(BaseModel): + cloud_account_count: int + cloud_accounts: List[str] + region_count: int + regions: List[str] + model_count: int + models: List[str] + alert_metrics: AlertMetrics + alerts: List[Alert] + +class AlertResponse(BaseModel): + cloud_provider: Dict[str, CloudProviderAlerts] + time_range: TimeRange diff --git a/javelin_sdk/services/aispm_service.py b/javelin_sdk/services/aispm_service.py new file mode 100644 index 0000000..ae09937 --- /dev/null +++ b/javelin_sdk/services/aispm_service.py @@ -0,0 +1,151 @@ +from typing import Dict, List, Optional, Union +from httpx import Response +import json + +from javelin_sdk.models import ( + Customer, CustomerResponse, + AWSConfig, AzureConfig, GCPConfig, CloudConfigResponse, + UsageResponse, AlertResponse, TimeRange, + HttpMethod, Request +) + +class AISPMService: + def __init__(self, client): + self.client = client + + def _handle_response(self, response: Response) -> None: + if response.status_code >= 400: + error = response.json().get("error", "Unknown error") + raise Exception(f"API error: {error}") + + # Customer Methods + def create_customer(self, customer: Customer) -> CustomerResponse: + request = Request( + method=HttpMethod.POST, + route="v1/admin/aispm/customer", + data=customer.dict() + ) + print(f"Sending request: {request.method} {request.route}") + response = self.client._send_request_sync(request) + print(f"Raw response: {response.text}") + self._handle_response(response) + return CustomerResponse(**response.json()) + + def get_customer(self) -> CustomerResponse: + request = Request( + method=HttpMethod.GET, + route="v1/admin/aispm/customer" + ) + response = self.client._send_request_sync(request) + self._handle_response(response) + return CustomerResponse(**response.json()) + + def update_customer(self, customer: Customer) -> CustomerResponse: + request = Request( + method=HttpMethod.PUT, + route="v1/admin/aispm/customer", + data=customer.dict() + ) + response = self.client._send_request_sync(request) + self._handle_response(response) + return CustomerResponse(**response.json()) + + # Cloud Config Methods + def configure_aws(self, configs: List[AWSConfig]) -> List[CloudConfigResponse]: + request = Request( + method=HttpMethod.POST, + route="v1/admin/aispm/config/aws", + data=[config.dict() for config in configs] + ) + response = self.client._send_request_sync(request) + self._handle_response(response) + return [CloudConfigResponse(**config) for config in response.json()] + + def configure_azure(self, configs: List[AzureConfig]) -> List[CloudConfigResponse]: + request = Request( + method=HttpMethod.POST, + route="v1/admin/aispm/config/azure", + data=[config.dict() for config in configs] + ) + response = self.client._send_request_sync(request) + self._handle_response(response) + return [CloudConfigResponse(**config) for config in response.json()] + + def configure_gcp(self, configs: List[GCPConfig]) -> List[CloudConfigResponse]: + request = Request( + method=HttpMethod.POST, + route="v1/admin/aispm/config/gcp", + data=[config.dict() for config in configs] + ) + response = self.client._send_request_sync(request) + self._handle_response(response) + return [CloudConfigResponse(**config) for config in response.json()] + + # Usage Methods + def get_usage(self, + provider: Optional[str] = None, + cloud_account: Optional[str] = None, + model: Optional[str] = None, + region: Optional[str] = None) -> UsageResponse: + + route = "v1/admin/aispm/usage" + if provider: + route += f"/{provider}" + if cloud_account: + route += f"/{cloud_account}" + + params = {} + if model: + params["model"] = model + if region: + params["region"] = region + + request = Request( + method=HttpMethod.GET, + route=route, + query_params=params + ) + response = self.client._send_request_sync(request) + self._handle_response(response) + return UsageResponse(**response.json()) + + # Alert Methods + def get_alerts(self, + provider: Optional[str] = None, + cloud_account: Optional[str] = None, + model: Optional[str] = None, + region: Optional[str] = None) -> AlertResponse: + + route = "v1/admin/aispm/alerts" + if provider: + route += f"/{provider}" + if cloud_account: + route += f"/{cloud_account}" + + params = {} + if model: + params["model"] = model + if region: + params["region"] = region + + request = Request( + method=HttpMethod.GET, + route=route, + query_params=params + ) + response = self.client._send_request_sync(request) + self._handle_response(response) + return AlertResponse(**response.json()) + + # Helpers + def _validate_provider(self, provider: str) -> None: + valid_providers = ["aws", "azure", "gcp", "openai"] + if provider.lower() not in valid_providers: + raise ValueError(f"Invalid provider. Must be one of: {valid_providers}") + + def _construct_error(self, response: Response) -> Dict: + try: + error = response.json() + return error.get("error", str(response.content)) + except json.JSONDecodeError: + return str(response.content) \ No newline at end of file