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. diff --git a/CHANGELOG.md b/CHANGELOG.md index 48c86a7..6feef44 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,6 +1,10 @@ # CHANGELOG ## 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 * Update dependencies 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 5c1b0ba..c1e955a 100644 --- a/assistant_core/builder/context.py +++ b/assistant_core/builder/context.py @@ -1,13 +1,41 @@ -from assistant_core.factories import BaseAgentFactory +import warnings +from typing import Callable, Self + +from assistant_core.factories import ( + AgentFactory, + BaseToolsFactory, + ContextFactory, + GraphFactory, + ModelFactory, + ResolverFactory, +) 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 + + # 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() + self.resolver_node = context_factory.create_resolver_node() + self.tools = context_factory.create_base_tools() self._entrypoint = None @@ -25,3 +53,65 @@ 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: + """Create a BuilderContext with a custom ContextFactory.""" + 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: + """Clone the context with optional overrides for component factories.""" + 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) + + +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) -> 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 + ) + + return entrypoint_selector 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/assistant_core/factories.py b/assistant_core/factories.py index 30cc660..5ba065d 100644 --- a/assistant_core/factories.py +++ b/assistant_core/factories.py @@ -1,23 +1,72 @@ -from abc import ABC, abstractmethod -from typing import TypedDict +import warnings +from typing import Callable, Self, 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 + + 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: + """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 + 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): + def model(self) -> BaseChatModel: """ Get the model instance. If not already created, it will create one using the factory method. @@ -26,34 +75,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/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 6640f15..98a2eac 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.builder import BuilderContext, MultiAgentContext +from assistant_core.factories import ContextFactory from assistant_core.nodes import AgentNode @@ -30,24 +29,23 @@ 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 context_factory(mock_model, agent_node, factory_config): + return ContextFactory( + factory_config, + agent_factory=lambda _: agent_node, + model_factory=lambda _: mock_model, + ) - def create_agent_node(self): - return agent_node - return MockAgentFactory(factory_config) +@pytest.fixture +def builder_context(context_factory): + return BuilderContext(context_factory=context_factory) @pytest.fixture -def builder_context(agent_factory): - return BuilderContext(agent_factory=agent_factory) +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 new file mode 100644 index 0000000..0d7cc2e --- /dev/null +++ b/tests/unit/test_builder_context.py @@ -0,0 +1,175 @@ +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, MultiAgentContext +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" + + +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" 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