1919from google .genai import types
2020from google .genai .types import Blob
2121
22+ from veadk .a2a .remote_ve_agent import RemoteVeAgent
2223from veadk .agent import Agent
2324from veadk .evaluation import EvalSetRecorder
2425from veadk .memory .short_term_memory import ShortTermMemory
4142class Runner :
4243 def __init__ (
4344 self ,
44- agent : Agent ,
45+ agent : Agent | RemoteVeAgent ,
4546 short_term_memory : ShortTermMemory ,
4647 app_name : str = "veadk_default_app" ,
4748 user_id : str = "veadk_default_user" ,
@@ -56,7 +57,10 @@ def __init__(
5657 self .short_term_memory = short_term_memory
5758 self .session_service = short_term_memory .session_service
5859
59- self .long_term_memory = self .agent .long_term_memory
60+ if isinstance (self .agent , Agent ):
61+ self .long_term_memory = self .agent .long_term_memory
62+ else :
63+ self .long_term_memory = None
6064
6165 # maintain a in-memory runner for fast inference
6266 self .runner = ADKRunner (
@@ -67,13 +71,6 @@ def __init__(
6771 )
6872
6973 def _convert_messages (self , messages ) -> list :
70- # RunnerMessage = Union[
71- # str, # single turn text-based prompt
72- # list[str], # multiple turn text-based prompt
73- # dict, # single turn prompt with media
74- # list[dict], # multiple turn prompt with media
75- # list[dict | str], # multiple turn prompt with media and text-based prompt
76- # ]
7774 if isinstance (messages , str ):
7875 messages = [types .Content (role = "user" , parts = [types .Part (text = messages )])]
7976 elif isinstance (messages , MediaMessage ):
@@ -159,7 +156,8 @@ async def run(
159156 final_output = await self ._run (session_id , message , stream )
160157
161158 # try to save tracing file
162- self .save_tracing_file (session_id )
159+ if isinstance (self .agent , Agent ):
160+ self .save_tracing_file (session_id )
163161
164162 return final_output
165163
0 commit comments