feat: megamind context providers (#687)

Author: jmatejcz

src/rai_core/rai/agents/langchain/core/__init__.py

@@ -14,7 +14,12 @@
 
 from .conversational_agent import State as ConversationalAgentState
 from .conversational_agent import create_conversational_agent
-from .megamind import Executor, create_megamind, get_initial_megamind_state
+from .megamind import (
+    ContextProvider,
+    Executor,
+    create_megamind,
+    get_initial_megamind_state,
+)
 from .react_agent import (
     ReActAgentState,
     create_react_runnable,
@@ -23,6 +28,7 @@
 from .tool_runner import SubAgentToolRunner, ToolRunner
 
 __all__ = [
+    "ContextProvider",
     "ConversationalAgentState",
     "Executor",
     "ReActAgentState",

src/rai_core/rai/agents/langchain/core/megamind.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 ### NOTE (jmatejcz) this agent is still in process of testing and refining
+from abc import ABC, abstractmethod
 from dataclasses import dataclass
 from functools import partial
 from typing import (
@@ -185,6 +186,14 @@ class Executor:
     system_prompt: str
 
 
+class ContextProvider(ABC):
+    """Context provider are meant to inject exteral info to megamind prompt"""
+
+    @abstractmethod
+    def get_context(self) -> str:
+        pass
+
+
 def get_initial_megamind_state(task: str):
     return MegamindState(
         {
@@ -198,7 +207,11 @@ def get_initial_megamind_state(task: str):
     )
 
 
-def plan_step(megamind_agent: BaseChatModel, state: MegamindState) -> MegamindState:
+def plan_step(
+    megamind_agent: BaseChatModel,
+    state: MegamindState,
+    context_providers: Optional[List[ContextProvider]] = None,
+) -> MegamindState:
     """Initial planning step."""
     if "original_task" not in state:
         state["original_task"] = state["messages"][0].content[0]["text"]
@@ -208,6 +221,9 @@ def plan_step(megamind_agent: BaseChatModel, state: MegamindState) -> MegamindSt
         state["step"] = None
 
     megamind_prompt = f"You are given objective to complete: {state['original_task']}"
+    for provider in context_providers:
+        megamind_prompt += provider.get_context()
+        megamind_prompt += "\n"
     if state["steps_done"]:
         megamind_prompt += "\n\n"
         megamind_prompt += "Steps that were already done successfully:\n"
@@ -244,17 +260,27 @@ def plan_step(megamind_agent: BaseChatModel, state: MegamindState) -> MegamindSt
 
 def create_megamind(
     megamind_llm: BaseChatModel,
-    megamind_system_prompt: str,
     executors: List[Executor],
+    megamind_system_prompt: Optional[str] = None,
     task_planning_prompt: Optional[str] = None,
+    context_providers: List[ContextProvider] = [],
 ) -> CompiledStateGraph:
     """Create a megamind langchain agent
 
     Args:
         executors (List[Executor]): Subagents for megamind, each can be a specialist with
         its own tools llm and system prompt
+
+        megamind_system_prompt (Optional[str]): Prompt for megamind node. If not provided
+        it will default to informing agent of the avaialble executors and listing their tools.
+
         task_planning_prompt (Optional[str]): Prompt that helps summarize the step in a way
         that helps planning task
+
+        context_providers (List[ContextProvider]): Each ContextProvider can inject external info
+        to prompt during planning phase
+
+
     """
     executor_agents = {}
     handoff_tools = []
@@ -295,7 +321,8 @@ def create_megamind(
     )
 
     graph = StateGraph(MegamindState).add_node(
-        "megamind", partial(plan_step, megamind_agent)
+        "megamind",
+        partial(plan_step, megamind_agent, context_providers=context_providers),
     )
     for agent_name, agent in executor_agents.items():
         graph.add_node(agent_name, agent)