diff --git a/tests/test_agent_loading.py b/tests/test_agent_loading.py new file mode 100644 index 0000000..c9821f3 --- /dev/null +++ b/tests/test_agent_loading.py @@ -0,0 +1,113 @@ +import os +import sys +import tempfile +import unittest +from unittest.mock import patch + +sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))) + +from ventis.stub_generator import _write_requirements + + +AGENT_MODULE = ''' +class Greeter(object): + """A class, the way agents have always been written.""" + + def hello(self, name): + return f"hello {name}" + + +class _Prebuilt(object): + """Stands in for a CompiledStateGraph / Runnable / Crew.""" + + def __init__(self, greeting): + self.greeting = greeting + + def invoke(self, name): + return f"{self.greeting} {name}" + + +# What LangGraph, LCEL and CrewAI actually export: an object, not a class. +graph = _Prebuilt("hi") +''' + + +def _load(agent_name, agent_path): + """The body of LocalController._load_agent, isolated from Redis and gRPC.""" + import importlib.util + + spec = importlib.util.spec_from_file_location("agent_under_test", agent_path) + module = importlib.util.module_from_spec(spec) + sys.modules["agent_under_test"] = module + spec.loader.exec_module(module) + + target = getattr(module, agent_name) + return target() if isinstance(target, type) else target + + +class LoadAgentTargetTests(unittest.TestCase): + """A yaml `agent.name` may point at a class or at an already-built object.""" + + def setUp(self): + self.dir = tempfile.mkdtemp() + self.path = os.path.join(self.dir, "agent_mod.py") + with open(self.path, "w") as f: + f.write(AGENT_MODULE) + + def test_class_is_instantiated(self): + agent = _load("Greeter", self.path) + self.assertEqual(agent.hello("world"), "hello world") + + def test_instance_is_used_as_is(self): + # Calling this one would raise TypeError, which _load_agent swallows, + # leaving the replica alive and answering "No agent loaded". + agent = _load("graph", self.path) + self.assertEqual(agent.invoke("world"), "hi world") + + def test_instance_keeps_its_construction_arguments(self): + # The reason this matters beyond convenience: `agent_class()` takes no + # arguments, so a configured agent had nowhere to receive its config. + agent = _load("graph", self.path) + self.assertEqual(agent.greeting, "hi") + + +class WriteRequirementsTests(unittest.TestCase): + BASE = "grpcio\nredis\npyyaml\n" + + def setUp(self): + self.path = os.path.join(tempfile.mkdtemp(), "requirements.txt") + + def _write(self, extra): + _write_requirements(self.path, self.BASE, extra) + with open(self.path) as f: + return f.read().split() + + def test_declared_requirements_are_appended(self): + self.assertEqual( + self._write(["langchain>=1.0.0", "langgraph>=1.0.0"]), + ["grpcio", "redis", "pyyaml", "langchain>=1.0.0", "langgraph>=1.0.0"], + ) + + def test_none_leaves_the_runtime_list_alone(self): + self.assertEqual(self._write(None), ["grpcio", "redis", "pyyaml"]) + + def test_a_bare_string_is_accepted(self): + self.assertEqual( + self._write("html2text"), ["grpcio", "redis", "pyyaml", "html2text"] + ) + + def test_runtime_dependencies_cannot_be_repinned(self): + # A second `redis` line would let a project pin a version the runtime + # cannot actually run against. + self.assertEqual(self._write(["redis>=5.0", "redis"]), ["grpcio", "redis", "pyyaml"]) + + def test_blank_entries_are_dropped(self): + self.assertEqual(self._write(["", " ", "html2text"]), + ["grpcio", "redis", "pyyaml", "html2text"]) + + def test_extras_marker_is_matched_against_the_base_name(self): + self.assertEqual(self._write(["redis[hiredis]"]), ["grpcio", "redis", "pyyaml"]) + + +if __name__ == "__main__": + unittest.main() diff --git a/ventis/cli.py b/ventis/cli.py index 3aceb18..e5d60dc 100644 --- a/ventis/cli.py +++ b/ventis/cli.py @@ -274,6 +274,7 @@ def cmd_build(args): output_dir=docker_context, grpc_stubs_dir=grpc_stubs_dir, api_port=agent_cfg.get("api_port", 8080), + requirements=agent_cfg.get("requirements"), ) else: @@ -316,6 +317,7 @@ def cmd_build(args): output_dir=docker_context, grpc_stubs_dir=grpc_stubs_dir, stub_files=stub_paths, + requirements=agent_cfg.get("requirements"), ) bake_targets.append( diff --git a/ventis/controller/local_controller.py b/ventis/controller/local_controller.py index 9f53a3e..68db42c 100644 --- a/ventis/controller/local_controller.py +++ b/ventis/controller/local_controller.py @@ -169,11 +169,13 @@ def _load_agent(self): sys.modules[agent_module_name] = module loader.exec_module(module) - agent_class = getattr(module, self.agent_name) - agent_instance = agent_class() - logger.info( - f"Successfully loaded and instantiated agent: {self.agent_name}" - ) + # The name may resolve to a class to instantiate, or to something + # that is already an object -- a compiled LangGraph, a LangChain + # Runnable, a CrewAI Crew. Calling those would be wrong, so only + # classes get called. + target = getattr(module, self.agent_name) + agent_instance = target() if isinstance(target, type) else target + logger.info(f"Successfully loaded agent: {self.agent_name}") return agent_instance except Exception as e: logger.error( diff --git a/ventis/stub_generator.py b/ventis/stub_generator.py index 8c956ef..ab91ff2 100644 --- a/ventis/stub_generator.py +++ b/ventis/stub_generator.py @@ -13,6 +13,7 @@ import argparse import ast import os +import re import shutil import yaml @@ -263,8 +264,40 @@ def _format_source(source): return "\n".join(formatted) + "\n" +def _write_requirements(path, base, extra): + """Write a build context's requirements.txt: the runtime's own, then any + the config declared for this agent or workflow. + + `extra` is whatever the yaml held, so it may be a list or a single string. + Entries already covered by the runtime's own are dropped, so declaring e.g. + `redis` cannot produce a duplicate pin that conflicts with it. + """ + if isinstance(extra, str): + extra = [extra] + + base_names = {line.split("[")[0].strip().lower() for line in base.split() if line} + lines = base.splitlines() + for item in extra or []: + item = str(item).strip() + if not item: + continue + name = re.split(r"[<>=!~\[]", item, 1)[0].strip().lower() + if name in base_names: + continue + base_names.add(name) + lines.append(item) + + with open(path, "w") as f: + f.write("\n".join(lines) + "\n") + + def generate_docker( - yaml_path, agent_file, output_dir=None, grpc_stubs_dir=None, stub_files=None + yaml_path, + agent_file, + output_dir=None, + grpc_stubs_dir=None, + stub_files=None, + requirements=None, ): """ Generate a minimal Docker build context for an agent. @@ -278,6 +311,9 @@ def generate_docker( 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. + requirements: Optional extra pip requirements for this agent, from the + `requirements` key on its global_controller entry. Appended + to the runtime's own, which the agent cannot go without. """ with open(yaml_path, "r") as f: config = yaml.safe_load(f) @@ -297,9 +333,8 @@ def generate_docker( # ---- requirements.txt ------------------------------------------------ # psutil is required unconditionally -- local_controller.py imports it at # module level for CPU/disk/memory metrics reporting on every agent. - requirements = "grpcio\ngrpcio-tools\nredis\npyyaml\nboto3\nyfinance\npsutil\nipdb\nipython\n" - with open(os.path.join(output_dir, "requirements.txt"), "w") as f: - f.write(requirements) + base = "grpcio\ngrpcio-tools\nredis\npyyaml\nboto3\nyfinance\npsutil\nipdb\nipython\n" + _write_requirements(os.path.join(output_dir, "requirements.txt"), base, requirements) # Copy general agent files files_to_copy = [ @@ -377,7 +412,12 @@ def generate_docker( def generate_workflow_docker( - workflow_file, stub_files, output_dir=None, grpc_stubs_dir=None, api_port=8080 + workflow_file, + stub_files, + output_dir=None, + grpc_stubs_dir=None, + api_port=8080, + requirements=None, ): """ Generate a Docker build context for a workflow. @@ -407,12 +447,11 @@ def generate_workflow_docker( # psutil is required unconditionally -- local_controller.py imports it at # module level for CPU/disk/memory metrics reporting on every controller, # including the Workflow's own embedded one. - requirements = ( + base = ( "grpcio\ngrpcio-tools\nredis\npyyaml\nflask\nboto3\nyfinance\npsutil\nipdb\nipython\n" "sqlalchemy\npsycopg[binary]\n" ) - with open(os.path.join(output_dir, "requirements.txt"), "w") as f: - f.write(requirements) + _write_requirements(os.path.join(output_dir, "requirements.txt"), base, requirements) # ---- Copy source files into the build context ------------------------ workflow_basename = os.path.basename(workflow_file)