From 98b0b8561aba1e88a8087a24fe03d46b3a893273 Mon Sep 17 00:00:00 2001 From: Andres Javier Lopez Date: Mon, 3 Nov 2025 18:05:46 -0600 Subject: [PATCH 1/6] Replace BaseAgentFactory for ContextFactory This class was actually used as an abstract factory for BuilderContext, the old name was misleading --- CHANGELOG.md | 1 + assistant_core/builder/context.py | 31 +++++++++++--- assistant_core/factories.py | 71 +++++++++++++++++++++++++------ tests/conftest.py | 27 +++++------- 4 files changed, 92 insertions(+), 38 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 48c86a7..2f1d513 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,6 +1,7 @@ # CHANGELOG ## NEXT +* Replace BaseAgentFactory for ContextFactory for pluggable components (agent/graph/model/tools/resolver) ## v0.7.0 * Update dependencies diff --git a/assistant_core/builder/context.py b/assistant_core/builder/context.py index 5c1b0ba..1ef3df5 100644 --- a/assistant_core/builder/context.py +++ b/assistant_core/builder/context.py @@ -1,13 +1,30 @@ -from assistant_core.factories import BaseAgentFactory +import warnings + +from assistant_core.factories import ContextFactory class BuilderContext: - def __init__(self, agent_factory: BaseAgentFactory): - self.model = agent_factory.model - self.graph_builder = agent_factory.create_graph_builder() - self.agent_node = agent_factory.create_agent_node() - self.resolver_node = agent_factory.create_resolver_node() - self.tools = agent_factory.create_base_tools() + def __init__( + self, + context_factory: ContextFactory = None, + agent_factory: ContextFactory = None, + ): + if context_factory is None: + if agent_factory is None: + raise ValueError("context_factory must be provided") + + warnings.warn( + "BuilderContext: 'agent_factory' is deprecated, " + "use 'context_factory' instead", + DeprecationWarning, + ) + context_factory = agent_factory + + self.model = context_factory.model + self.graph_builder = context_factory.create_graph_builder() + self.agent_node = context_factory.create_agent_node() + self.resolver_node = context_factory.create_resolver_node() + self.tools = context_factory.create_base_tools() self._entrypoint = None diff --git a/assistant_core/factories.py b/assistant_core/factories.py index 30cc660..21e3bb1 100644 --- a/assistant_core/factories.py +++ b/assistant_core/factories.py @@ -1,23 +1,44 @@ -from abc import ABC, abstractmethod -from typing import TypedDict +import warnings +from typing import Callable, TypedDict -from langgraph.graph import StateGraph +from langchain_core.language_models.chat_models import BaseChatModel +from langchain_core.tools import BaseTool +from langgraph.graph import MessagesState, StateGraph from assistant_core.models import load_default_model from assistant_core.nodes import AgentNode, ResolverNode +AgentFactory = Callable[[BaseChatModel], AgentNode] | None +GraphFactory = Callable[[], StateGraph] | None +ModelFactory = Callable[["ContextFactory.FactoryConfig"], BaseChatModel] | None +ResolverFactory = Callable[[], ResolverNode] | None +BaseToolsFactory = Callable[[], list[BaseTool]] | None -class BaseAgentFactory(ABC): + +class ContextFactory: class FactoryConfig(TypedDict): OPENAI_API_KEY: str - def __init__(self, config: FactoryConfig): - super().__init__() - self._model = None + def __init__( + self, + config: FactoryConfig, + *, + agent_factory: AgentFactory = None, + graph_factory: GraphFactory = None, + model_factory: ModelFactory = None, + resolver_factory: ResolverFactory = None, + base_tools_factory: BaseToolsFactory = None, + ): + self._model: BaseChatModel | None = None self.config = config + self._agent_factory = agent_factory + self._graph_factory = graph_factory + self._model_factory = model_factory + self._resolver_factory = resolver_factory + self._base_tools_factory = base_tools_factory @property - def model(self): + def model(self) -> BaseChatModel: """ Get the model instance. If not already created, it will create one using the factory method. @@ -26,34 +47,56 @@ def model(self): self._model = self.create_model() return self._model - @abstractmethod def create_graph_builder(self) -> StateGraph: """ Create a graph builder instance. """ - raise NotImplementedError("Subclasses must implement this method") + if self._graph_factory: + return self._graph_factory() + + return StateGraph(MessagesState) - @abstractmethod def create_agent_node(self) -> AgentNode: """ Create an agent instance. """ - raise NotImplementedError("Subclasses must implement this method") + if self._agent_factory: + return self._agent_factory(self.model) - def create_model(self): + raise NotImplementedError("Agent factory must be provided") + + def create_model(self) -> BaseChatModel: """ Create a model instance. """ + if self._model_factory: + return self._model_factory(self.config) + return load_default_model(openai_api_key=self.config["OPENAI_API_KEY"]) def create_resolver_node(self) -> ResolverNode: """ Create a resolver instance. """ + if self._resolver_factory: + return self._resolver_factory() + return ResolverNode(name="resolver") - def create_base_tools(self) -> list: + def create_base_tools(self) -> list[BaseTool]: """ Create a set of base tools for the agent. """ + if self._base_tools_factory: + return self._base_tools_factory() + return [] + + +class BaseAgentFactory(ContextFactory): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + warnings.warn( + "BaseAgentFactory is deprecated, use ContextFactory instead.", + DeprecationWarning, + ) diff --git a/tests/conftest.py b/tests/conftest.py index 6640f15..ae2ac97 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -2,10 +2,9 @@ import pytest from langchain_core.language_models.chat_models import BaseChatModel -from langgraph.graph import MessagesState, StateGraph from assistant_core.builder import BuilderContext -from assistant_core.factories import BaseAgentFactory +from assistant_core.factories import ContextFactory from assistant_core.nodes import AgentNode @@ -30,24 +29,18 @@ def agent_node(mock_model): @pytest.fixture def factory_config(): - return BaseAgentFactory.FactoryConfig(OPENAI_API_KEY="test_key") + return ContextFactory.FactoryConfig(OPENAI_API_KEY="test_key") @pytest.fixture -def agent_factory(mock_model, agent_node, factory_config): - class MockAgentFactory(BaseAgentFactory): - def create_model(self): - return mock_model - - def create_graph_builder(self): - return StateGraph(MessagesState) - - def create_agent_node(self): - return agent_node - - return MockAgentFactory(factory_config) +def context_factory(mock_model, agent_node, factory_config): + return ContextFactory( + factory_config, + agent_factory=lambda _: agent_node, + model_factory=lambda _: mock_model, + ) @pytest.fixture -def builder_context(agent_factory): - return BuilderContext(agent_factory=agent_factory) +def builder_context(context_factory): + return BuilderContext(context_factory=context_factory) From 5fd28667487a1d998c2f29bafaf93c86f2c52444 Mon Sep 17 00:00:00 2001 From: Andres Javier Lopez Date: Mon, 3 Nov 2025 20:46:36 -0600 Subject: [PATCH 2/6] Add helper to create BuilderContext and clone method --- assistant_core/builder/context.py | 54 +++++++++- assistant_core/factories.py | 29 +++++- tests/unit/test_builder_context.py | 155 +++++++++++++++++++++++++++++ 3 files changed, 236 insertions(+), 2 deletions(-) create mode 100644 tests/unit/test_builder_context.py diff --git a/assistant_core/builder/context.py b/assistant_core/builder/context.py index 1ef3df5..8d81372 100644 --- a/assistant_core/builder/context.py +++ b/assistant_core/builder/context.py @@ -1,6 +1,14 @@ import warnings +from typing import Self -from assistant_core.factories import ContextFactory +from assistant_core.factories import ( + AgentFactory, + BaseToolsFactory, + ContextFactory, + GraphFactory, + ModelFactory, + ResolverFactory, +) class BuilderContext: @@ -20,6 +28,9 @@ def __init__( ) context_factory = agent_factory + # Keep a reference to enable cloning/customization + self._factory = context_factory + self.model = context_factory.model self.graph_builder = context_factory.create_graph_builder() self.agent_node = context_factory.create_agent_node() @@ -42,3 +53,44 @@ def entrypoint(self, value: str): self.graph_builder.add_edge(value, self.entrypoint) self._entrypoint = value + + @classmethod + def create( + cls, + config: ContextFactory.FactoryConfig, + *, + agent_factory: AgentFactory, + graph_factory: GraphFactory = None, + model_factory: ModelFactory = None, + resolver_factory: ResolverFactory = None, + base_tools_factory: BaseToolsFactory = None, + ) -> Self: + factory = ContextFactory( + config, + agent_factory=agent_factory, + graph_factory=graph_factory, + model_factory=model_factory, + resolver_factory=resolver_factory, + base_tools_factory=base_tools_factory, + ) + return cls(factory) + + def clone( + self, + *, + agent_factory: AgentFactory = None, + graph_factory: GraphFactory = None, + model_factory: ModelFactory = None, + resolver_factory: ResolverFactory = None, + base_tools_factory: BaseToolsFactory = None, + ) -> Self: + + new_factory = self._factory.clone( + agent_factory=agent_factory, + graph_factory=graph_factory, + model_factory=model_factory, + resolver_factory=resolver_factory, + base_tools_factory=base_tools_factory, + ) + + return self.__class__(new_factory) diff --git a/assistant_core/factories.py b/assistant_core/factories.py index 21e3bb1..46a9a11 100644 --- a/assistant_core/factories.py +++ b/assistant_core/factories.py @@ -1,5 +1,5 @@ import warnings -from typing import Callable, TypedDict +from typing import Callable, Self, TypedDict from langchain_core.language_models.chat_models import BaseChatModel from langchain_core.tools import BaseTool @@ -37,6 +37,33 @@ def __init__( self._resolver_factory = resolver_factory self._base_tools_factory = base_tools_factory + def clone( + self, + *, + agent_factory: AgentFactory = None, + graph_factory: GraphFactory = None, + model_factory: ModelFactory = None, + resolver_factory: ResolverFactory = None, + base_tools_factory: BaseToolsFactory = None, + ) -> Self: + + agent_factory = agent_factory or self._agent_factory + graph_factory = graph_factory or self._graph_factory + model_factory = model_factory or self._model_factory + resolver_factory = resolver_factory or self._resolver_factory + base_tools_factory = base_tools_factory or self._base_tools_factory + + new_factory = self.__class__( + self.config, + agent_factory=agent_factory, + graph_factory=graph_factory, + model_factory=model_factory, + resolver_factory=resolver_factory, + base_tools_factory=base_tools_factory, + ) + + return new_factory + @property def model(self) -> BaseChatModel: """ diff --git a/tests/unit/test_builder_context.py b/tests/unit/test_builder_context.py new file mode 100644 index 0000000..7419a8b --- /dev/null +++ b/tests/unit/test_builder_context.py @@ -0,0 +1,155 @@ +import warnings +from unittest import mock + +from langchain_core.language_models.chat_models import BaseChatModel +from langgraph.graph import MessagesState, StateGraph + +from assistant_core.builder.context import BuilderContext +from assistant_core.nodes import AgentNode, PromptNode, ResolverNode + + +def test_builder_context_initialization_defaults(builder_context): + # Provided via fixture with ContextFactory using agent/model factories + ctx = builder_context + + assert isinstance(ctx.graph_builder, StateGraph) + assert ctx.agent_node.name == "test_agent" + assert ctx.resolver_node.name == "resolver" + assert isinstance(ctx.tools, list) + # Default entrypoint is the agent node + assert ctx.entrypoint == ctx.agent_node.name + + +def test_entrypoint_setter_creates_lifo_edges(builder_context): + ctx = builder_context + + # Register two custom nodes + n1 = PromptNode(name="node_1", prompt="n1") + n2 = PromptNode(name="node_2", prompt="n2") + ctx.graph_builder.add_node(n1.name, n1) + ctx.graph_builder.add_node(n2.name, n2) + + # Initial entrypoint should be the agent + assert ctx.entrypoint == "test_agent" + + # Setting entrypoint should add an edge to previous entrypoint + ctx.entrypoint = n1.name + ctx.entrypoint = n2.name + + assert ("node_1", "test_agent") in ctx.graph_builder.edges + assert ("node_2", "node_1") in ctx.graph_builder.edges + # Current entrypoint is the last set value + assert ctx.entrypoint == "node_2" + + +def test_create_classmethod_uses_factories(mock_model, factory_config): + def model_factory(cfg): + assert cfg == factory_config + return mock_model + + def agent_factory(model): + assert model is mock_model + return AgentNode(name="custom_agent", model=model, prompts=["Hi"]) + + def graph_factory(): + return StateGraph(MessagesState) + + def resolver_factory(): + return ResolverNode(name="custom_resolver") + + tools = [mock.Mock(name="tool1"), mock.Mock(name="tool2")] + + ctx = BuilderContext.create( + factory_config, + agent_factory=agent_factory, + graph_factory=graph_factory, + model_factory=model_factory, + resolver_factory=resolver_factory, + base_tools_factory=lambda: tools, + ) + + assert isinstance(ctx, BuilderContext) + assert isinstance(ctx.graph_builder, StateGraph) + assert ctx.agent_node.name == "custom_agent" + assert ctx.resolver_node.name == "custom_resolver" + assert ctx.model is mock_model + assert ctx.tools is tools + assert ctx.entrypoint == "custom_agent" + + +def test_clone_same_factories(builder_context): + # Clone without overrides should produce an independent but equivalent context + original = builder_context + + cloned = original.clone() + + assert cloned is not original + assert cloned.agent_node.name == original.agent_node.name + assert cloned.resolver_node.name == original.resolver_node.name + assert cloned.model is original.model + assert cloned.tools == original.tools + assert isinstance(cloned.graph_builder, StateGraph) + # Original remains unchanged + assert original.agent_node.name == "test_agent" + assert original.resolver_node.name == "resolver" + + +def test_clone_overrides_factories(context_factory, builder_context): + # Original context + original = builder_context + assert original.agent_node.name == "test_agent" + assert original.resolver_node.name == "resolver" + + # New model and overrides + new_model = mock.Mock(spec=BaseChatModel) + + def model_factory(_cfg): + return new_model + + received_models: list[BaseChatModel] = [] + + def agent_factory(model): + # Record the model used to build the agent; cloning may keep cached model + received_models.append(model) + return AgentNode(name="agent_clone", model=model, prompts=["clone"]) + + def graph_factory(): + return StateGraph(MessagesState) + + def resolver_factory(): + return ResolverNode(name="resolver_clone") + + tools = [mock.Mock(name="toolX")] + + cloned = original.clone( + agent_factory=agent_factory, + graph_factory=graph_factory, + model_factory=model_factory, + resolver_factory=resolver_factory, + base_tools_factory=lambda: tools, + ) + + # Ensure the clone uses overrides and is independent from original + assert cloned is not original + assert cloned.agent_node.name == "agent_clone" + assert cloned.resolver_node.name == "resolver_clone" + # Model may be cached from the original factory; ensure consistency + assert cloned.model is new_model + assert received_models == [cloned.model] + assert cloned.tools is tools + assert isinstance(cloned.graph_builder, StateGraph) + # Original remains unchanged + assert original.agent_node.name == "test_agent" + assert original.resolver_node.name == "resolver" + + +def test_init_with_deprecated_agent_factory_param(context_factory): + # Passing a ContextFactory via deprecated param should warn and work + with warnings.catch_warnings(record=True) as rec: + warnings.simplefilter("always") + ctx = BuilderContext(agent_factory=context_factory) + assert any( + issubclass(w.category, DeprecationWarning) for w in rec + ), "Expected a DeprecationWarning when using agent_factory param" + assert ctx.agent_node.name == "test_agent" + assert ctx.entrypoint == "test_agent" From ab5f140a96144969c4a6326b6ca14e98518e736b Mon Sep 17 00:00:00 2001 From: Andres Javier Lopez Date: Mon, 3 Nov 2025 21:00:05 -0600 Subject: [PATCH 3/6] Improve prompt for datetime builder --- CHANGELOG.md | 1 + assistant_core/builder/datetime.py | 5 ++++- tests/unit/test_datetime_builder.py | 6 ++++-- 3 files changed, 9 insertions(+), 3 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 2f1d513..22ef2f8 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,7 @@ ## NEXT * Replace BaseAgentFactory for ContextFactory for pluggable components (agent/graph/model/tools/resolver) +* Add an internal-use notice to datetime builder ## v0.7.0 * Update dependencies diff --git a/assistant_core/builder/datetime.py b/assistant_core/builder/datetime.py index 53e2112..42329d5 100644 --- a/assistant_core/builder/datetime.py +++ b/assistant_core/builder/datetime.py @@ -21,7 +21,10 @@ def get_current_time(TZ: ZoneInfo) -> str: def get_current_date_prompt(TZ: ZoneInfo) -> str: """Return the current date prompt.""" - return f"Today is {get_weekday(TZ)} and the time is {get_current_time(TZ)}." + return ( + f"CONTEXT: Today is {get_weekday(TZ)} and the time is {get_current_time(TZ)}.\n" + "This is for internal use only. Do not provide date or time unless asked." + ) class DateTimeNode(PromptNode): diff --git a/tests/unit/test_datetime_builder.py b/tests/unit/test_datetime_builder.py index 4ac0476..1e7a3ca 100644 --- a/tests/unit/test_datetime_builder.py +++ b/tests/unit/test_datetime_builder.py @@ -30,10 +30,12 @@ def now(cls, tz=None): def test_helpers_return_expected_strings(): # Patch the datetime used in the module to our fixed datetime class with mock.patch("assistant_core.builder.datetime.datetime.datetime", FixedDateTime): + date_isoformat = _make_fixed_dt().isoformat() assert get_weekday(UTC) == "Friday" - assert get_current_time(UTC) == _make_fixed_dt().isoformat() + assert get_current_time(UTC) == date_isoformat expected_prompt = ( - f"Today is Friday and the time is {_make_fixed_dt().isoformat()}." + f"CONTEXT: Today is Friday and the time is {date_isoformat}.\n" + "This is for internal use only. Do not provide date or time unless asked." ) assert get_current_date_prompt(UTC) == expected_prompt From 05f7773c5a9e2e8cbbe4b518d1a630f00a9769a6 Mon Sep 17 00:00:00 2001 From: Andres Javier Lopez Date: Mon, 3 Nov 2025 21:01:15 -0600 Subject: [PATCH 4/6] Add poetry instructions to copilot --- .github/copilot-instructions.md | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/.github/copilot-instructions.md b/.github/copilot-instructions.md index 46373fd..29a28d9 100644 --- a/.github/copilot-instructions.md +++ b/.github/copilot-instructions.md @@ -16,3 +16,7 @@ instructions: - If the suggestion touches public APIs or behavior, encourage adding typed signatures and small unit tests. + + - When suggesting commands to run tests or test-related scripts, prefer + invoking them via Poetry in this repository. Example: `poetry run pytest -q`. + This ensures the project's virtual environment and dev-dependencies are used. From 9ab94694eba82cc2060b49276f6dbc8e95cc0198 Mon Sep 17 00:00:00 2001 From: Andres Javier Lopez Date: Mon, 3 Nov 2025 21:26:28 -0600 Subject: [PATCH 5/6] Add MultiAgent builder support --- CHANGELOG.md | 2 + assistant_core/builder/__init__.py | 4 +- assistant_core/builder/agent.py | 25 +++++++++++- assistant_core/builder/context.py | 16 ++++++++ assistant_core/state.py | 10 +++++ tests/conftest.py | 7 +++- tests/unit/test_builder.py | 65 +++++++++++++++++++++++++++++- tests/unit/test_builder_context.py | 22 +++++++++- 8 files changed, 145 insertions(+), 6 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 22ef2f8..6feef44 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,8 @@ ## NEXT * Replace BaseAgentFactory for ContextFactory for pluggable components (agent/graph/model/tools/resolver) +* Add MultiAgent context and director with conditional entry points +* Add MultiAgentState with `active_agent` selector * Add an internal-use notice to datetime builder ## v0.7.0 diff --git a/assistant_core/builder/__init__.py b/assistant_core/builder/__init__.py index 1b1fe4f..4fc54d7 100644 --- a/assistant_core/builder/__init__.py +++ b/assistant_core/builder/__init__.py @@ -1,4 +1,4 @@ -from .agent import AgentBuilder, SingleAgent +from .agent import AgentBuilder, MultiAgent, SingleAgent from .base import BaseBuilder, BaseDirector -from .context import BuilderContext +from .context import BuilderContext, MultiAgentContext from .exceptions import BuilderError diff --git a/assistant_core/builder/agent.py b/assistant_core/builder/agent.py index 134fa8e..5013926 100644 --- a/assistant_core/builder/agent.py +++ b/assistant_core/builder/agent.py @@ -2,7 +2,7 @@ from langgraph.prebuilt import ToolNode, tools_condition from .base import BaseBuilder, BaseDirector -from .context import BuilderContext +from .context import BuilderContext, MultiAgentContext class AgentBuilder(BaseBuilder): @@ -48,3 +48,26 @@ def make(self, context: BuilderContext) -> StateGraph: SingleAgent = SingleAgentDirector + + +class MultiAgentDirector(BaseDirector): + def make(self, context: MultiAgentContext) -> StateGraph: + """ + Execute the build process for multiple agents. + """ + for builder in self.builders: + builder.build(context) + + if "default" not in context.entrypoint_mapping: + context.entrypoint_mapping["default"] = context.entrypoint + + # Use conditional entry depending on state["active_agent"] + context.graph_builder.set_conditional_entry_point( + context.conditional_entrypoint_factory() + ) + AgentBuilder().build(context) + + return context.graph_builder + + +MultiAgent = MultiAgentDirector diff --git a/assistant_core/builder/context.py b/assistant_core/builder/context.py index 8d81372..c317e05 100644 --- a/assistant_core/builder/context.py +++ b/assistant_core/builder/context.py @@ -94,3 +94,19 @@ def clone( ) return self.__class__(new_factory) + + +class MultiAgentContext(BuilderContext): + """Context for multi-agent workflows with conditional entrypoint selection.""" + + def __init__(self, context_factory: ContextFactory): + super().__init__(context_factory) + self.entrypoint_mapping: dict[str, str] = {} + + def conditional_entrypoint_factory(self): + def entrypoint_selector(state: dict): + return self.entrypoint_mapping.get( + state.get("active_agent", "default"), self.entrypoint + ) + + return entrypoint_selector diff --git a/assistant_core/state.py b/assistant_core/state.py index 9dddec2..3288f7d 100644 --- a/assistant_core/state.py +++ b/assistant_core/state.py @@ -14,3 +14,13 @@ class QuestionState(MessagesState): class NextProcessState(MessagesState): next: Annotated[str | None, replace] = None + + +class MultiAgentState(MessagesState): + """State for multi-agent workflows. + + When building graphs with conditional entry points, the value of + `active_agent` selects which entrypoint to use. + """ + + active_agent: str | None = None diff --git a/tests/conftest.py b/tests/conftest.py index ae2ac97..98a2eac 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -3,7 +3,7 @@ import pytest from langchain_core.language_models.chat_models import BaseChatModel -from assistant_core.builder import BuilderContext +from assistant_core.builder import BuilderContext, MultiAgentContext from assistant_core.factories import ContextFactory from assistant_core.nodes import AgentNode @@ -44,3 +44,8 @@ def context_factory(mock_model, agent_node, factory_config): @pytest.fixture def builder_context(context_factory): return BuilderContext(context_factory=context_factory) + + +@pytest.fixture +def multi_agent_context(context_factory): + return MultiAgentContext(context_factory=context_factory) diff --git a/tests/unit/test_builder.py b/tests/unit/test_builder.py index 63747cd..d8b6b6e 100644 --- a/tests/unit/test_builder.py +++ b/tests/unit/test_builder.py @@ -1,4 +1,4 @@ -from assistant_core.builder import BaseBuilder, SingleAgent +from assistant_core.builder import BaseBuilder, MultiAgent, SingleAgent from assistant_core.nodes import PromptNode @@ -45,3 +45,66 @@ def build(self, context): ("custom_node_1", "test_agent"), ("tools_test_agent", "resolver"), } == workflow.edges + + +def test_multi_agent_entrypoint_chain_and_default_mapping(multi_agent_context): + class CustomBuilderWithEntrypoint(BaseBuilder): + def __init__(self, custom_node_name: str): + self.custom_node_name = custom_node_name + + def build(self, context): + node = PromptNode( + name=self.custom_node_name, + prompt=f"Testing the custom node {self.custom_node_name}", + ) + context.graph_builder.add_node(node.name, node) + # Push this node as the new entrypoint (LIFO), linking to previous + context.entrypoint = node.name + + ctx = multi_agent_context + + agent = MultiAgent() + agent.add_builder(CustomBuilderWithEntrypoint("custom_node_1")) + agent.add_builder(CustomBuilderWithEntrypoint("custom_node_2")) + agent.add_builder(CustomBuilderWithEntrypoint("custom_node_3")) + + workflow = agent.make(ctx) + assert workflow is not None + + # Default mapping should be set to last entrypoint when not provided + assert ctx.entrypoint_mapping.get("default") == "custom_node_3" + + # Edges should reflect the LIFO chain and tools -> resolver wiring + expected_edges = { + ("custom_node_3", "custom_node_2"), + ("custom_node_2", "custom_node_1"), + ("custom_node_1", "test_agent"), + ("tools_test_agent", "resolver"), + } + assert workflow.edges == expected_edges + + +def test_multi_agent_preserves_existing_default_mapping(multi_agent_context): + class SimpleEntrypointBuilder(BaseBuilder): + def build(self, context): + node = PromptNode(name="ep_node", prompt="n") + context.graph_builder.add_node(node.name, node) + context.entrypoint = node.name + context.entrypoint_mapping["ep_agent"] = node.name + + ctx = multi_agent_context + # Predefine default mapping; director should not override it + ctx.entrypoint_mapping["default"] = "predefined_node" + + agent = MultiAgent() + agent.add_builder(SimpleEntrypointBuilder()) + + workflow = agent.make(ctx) + assert workflow is not None + + # Ensure mapping wasn't overwritten + assert ctx.entrypoint_mapping["default"] == "predefined_node" + assert ctx.entrypoint_mapping["ep_agent"] == "ep_node" + # LIFO edge and tools wiring should still be present + assert ("ep_node", "test_agent") in workflow.edges + assert ("tools_test_agent", "resolver") in workflow.edges diff --git a/tests/unit/test_builder_context.py b/tests/unit/test_builder_context.py index 7419a8b..0d7cc2e 100644 --- a/tests/unit/test_builder_context.py +++ b/tests/unit/test_builder_context.py @@ -4,7 +4,7 @@ from langchain_core.language_models.chat_models import BaseChatModel from langgraph.graph import MessagesState, StateGraph -from assistant_core.builder.context import BuilderContext +from assistant_core.builder.context import BuilderContext, MultiAgentContext from assistant_core.nodes import AgentNode, PromptNode, ResolverNode @@ -153,3 +153,23 @@ def test_init_with_deprecated_agent_factory_param(context_factory): ), "Expected a DeprecationWarning when using agent_factory param" assert ctx.agent_node.name == "test_agent" assert ctx.entrypoint == "test_agent" + + +def test_multi_agent_context_selector_behavior(context_factory): + ctx = MultiAgentContext(context_factory) + + # Default entrypoint is agent name + assert ctx.entrypoint == "test_agent" + + # Map some agents to custom entrypoints + ctx.entrypoint_mapping["a1"] = "node_a1" + ctx.entrypoint_mapping["default"] = "node_default" + + selector = ctx.conditional_entrypoint_factory() + + # When no active_agent present, use mapping['default'] + assert selector({}) == "node_default" + # Known active_agent → mapped entrypoint + assert selector({"active_agent": "a1"}) == "node_a1" + # Unknown active_agent → fallback to context.entrypoint (agent node) + assert selector({"active_agent": "unknown"}) == "test_agent" From af1485533cdc75214e4fbcf5e41542dc1052741a Mon Sep 17 00:00:00 2001 From: Andres Javier Lopez Date: Tue, 4 Nov 2025 09:25:59 -0600 Subject: [PATCH 6/6] Add missing docstrings --- assistant_core/builder/context.py | 13 +++++++++---- assistant_core/factories.py | 1 + 2 files changed, 10 insertions(+), 4 deletions(-) diff --git a/assistant_core/builder/context.py b/assistant_core/builder/context.py index c317e05..c1e955a 100644 --- a/assistant_core/builder/context.py +++ b/assistant_core/builder/context.py @@ -1,5 +1,5 @@ import warnings -from typing import Self +from typing import Callable, Self from assistant_core.factories import ( AgentFactory, @@ -65,6 +65,7 @@ def create( resolver_factory: ResolverFactory = None, base_tools_factory: BaseToolsFactory = None, ) -> Self: + """Create a BuilderContext with a custom ContextFactory.""" factory = ContextFactory( config, agent_factory=agent_factory, @@ -84,7 +85,7 @@ def clone( resolver_factory: ResolverFactory = None, base_tools_factory: BaseToolsFactory = None, ) -> Self: - + """Clone the context with optional overrides for component factories.""" new_factory = self._factory.clone( agent_factory=agent_factory, graph_factory=graph_factory, @@ -103,8 +104,12 @@ def __init__(self, context_factory: ContextFactory): super().__init__(context_factory) self.entrypoint_mapping: dict[str, str] = {} - def conditional_entrypoint_factory(self): - def entrypoint_selector(state: dict): + def conditional_entrypoint_factory(self) -> Callable[[dict], str]: + """ + Return a selector callable that chooses an entrypoint based on state. + """ + + def entrypoint_selector(state: dict) -> str: return self.entrypoint_mapping.get( state.get("active_agent", "default"), self.entrypoint ) diff --git a/assistant_core/factories.py b/assistant_core/factories.py index 46a9a11..5ba065d 100644 --- a/assistant_core/factories.py +++ b/assistant_core/factories.py @@ -46,6 +46,7 @@ def clone( resolver_factory: ResolverFactory = None, base_tools_factory: BaseToolsFactory = None, ) -> Self: + """Clone the factory with optional overrides for component factories.""" agent_factory = agent_factory or self._agent_factory graph_factory = graph_factory or self._graph_factory