""" Comprehensive tests for src/builder/graph_builder.py. Covers BuilderConfig, GraphBuilder, build_from_schema, build_from_adjacency, build_property_graph, default_edges, default_sequence helper functions. """ import pytest import torch from builder.graph_builder import ( BuilderConfig, GraphBuilder, build_from_schema, build_property_graph, default_edges, default_sequence, ) from core.schema import ( SCHEMA_VERSION, AgentNodeSchema, EdgeType, GraphSchema, LLMConfig, NodeType, TaskNodeSchema, WorkflowEdgeSchema, ) # ─────────────────────────── BuilderConfig ──────────────────────────────────── class TestBuilderConfig: def test_defaults(self): cfg = BuilderConfig() assert cfg.validate is True assert cfg.check_cycles is True assert cfg.check_duplicates is True assert cfg.allow_self_loops is False assert cfg.node_feature_names == [] assert cfg.edge_feature_names == [] assert cfg.default_edge_dim is None assert cfg.weight_fn is None assert cfg.default_weight == 1.0 assert cfg.include_task_node is True assert cfg.task_edge_weight == 1.0 def test_custom_values(self): cfg = BuilderConfig( validate=False, check_cycles=False, check_duplicates=False, allow_self_loops=True, node_feature_names=["trust", "latency"], edge_feature_names=["weight"], default_edge_dim=4, default_weight=0.5, include_task_node=False, task_edge_weight=2.0, ) assert cfg.validate is False assert cfg.check_cycles is False assert cfg.allow_self_loops is True assert cfg.node_feature_names == ["trust", "latency"] assert cfg.default_weight == 0.5 assert cfg.task_edge_weight == 2.0 def test_weight_fn(self): def fn(s, t, m): return 0.42 cfg = BuilderConfig(weight_fn=fn) assert cfg.weight_fn is fn def test_none_feature_names_become_empty_list(self): cfg = BuilderConfig(node_feature_names=None, edge_feature_names=None) assert cfg.node_feature_names == [] assert cfg.edge_feature_names == [] # ─────────────────────────── GraphBuilder ───────────────────────────────────── class TestGraphBuilderAddAgent: def test_add_single_agent(self): builder = GraphBuilder() builder.add_agent("agent1", description="Test agent") schema = builder.schema assert "agent1" in schema.nodes assert schema.nodes["agent1"].type == NodeType.AGENT def test_add_agent_with_llm_params(self): builder = GraphBuilder() builder.add_agent( "agent1", llm_backbone="gpt-4", base_url="https://api.openai.com/v1", api_key="test-key", max_tokens=2000, temperature=0.7, timeout=30.0, top_p=0.9, stop_sequences=[""], ) node = builder.schema.nodes["agent1"] assert isinstance(node, AgentNodeSchema) assert node.llm_backbone == "gpt-4" assert node.temperature == 0.7 assert node.max_tokens == 2000 assert node.stop_sequences == [""] def test_add_agent_with_llm_config(self): llm_cfg = LLMConfig( model_name="gpt-3.5-turbo", base_url="https://api.openai.com/v1", api_key="key", temperature=0.5, max_tokens=1000, ) builder = GraphBuilder() builder.add_agent("agent1", llm_config=llm_cfg) node = builder.schema.nodes["agent1"] assert isinstance(node, AgentNodeSchema) assert node.llm_backbone == "gpt-3.5-turbo" assert node.temperature == 0.5 def test_add_agent_llm_params_override_config(self): """Direct params override llm_config values.""" llm_cfg = LLMConfig(model_name="gpt-3.5-turbo", temperature=0.3) builder = GraphBuilder() builder.add_agent("agent1", llm_config=llm_cfg, temperature=0.9) node = builder.schema.nodes["agent1"] assert isinstance(node, AgentNodeSchema) assert node.temperature == 0.9 def test_add_agent_with_tools(self): builder = GraphBuilder() builder.add_agent("agent1", tools=["code_interpreter", "web_search"]) node = builder.schema.nodes["agent1"] assert isinstance(node, AgentNodeSchema) assert "code_interpreter" in node.tools def test_add_agent_chaining(self): result = GraphBuilder().add_agent("a").add_agent("b").add_agent("c") assert isinstance(result, GraphBuilder) assert len(result.schema.nodes) == 3 def test_add_agent_with_embedding(self): emb = [0.1, 0.2, 0.3] builder = GraphBuilder() builder.add_agent("agent1", embedding=emb) node = builder.schema.nodes["agent1"] assert node.embedding == emb def test_add_agent_display_name_default(self): builder = GraphBuilder() builder.add_agent("my_agent") node = builder.schema.nodes["my_agent"] assert node.display_name == "my_agent" def test_add_agent_with_metadata(self): builder = GraphBuilder() builder.add_agent("agent1", role="planner", priority=1) node = builder.schema.nodes["agent1"] assert node.metadata.get("role") == "planner" class TestGraphBuilderAddTask: def test_add_task_node(self): builder = GraphBuilder() builder.add_task(query="Solve this problem", description="Task description") assert "__task__" in builder.schema.nodes def test_add_task_custom_id(self): builder = GraphBuilder() builder.add_task(task_id="my_task", query="query") assert "my_task" in builder.schema.nodes def test_task_type(self): builder = GraphBuilder() builder.add_task() node = builder.schema.nodes["__task__"] assert node.type == NodeType.TASK class TestGraphBuilderAddNode: def test_add_custom_node(self): builder = GraphBuilder() builder.add_node("custom1", node_type=NodeType.CUSTOM) assert "custom1" in builder.schema.nodes def test_add_tool_node(self): builder = GraphBuilder() builder.add_node("tool1", node_type=NodeType.TOOL) node = builder.schema.nodes["tool1"] assert node.type == NodeType.TOOL class TestGraphBuilderAddEdges: def test_add_edge_basic(self): builder = GraphBuilder() builder.add_agent("a").add_agent("b") builder.add_edge("a", "b") assert len(builder.schema.edges) == 1 def test_add_edge_with_weight(self): builder = GraphBuilder() builder.add_agent("a").add_agent("b") builder.add_edge("a", "b", weight=0.75) edge = builder.schema.edges[0] assert edge.weight == 0.75 def test_add_edge_with_weight_fn(self): def fn(s, t, m): return 0.33 cfg = BuilderConfig(weight_fn=fn) builder = GraphBuilder(config=cfg) builder.add_agent("a").add_agent("b") builder.add_edge("a", "b") edge = builder.schema.edges[0] assert abs(edge.weight - 0.33) < 1e-6 def test_add_edge_self_loop_raises(self): builder = GraphBuilder() builder.add_agent("a") with pytest.raises(ValueError, match="Self-loops not allowed"): builder.add_edge("a", "a") def test_add_edge_self_loop_allowed(self): cfg = BuilderConfig(allow_self_loops=True) builder = GraphBuilder(config=cfg) builder.add_agent("a") builder.add_edge("a", "a") assert len(builder.schema.edges) == 1 def test_add_workflow_edge(self): builder = GraphBuilder() builder.add_agent("a").add_agent("b") builder.add_workflow_edge("a", "b", condition="source_success") edge = builder.schema.edges[0] assert edge.type == EdgeType.WORKFLOW def test_add_workflow_edge_with_weight_fn(self): def fn(s, t, m): return 2.0 cfg = BuilderConfig(weight_fn=fn) builder = GraphBuilder(config=cfg) builder.add_agent("a").add_agent("b") builder.add_workflow_edge("a", "b") assert builder.schema.edges[0].weight == 2.0 def test_add_conditional_edge_callable(self): builder = GraphBuilder() builder.add_agent("a").add_agent("b") def cond(ctx): return True builder.add_conditional_edge("a", "b", condition=cond) assert ("a", "b") in builder.edge_conditions assert builder.edge_conditions[("a", "b")] is cond def test_add_conditional_edge_string(self): builder = GraphBuilder() builder.add_agent("a").add_agent("b") builder.add_conditional_edge("a", "b", condition="source_success") # String condition goes into schema, not _edge_conditions assert ("a", "b") not in builder.edge_conditions edge = builder.schema.edges[0] assert isinstance(edge, WorkflowEdgeSchema) assert edge.is_conditional is True def test_add_conditional_edge_with_weight_fn(self): def fn(s, t, m): return 3.0 cfg = BuilderConfig(weight_fn=fn) builder = GraphBuilder(config=cfg) builder.add_agent("a").add_agent("b") builder.add_conditional_edge("a", "b", condition="source_success") assert builder.schema.edges[0].weight == 3.0 def test_add_conditional_edges_dict(self): builder = GraphBuilder() builder.add_agent("a").add_agent("b").add_agent("c") builder.add_conditional_edges( "a", path_map={"b": lambda _ctx: True, "c": None}, ) assert len(builder.schema.edges) == 2 def test_add_conditional_edges_with_default(self): builder = GraphBuilder() builder.add_agent("a").add_agent("b").add_agent("c").add_agent("d") builder.add_conditional_edges( "a", path_map={"b": lambda _ctx: True, "c": lambda _ctx: False}, default="d", ) # d is not in path_map, so it gets an unconditional edge targets = [e.target for e in builder.schema.edges] assert "d" in targets def test_from_edges(self): builder = GraphBuilder() builder.add_agent("a").add_agent("b").add_agent("c") builder.from_edges([("a", "b"), ("b", "c")]) assert len(builder.schema.edges) == 2 class TestGraphBuilderExecutionBounds: def test_set_start_node(self): builder = GraphBuilder() builder.add_agent("a") result = builder.set_start_node("a") assert builder.start_node == "a" assert result is builder # chaining def test_set_end_node(self): builder = GraphBuilder() builder.add_agent("b") result = builder.set_end_node("b") assert builder.end_node == "b" assert result is builder def test_set_execution_bounds(self): builder = GraphBuilder() builder.add_agent("a").add_agent("b") builder.set_execution_bounds("a", "b") assert builder.start_node == "a" assert builder.end_node == "b" def test_set_execution_bounds_none(self): builder = GraphBuilder() builder.set_execution_bounds(None, None) assert builder.start_node is None assert builder.end_node is None def test_start_end_node_defaults(self): builder = GraphBuilder() assert builder.start_node is None assert builder.end_node is None class TestGraphBuilderConnectTask: def test_connect_task_to_agents(self): builder = GraphBuilder() builder.add_task() builder.add_agent("a").add_agent("b") builder.connect_task_to_agents() # Should add edges from task to each agent and back edge_pairs = {(e.source, e.target) for e in builder.schema.edges} assert ("__task__", "a") in edge_pairs assert ("__task__", "b") in edge_pairs assert ("a", "__task__") in edge_pairs # bidirectional def test_connect_task_unidirectional(self): builder = GraphBuilder() builder.add_task() builder.add_agent("a") builder.connect_task_to_agents(bidirectional=False) edge_pairs = {(e.source, e.target) for e in builder.schema.edges} assert ("__task__", "a") in edge_pairs assert ("a", "__task__") not in edge_pairs def test_connect_task_specific_agents(self): builder = GraphBuilder() builder.add_task() builder.add_agent("a").add_agent("b").add_agent("c") builder.connect_task_to_agents(agent_ids=["a", "c"]) edge_pairs = {(e.source, e.target) for e in builder.schema.edges} assert ("__task__", "a") in edge_pairs assert ("__task__", "c") in edge_pairs assert ("__task__", "b") not in edge_pairs class TestGraphBuilderValidate: def test_validate_valid_schema(self): builder = GraphBuilder(config=BuilderConfig(check_cycles=False)) builder.add_agent("a").add_agent("b") builder.add_workflow_edge("a", "b") result = builder.validate() assert result.valid def test_validate_with_cycles_disabled(self): cfg = BuilderConfig(check_cycles=False, check_duplicates=False) builder = GraphBuilder(config=cfg) builder.add_agent("a").add_agent("b").add_agent("c") builder.add_workflow_edge("a", "b") builder.add_workflow_edge("b", "c") builder.add_workflow_edge("c", "a") # cycle result = builder.validate() assert result.valid class TestGraphBuilderBuild: def test_build_simple_graph(self): cfg = BuilderConfig(check_cycles=False) builder = GraphBuilder(config=cfg) builder.add_agent("a", description="Agent A") builder.add_agent("b", description="Agent B") builder.add_workflow_edge("a", "b") graph = builder.build() assert graph.num_nodes == 2 assert "a" in graph.node_ids assert "b" in graph.node_ids def test_build_with_validation_disabled(self): cfg = BuilderConfig(validate=False) builder = GraphBuilder(config=cfg) builder.add_agent("x") graph = builder.build() assert graph.num_nodes == 1 def test_build_with_start_end_nodes(self): cfg = BuilderConfig(check_cycles=False) builder = GraphBuilder(config=cfg) builder.add_agent("start_agent").add_agent("end_agent") builder.add_workflow_edge("start_agent", "end_agent") builder.set_start_node("start_agent") builder.set_end_node("end_agent") graph = builder.build() assert graph.start_node == "start_agent" assert graph.end_node == "end_agent" def test_build_validation_fails_on_cycles(self): cfg = BuilderConfig(check_cycles=True, check_duplicates=False) builder = GraphBuilder(config=cfg) builder.add_agent("a").add_agent("b").add_agent("c") builder.add_workflow_edge("a", "b") builder.add_workflow_edge("b", "c") builder.add_workflow_edge("c", "a") # Cycle detection may or may not fail depending on validator implementation result = builder.validate() if not result.valid: with pytest.raises(ValueError, match="Schema validation failed"): builder.build() else: # If cycle check is lenient, just verify build works graph = builder.build() assert graph.num_nodes == 3 def test_build_with_task_node(self): cfg = BuilderConfig(check_cycles=False) builder = GraphBuilder(config=cfg) builder.add_task(query="Test query") builder.add_agent("solver") builder.add_edge("__task__", "solver", edge_type=EdgeType.TASK_CONTEXT) graph = builder.build() assert graph.task_node == "__task__" def test_build_with_llm_config_propagation(self): cfg = BuilderConfig(check_cycles=False) builder = GraphBuilder(config=cfg) builder.add_agent( "agent1", llm_backbone="gpt-4", base_url="https://api.openai.com/v1", api_key="test-key", ) graph = builder.build() agent = next(a for a in graph.agents if a.agent_id == "agent1") assert agent.llm_config is not None assert agent.llm_config.model_name == "gpt-4" def test_build_with_conditional_edges(self): cfg = BuilderConfig(check_cycles=False) builder = GraphBuilder(config=cfg) builder.add_agent("a").add_agent("b") def cond(ctx): return True builder.add_conditional_edge("a", "b", condition=cond) graph = builder.build() assert ("a", "b") in graph.edge_conditions def test_schema_property(self): builder = GraphBuilder() builder.add_agent("a") schema = builder.schema assert isinstance(schema, GraphSchema) assert "a" in schema.nodes class TestGraphBuilderAddAgentProfile: def test_add_agent_profile(self): from core.agent import AgentProfile builder = GraphBuilder() profile = AgentProfile(agent_id="profiled", display_name="Profiled Agent", persona="test") builder.add_agent_profile(profile) assert "profiled" in builder.schema.nodes def test_add_agent_profile_with_tools(self): from core.agent import AgentProfile builder = GraphBuilder() profile = AgentProfile(agent_id="agent1", display_name="Agent", tools=["tool_a"]) builder.add_agent_profile(profile) node = builder.schema.nodes["agent1"] assert isinstance(node, AgentNodeSchema) assert "tool_a" in node.tools def test_add_agent_profile_with_llm_config(self): from core.agent import AgentLLMConfig, AgentProfile llm_cfg = AgentLLMConfig(model_name="llama3", base_url="http://localhost:11434/v1") builder = GraphBuilder() profile = AgentProfile(agent_id="local_agent", display_name="Local", llm_config=llm_cfg) builder.add_agent_profile(profile) node = builder.schema.nodes["local_agent"] assert isinstance(node, AgentNodeSchema) assert node.base_url == "http://localhost:11434/v1" # ─────────────────────────── build_from_schema ──────────────────────────────── class TestBuildFromSchema: def _make_schema_with_agents(self, agent_ids: list[str]) -> GraphSchema: schema = GraphSchema(schema_version=SCHEMA_VERSION) for aid in agent_ids: schema.add_node(AgentNodeSchema(id=aid, display_name=aid)) return schema def test_basic_build(self): schema = self._make_schema_with_agents(["a", "b"]) schema.add_edge(WorkflowEdgeSchema(source="a", target="b")) graph = build_from_schema(schema) assert "a" in graph.node_ids assert "b" in graph.node_ids def test_task_node_in_schema(self): schema = GraphSchema(schema_version=SCHEMA_VERSION) schema.add_node(TaskNodeSchema(id="task", query="test query")) schema.add_node(AgentNodeSchema(id="agent1", display_name="Agent")) schema.add_edge(WorkflowEdgeSchema(source="task", target="agent1")) graph = build_from_schema(schema) assert graph.task_node == "task" def test_with_edge_conditions(self): schema = self._make_schema_with_agents(["a", "b"]) schema.add_edge(WorkflowEdgeSchema(source="a", target="b")) def cond(ctx): return True graph = build_from_schema(schema, edge_conditions={("a", "b"): cond}) assert ("a", "b") in graph.edge_conditions def test_start_end_nodes_passed_to_graph(self): schema = self._make_schema_with_agents(["start", "end"]) schema.add_edge(WorkflowEdgeSchema(source="start", target="end")) graph = build_from_schema(schema, start_node="start", end_node="end") assert graph.start_node == "start" assert graph.end_node == "end" def test_edge_condition_name_extracted(self): schema = self._make_schema_with_agents(["a", "b"]) edge = WorkflowEdgeSchema(source="a", target="b", condition="source_success") schema.add_edge(edge) graph = build_from_schema(schema) # edge_condition_names should contain "source_success" assert ("a", "b") in graph.edge_condition_names assert graph.edge_condition_names[("a", "b")] == "source_success" def test_agent_with_embedding_converted_to_tensor(self): schema = GraphSchema(schema_version=SCHEMA_VERSION) schema.add_node( AgentNodeSchema(id="agent1", display_name="Agent", embedding=[0.1, 0.2, 0.3]) ) graph = build_from_schema(schema) agent = next(a for a in graph.agents if a.agent_id == "agent1") assert agent.embedding is not None assert isinstance(agent.embedding, torch.Tensor) # (build_from_adjacency does not exist in the module, skipped) # ─────────────────────────── build_property_graph ───────────────────────────── class TestBuildPropertyGraph: def _make_agents(self, ids): from core.agent import AgentProfile return [AgentProfile(agent_id=aid, display_name=aid) for aid in ids] def test_basic_property_graph(self): agents = self._make_agents(["a", "b"]) graph = build_property_graph( agents=agents, workflow_edges=[("a", "b")], include_task_node=False, ) assert graph.num_nodes >= 2 def test_property_graph_with_task(self): agents = self._make_agents(["solver"]) graph = build_property_graph( agents=agents, workflow_edges=[], query="Solve X", include_task_node=True, ) assert "__task__" in graph.node_ids assert "solver" in graph.node_ids def test_property_graph_empty(self): graph = build_property_graph(agents=[], workflow_edges=[], include_task_node=False) assert graph.num_nodes == 0 def test_property_graph_with_anchor(self): agents = self._make_agents(["leader", "worker"]) graph = build_property_graph( agents=agents, workflow_edges=[("leader", "worker")], anchor="leader", include_task_node=False, ) assert graph.num_nodes >= 2 # ─────────────────────────── helper functions ───────────────────────────────── class TestDefaultEdges: def test_default_edges_creates_chain(self): """default_edges should create fully-connected or chain edges.""" node_ids = ["a", "b", "c"] edges = default_edges(node_ids) assert isinstance(edges, list) def test_default_edges_empty(self): edges = default_edges([]) assert edges == [] def test_default_edges_single_node(self): edges = default_edges(["solo"]) assert isinstance(edges, list) class TestDefaultSequence: def test_default_sequence_with_anchor(self): """default_sequence(roles, anchor) returns anchor-first ordering.""" result = default_sequence(["b", "c", "a"], "a") assert result[0] == "a" assert set(result) == {"a", "b", "c"} def test_default_sequence_anchor_not_in_roles(self): result = default_sequence(["x", "y"], "z") # anchor not in roles → just return roles in order assert set(result) == {"x", "y"} def test_default_sequence_empty(self): result = default_sequence([], "a") assert result == [] # ─────────────────────────── Integration tests ──────────────────────────────── class TestBuilderIntegration: def test_full_workflow_graph(self): """Build a complete workflow graph and verify all properties.""" cfg = BuilderConfig(check_cycles=False) builder = GraphBuilder(config=cfg) ( builder .add_agent("coordinator", description="Manages workflow", llm_backbone="gpt-4") .add_agent("researcher", description="Researches topics") .add_agent("writer", description="Writes content") .add_agent("reviewer", description="Reviews output") ) ( builder .add_workflow_edge("coordinator", "researcher") .add_workflow_edge("coordinator", "writer") .add_workflow_edge("researcher", "reviewer") .add_workflow_edge("writer", "reviewer") ) builder.set_start_node("coordinator").set_end_node("reviewer") graph = builder.build() assert graph.num_nodes == 4 assert graph.start_node == "coordinator" assert graph.end_node == "reviewer" def test_graph_with_task_and_agents(self): cfg = BuilderConfig(check_cycles=False) builder = GraphBuilder(config=cfg) builder.add_task(query="Analyze data") builder.add_agent("analyzer").add_agent("reporter") builder.connect_task_to_agents() builder.add_workflow_edge("analyzer", "reporter") graph = builder.build() assert graph.task_node == "__task__" assert graph.num_nodes == 3 def test_conditional_routing_pattern(self): cfg = BuilderConfig(check_cycles=False) builder = GraphBuilder(config=cfg) builder.add_agent("classifier") builder.add_agent("math_agent") builder.add_agent("code_agent") builder.add_agent("general_agent") builder.add_conditional_edges( "classifier", path_map={ "math_agent": lambda ctx: "math" in str(ctx), "code_agent": lambda ctx: "code" in str(ctx), }, default="general_agent", ) graph = builder.build() assert graph.num_nodes == 4 def test_pyg_format_export(self): """Test that the built graph can export node features.""" cfg = BuilderConfig(check_cycles=False, node_feature_names=["trust"]) builder = GraphBuilder(config=cfg) builder.add_agent("a", trust_score=0.9) builder.add_agent("b", trust_score=0.7) builder.add_workflow_edge("a", "b") # Build shouldn't raise graph = builder.build() assert graph.num_nodes == 2