kalavai-client 0.6.17__py3-none-any.whl → 0.6.18__py3-none-any.whl

kalavai_client/__init__.py

@@ -1,2 +1,2 @@
 
-__version__ = "0.6.17"
+__version__ = "0.6.18"

kalavai_client/assets/apps.yaml

@@ -113,7 +113,7 @@ releases:
   - name: kuberay
     namespace: kuberay
     chart: kuberay/kuberay-operator
-    installed: true
+    installed: {{deploy_kuberay}}
     version: "1.2.2"
   - name: kuberay-apiserver
     namespace: kuberay
@@ -154,7 +154,7 @@ releases:
     - name: replicas
       value: 1
     - name: image_tag
-      value: "v2025.06.15"
+      value: "v2025.07.31"
     - name: deployment.in_cluster
       value: "True"
     - name: deployment.kalavai_username_key
@@ -184,7 +184,7 @@ releases:
   - name: nvidia-gpu-operator
     namespace: kalavai
     chart: kalavai/gpu
-    installed: true    
+    installed: false    
   - name: hami-vgpu
     namespace: kalavai
     chart: kalavai/hami

kalavai_client/assets/apps_values.yaml

@@ -8,13 +8,17 @@
   description: "Deploy Opencost cost monitoring system"
 
 - name: deploy_prometheus
-  default: "True"
+  default: "False"
   description: "Deploy Prometheus system monitoring system"
 
 - name: deploy_langfuse
   default: "False"
   description: "Deploy Langfuse LLM tracing system"
 
+- name: deploy_kuberay
+  default: "False"
+  description: "Deploy Langfuse LLM tracing system"
+
 ######
 
 ### VARIABLES ###

kalavai_client/bridge_api.py

@@ -2,10 +2,13 @@
 Core kalavai service.
 Used as a bridge between the kalavai-client app and the reflex frontend
 """
-from fastapi import FastAPI, HTTPException, Depends
+from fastapi import FastAPI, HTTPException, Depends, Query, Body
+from typing import Optional, List
+from fastapi_mcp import FastApiMCP
 from starlette.requests import Request
 import uvicorn
 
+from kalavai_client.core import Job
 from kalavai_client.bridge_models import (
     CreatePoolRequest,
     InvitesRequest,
@@ -15,8 +18,7 @@ from kalavai_client.bridge_models import (
     DeleteJobRequest,
     JobDetailsRequest,
     NodesActionRequest,
-    NodeLabelsRequest,
-    GetNodeLabelsRequest
+    NodeLabelsRequest
 )
 from kalavai_client.core import (
     create_pool,
@@ -52,7 +54,10 @@ from kalavai_client.core import (
     get_node_labels,
     TokenType
 )
-from kalavai_client.utils import load_user_id
+from kalavai_client.utils import (
+    load_user_id,
+    extract_auth_token
+)
 
 app = FastAPI(
     title="Kalavai Bridge API",
@@ -73,14 +78,16 @@ async def verify_api_key(request: Request):
     user_id = load_user_id()
     if user_id is None:
         return None
-    api_key = request.headers.get("X-API-KEY")
+    api_key = extract_auth_token(headers=request.headers)
     if api_key != user_id:
         raise HTTPException(status_code=401, detail="Request requires API Key")
     return api_key
 
 @app.post("/create_pool", 
-    summary="Create a new pool",
-    description="Creates a new pool with the specified configuration",
+    operation_id="create_pool",
+    summary="Create a new Kalavai compute pool",
+    tags=["pool_management"],
+    description="Creates a new distributed compute pool that allows multiple nodes to join and share GPU resources. The pool acts as a Kubernetes cluster where users can deploy and manage machine learning jobs across multiple devices.",
     response_description="Result of pool creation")
 def pool_create(request: CreatePoolRequest, api_key: str = Depends(verify_api_key)):
     """
@@ -109,8 +116,10 @@ def pool_create(request: CreatePoolRequest, api_key: str = Depends(verify_api_ke
     return result
 
 @app.post("/join_pool",
-    summary="Join an existing pool",
-    description="Join a pool using a token",
+    operation_id="join_pool",
+    summary="Join an existing Kalavai pool as a compute node",
+    description="Joins a running Kalavai pool by providing a valid join token. This endpoint registers the current machine as a compute node in the pool, making its GPU resources available for job scheduling. The node will receive workloads based on the pool's scheduling policy.",
+    tags=["pool_management"],
     response_description="Result of joining the pool")
 def pool_join(request: JoinPoolRequest, api_key: str = Depends(verify_api_key)):
     """
@@ -130,8 +139,10 @@ def pool_join(request: JoinPoolRequest, api_key: str = Depends(verify_api_key)):
     return result
 
 @app.post("/attach_to_pool",
-    summary="Attach to an existing pool",
-    description="Attach to a pool using a token",
+    operation_id="attach_to_pool",
+    summary="Attach to a pool for management purposes",
+    description="Attaches to an existing Kalavai pool for administrative and monitoring purposes without contributing compute resources. This is typically used by frontend applications or management tools that need to interact with the pool but don't provide GPU resources.",
+    tags=["pool_management"],
     response_description="Result of attaching to the pool")
 def pool_attach(request: JoinPoolRequest, api_key: str = Depends(verify_api_key)):
     """
@@ -149,8 +160,10 @@ def pool_attach(request: JoinPoolRequest, api_key: str = Depends(verify_api_key)
     return result
 
 @app.post("/stop_pool",
-    summary="Stop a pool",
-    description="Stop the current pool",
+    operation_id="stop_pool",
+    summary="Stop and clean up the current Kalavai pool",
+    description="Gracefully shuts down the current Kalavai pool, terminating all running jobs and optionally removing all compute nodes from the cluster. This operation is irreversible and will disconnect all nodes from the pool.",
+    tags=["pool_management"],
     response_description="Result of stopping the pool")
 def pool_stop(request: StopPoolRequest, api_key: str = Depends(verify_api_key)):
     """
@@ -164,8 +177,10 @@ def pool_stop(request: StopPoolRequest, api_key: str = Depends(verify_api_key)):
     return result
 
 @app.post("/delete_nodes",
-    summary="Delete nodes",
-    description="Delete specified nodes from the pool",
+    operation_id="delete_nodes",
+    summary="Remove specific nodes from the pool",
+    description="Removes specified compute nodes from the Kalavai pool. This operation will terminate any jobs running on the target nodes and clean up their resources. Use with caution as it may interrupt running workloads.",
+    tags=["pool_management"],
     response_description="Result of node deletion")
 def device_delete(request: NodesActionRequest, api_key: str = Depends(verify_api_key)):
     """
@@ -179,8 +194,10 @@ def device_delete(request: NodesActionRequest, api_key: str = Depends(verify_api
     return result
 
 @app.post("/cordon_nodes",
-    summary="Cordon nodes",
-    description="Mark nodes as unschedulable",
+    operation_id="cordon_nodes",
+    summary="Mark nodes as unschedulable",
+    description="Marks specified nodes as unschedulable, preventing new jobs from being assigned to them while allowing existing jobs to complete. This is useful for maintenance operations or when you want to gradually remove nodes from the pool.",
+    tags=["pool_management"],
     response_description="Result of cordoning nodes")
 def device_cordon(request: NodesActionRequest, api_key: str = Depends(verify_api_key)):
     """
@@ -194,8 +211,10 @@ def device_cordon(request: NodesActionRequest, api_key: str = Depends(verify_api
     return result
 
 @app.post("/uncordon_nodes",
-    summary="Uncordon nodes",
-    description="Mark nodes as schedulable",
+    operation_id="uncordon_nodes",
+    summary="Mark nodes as schedulable again",
+    description="Re-enables job scheduling on previously cordoned nodes, allowing them to receive new workloads. This reverses the effect of the cordon operation.",
+    tags=["pool_management"],
     response_description="Result of uncordoning nodes")
 def device_uncordon(request: NodesActionRequest, api_key: str = Depends(verify_api_key)):
     """
@@ -209,8 +228,10 @@ def device_uncordon(request: NodesActionRequest, api_key: str = Depends(verify_a
     return result
 
 @app.get("/get_pool_token",
-    summary="Get pool token",
-    description="Get a token for the pool",
+    operation_id="get_pool_token",
+    summary="Generate a token for pool access",
+    description="Generates a secure token that can be used to join or attach to the current Kalavai pool. Different token types provide different levels of access - join tokens allow nodes to contribute resources, while attach tokens allow management access.",
+    tags=["auth"],
     response_description="Pool token")
 def get_token(mode: int, api_key: str = Depends(verify_api_key)):
     """
@@ -221,16 +242,20 @@ def get_token(mode: int, api_key: str = Depends(verify_api_key)):
     return get_pool_token(mode=TokenType(mode))
 
 @app.get("/fetch_devices",
-    summary="Fetch devices",
-    description="Get list of available devices",
+    operation_id="fetch_devices",
+    summary="Get list of all compute devices in the pool",
+    description="Retrieves information about all compute devices (nodes) currently connected to the Kalavai pool, including their status, available resources, and current workload distribution.",
+    tags=["info"],
     response_description="List of devices")
 def get_devices(api_key: str = Depends(verify_api_key)):
     """Get list of available devices"""
     return fetch_devices()
 
 @app.post("/send_pool_invites",
-    summary="Send pool invites",
-    description="Send invites to join the pool",
+    operation_id="send_pool_invites",
+    summary="Send invitations to join the pool",
+    description="Sends invitations to potential users or nodes to join the current Kalavai pool. Invitees will receive tokens that allow them to connect to the pool and contribute their resources.",
+    tags=["avoid"],
     response_description="Result of sending invites")
 def send_pool_invites(request: InvitesRequest, api_key: str = Depends(verify_api_key)):
     """
@@ -241,24 +266,30 @@ def send_pool_invites(request: InvitesRequest, api_key: str = Depends(verify_api
     return send_invites(invitees=request.invitees)
 
 @app.get("/fetch_resources",
-    summary="Fetch resources",
-    description="Get available resources",
+    operation_id="fetch_resources",
+    summary="Get resource utilization for specific nodes",
+    description="Retrieves detailed resource information (CPU, memory, GPU usage) for the pool; optionally for a list of specified nodes in the pool (as {'nodes': node_list}). This helps monitor resource utilization and plan workload distribution.",
+    tags=["info"],
     response_description="Resource information")
-def resources(request: NodesActionRequest, api_key: str = Depends(verify_api_key)):
+def resources(request: Optional[NodesActionRequest]=NodesActionRequest(), api_key: str = Depends(verify_api_key)):
     """Get available resources"""
     return fetch_resources(node_names=request.nodes)
 
 @app.get("/fetch_job_names",
-    summary="Fetch job names",
-    description="Get list of job names",
+    operation_id="fetch_job_names",
+    summary="Get list of all jobs (model deployments) in the pool",
+    description="Retrieves the names of all jobs and models currently deployed or scheduled in the Kalavai pool. This provides an overview of all workloads in the system.",
+    tags=["info"],
     response_description="List of job names")
 def job_names(api_key: str = Depends(verify_api_key)):
     """Get list of job names"""
     return fetch_job_names()
 
 @app.get("/fetch_gpus",
-    summary="Fetch GPUs",
-    description="Get list of available GPUs",
+    operation_id="fetch_gpus",
+    summary="Get GPU information across the pool",
+    description="Retrieves detailed information about all GPUs in the Kalavai pool, including their availability status, current utilization, and which jobs are using them. Can filter to show only available GPUs.",
+    tags=["info"],
     response_description="List of GPUs")
 def gpus(available: bool = False, api_key: str = Depends(verify_api_key)):
     """
@@ -269,26 +300,26 @@ def gpus(available: bool = False, api_key: str = Depends(verify_api_key)):
     return fetch_gpus(available=available)
 
 @app.post("/fetch_job_details",
-    summary="Fetch job details",
-    description="Get details for specified jobs",
+    operation_id="fetch_job_details",
+    summary="Get detailed information about specific job and model deployments",
+    description="Given a list of jobs (as {'jobs': [{'name': job_name}]}'), retrieves comprehensive information about specified jobs or models including their status, resource usage, runtime, and configuration. Useful for monitoring and debugging job execution.",
+    tags=["info"],
     response_description="Job details")
 def job_details(request: JobDetailsRequest, api_key: str = Depends(verify_api_key)):
-    """
-    Get job details with the following parameters:
-    
-    - **jobs**: List of jobs to get details for
-    """
+    """Get job details"""
     return fetch_job_details(jobs=request.jobs)
 
 @app.get("/fetch_job_logs",
-    summary="Fetch job logs",
-    description="Get logs for a specific job",
+    operation_id="fetch_job_logs",
+    summary="Get execution logs for a specific job",
+    description="Retrieves the execution logs for a specified job, providing real-time or historical output from the job's containers. Useful for debugging, monitoring progress, and understanding job behavior.",
+    tags=["info", "avoid"],
     response_description="Job logs")
 def job_logs(
     job_name: str,
-    force_namespace: str = None,
-    pod_name: str = None,
-    tail: int = 100,
+    force_namespace: str = Query(None),
+    pod_name: str = Query(None),
+    tail: int = Query(100),
     api_key: str = Depends(verify_api_key)
 ):
     """
@@ -307,28 +338,68 @@ def job_logs(
     )
 
 @app.get("/fetch_job_templates",
-    summary="Fetch job templates",
-    description="Get available job templates",
+    operation_id="fetch_job_templates",
+    summary="Get available job templates",
+    description="Retrieves a list of all available job templates that can be used to deploy workloads. Templates provide predefined configurations for frameworks.",
+    tags=["info"],
     response_description="List of job templates")
 def job_templates(api_key: str = Depends(verify_api_key)):
-    """Get available job templates"""
     return fetch_job_templates()
 
+@app.get("/fetch_model_templates",
+    operation_id="fetch_model_templates",
+    summary="Get available model engines templates",
+    description="Retrieves a list of all available model engine templates that can be used to deploy models. Templates provide predefined configurations for model engine frameworks.",
+    tags=["info"],
+    response_description="List of model engine templates")
+def model_templates(api_key: str = Depends(verify_api_key)):
+    return fetch_job_templates(type="model")
+
 @app.get("/fetch_job_defaults",
-    summary="Fetch job defaults",
-    description="Get default values for a job template",
-    response_description="Job metadata values")
-def job_templates(name: str, api_key: str = Depends(verify_api_key)):
-    """
-    Get job defaults with the following parameters:
-    
-    - **name**: Name of the job template
-    """
-    return fetch_job_defaults(name=name)
+    operation_id="fetch_job_defaults",
+    summary="Get default values for a job or model engine template deployment",
+    description="Retrieves the default values for a specific job or model engine template deployment. This helps users understand what parameters are required and what their default values are before deploying a job.",
+    tags=["info"],
+    response_description="Job and model engine default values")
+def job_defaults(name: str, api_key: str = Depends(verify_api_key)):
+    result = fetch_job_defaults(name=name)
+    return result["defaults"]
+
+@app.get("/fetch_job_metadata",
+    operation_id="fetch_job_metadata",
+    summary="Get metadata with information about a given job or model engine template deployment",
+    description="Retrieves the metadata associated with a specific job or model engine template deployment. This helps users understand what the template can be used for.",
+    tags=["info"],
+    response_description="Job and model engine metadata values")
+def job_metadata(name: str, api_key: str = Depends(verify_api_key)):
+    result = fetch_job_defaults(name=name)
+    return result["metadata"]
+
+@app.get("/fetch_job_rules",
+    operation_id="fetch_job_rules",
+    summary="Get the rules associated with the use of a given job or model engine template",
+    description="Retrieves the rules associated with a specific job or model engine template deployment. This helps users and AI agents determine if a given model engine template is adequate for the task.",
+    tags=["info"],
+    response_description="Job and model engine rules")
+def job_rules(name: str, api_key: str = Depends(verify_api_key)):
+    result = job_metadata(name=name)
+    return result["template_rules"]
+
+@app.get("/fetch_job_values_rules",
+    operation_id="fetch_job_values_rules",
+    summary="Get information on how to provide values to the parameters of a specific job or model engine template",
+    description="Retrieves information necessary to fill up the values required to deploy a specific job or model engine template. This helps users and AI agents generate the values dictionary for a job or model engine template deployment.",
+    tags=["info"],
+    response_description="Job and model engine info for values")
+def job_values_rules(name: str, api_key: str = Depends(verify_api_key)):
+    result = job_metadata(name=name)
+    return result["values_rules"]
 
 @app.post("/deploy_job",
-    summary="Deploy job",
-    description="Deploy a new job",
+    operation_id="deploy_job",
+    summary="Deploy a new job to the pool",
+    description="Deploys a new job to the Kalavai pool using a specified template and configuration. The job will be scheduled on appropriate nodes based on resource availability and any specified target labels.",
+    tags=["job_management"],
     response_description="Result of job deployment")
 def job_deploy(request: DeployJobRequest, api_key: str = Depends(verify_api_key)):
     """
@@ -348,8 +419,10 @@ def job_deploy(request: DeployJobRequest, api_key: str = Depends(verify_api_key)
     return result
 
 @app.post("/delete_job",
-    summary="Delete job",
-    description="Delete a job",
+    operation_id="delete_job",
+    summary="Terminate and remove a job from the pool",
+    description="Terminates a running job and removes it from the Kalavai pool. This will stop all containers associated with the job and free up the resources they were using.",
+    tags=["job_management"],
     response_description="Result of job deletion")
 def job_delete(request: DeleteJobRequest, api_key: str = Depends(verify_api_key)):
     """
@@ -365,8 +438,10 @@ def job_delete(request: DeleteJobRequest, api_key: str = Depends(verify_api_key)
     return result
 
 @app.get("/authenticate_user",
-    summary="Authenticate user",
-    description="Authenticate a user",
+    operation_id="authenticate_user",
+    summary="Authenticate a user with the Kalavai system",
+    description="Authenticates a user against the Kalavai system, establishing their identity and permissions. This is required for accessing pool management features and deploying jobs.",
+    tags=["info", "auth"],
     response_description="Authentication result")
 def user_authenticate(user_id: str, api_key: str = Depends(verify_api_key)):
     """
@@ -380,8 +455,10 @@ def user_authenticate(user_id: str, api_key: str = Depends(verify_api_key)):
     return result
 
 @app.get("/load_user_session",
-    summary="Load user session",
-    description="Load the current user session",
+    operation_id="load_user_session",
+    summary="Load current user session information",
+    description="Retrieves information about the currently authenticated user's session, including their identity, permissions, and any active connections to pools.",
+    tags=["info", "auth"],
     response_description="User session information")
 def user_session(api_key: str = Depends(verify_api_key)):
     """Load the current user session"""
@@ -389,8 +466,10 @@ def user_session(api_key: str = Depends(verify_api_key)):
     return result
 
 @app.get("/user_logout",
-    summary="User logout",
-    description="Log out the current user",
+    operation_id="user_logout",
+    summary="Log out the current user",
+    description="Terminates the current user's session and clears authentication credentials. This should be called when the user is done using the system to ensure proper cleanup.",
+    tags=["auth"],
     response_description="Logout result")
 def logout_user():
     """Log out the current user"""
@@ -398,8 +477,10 @@ def logout_user():
     return result
 
 @app.get("/is_connected",
-    summary="Check connection",
-    description="Check if connected to a pool",
+    operation_id="is_connected",
+    summary="Check if connected to a Kalavai pool",
+    description="Verifies whether the current instance is connected to a Kalavai pool. Returns connection status and pool information if connected.",
+    tags=["agent_management"],
     response_description="Connection status")
 def pool_connected():
     """Check if connected to a pool"""
@@ -407,8 +488,10 @@ def pool_connected():
     return result
 
 @app.get("/is_agent_running",
-    summary="Check agent status",
-    description="Check if the agent is running",
+    operation_id="is_agent_running",
+    summary="Check if the Kalavai agent is running",
+    description="Verifies whether the Kalavai agent service is currently running on this machine. The agent is responsible for managing pool connections and job execution.",
+    tags=["agent_management"],
     response_description="Agent status")
 def agent_running():
     """Check if the agent is running"""
@@ -416,8 +499,10 @@ def agent_running():
     return result
 
 @app.get("/is_server",
-    summary="Check server status",
-    description="Check if running as server",
+    operation_id="is_server",
+    summary="Check if running as a pool server",
+    description="Determines whether this instance is running as a Kalavai pool server (coordinator) or as a client node. Server instances manage the pool while client instances contribute resources.",
+    tags=["agent_management"],
     response_description="Server status")
 def server():
     """Check if running as server"""
@@ -425,8 +510,10 @@ def server():
     return result
 
 @app.post("/pause_agent",
-    summary="Pause agent",
-    description="Pause the agent",
+    operation_id="pause_agent",
+    summary="Pause the Kalavai agent service",
+    description="Temporarily pauses the Kalavai agent, stopping it from accepting new jobs or participating in pool operations. Existing jobs will continue running until completion.",
+    tags=["agent_management"],
     response_description="Result of pausing agent")
 def agent_pause():
     """Pause the agent"""
@@ -434,8 +521,10 @@ def agent_pause():
     return result
 
 @app.post("/resume_agent",
-    summary="Resume agent",
-    description="Resume the agent",
+    operation_id="resume_agent",
+    summary="Resume the Kalavai agent service",
+    description="Resumes the previously paused Kalavai agent, allowing it to accept new jobs and participate in pool operations again.",
+    tags=["agent_management"],
     response_description="Result of resuming agent")
 def agent_resume():
     """Resume the agent"""
@@ -443,8 +532,10 @@ def agent_resume():
     return result
 
 @app.get("/get_ip_addresses",
-    summary="Get IP addresses",
-    description="Get available IP addresses",
+    operation_id="get_ip_addresses",
+    summary="Get available IP addresses for pool configuration",
+    description="Retrieves a list of available IP addresses that can be used for pool configuration. Optionally filters by subnet to help with network planning and pool setup.",
+    tags=["agent_management"],
     response_description="List of IP addresses")
 def ip_addresses(subnet: str = None, api_key: str = Depends(verify_api_key)):
     """
@@ -456,8 +547,10 @@ def ip_addresses(subnet: str = None, api_key: str = Depends(verify_api_key)):
     return result
 
 @app.get("/list_available_pools",
-    summary="List available pools",
-    description="Get list of available pools",
+    operation_id="list_available_pools",
+    summary="List all available Kalavai pools",
+    description="Retrieves a list of all Kalavai pools that are currently available for connection. Can filter to show only pools owned by the current user or all public pools.",
+    tags=["agent_management"],
     response_description="List of available pools")
 def pool_connected(user_only: bool = False, api_key: str = Depends(verify_api_key)):
     """
@@ -469,8 +562,10 @@ def pool_connected(user_only: bool = False, api_key: str = Depends(verify_api_ke
     return result
 
 @app.post("/add_node_labels",
-    summary="Add node labels",
-    description="Add labels to a node",
+    operation_id="add_node_labels",
+    summary="Add custom labels to a compute node",
+    description="Adds custom labels to a specific compute node in the pool. Labels can be used for job scheduling, resource allocation, and organizational purposes. Labels are key-value pairs that help categorize and identify nodes.",
+    tags=["pool_management"],
     response_description="Result of adding labels")
 def node_labels(request: NodeLabelsRequest, api_key: str = Depends(verify_api_key)):
     """
@@ -485,21 +580,40 @@ def node_labels(request: NodeLabelsRequest, api_key: str = Depends(verify_api_ke
     )
     return result
 
-@app.post("/get_node_labels",
-    summary="Get node labels",
-    description="Get labels for specified nodes",
+@app.get("/get_node_labels",
+    operation_id="get_node_labels",
+    summary="Get labels for specified compute nodes",
+    description="Retrieves all labels associated with specified compute nodes in the pool. Labels provide metadata about nodes and can be used for filtering and scheduling decisions.",
+    tags=["info"],
     response_description="Node labels")
-def node_labels_get(request: GetNodeLabelsRequest, api_key: str = Depends(verify_api_key)):
+def node_labels_get(request: Optional[NodesActionRequest]=NodesActionRequest(), api_key: str = Depends(verify_api_key)):
     """
     Get node labels with the following parameters:
     
     - **node_names**: List of node names to get labels for
     """
     result = get_node_labels(
-        node_names=request.node_names
+        node_names=request.nodes
     )
     return result
 
+### BUILD MCP WRAPPER ###
+mcp = FastApiMCP(
+    app,
+    name="Protected MCP",
+    #exclude_operations=[],
+    exclude_tags=[
+        "auth",
+        "agent_management",
+        "job_management",
+        "pool_management",
+        "avoid"
+    ]
+)
+mcp.mount()
+##########################
+
+
 def run_api(host="0.0.0.0", port=8001, log_level="critical"):
     uvicorn.run(
         app,

kalavai_client/bridge_models.py

@@ -18,7 +18,7 @@ class CreatePoolRequest(BaseModel):
     description: str = Field("", description="Description of the pool")
 
 class NodesActionRequest(BaseModel):
-    nodes: list[str] = Field(None, description="List of node names to perform the action on")
+    nodes: list[str] = Field(None, description="List of node names to perform the action on, defaults to None")
 
 class JoinPoolRequest(BaseModel):
     token: str = Field(description="Token to join the pool")
@@ -26,10 +26,10 @@ class JoinPoolRequest(BaseModel):
     node_name: str = Field(None, description="Name of the node")
     num_gpus: int = Field(None, description="Number of GPUs to allocate")
     frontend: bool = Field(False, description="Whether this is a frontend request")
+
 class JobDetailsRequest(BaseModel):
     jobs: list[Job] = Field(description="List of jobs to get details for")
 
-
 class StopPoolRequest(BaseModel):
     skip_node_deletion: bool = Field(False, description="Whether to skip node deletion when stopping the pool")
 
@@ -46,6 +46,3 @@ class DeleteJobRequest(BaseModel):
 class NodeLabelsRequest(BaseModel):
     node_name: str = Field(description="Name of the node to add labels to")
     labels: Dict[str, str] = Field(description="Dictionary of labels to add to the node")
-
-class GetNodeLabelsRequest(BaseModel):
-    node_names: List[str] = Field(description="List of node names to get labels for")

kalavai_client/core.py

@@ -74,7 +74,7 @@ from kalavai_client.env import (
 )
 
 class Job(BaseModel):
-    owner: Optional[str] = None
+    owner: Optional[str] = "default"
     name: Optional[str] = None
     workers: Optional[str] = None
     endpoint: Optional[str] = None
@@ -238,13 +238,17 @@ def fetch_job_defaults(name):
     except Exception as e:
         return {"error": str(e)}
     
-def fetch_job_templates():
+def fetch_job_templates(type: str=None):
+    data = None
+    if type is not None:
+        data = {"type": type}
     try:
         templates = request_to_server(
             method="get",
             endpoint="/v1/get_job_templates",
             server_creds=USER_LOCAL_SERVER_FILE,
             data=None,
+            params=data,
             user_cookie=USER_COOKIE
         )
         return templates

kalavai_client/utils.py

@@ -105,6 +105,29 @@ def is_storage_compatible():
         return False
 ################
 
+def extract_auth_token(headers):
+    """
+    Extract auth token. Valid headers:
+        X-API-KEY: token
+        X-API-Key: token
+        Authorization: Bearer token
+        authorization: Bearer token
+    """
+    #return headers.get("X-API-Key")
+    bearer = None
+    try:
+        for header in ["Authorization", "authorization", "X-API-KEY", "X-API-Key"]:
+            bearer = headers.get(header, None)
+            if bearer is not None:
+                break
+        if bearer is not None and " " in bearer:
+            return bearer.split()[-1]
+        else:
+            return bearer
+    except Exception as e:
+        return {"error": str(e)}
+
+
 def generate_compose_config(role, node_name, target_platform="amd64", write_to_file=True, node_ip_address="0.0.0.0", num_gpus=0, node_labels=None, pool_ip=None, vpn_token=None, pool_token=None):
     
     if node_labels is not None:
@@ -268,6 +291,7 @@ def request_to_server(
         endpoint,
         data,
         server_creds,
+        params=None,
         force_url=None,
         force_key=None,
         user_cookie=None,
@@ -296,6 +320,7 @@ def request_to_server(
         method=method,
         url=f"http://{service_url}{endpoint}",
         json=data,
+        params=params,
         headers=headers,
         timeout=timeout
     )

{kalavai_client-0.6.17.dist-info → kalavai_client-0.6.18.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: kalavai-client
-Version: 0.6.17
+Version: 0.6.18
 Summary: Client app for kalavai platform
 License: Apache-2.0
 Keywords: LLM,platform
@@ -8,15 +8,9 @@ Author: Carlos Fernandez Musoles
 Author-email: carlos@kalavai.net
 Maintainer: Carlos Fernandez Musoles
 Maintainer-email: carlos@kalavai.net
-Requires-Python: >=3.4
+Requires-Python: >=3.10
 Classifier: License :: OSI Approved :: Apache Software License
 Classifier: Programming Language :: Python :: 3
-Classifier: Programming Language :: Python :: 3.4
-Classifier: Programming Language :: Python :: 3.5
-Classifier: Programming Language :: Python :: 3.6
-Classifier: Programming Language :: Python :: 3.7
-Classifier: Programming Language :: Python :: 3.8
-Classifier: Programming Language :: Python :: 3.9
 Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
@@ -26,6 +20,7 @@ Requires-Dist: Pillow (==10.3.0)
 Requires-Dist: arguably (>=1.2.5)
 Requires-Dist: build ; extra == "dev"
 Requires-Dist: fastapi (==0.115.8)
+Requires-Dist: fastapi-mcp (==0.3.0)
 Requires-Dist: importlib_resources (==6.5.2)
 Requires-Dist: jinja2 (==3.1.4)
 Requires-Dist: netifaces (==0.11.0)

{kalavai_client-0.6.17.dist-info → kalavai_client-0.6.18.dist-info}/RECORD

@@ -1,8 +1,8 @@
-kalavai_client/__init__.py,sha256=6D1LEuHzwFI-n2PmIawhX9C0K4e_BwIIJUAUnpkzCWE,23
+kalavai_client/__init__.py,sha256=ErdtY8HWYl_n6MmTR8hlH878NUP9glayQd1egl02vKY,23
 kalavai_client/__main__.py,sha256=WQUfxvRsBJH5gsCJg8pLz95QnZIj7Ol8psTO77m0QE0,73
 kalavai_client/assets/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-kalavai_client/assets/apps.yaml,sha256=zVtfPqesNhoBLpNlhIPAVtxgXLqEQU2pK1GTzKGEqiQ,6395
-kalavai_client/assets/apps_values.yaml,sha256=dvsAnMC1uk9oDsnITIYJc5CIg9LTwGzXldjPZTwRQyE,2069
+kalavai_client/assets/apps.yaml,sha256=17JuXSv-Qj5Az6ZTRyiEaQXVbI325uTrZzKk2irts2g,6410
+kalavai_client/assets/apps_values.yaml,sha256=LeSNd3PwkIx0wkTIlEk2KNz3Yy4sXSaHALQEkopdhKE,2165
 kalavai_client/assets/docker-compose-gui.yaml,sha256=shqN78YLw0QP7bqTKveI4ppz5E-5b1JowmsSB4OG3nA,778
 kalavai_client/assets/docker-compose-template.yaml,sha256=KHIwJ2WWX7Y7wQKiXRr82Jqd3IKRyls5zhTyl8mSmrc,1805
 kalavai_client/assets/nginx.conf,sha256=drVVCg8GHucz7hmt_BI6giAhK92OV71257NTs3LthwM,225
@@ -11,15 +11,15 @@ kalavai_client/assets/pool_config_values.yaml,sha256=_iAnugramLiwJaaDcPSetThvOdR
 kalavai_client/assets/user_workspace.yaml,sha256=wDvlMYknOPABAEo0dsQwU7bac8iubjAG9tdkFbJZ5Go,476
 kalavai_client/assets/user_workspace_values.yaml,sha256=G0HOzQUxrDMCwuW9kbWUZaKMzDDPVwDwzBHCL2Xi2ZM,542
 kalavai_client/auth.py,sha256=EB3PMvKUn5_KAQkezkEHEt-OMZXyfkZguIQlUFkEHcA,3243
-kalavai_client/bridge_api.py,sha256=5tYqI8UdG7K1Qskywk97kC0TpvYruUZxqxvbn-2nve4,15405
-kalavai_client/bridge_models.py,sha256=t1fJGaF6YDMQdOnEU3XT8zTBHU8eUWJ1yhM5s7c6vMo,2546
+kalavai_client/bridge_api.py,sha256=Hd7whTX2TAiNYX1G237hv2rqtKUBGRJkzUoWOMZm44A,25562
+kalavai_client/bridge_models.py,sha256=3mHCqIHVysLLkQvGT-DKqKOrtAlQSfEOdrwSq2yTRRU,2439
 kalavai_client/cli.py,sha256=SzKG7_ZG0ehMQsECQRWSvqj2Fju2Gd5O7uBa60bFBAY,47830
 kalavai_client/cluster.py,sha256=Z2PIXbZuSAv9xmw-MyZP1M41BpVMpirLzG51bqGA-zc,13548
-kalavai_client/core.py,sha256=dJVX5mhFzIshazCfAzb-AqpqWjkp_djgbMyNXzuAF48,34650
+kalavai_client/core.py,sha256=weg54lc03gp2qGwEXl90XEnXGdwFFlaTqZjxyKsngj4,34765
 kalavai_client/env.py,sha256=YsfZj7LWf6ABquDsoIFFkXCFYwenpDk8zVnGsf7qv98,2823
-kalavai_client/utils.py,sha256=bhvQzF12q7L2hGVrbcmXRDXXIsAdlzcsTms6RQRxGU4,12733
-kalavai_client-0.6.17.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-kalavai_client-0.6.17.dist-info/METADATA,sha256=J5TEqnwPm3ZopGEi2MRn_ddy_VV101JcB8sl6MGL4iI,12655
-kalavai_client-0.6.17.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
-kalavai_client-0.6.17.dist-info/entry_points.txt,sha256=9T6D45gxwzfVbglMm1r6XPdXuuZdHfy_7fCeu2jUphc,50
-kalavai_client-0.6.17.dist-info/RECORD,,
+kalavai_client/utils.py,sha256=GeX1rKUdlQoOW_K2relER8jRQEN1M0UdhsLKOkv5D_g,13428
+kalavai_client-0.6.18.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+kalavai_client-0.6.18.dist-info/METADATA,sha256=nCvnC5f8QM1sHV4wk3HI9YjH0c_vkpOIkcVpMFIKEx0,12393
+kalavai_client-0.6.18.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
+kalavai_client-0.6.18.dist-info/entry_points.txt,sha256=9T6D45gxwzfVbglMm1r6XPdXuuZdHfy_7fCeu2jUphc,50
+kalavai_client-0.6.18.dist-info/RECORD,,