Conversation
📝 WalkthroughWalkthroughThe PR refactors the chat model's tool execution from LangChain's Changes
Estimated code review effort🎯 4 (Complex) | ⏱️ ~75 minutes Poem
🚥 Pre-merge checks | ✅ 1 | ❌ 2❌ Failed checks (1 warning, 1 inconclusive)
✅ Passed checks (1 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
There was a problem hiding this comment.
Pull request overview
This PR attempts to enable/repair web search tool-calling by replacing the previous LangChain agent construction with an explicit tool-calling loop over ChatBedrock.bind_tools(...), and adjusts an integration script to import Django-dependent modules only after django.setup().
Changes:
- Migrate web-search execution from
create_agent(...)to a manual tool-call loop usingbind_tools+ToolMessage. - Add
ToolMessagehandling and new logic for extracting the final AI response from the message sequence. - Reorder imports in
back/test_integration.pyto occur after Django initialization.
Reviewed changes
Copilot reviewed 2 out of 2 changed files in this pull request and generated 3 comments.
| File | Description |
|---|---|
| back/bots/models/chat.py | Reworks the web-search “agent” execution to use Bedrock tool binding and a manual tool loop. |
| back/test_integration.py | Moves Django-dependent imports to after django.setup() for safer script execution. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| # Agentagent loop - keep invoking until no more tool calls | ||
| max_iterations = 5 | ||
| iteration = 0 | ||
|
|
||
| while iteration < max_iterations: | ||
| iteration += 1 | ||
| logger.info(f"🤖 AGENT_LOOP_ITERATION: {iteration}") | ||
|
|
||
| # Call the model | ||
| response = model_with_tools.invoke(messages) | ||
| messages.append(response) | ||
|
|
||
| # Check if model wants to call a tool | ||
| if not response.tool_calls: | ||
| logger.info(f"🤖 AGENT_LOOP_COMPLETE: no more tool calls after {iteration} iterations") | ||
| break |
There was a problem hiding this comment.
The manual tool loop stops after max_iterations, but if the model is still returning tool calls at that point there may be no final AIMessage without tool_calls. In that case response_text falls back to the last AIMessage which is likely a tool-call message (often empty), resulting in a blank or incorrect assistant reply. Handle the max-iteration exit explicitly (e.g., raise/log an error and return a safe message, or force one final model call after the last ToolMessage) so callers always get a real final response.
| # Build and run the agent loop manually | ||
| messages = [SystemMessage(content=self.get_system_message()), HumanMessage(content=agent_input)] | ||
|
|
||
| # Agentagent loop - keep invoking until no more tool calls |
There was a problem hiding this comment.
Typo in comment: # Agentagent loop should be corrected (e.g., # Agent loop).
| # Agentagent loop - keep invoking until no more tool calls | |
| # Agent loop - keep invoking until no more tool calls |
| from langchain_aws import ChatBedrock | ||
| from langchain_core.messages import HumanMessage, SystemMessage, AIMessage | ||
| from langchain_core.messages import HumanMessage, SystemMessage, AIMessage, ToolMessage | ||
| from langchain_core.tools import tool | ||
| from langchain_core.callbacks.base import BaseCallbackHandler | ||
| from langchain.agents import create_agent | ||
| from tavily import TavilyClient |
There was a problem hiding this comment.
create_agent is removed in this PR, but the existing unit test back/bots/tests/test_chat.py patches bots.models.chat.create_agent. With this change, that test will start failing (patch target no longer exists). Update the tests to patch/mock the new bind_tools(...).invoke(...) flow instead.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 4
🧹 Nitpick comments (1)
back/bots/models/chat.py (1)
118-211: Extract the tool orchestration into helper methods.This branch now combines model setup, tool execution, loop control, response parsing, and persistence in one place. Pulling those pieces apart would make the web-search path much easier to test and reason about.
As per coding guidelines,
back/bots/**/*.py: "Keep methods focused and concise, under 10 lines when possible, extracting complex logic into separate methods".🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@back/bots/models/chat.py` around lines 118 - 211, The current method mixes model setup, loop control, tool execution and response parsing; split it into focused helpers: implement _create_chat_with_tools(self, web_search) to construct ChatBedrock(model_id=...) and call bind_tools, _run_agent_loop(self, model_with_tools, messages, max_iterations=5) to perform the while loop and return the final messages list, _process_tool_call(self, tool_call, web_search) to execute a single tool call (handling "web_search" via web_search.invoke and returning a ToolMessage), and _extract_final_response(self, messages) to locate the last AIMessage without tool_calls and return (response_text, usage_metadata); update the main flow to call _extract_agent_input, _create_chat_with_tools, build initial messages, call _run_agent_loop which uses _process_tool_call internally, then call _extract_final_response to populate response_text and usage_metadata.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@back/bots/models/chat.py`:
- Around line 143-144: The code calls model_with_tools.invoke(messages) which
may append multiple AIMessage responses but only captures usage_metadata from
the last AIMessage, causing under-reported token usage; update the handling
after each invoke to accumulate token counts from every AIMessage's
usage_metadata (e.g., sum input_tokens/output_tokens from each AIMessage in
messages or from the returned response objects) after each call to
model_with_tools.invoke and before appending additional messages so that total
input/output/token usage reflects all tool-enabled invocations; apply the same
accumulation fix where model_with_tools.invoke is used (including the block
around lines 176-197) and ensure you update whatever counters or usage fields
(input_tokens, output_tokens, total_tokens, usage_metadata) used by
billing/quota logic.
- Around line 118-123: The branch that enables tools incorrectly constructs a
new ChatBedrock instance (ChatBedrock(model_id=...)) instead of reusing the
injected client self.ai, causing loss of injected callbacks and differing
behavior; change the code to use self.ai (which already implements bind_tools)
when creating model_with_tools (e.g., call self.ai.bind_tools([web_search]) or
equivalent) so the tool-enabled path reuses the same client and preserves the
seams used by get_response_standard and tests.
- Around line 135-149: The loop can exit on max_iterations while the last
response still contains a tool_call stub, so after the while using
max_iterations/iteration and the last response from model_with_tools.invoke,
detect if iteration >= max_iterations and response.tool_calls is truthy; in that
case produce and append a proper terminal assistant message (e.g., an AIMessage
summarizing that the agent could not finish the task or asking to try again) or
re-invoke model_with_tools.invoke with instructions to return a final assistant
answer (no tool_calls) before persisting — update the code paths around the loop
and the fallback persistence (the messages list, response variable, and the code
that saves the last AIMessage) to ensure a non-tool_call terminal assistant
message is saved and returned.
- Around line 125-133: The code currently discards prior turns by building
messages as only SystemMessage + last HumanMessage (using agent_input); instead,
preserve the full prior conversation/multimodal context by using the original
message list (or the output of the existing get_input/_extract_agent_input flow)
when creating the agent messages: convert message_list (or the previously
assembled messages from get_input) into LLM message objects, prepend the
SystemMessage from get_system_message(), and then append the new
HumanMessage(content=agent_input) rather than replacing history; update the code
around agent_input, _extract_agent_input, get_system_message, and the messages
list construction so prior turns remain when web search is enabled.
---
Nitpick comments:
In `@back/bots/models/chat.py`:
- Around line 118-211: The current method mixes model setup, loop control, tool
execution and response parsing; split it into focused helpers: implement
_create_chat_with_tools(self, web_search) to construct ChatBedrock(model_id=...)
and call bind_tools, _run_agent_loop(self, model_with_tools, messages,
max_iterations=5) to perform the while loop and return the final messages list,
_process_tool_call(self, tool_call, web_search) to execute a single tool call
(handling "web_search" via web_search.invoke and returning a ToolMessage), and
_extract_final_response(self, messages) to locate the last AIMessage without
tool_calls and return (response_text, usage_metadata); update the main flow to
call _extract_agent_input, _create_chat_with_tools, build initial messages, call
_run_agent_loop which uses _process_tool_call internally, then call
_extract_final_response to populate response_text and usage_metadata.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 62382766-50a0-4fee-9fa1-88105f5223a6
📒 Files selected for processing (2)
back/bots/models/chat.pyback/test_integration.py
| # Create chat model with tool binding | ||
| chat_model = ChatBedrock(model_id=self.ai.model_id) | ||
| tools = [web_search] | ||
|
|
||
| # Create modern agent with tool calling support | ||
| # This is the recommended approach per LangChain docs | ||
| agent = create_agent( | ||
| model=chat_model, | ||
| tools=tools, | ||
| system_prompt=self.get_system_message(), | ||
| debug=settings.DEBUG | ||
| ) | ||
| # Bind tools to the model for proper tool calling | ||
| model_with_tools = chat_model.bind_tools(tools) |
There was a problem hiding this comment.
Don't bypass ai in the tool-enabled path.
Line 119 creates a new ChatBedrock instead of reusing self.ai, so this branch ignores the injected client/callbacks used by get_response_standard(). That makes tool-enabled chats behave differently from the non-tool path and leaves this code path without the same seam the current tests were mocking.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@back/bots/models/chat.py` around lines 118 - 123, The branch that enables
tools incorrectly constructs a new ChatBedrock instance
(ChatBedrock(model_id=...)) instead of reusing the injected client self.ai,
causing loss of injected callbacks and differing behavior; change the code to
use self.ai (which already implements bind_tools) when creating model_with_tools
(e.g., call self.ai.bind_tools([web_search]) or equivalent) so the tool-enabled
path reuses the same client and preserves the seams used by
get_response_standard and tests.
| # Extract text input from message_list for agent | ||
| agent_input = self._extract_agent_input(message_list) | ||
|
|
||
| logger.info(f"Invoking agent with input: {agent_input[:100]}...") | ||
| logger.info("🤖 AGENT_INVOKE_START: web_search tool available") | ||
|
|
||
| # Invoke agent - the CompiledStateGraph handles tool loop internally | ||
| response = agent.invoke({"messages": [HumanMessage(content=agent_input)]}) | ||
| # Build and run the agent loop manually | ||
| messages = [SystemMessage(content=self.get_system_message()), HumanMessage(content=agent_input)] | ||
|
|
There was a problem hiding this comment.
Preserve prior turns when web search is enabled.
Lines 125-132 throw away the history that get_input() just assembled and restart the model with only SystemMessage + last HumanMessage. Any follow-up question, and any multimodal context already present in message_list, is lost as soon as enable_web_search is on.
🧩 Minimal fix
- # Build and run the agent loop manually
- messages = [SystemMessage(content=self.get_system_message()), HumanMessage(content=agent_input)]
+ # Build and run the agent loop manually with the existing chat history
+ messages = list(message_list)🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@back/bots/models/chat.py` around lines 125 - 133, The code currently discards
prior turns by building messages as only SystemMessage + last HumanMessage
(using agent_input); instead, preserve the full prior conversation/multimodal
context by using the original message list (or the output of the existing
get_input/_extract_agent_input flow) when creating the agent messages: convert
message_list (or the previously assembled messages from get_input) into LLM
message objects, prepend the SystemMessage from get_system_message(), and then
append the new HumanMessage(content=agent_input) rather than replacing history;
update the code around agent_input, _extract_agent_input, get_system_message,
and the messages list construction so prior turns remain when web search is
enabled.
| max_iterations = 5 | ||
| iteration = 0 | ||
|
|
||
| while iteration < max_iterations: | ||
| iteration += 1 | ||
| logger.info(f"🤖 AGENT_LOOP_ITERATION: {iteration}") | ||
|
|
||
| # Call the model | ||
| response = model_with_tools.invoke(messages) | ||
| messages.append(response) | ||
|
|
||
| # Check if model wants to call a tool | ||
| if not response.tool_calls: | ||
| logger.info(f"🤖 AGENT_LOOP_COMPLETE: no more tool calls after {iteration} iterations") | ||
| break |
There was a problem hiding this comment.
Handle loop exhaustion explicitly.
If the model is still returning tool_calls on the fifth pass, the loop exits without a terminal assistant answer. The fallback at Lines 202-211 then persists the last AIMessage, which in that case is usually an empty/partial tool-call stub, and back/bots/views/get_chat_response.py:82-83 returns that straight to the client.
🛑 Safer fallback
max_iterations = 5
iteration = 0
+ exhausted = True
while iteration < max_iterations:
iteration += 1
logger.info(f"🤖 AGENT_LOOP_ITERATION: {iteration}")
@@
# Check if model wants to call a tool
if not response.tool_calls:
+ exhausted = False
logger.info(f"🤖 AGENT_LOOP_COMPLETE: no more tool calls after {iteration} iterations")
break
@@
- if not response_text:
+ if exhausted:
+ logger.error("🤖 AGENT_LOOP_ABORTED: max iterations reached without a final response")
+ response_text = "I'm sorry, I couldn't complete the web search. Please try again."
+ elif not response_text:
# Fallback: get the last message
for msg in reversed(messages):
if isinstance(msg, AIMessage):Also applies to: 202-211
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@back/bots/models/chat.py` around lines 135 - 149, The loop can exit on
max_iterations while the last response still contains a tool_call stub, so after
the while using max_iterations/iteration and the last response from
model_with_tools.invoke, detect if iteration >= max_iterations and
response.tool_calls is truthy; in that case produce and append a proper terminal
assistant message (e.g., an AIMessage summarizing that the agent could not
finish the task or asking to try again) or re-invoke model_with_tools.invoke
with instructions to return a final assistant answer (no tool_calls) before
persisting — update the code paths around the loop and the fallback persistence
(the messages list, response variable, and the code that saves the last
AIMessage) to ensure a non-tool_call terminal assistant message is saved and
returned.
| response = model_with_tools.invoke(messages) | ||
| messages.append(response) |
There was a problem hiding this comment.
Accumulate token usage for every model invoke.
model_with_tools.invoke(messages) can run multiple times here, but usage_metadata is taken only from the last AIMessage. Tool-enabled chats will under-report tokens whenever a tool is used, which skews any quota/billing logic that reads input_tokens / output_tokens.
📊 Suggested fix
response_text = ""
usage_metadata = {"input_tokens": 0, "output_tokens": 0}
@@
# Call the model
response = model_with_tools.invoke(messages)
+ if getattr(response, "usage_metadata", None):
+ usage_metadata["input_tokens"] += response.usage_metadata.get("input_tokens", 0)
+ usage_metadata["output_tokens"] += response.usage_metadata.get("output_tokens", 0)
messages.append(response)
@@
- # Extract token usage
- if hasattr(msg, 'usage_metadata') and msg.usage_metadata:
- usage_metadata = msg.usage_metadata
-
logger.info(f"🤖 FINAL_RESPONSE: {len(response_text)} chars")
breakAlso applies to: 176-197
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@back/bots/models/chat.py` around lines 143 - 144, The code calls
model_with_tools.invoke(messages) which may append multiple AIMessage responses
but only captures usage_metadata from the last AIMessage, causing under-reported
token usage; update the handling after each invoke to accumulate token counts
from every AIMessage's usage_metadata (e.g., sum input_tokens/output_tokens from
each AIMessage in messages or from the returned response objects) after each
call to model_with_tools.invoke and before appending additional messages so that
total input/output/token usage reflects all tool-enabled invocations; apply the
same accumulation fix where model_with_tools.invoke is used (including the block
around lines 176-197) and ensure you update whatever counters or usage fields
(input_tokens, output_tokens, total_tokens, usage_metadata) used by
billing/quota logic.
Summary by CodeRabbit