klbr/khoj
1
0
Fork 0
mirror of https://github.com/khoaliber/khoj.git synced 2026-03-10 13:26:13 +00:00
Code Issues Packages Projects Releases Wiki Activity

Add cancellation support to browser operator via asyncio.Event

Browse Source
This commit is contained in:
Debanjum
2025-04-07 21:22:44 +05:30
parent 9f75622346
commit 20f87542e5
3 changed files with 17 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

15
src/khoj/processor/operator/browser_operator.py
View File

@@ -38,6 +38,7 @@ async def operate_browser(
query_images: Optional[List[str]] = None, query_images: Optional[List[str]] = None,
agent: Agent = None, agent: Agent = None,
query_files: str = None, query_files: str = None,
cancellation_event: Optional[asyncio.Event] = None,
tracer: dict = {}, tracer: dict = {},
): ):
response, safety_check_message = None, None response, safety_check_message = None, None
@@ -75,6 +76,7 @@ async def operate_browser(
send_status_func=send_status_func, send_status_func=send_status_func,
user=user, user=user,
agent=agent, agent=agent,
cancellation_event=cancellation_event,
tracer=tracer, tracer=tracer,
): ):
if isinstance(result, dict) and ChatEvent.STATUS in result: if isinstance(result, dict) and ChatEvent.STATUS in result:
@@ -92,6 +94,7 @@ async def operate_browser(
send_status_func=send_status_func, send_status_func=send_status_func,
user=user, user=user,
agent=agent, agent=agent,
cancellation_event=cancellation_event,
tracer=tracer, tracer=tracer,
): ):
if isinstance(result, dict) and ChatEvent.STATUS in result: if isinstance(result, dict) and ChatEvent.STATUS in result:
@@ -142,6 +145,7 @@ async def browser_use_openai(
send_status_func: Optional[Callable] = None, send_status_func: Optional[Callable] = None,
user: KhojUser = None, user: KhojUser = None,
agent: Agent = None, agent: Agent = None,
cancellation_event: Optional[asyncio.Event] = None,
tracer: dict = {}, tracer: dict = {},
): ):
""" """
@@ -212,6 +216,11 @@ async def browser_use_openai(
last_call_id = None last_call_id = None
iterations = 0 iterations = 0
while iterations < max_iterations: while iterations < max_iterations:
# Check for cancellation at the start of each iteration
if cancellation_event and cancellation_event.is_set():
logger.info(f"Browser operator cancelled by client disconnect")
break
iterations += 1 iterations += 1
# Send the screenshot back as a computer_call_output # Send the screenshot back as a computer_call_output
@@ -345,6 +354,7 @@ async def browser_use_anthropic(
send_status_func: Optional[Callable] = None, send_status_func: Optional[Callable] = None,
user: KhojUser = None, user: KhojUser = None,
agent: Agent = None, agent: Agent = None,
cancellation_event: Optional[asyncio.Event] = None,
tracer: dict = {}, tracer: dict = {},
): ):
""" """
@@ -420,6 +430,11 @@ async def browser_use_anthropic(
task_completed = False task_completed = False
iterations = 0 iterations = 0
while iterations < max_iterations: while iterations < max_iterations:
# Check for cancellation at the start of each iteration
if cancellation_event and cancellation_event.is_set():
logger.info(f"Browser operator cancelled by client disconnect")
break
iterations += 1 iterations += 1
# Set up optional thinking parameter (for Claude 3.7 Sonnet) # Set up optional thinking parameter (for Claude 3.7 Sonnet)
thinking = {"type": "disabled"} thinking = {"type": "disabled"}

1
src/khoj/routers/api_chat.py
View File

@@ -1228,6 +1228,7 @@ async def chat(
query_files=attached_file_context, query_files=attached_file_context,
send_status_func=partial(send_event, ChatEvent.STATUS), send_status_func=partial(send_event, ChatEvent.STATUS),
agent=agent, agent=agent,
cancellation_event=cancellation_event,
tracer=tracer, tracer=tracer,
): ):
if isinstance(result, dict) and ChatEvent.STATUS in result: if isinstance(result, dict) and ChatEvent.STATUS in result:

1
src/khoj/routers/research.py
View File

@@ -411,6 +411,7 @@ async def execute_information_collection(
query_images=query_images, query_images=query_images,
agent=agent, agent=agent,
query_files=query_files, query_files=query_files,
cancellation_event=cancellation_event,
tracer=tracer, tracer=tracer,
): ):
if isinstance(result, dict) and ChatEvent.STATUS in result: if isinstance(result, dict) and ChatEvent.STATUS in result: