diff --git a/naptha_sdk/configs.py b/naptha_sdk/configs.py index b706fad..84ecec7 100644 --- a/naptha_sdk/configs.py +++ b/naptha_sdk/configs.py @@ -1,67 +1,98 @@ import json from pathlib import Path from naptha_sdk.package_manager import load_persona -from naptha_sdk.schemas import AgentDeployment, EnvironmentDeployment, LLMConfig, OrchestratorDeployment, ToolDeployment +from naptha_sdk.schemas import AgentDeployment, EnvironmentDeployment, LLMConfig, OrchestratorDeployment, ToolDeployment, KBDeployment +from naptha_sdk.utils import url_to_node def load_llm_configs(llm_configs_path): with open(llm_configs_path, "r") as file: llm_configs = json.loads(file.read()) return [LLMConfig(**config) for config in llm_configs] -def load_agent_deployments(agent_deployments_path, load_persona_data=False, load_persona_schema=False): - with open(agent_deployments_path, "r") as file: - agent_deployments = json.loads(file.read()) +def load_node_metadata(deployment, node_url): + print(f"Loading node metadata for {deployment['node']['ip']}") + print(f"QQQQQQ: {node_url}") + deployment["node"] = url_to_node(node_url) + print(f"Node metadata loaded {deployment['node']}") + return deployment - for deployment in agent_deployments: - # Load LLM config - config_name = deployment["agent_config"]["llm_config"]["config_name"] +def load_module_config_data(deployment, load_persona_data=False, load_persona_schema=False): + + if "llm_config" in deployment["config"] and deployment["config"]["llm_config"] is not None: + config_name = deployment["config"]["llm_config"]["config_name"] config_path = f"{Path.cwd().name}/configs/llm_configs.json" llm_configs = load_llm_configs(config_path) llm_config = next(config for config in llm_configs if config.config_name == config_name) - deployment["agent_config"]["llm_config"] = llm_config + deployment["config"]["llm_config"] = llm_config + if load_persona_data: + persona_data, input_schema = load_persona(deployment["agent_config"]["persona_module"]["module_url"]) + deployment["agent_config"]["persona_module"]["data"] = persona_data + if load_persona_schema: + deployment["agent_config"]["persona_module"]["data"] = input_schema(**persona_data) + + return deployment - # Load tool deployments if they exist - if "tool_deployments" in deployment and deployment["tool_deployments"]: - tool_deployment_name = deployment["tool_deployments"][0]["name"] - tool_deployment_path = f"{Path.cwd().name}/configs/tool_deployments.json" - tool_deployments = load_tool_deployments(tool_deployment_path) - tool_deployment = next(deployment for deployment in tool_deployments if deployment.name == tool_deployment_name) - deployment["tool_deployments"][0] = tool_deployment +def load_subdeployments(deployment): - if load_persona_data: - persona_data, input_schema = load_persona(deployment["agent_config"]["persona_module"]["module_url"]) - deployment["agent_config"]["persona_module"]["data"] = persona_data - if load_persona_schema: - deployment["agent_config"]["persona_module"]["data"] = input_schema(**persona_data) + configs_path = Path(f"{Path.cwd().name}/configs") - return [AgentDeployment(**deployment) for deployment in agent_deployments] + if "agent_deployments" in deployment and deployment["agent_deployments"]: + # Update defaults with non-None values from input + agent_deployments = [] + for i, agent_deployment in enumerate(deployment["agent_deployments"]): + deployment_name = deployment["agent_deployments"][i]["name"] + agent_deployment = setup_module_deployment("agent", configs_path / "agent_deployments.json", deployment_name) + agent_deployments.append(agent_deployment) + deployment["agent_deployments"] = agent_deployments + if "tool_deployments" in deployment and deployment["tool_deployments"]: + tool_deployments = [] + for i, tool_deployment in enumerate(deployment["tool_deployments"]): + deployment_name = deployment["tool_deployments"][i]["name"] + tool_deployment = setup_module_deployment("tool", configs_path / "tool_deployments.json", deployment_name) + tool_deployments.append(tool_deployment) + deployment["tool_deployments"] = tool_deployments + if "environment_deployments" in deployment and deployment["environment_deployments"]: + environment_deployments = [] + for i, environment_deployment in enumerate(deployment["environment_deployments"]): + deployment_name = deployment["environment_deployments"][i]["name"] + environment_deployment = setup_module_deployment("environment", configs_path / "environment_deployments.json", deployment_name) + environment_deployments.append(environment_deployment) + deployment["environment_deployments"] = environment_deployments + if "kb_deployments" in deployment and deployment["kb_deployments"]: + kb_deployments = [] + for i, kb_deployment in enumerate(deployment["kb_deployments"]): + deployment_name = deployment["kb_deployments"][i]["name"] + kb_deployment = setup_module_deployment("kb", configs_path / "kb_deployments.json", deployment_name) + kb_deployments.append(kb_deployment) + deployment["kb_deployments"] = kb_deployments + print(f"Subdeployments loaded {deployment}") + return deployment -def load_tool_deployments(tool_deployments_path): - with open(tool_deployments_path, "r") as file: - tool_deployments = json.loads(file.read()) +def setup_module_deployment(module_type: str, deployment_path: str, node_url: str, deployment_name: str = None): - for deployment in tool_deployments: - # Load LLM config if present - if "tool_config" in deployment and "llm_config" in deployment["tool_config"]: - config_name = deployment["tool_config"]["llm_config"]["config_name"] - config_path = f"{Path.cwd().name}/configs/llm_configs.json" - llm_configs = load_llm_configs(config_path) - llm_config = next(config for config in llm_configs if config.config_name == config_name) - deployment["tool_config"]["llm_config"] = llm_config + # Map deployment types to their corresponding classes + deployment_map = { + "agent": AgentDeployment, + "tool": ToolDeployment, + "environment": EnvironmentDeployment, + "kb": KBDeployment, + "orchestrator": OrchestratorDeployment + } - return [ToolDeployment(**deployment) for deployment in tool_deployments] + # Load default deployment config from module + with open(deployment_path, "r") as file: + deployment = json.loads(file.read()) -def load_orchestrator_deployments(orchestrator_deployments_path): - with open(orchestrator_deployments_path, "r") as file: - orchestrator_deployments = json.loads(file.read()) - return [OrchestratorDeployment(**deployment) for deployment in orchestrator_deployments] + if deployment_name is None: + deployment = deployment[0] + else: + # Get the first deployment with matching name + deployment = next((d for d in deployment if d["name"] == deployment_name), None) + if deployment is None: + raise ValueError(f"No deployment found with name {deployment_name}") -def load_environment_deployments(environment_deployments_path, config_schema=None): - with open(environment_deployments_path, "r") as file: - environment_deployments = json.loads(file.read()) - - if config_schema: - for deployment in environment_deployments: - deployment["environment_config"] = config_schema + deployment = load_node_metadata(deployment, node_url) + deployment = load_module_config_data(deployment) + deployment = load_subdeployments(deployment) - return [EnvironmentDeployment(**deployment) for deployment in environment_deployments] + return deployment_map[module_type](**deployment) diff --git a/naptha_sdk/schemas.py b/naptha_sdk/schemas.py index 7dedd69..0445c22 100644 --- a/naptha_sdk/schemas.py +++ b/naptha_sdk/schemas.py @@ -62,26 +62,26 @@ class DataGenerationConfig(BaseModel): default_filename: Optional[str] = None class ToolDeployment(BaseModel): - node: NodeSchema + node: Union[NodeSchema, Dict] name: Optional[str] = None module: Optional[Dict] = None config: Optional[ToolConfig] = None data_generation_config: Optional[DataGenerationConfig] = None class KBDeployment(BaseModel): - node: NodeSchema + node: Union[NodeSchema, Dict] name: Optional[str] = None module: Optional[Dict] = None config: Optional[Dict] = None class EnvironmentDeployment(BaseModel): - node: NodeSchema + node: Union[NodeSchema, Dict] name: Optional[str] = None module: Optional[Dict] = None config: Optional[EnvironmentConfig] = None class AgentDeployment(BaseModel): - node: NodeSchema + node: Union[NodeSchema, Dict] name: Optional[str] = None module: Optional[Dict] = None config: Optional[AgentConfig] = None @@ -91,7 +91,7 @@ class AgentDeployment(BaseModel): kb_deployments: Optional[List[KBDeployment]] = None class OrchestratorDeployment(BaseModel): - node: NodeSchema + node: Union[NodeSchema, Dict] name: Optional[str] = None module: Optional[Dict] = None config: Optional[OrchestratorConfig] = None