andre/crewAI
1
0
Fork 0
mirror of https://github.com/crewAIInc/crewAI.git synced 2026-05-01 15:22:37 +00:00
Code Issues Actions 6 Packages Projects Releases Wiki Activity

improve logic

Browse Source
This commit is contained in:
Brandon Hancock
2025-03-06 11:58:10 -05:00
parent 0aea27bcf0
commit e8707e15ef
1 changed files with 66 additions and 40 deletions
Download Patch File Download Diff File Expand all files Collapse all files

106
src/crewai/llm.py
View File

@@ -439,6 +439,7 @@ class LLM:
) )
# --- 7) Check for tool calls in the final response # --- 7) Check for tool calls in the final response
tool_calls = None
try: try:
if last_chunk: if last_chunk:
choices = None choices = None
@@ -458,53 +459,30 @@ class LLM:
message = getattr(choice, "message") message = getattr(choice, "message")
if message: if message:
tool_calls = None
if isinstance(message, dict) and "tool_calls" in message: if isinstance(message, dict) and "tool_calls" in message:
tool_calls = message["tool_calls"] tool_calls = message["tool_calls"]
elif hasattr(message, "tool_calls"): elif hasattr(message, "tool_calls"):
tool_calls = getattr(message, "tool_calls") tool_calls = getattr(message, "tool_calls")
if tool_calls:
tool_result = self._handle_tool_call(
tool_calls, available_functions
)
if tool_result is not None:
return tool_result
except Exception as e: except Exception as e:
logging.debug(f"Error checking for tool calls: {e}") logging.debug(f"Error checking for tool calls: {e}")
# --- 8) Log token usage if available in streaming mode # --- 8) If no tool calls or no available functions, return the text response directly
# Safely handle callbacks with usage info if not tool_calls or not available_functions:
if callbacks and len(callbacks) > 0: # Log token usage if available in streaming mode
for callback in callbacks: self._handle_streaming_callbacks(callbacks, usage_info, last_chunk)
if hasattr(callback, "log_success_event"): # Emit completion event and return response
# Use the usage_info we've been tracking self._handle_emit_call_events(full_response, LLMCallType.LLM_CALL)
if not usage_info: return full_response
# Try to get usage from the last chunk if we haven't already
try:
if last_chunk:
if (
isinstance(last_chunk, dict)
and "usage" in last_chunk
):
usage_info = last_chunk["usage"]
elif hasattr(last_chunk, "usage"):
if not isinstance(
getattr(last_chunk, "usage"), type
):
usage_info = getattr(last_chunk, "usage")
except Exception as e:
logging.debug(f"Error extracting usage info: {e}")
if usage_info: # --- 9) Handle tool calls if present
callback.log_success_event( tool_result = self._handle_tool_call(tool_calls, available_functions)
kwargs=params, if tool_result is not None:
response_obj={"usage": usage_info}, return tool_result
start_time=0,
end_time=0,
)
# --- 9) Emit completion event and return response # --- 10) Log token usage if available in streaming mode
self._handle_streaming_callbacks(callbacks, usage_info, last_chunk)
# --- 11) Emit completion event and return response
self._handle_emit_call_events(full_response, LLMCallType.LLM_CALL) self._handle_emit_call_events(full_response, LLMCallType.LLM_CALL)
return full_response return full_response
@@ -522,6 +500,48 @@ class LLM:
) )
raise Exception(f"Failed to get streaming response: {str(e)}") raise Exception(f"Failed to get streaming response: {str(e)}")
def _handle_streaming_callbacks(
self,
callbacks: Optional[List[Any]],
usage_info: Optional[Dict[str, Any]],
last_chunk: Optional[Any],
) -> None:
"""Handle callbacks with usage info for streaming responses.
Args:
callbacks: Optional list of callback functions
usage_info: Usage information collected during streaming
last_chunk: The last chunk received from the streaming response
"""
if callbacks and len(callbacks) > 0:
for callback in callbacks:
if hasattr(callback, "log_success_event"):
# Use the usage_info we've been tracking
if not usage_info:
# Try to get usage from the last chunk if we haven't already
try:
if last_chunk:
if (
isinstance(last_chunk, dict)
and "usage" in last_chunk
):
usage_info = last_chunk["usage"]
elif hasattr(last_chunk, "usage"):
if not isinstance(
getattr(last_chunk, "usage"), type
):
usage_info = getattr(last_chunk, "usage")
except Exception as e:
logging.debug(f"Error extracting usage info: {e}")
if usage_info:
callback.log_success_event(
kwargs={}, # We don't have the original params here
response_obj={"usage": usage_info},
start_time=0,
end_time=0,
)
def _handle_non_streaming_response( def _handle_non_streaming_response(
self, self,
params: Dict[str, Any], params: Dict[str, Any],
@@ -532,6 +552,7 @@ class LLM:
Args: Args:
params: Parameters for the completion call params: Parameters for the completion call
callbacks: Optional list of callback functions
available_functions: Dict of available functions available_functions: Dict of available functions
Returns: Returns:
@@ -562,12 +583,17 @@ class LLM:
# --- 4) Check for tool calls # --- 4) Check for tool calls
tool_calls = getattr(response_message, "tool_calls", []) tool_calls = getattr(response_message, "tool_calls", [])
# --- 5) Handle tool calls if present # --- 5) If no tool calls or no available functions, return the text response directly
if not tool_calls or not available_functions:
self._handle_emit_call_events(text_response, LLMCallType.LLM_CALL)
return text_response
# --- 6) Handle tool calls if present
tool_result = self._handle_tool_call(tool_calls, available_functions) tool_result = self._handle_tool_call(tool_calls, available_functions)
if tool_result is not None: if tool_result is not None:
return tool_result return tool_result
# --- 6) Emit completion event and return response # --- 7) If tool call handling didn't return a result, emit completion event and return text response
self._handle_emit_call_events(text_response, LLMCallType.LLM_CALL) self._handle_emit_call_events(text_response, LLMCallType.LLM_CALL)
return text_response return text_response