-
Notifications
You must be signed in to change notification settings - Fork 4.5k
Fix: Enable Gemini support via GOOGLE_API_KEY fallback and explicit api_key injection #2805
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Open
MathavanSG
wants to merge
4
commits into
crewAIInc:main
Choose a base branch
from
MathavanSG:fix/gemini-langchain-support
base: main
Could not load branches
Branch not found: {{ refName }}
Loading
Could not load tags
Nothing to show
Loading
Are you sure you want to change the base?
Some commits from the old base branch may be removed from the timeline,
and old review comments may become outdated.
+45
−9
Open
Changes from all commits
Commits
Show all changes
4 commits
Select commit
Hold shift + click to select a range
21961e7
Fix: Enable Gemini support via GOOGLE_API_KEY fallback and explicit a…
MathavanSG e37dd98
Fix: finalized Gemini API key fallback and completion param injection…
MathavanSG 064c550
Refactor: remove explicit api_key enforcement and redundant injection…
MathavanSG ec065a3
Merge branch 'main' into fix/gemini-langchain-support
lucasgomide File filter
Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -56,9 +56,16 @@ | |
from crewai.utilities.exceptions.context_window_exceeding_exception import ( | ||
LLMContextLengthExceededException, | ||
) | ||
import logging | ||
|
||
|
||
load_dotenv() | ||
|
||
# Patch: Gemini key fallback | ||
if "GEMINI_API_KEY" not in os.environ and "GOOGLE_API_KEY" in os.environ: | ||
os.environ["GEMINI_API_KEY"] = os.environ["GOOGLE_API_KEY"] | ||
logging.info("[CrewAI Gemini Patch] Set GEMINI_API_KEY from GOOGLE_API_KEY") | ||
|
||
|
||
class FilteredStream(io.TextIOBase): | ||
_lock = None | ||
|
@@ -287,6 +294,17 @@ class AccumulatedToolArgs(BaseModel): | |
|
||
|
||
class LLM(BaseLLM): | ||
""" | ||
LLM class for handling language model interactions via LiteLLM. | ||
|
||
Features: | ||
- Supports multiple model providers (e.g., OpenAI, Gemini, Anthropic) | ||
- Automatically uses GOOGLE_API_KEY if GEMINI_API_KEY is not explicitly set | ||
- Injects the resolved API key directly into the LLM completion parameters | ||
- Ensures compatibility with both legacy and AI Studio-style key environments | ||
- Designed for use in CrewAI agent workflows and tool-based LLM interactions | ||
""" | ||
|
||
def __init__( | ||
self, | ||
model: str, | ||
|
@@ -348,6 +366,15 @@ def __init__( | |
else: | ||
self.stop = stop | ||
|
||
# Fallback logic | ||
if "GEMINI_API_KEY" in os.environ: | ||
api_key = os.environ["GEMINI_API_KEY"] | ||
elif "GOOGLE_API_KEY" in os.environ: | ||
api_key = os.environ["GOOGLE_API_KEY"] | ||
os.environ["GEMINI_API_KEY"] = api_key | ||
|
||
self.api_key = api_key | ||
|
||
self.set_callbacks(callbacks) | ||
self.set_env_callbacks() | ||
|
||
|
@@ -367,19 +394,26 @@ def _prepare_completion_params( | |
self, | ||
messages: Union[str, List[Dict[str, str]]], | ||
tools: Optional[List[dict]] = None, | ||
**kwargs: Any, | ||
) -> Dict[str, Any]: | ||
"""Prepare parameters for the completion call. | ||
""" | ||
Prepare parameters for the LLM completion API call. | ||
|
||
This method: | ||
- Formats input messages for the model provider | ||
- Accepts optional tool definitions | ||
- Injects API key into the request (fallback to GOOGLE_API_KEY if GEMINI_API_KEY is not set) | ||
- Merges additional keyword arguments passed to support flexibility | ||
|
||
Args: | ||
messages: Input messages for the LLM | ||
tools: Optional list of tool schemas | ||
callbacks: Optional list of callback functions | ||
available_functions: Optional dict of available functions | ||
messages (Union[str, List[Dict[str, str]]]): Prompt or structured messages to send to the LLM. | ||
tools (Optional[List[dict]]): Optional tool definitions (for function calling). | ||
**kwargs (Any): Additional optional parameters for the completion call. | ||
|
||
Returns: | ||
Dict[str, Any]: Parameters for the completion call | ||
Dict[str, Any]: Final parameters dictionary to be passed to `litellm.completion(...)`. | ||
""" | ||
# --- 1) Format messages according to provider requirements | ||
# Format messages | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. You removed 1), but left the other section numbers like 2). |
||
if isinstance(messages, str): | ||
messages = [{"role": "user", "content": messages}] | ||
formatted_messages = self._format_messages_for_provider(messages) | ||
|
@@ -411,8 +445,10 @@ def _prepare_completion_params( | |
**self.additional_params, | ||
} | ||
|
||
# Remove None values from params | ||
return {k: v for k, v in params.items() if v is not None} | ||
# Remove None values | ||
params = {k: v for k, v in params.items() if v is not None} | ||
|
||
return params | ||
|
||
def _handle_streaming_response( | ||
self, | ||
|
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
You have couple of issues here
api_key
was provided you have to use it even those GOOGLE/GEMINI_API_KEY have been set.