From 0a05a5709ee35cb94152f5f0f82314ab372e4548 Mon Sep 17 00:00:00 2001 From: Debanjum Date: Tue, 15 Jul 2025 16:21:06 -0700 Subject: [PATCH] Use agent chat model to generate code instead of default chat model This is consistent with chat model preference order for other tools --- src/khoj/processor/tools/run_code.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/src/khoj/processor/tools/run_code.py b/src/khoj/processor/tools/run_code.py index 8232c6ca..448ed699 100644 --- a/src/khoj/processor/tools/run_code.py +++ b/src/khoj/processor/tools/run_code.py @@ -19,7 +19,7 @@ from tenacity import ( wait_random_exponential, ) -from khoj.database.adapters import FileObjectAdapters +from khoj.database.adapters import AgentAdapters, FileObjectAdapters from khoj.database.models import Agent, ChatMessageModel, FileObject, KhojUser from khoj.processor.conversation import prompts from khoj.processor.conversation.utils import ( @@ -152,12 +152,15 @@ async def generate_python_code( personality_context=personality_context, ) + agent_chat_model = AgentAdapters.get_agent_chat_model(agent, user) if agent else None + response = await send_message_to_model_wrapper( code_generation_prompt, query_images=query_images, - user=user, - tracer=tracer, query_files=query_files, + user=user, + agent_chat_model=agent_chat_model, + tracer=tracer, ) # Extract python code wrapped in markdown code blocks from the response