update flow

Author: yaozheng-fang

veadk/agent.py

@@ -121,7 +121,7 @@ class Agent(LlmAgent):
 
     enable_responses: bool = False
 
-    enable_shadow_agent: bool = False
+    enable_supervisor_flow: bool = False
 
     context_cache_config: Optional[ContextCacheConfig] = None
 
@@ -306,15 +306,15 @@ def _llm_flow(self) -> BaseLlmFlow:
         ):
             from veadk.flows.supervisor_single_flow import SupervisorSingleFlow
 
-            if self.enable_shadow_agent:
+            if self.enable_supervisor_flow:
                 logger.debug(f"Enable supervisor flow for agent: {self.name}")
                 return SupervisorSingleFlow(supervised_agent=self)
             else:
                 return SingleFlow()
         else:
             from veadk.flows.supervisor_auto_flow import SupervisorAutoFlow
 
-            if self.enable_shadow_agent:
+            if self.enable_supervisor_flow:
                 logger.debug(f"Enable supervisor flow for agent: {self.name}")
                 return SupervisorAutoFlow(supervised_agent=self)
             return AutoFlow()

veadk/agents/supervise_agent.py

@@ -1,41 +1,27 @@
 from google.adk.models.llm_request import LlmRequest
 from jinja2 import Template
-from pydantic import BaseModel
 
 from veadk import Agent, Runner
+from veadk.utils.logger import get_logger
 
-
-class SupervisorAgentOutput(BaseModel):
-    advice: str = ""
-    """
-    Advices for the worker agent.
-    For example, suggested function call / actions / responses.
-    """
-
+logger = get_logger(__name__)
 
 instruction = Template("""You are a supervisor of an agent system. The system prompt of worker agent is:
 
 ```system prompt
 {{ system_prompt }}
 ```
                        
-```worker agent tools
-{{ agent_tools }}
-```
-
 You should guide the agent to finish task. If you think the history execution is not correct, you should give your advice to the worker agent. If you think the history execution is correct, you should output an empty string.
-
-Your final response should be in `json` format.
 """)
 
 
 def build_supervisor(supervised_agent: Agent) -> Agent:
     custom_instruction = instruction.render(system_prompt=supervised_agent.instruction)
     agent = Agent(
         name="supervisor",
-        description="",
+        description="A supervisor for agent execution",
         instruction=custom_instruction,
-        output_schema=SupervisorAgentOutput,
     )
 
     return agent
@@ -55,4 +41,11 @@ async def generate_advice(agent: Agent, llm_request: LlmRequest) -> str:
                 if part.function_response:
                     messages += f"{content.role}: {part.function_response}"
 
-    return await runner.run(messages="History trajectory is: " + messages)
+    prompt = (
+        f"Tools of agent is {llm_request.tools_dict}. History trajectory is: "
+        + messages
+    )
+
+    logger.debug(f"Prompt for supervisor: {prompt}")
+
+    return await runner.run(messages=prompt)

veadk/flows/supervisor_auto_flow.py

@@ -10,6 +10,9 @@
 from veadk import Agent
 from veadk.agents.supervise_agent import generate_advice
 from veadk.flows.supervisor_single_flow import SupervisorSingleFlow
+from veadk.utils.logger import get_logger
+
+logger = get_logger(__name__)
 
 
 class SupervisorAutoFlow(SupervisorSingleFlow):
@@ -24,9 +27,17 @@ async def _call_llm_async(
         model_response_event: Event,
     ) -> AsyncGenerator[LlmResponse, None]:
         advice = await generate_advice(self._supervisor, llm_request)
-        print(f"Advice: {advice}")
+        logger.debug(f"Advice from supervisor: {advice}")
+
+        llm_request.contents.append(
+            Content(
+                parts=[Part(text=f"Message from your supervisor: {advice}")],
+                role="user",
+            )
+        )
 
-        llm_request.contents.append(Content(parts=[Part(text=advice)], role="model"))
+        print("====")
+        print(llm_request)
 
         async for llm_response in super()._call_llm_async(
             invocation_context, llm_request, model_response_event