Multi agent support

.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.

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

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

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

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

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):

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,
+        )

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

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)