import time
from google.adk.agents import Agent
from google.adk.runners import Runner
from google.adk.tools import LongRunningFunctionTool
from google.adk.sessions import InMemorySessionService
from google.genai import types
from google.adk.models.lite_llm import LiteLlm


# 1. Define the generator function
def process_large_file(file_path: str) -> dict:
    """
    Simulates processing a large file, yielding progress updates.

    Args:
      file_path: Path to the file being processed.

    Returns: 
      A final status dictionary.
    """
    total_steps = 5

    # This dict will be sent in the first FunctionResponse
    yield {"status": "pending", "message": f"Starting processing for {file_path}..."}

    for i in range(total_steps):
        time.sleep(1)  # Simulate work for one step
        progress = (i + 1) / total_steps
        # Each yielded dict is sent in a subsequent FunctionResponse
        yield {
            "status": "pending",
            "progress": f"{int(progress * 100)}%",
            "estimated_completion_time": f"~{total_steps - (i + 1)} seconds remaining"
        }

    # This returned dict will be sent in the final FunctionResponse
    return {"status": "completed", "result": f"Successfully processed file: {file_path}"}


# 2. Wrap the function with LongRunningFunctionTool
long_running_tool = LongRunningFunctionTool(func=process_large_file)

# 3. Use the tool in an Agent
file_processor_agent = Agent(
    # Use a model compatible with function calling
    model=LiteLlm(name="deepseek-v3", model="openai/deepseek-chat", api_base="https://api.deepseek.com",
                  api_key="sk-xxx"),
    name='file_processor_agent',
    instruction="""You are an agent that processes large files. When the user provides a file path, use the 'process_large_file' tool. Keep the user informed about the progress based on the tool's updates (which arrive as function responses). Only provide the final result when the tool indicates completion in its final function response.""",
    tools=[long_running_tool]
)

APP_NAME = "file_processor"
USER_ID = "1234"
SESSION_ID = "session1234"

# Session and Runner
session_service = InMemorySessionService()
session = session_service.create_session(app_name=APP_NAME, user_id=USER_ID, session_id=SESSION_ID)
runner = Runner(agent=file_processor_agent, app_name=APP_NAME, session_service=session_service)


# Agent Interaction
def call_agent(query):
    content = types.Content(role='user', parts=[types.Part(text=query)])
    events = runner.run(user_id=USER_ID, session_id=SESSION_ID, new_message=content)

    for event in events:
        print(event)
        if event.is_final_response():
            final_response = event.content.parts[0].text
            print("Agent Response: ", final_response)


call_agent("/Users/projects/demo//browser.py")

content=Content(parts=[Part(video_metadata=None, thought=None, code_execution_result=None, executable_code=None, file_data=None, function_call=FunctionCall(id='call_0_8d5face4-4963-42e6-9236-091a8b4bab71', args={'file_path': '/Users/projects/demo//browser.py'}, name='process_large_file'), function_response=None, inline_data=None, text=None)], role='model') grounding_metadata=None partial=False turn_complete=None error_code=None error_message=None interrupted=None invocation_id='e-dd2c6b0a-d014-4abb-8826-20a98c696f49' author='file_processor_agent' actions=EventActions(skip_summarization=None, state_delta={}, artifact_delta={}, transfer_to_agent=None, escalate=None, requested_auth_configs={}) long_running_tool_ids={'call_0_8d5face4-4963-42e6-9236-091a8b4bab71'} branch=None id='vUtUSIps' timestamp=1744341855.652638
Agent Response: None
Exception in thread Thread-1 (_asyncio_thread_main):
content=Content(parts=[Part(video_metadata=None, thought=None, code_execution_result=None, executable_code=None, file_data=None, function_call=None, function_response=FunctionResponse(id='call_0_8d5face4-4963-42e6-9236-091a8b4bab71', name='process_large_file', response={'result': <generator object process_large_file at 0x12c2bb340>}), inline_data=None, text=None)], role='user') grounding_metadata=None partial=None turn_complete=None error_code=None error_message=None interrupted=None invocation_id='e-dd2c6b0a-d014-4abb-8826-20a98c696f49' author='file_processor_agent' actions=EventActions(skip_summarization=None, state_delta={}, artifact_delta={}, transfer_to_agent=None, escalate=None, requested_auth_configs={}) long_running_tool_ids=None branch=None id='inD1Z1Va' timestamp=1744341861.393286
Traceback (most recent call last):
File "/opt/homebrew/Cellar/[email protected]/3.13.2/Frameworks/Python.framework/Versions/3.13/lib/python3.13/threading.py", line 1041, in _bootstrap_inner
self.run()
~~~~~~~~^^
File "/opt/homebrew/Cellar/[email protected]/3.13.2/Frameworks/Python.framework/Versions/3.13/lib/python3.13/threading.py", line 992, in run
self._target(*self._args, **self._kwargs)
~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/Users/zhen/projects/demo/mcp-agent-auth/.venv/lib/python3.13/site-packages/google/adk/runners.py", line 138, in _asyncio_thread_main
asyncio.run(_invoke_run_async())
~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^
File "/opt/homebrew/Cellar/[email protected]/3.13.2/Frameworks/Python.framework/Versions/3.13/lib/python3.13/asyncio/runners.py", line 195, in run
return runner.run(main)
~~~~~~~~~~^^^^^^
File "/opt/homebrew/Cellar/[email protected]/3.13.2/Frameworks/Python.framework/Versions/3.13/lib/python3.13/asyncio/runners.py", line 118, in run
return self._loop.run_until_complete(task)
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~^^^^^^
File "/opt/homebrew/Cellar/[email protected]/3.13.2/Frameworks/Python.framework/Versions/3.13/lib/python3.13/asyncio/base_events.py", line 725, in run_until_complete
return future.result()
~~~~~~~~~~~~~^^
File "/Users/zhen/projects/demo/mcp-agent-auth/.venv/lib/python3.13/site-packages/google/adk/runners.py", line 126, in _invoke_run_async
async for event in self.run_async(
...<5 lines>...
event_queue.put(event)
File "/Users/zhen/projects/demo/mcp-agent-auth/.venv/lib/python3.13/site-packages/google/adk/runners.py", line 197, in run_async
async for event in invocation_context.agent.run_async(invocation_context):
...<2 lines>...
yield event
File "/Users/zhen/projects/demo/mcp-agent-auth/.venv/lib/python3.13/site-packages/google/adk/agents/base_agent.py", line 141, in run_async
async for event in self._run_async_impl(ctx):
yield event
File "/Users/zhen/projects/demo/mcp-agent-auth/.venv/lib/python3.13/site-packages/google/adk/agents/llm_agent.py", line 232, in _run_async_impl
async for event in self._llm_flow.run_async(ctx):
self.__maybe_save_output_to_state(event)
yield event
File "/Users/zhen/projects/demo/mcp-agent-auth/.venv/lib/python3.13/site-packages/google/adk/flows/llm_flows/base_llm_flow.py", line 231, in run_async
async for event in self._run_one_step_async(invocation_context):
last_event = event
yield event
File "/Users/zhen/projects/demo/mcp-agent-auth/.venv/lib/python3.13/site-packages/google/adk/flows/llm_flows/base_llm_flow.py", line 245, in _run_one_step_async
async for event in self._preprocess_async(invocation_context, llm_request):
yield event
File "/Users/zhen/projects/demo/mcp-agent-auth/.venv/lib/python3.13/site-packages/google/adk/flows/llm_flows/base_llm_flow.py", line 277, in _preprocess_async
async for event in processor.run_async(invocation_context, llm_request):
yield event
File "/Users/zhen/projects/demo/mcp-agent-auth/.venv/lib/python3.13/site-packages/google/adk/flows/llm_flows/contents.py", line 47, in run_async
llm_request.contents = _get_contents(
~~~~~~~~~~~~~^
invocation_context.branch,
^^^^^^^^^^^^^^^^^^^^^^^^^^
invocation_context.session.events,
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
agent.name,
^^^^^^^^^^^
)
^
File "/Users/zhen/projects/demo/mcp-agent-auth/.venv/lib/python3.13/site-packages/google/adk/flows/llm_flows/contents.py", line 229, in _get_contents
content = copy.deepcopy(event.content)
File "/opt/homebrew/Cellar/[email protected]/3.13.2/Frameworks/Python.framework/Versions/3.13/lib/python3.13/copy.py", line 144, in deepcopy
y = copier(memo)
File "/Users/zhen/projects/demo/mcp-agent-auth/.venv/lib/python3.13/site-packages/pydantic/main.py", line 840, in deepcopy
_object_setattr(m, 'dict', deepcopy(self.dict, memo=memo))
~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/opt/homebrew/Cellar/[email protected]/3.13.2/Frameworks/Python.framework/Versions/3.13/lib/python3.13/copy.py", line 137, in deepcopy
y = copier(x, memo)
File "/opt/homebrew/Cellar/[email protected]/3.13.2/Frameworks/Python.framework/Versions/3.13/lib/python3.13/copy.py", line 222, in _deepcopy_dict
y[deepcopy(key, memo)] = deepcopy(value, memo)
~~~~~~~~^^^^^^^^^^^^^
File "/opt/homebrew/Cellar/[email protected]/3.13.2/Frameworks/Python.framework/Versions/3.13/lib/python3.13/copy.py", line 137, in deepcopy
y = copier(x, memo)
File "/opt/homebrew/Cellar/[email protected]/3.13.2/Frameworks/Python.framework/Versions/3.13/lib/python3.13/copy.py", line 197, in _deepcopy_list
append(deepcopy(a, memo))
~~~~~~~~^^^^^^^^^
File "/opt/homebrew/Cellar/[email protected]/3.13.2/Frameworks/Python.framework/Versions/3.13/lib/python3.13/copy.py", line 144, in deepcopy
y = copier(memo)
File "/Users/zhen/projects/demo/mcp-agent-auth/.venv/lib/python3.13/site-packages/pydantic/main.py", line 840, in deepcopy
_object_setattr(m, 'dict', deepcopy(self.dict, memo=memo))
~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/opt/homebrew/Cellar/[email protected]/3.13.2/Frameworks/Python.framework/Versions/3.13/lib/python3.13/copy.py", line 137, in deepcopy
y = copier(x, memo)
File "/opt/homebrew/Cellar/[email protected]/3.13.2/Frameworks/Python.framework/Versions/3.13/lib/python3.13/copy.py", line 222, in _deepcopy_dict
y[deepcopy(key, memo)] = deepcopy(value, memo)
~~~~~~~~^^^^^^^^^^^^^
File "/opt/homebrew/Cellar/[email protected]/3.13.2/Frameworks/Python.framework/Versions/3.13/lib/python3.13/copy.py", line 144, in deepcopy
y = copier(memo)
File "/Users/zhen/projects/demo/mcp-agent-auth/.venv/lib/python3.13/site-packages/pydantic/main.py", line 840, in deepcopy
_object_setattr(m, 'dict', deepcopy(self.dict, memo=memo))
~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/opt/homebrew/Cellar/[email protected]/3.13.2/Frameworks/Python.framework/Versions/3.13/lib/python3.13/copy.py", line 137, in deepcopy
y = copier(x, memo)
File "/opt/homebrew/Cellar/[email protected]/3.13.2/Frameworks/Python.framework/Versions/3.13/lib/python3.13/copy.py", line 222, in _deepcopy_dict
y[deepcopy(key, memo)] = deepcopy(value, memo)
~~~~~~~~^^^^^^^^^^^^^
File "/opt/homebrew/Cellar/[email protected]/3.13.2/Frameworks/Python.framework/Versions/3.13/lib/python3.13/copy.py", line 137, in deepcopy
y = copier(x, memo)
File "/opt/homebrew/Cellar/[email protected]/3.13.2/Frameworks/Python.framework/Versions/3.13/lib/python3.13/copy.py", line 222, in _deepcopy_dict
y[deepcopy(key, memo)] = deepcopy(value, memo)
~~~~~~~~^^^^^^^^^^^^^
File "/opt/homebrew/Cellar/[email protected]/3.13.2/Frameworks/Python.framework/Versions/3.13/lib/python3.13/copy.py", line 152, in deepcopy
rv = reductor(4)
TypeError: cannot pickle 'generator' object