diff --git a/examples/portfolio/agents/metrics_agent.py b/examples/portfolio/agents/metrics_agent.py index 6fb375f..d5f722e 100644 --- a/examples/portfolio/agents/metrics_agent.py +++ b/examples/portfolio/agents/metrics_agent.py @@ -8,6 +8,7 @@ # downstream RiskAgent can build the portfolio covariance. # # Resource profile: cheap CPU, high fan-out — one compute() call per holding. +import os import sys import os diff --git a/examples/portfolio/workflow/portfolio_workflow.py b/examples/portfolio/workflow/portfolio_workflow.py index 299a3f5..73e94d2 100644 --- a/examples/portfolio/workflow/portfolio_workflow.py +++ b/examples/portfolio/workflow/portfolio_workflow.py @@ -46,7 +46,7 @@ def main( advisor = AdvisorAgent() # Stage 0: parse the free-text request into structured holdings + window. - intent = intent_agent.parse(query=query).value() + intent = intent_agent.parse(query=query) # parse() returns a dict, but a Future's .value() only ever gives back the # raw string ventis stored in Redis -- same deserialization requirement as # every other dict-returning agent call below. diff --git a/ventis/cli.py b/ventis/cli.py index b43a6b3..4b5c95e 100644 --- a/ventis/cli.py +++ b/ventis/cli.py @@ -226,6 +226,23 @@ def cmd_build(args): if not yaml_files: logger.warning("No agent YAML files found in %s", agents_dir) + import yaml + + # Looks up a config entry's YAML and to map stubs to entrypoints. + yaml_by_name = {} + for yaml_path in yaml_files: + with open(yaml_path) as f: + name = yaml.safe_load(f).get("agent", {}).get("name") + if name: + yaml_by_name[name] = yaml_path + + entrypoints_by_name = {a["name"]: a.get("entrypoint") for a in agents} + stub_entrypoints = { + f"{os.path.splitext(os.path.basename(p))[0]}.py": entrypoints_by_name[n] + for n, p in yaml_by_name.items() + if entrypoints_by_name.get(n) + } + stub_paths = [] for yaml_path in yaml_files: base_name = os.path.splitext(os.path.basename(yaml_path))[0] @@ -288,6 +305,8 @@ def cmd_build(args): output_dir=docker_context, grpc_stubs_dir=grpc_stubs_dir, api_port=agent_cfg.get("api_port", 8080), + project_dir=project_dir, + stub_entrypoints=stub_entrypoints, requirements=_normalize_requirements(agent_cfg), ) @@ -306,16 +325,7 @@ def cmd_build(args): continue # Find matching YAML by agent name - matching_yaml = None - for yaml_path in yaml_files: - import yaml - - with open(yaml_path) as f: - ydata = yaml.safe_load(f) - if ydata.get("agent", {}).get("name") == agent_name: - matching_yaml = yaml_path - break - + matching_yaml = yaml_by_name.get(agent_name) if not matching_yaml: logger.warning( "No YAML definition found for agent '%s', skipping Docker", @@ -331,6 +341,8 @@ def cmd_build(args): output_dir=docker_context, grpc_stubs_dir=grpc_stubs_dir, stub_files=stub_paths, + project_dir=project_dir, + stub_entrypoints=stub_entrypoints, requirements=_normalize_requirements(agent_cfg), ) diff --git a/ventis/stub_generator.py b/ventis/stub_generator.py index e408553..54a46b9 100644 --- a/ventis/stub_generator.py +++ b/ventis/stub_generator.py @@ -271,12 +271,66 @@ def _format_source(source): return "\n".join(formatted) + "\n" +# Directories ventis build itself generates inside a project -- never swept. +_GENERATED_DIRS = {"docker_container", "stubs", "grpc_stubs"} + + +def _sweep_py_files(project_dir): + """Recursively collect (abs_src, rel_dst) for every .py file under project_dir, preserving its directory structure.""" + swept = [] + for root, dirs, files in os.walk(project_dir): + dirs[:] = [ + d + for d in dirs + if not d.startswith(".") + and not (root == project_dir and d in _GENERATED_DIRS) + ] + for fname in files: + abs_src = os.path.join(root, fname) + if fname.endswith(".py") and not os.path.islink(abs_src): + rel_dst = os.path.relpath(abs_src, project_dir) + swept.append((abs_src, rel_dst)) + return swept + + +def _stub_destination(stub_file, stub_entrypoints): + """Where to copy a stub so it overwrites the real file it replaces, falling back to flat if that's unsafe.""" + basename = os.path.basename(stub_file) + entrypoint = stub_entrypoints.get(basename) + if entrypoint: + normalized = entrypoint.replace("\\", "/") + if not normalized.startswith("/") and ".." not in normalized.split("/"): + return normalized + print(f" Warning: unsafe entrypoint '{entrypoint}' for stub {basename}, placing flat instead") + elif stub_entrypoints: + print(f" Warning: no entrypoint mapping for stub {basename}, placing flat instead") + return basename + + +def _copy_files(output_dir, files_to_copy): + """Copy each (src, dst) pair into output_dir, refusing to write outside it (e.g. via a symlinked destination parent).""" + real_output_dir = os.path.realpath(output_dir) + for src, dst in files_to_copy: + if not os.path.isfile(src): + print(f" Warning: source file not found, skipping: {src}") + continue + dest_path = os.path.join(output_dir, dst) + real_dest = os.path.realpath(dest_path) + if os.path.commonpath([real_output_dir, real_dest]) != real_output_dir: + print(f" Warning: destination escapes build context, skipping: {dst}") + continue + os.makedirs(os.path.dirname(dest_path), exist_ok=True) + shutil.copy2(src, dest_path) + + def generate_docker( yaml_path, agent_file, output_dir=None, grpc_stubs_dir=None, stub_files=None, + project_dir=None, + stub_entrypoints=None, requirements=None, ): """ @@ -286,11 +340,13 @@ def generate_docker( source files needed to run the agent with its own local controller. Args: - yaml_path: Path to the YAML agent definition. - agent_file: Path to the original Python agent implementation. - output_dir: Optional output directory (default: docker_container//). - grpc_stubs_dir: Optional path to compiled gRPC stubs (default: /grpc_stubs). - stub_files: Optional list of agent stub files to copy into the context. + yaml_path: Path to the YAML agent definition. + agent_file: Path to the original Python agent implementation. + output_dir: Optional output directory (default: docker_container//). + grpc_stubs_dir: Optional path to compiled gRPC stubs (default: /grpc_stubs). + stub_files: Optional list of agent stub files to copy into the context. + project_dir: Optional project root to sweep for extra .py helper files. + stub_entrypoints: Optional {stub_basename: entrypoint} map for exact stub placement. requirements: Optional list of extra pip packages this agent needs. """ with open(yaml_path, "r") as f: @@ -314,8 +370,13 @@ def generate_docker( with open(os.path.join(output_dir, "requirements.txt"), "w") as f: f.write(requirements_txt) + # Sweep the project for extra .py helper files not on the explicit list below. + files_to_copy = [] + if project_dir: + files_to_copy += _sweep_py_files(project_dir) + # Copy general agent files - files_to_copy = [ + files_to_copy += [ # (source_path, destination_filename) (os.path.join(script_dir, "future.py"), "future.py"), (os.path.join(script_dir, "ventis_context.py"), "ventis_context.py"), @@ -336,13 +397,16 @@ def generate_docker( (os.path.join(script_dir, "llm", "bedrock.py"), "bedrock.py"), ] - # Copy provided agent stubs + # Copy provided agent stubs, overwriting the swept real file at the same path if stub_files: for stub_file in stub_files: files_to_copy.append( - (os.path.abspath(stub_file), os.path.basename(stub_file)) + ( + os.path.abspath(stub_file), + _stub_destination(stub_file, stub_entrypoints or {}), + ) ) - + files_to_copy.append((os.path.abspath(agent_file), os.path.basename(agent_file))) # Copy gRPC generated stubs if they exist @@ -351,11 +415,7 @@ def generate_docker( if fname.endswith(".py"): files_to_copy.append((os.path.join(grpc_stubs_dir, fname), fname)) - for src, dst in files_to_copy: - if os.path.isfile(src): - shutil.copy2(src, os.path.join(output_dir, dst)) - else: - print(f" Warning: source file not found, skipping: {src}") + _copy_files(output_dir, files_to_copy) # Copy the YAML definition too shutil.copy2( @@ -396,6 +456,8 @@ def generate_workflow_docker( output_dir=None, grpc_stubs_dir=None, api_port=8080, + project_dir=None, + stub_entrypoints=None, requirements=None, ): """ @@ -406,10 +468,12 @@ def generate_workflow_docker( with its own local controller. Args: - workflow_file: Path to the workflow Python file. - stub_files: List of stub file paths to include. - output_dir: Optional output directory (default: docker_container/Workflow/). - grpc_stubs_dir: Optional path to compiled gRPC stubs (default: /grpc_stubs). + workflow_file: Path to the workflow Python file. + stub_files: List of stub file paths to include. + output_dir: Optional output directory (default: docker_container/Workflow/). + grpc_stubs_dir: Optional path to compiled gRPC stubs (default: /grpc_stubs). + project_dir: Optional project root to sweep for extra .py helper files. + stub_entrypoints: Optional {stub_basename: entrypoint} map for exact stub placement. requirements: Optional list of extra pip packages this workflow needs. """ script_dir = os.path.dirname(os.path.abspath(__file__)) @@ -432,7 +496,10 @@ def generate_workflow_docker( # ---- Copy source files into the build context ------------------------ workflow_basename = os.path.basename(workflow_file) - files_to_copy = [ + # Sweep the project for extra .py helper files not on the explicit list below. + files_to_copy = _sweep_py_files(project_dir) if project_dir else [] + + files_to_copy += [ (os.path.abspath(workflow_file), workflow_basename), (os.path.join(script_dir, "future.py"), "future.py"), (os.path.join(script_dir, "ventis_context.py"), "ventis_context.py"), @@ -452,10 +519,15 @@ def generate_workflow_docker( for name in ("gpu_metrics.py", "session_logging.py") ], ] - - # Copy stub files + + # Copy stub files, overwriting the swept real file at the same path for stub_file in stub_files: - files_to_copy.append((os.path.abspath(stub_file), os.path.basename(stub_file))) + files_to_copy.append( + ( + os.path.abspath(stub_file), + _stub_destination(stub_file, stub_entrypoints or {}), + ) + ) # Copy gRPC generated stubs if they exist if os.path.isdir(grpc_stubs_dir): @@ -463,11 +535,7 @@ def generate_workflow_docker( if fname.endswith(".py"): files_to_copy.append((os.path.join(grpc_stubs_dir, fname), fname)) - for src, dst in files_to_copy: - if os.path.isfile(src): - shutil.copy2(src, os.path.join(output_dir, dst)) - else: - print(f" Warning: source file not found, skipping: {src}") + _copy_files(output_dir, files_to_copy) # ---- workflow_launcher.py -------------------------------------------- launcher = f"""import threading