-
Notifications
You must be signed in to change notification settings - Fork 150
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
vertexai[patch]: ChatVertexAI.generate, stream, bind_tools #166
Conversation
baskaryan
commented
Apr 18, 2024
•
edited
Loading
edited
- Update bind_tools to construct a dict that matches VertexAI Tool proto and bind that to the model
- Update bind_tools to accept FunctionDeclaration and Vertex Tool directly
- Update bind_tools to accept a tool_choice param for easier tool_config setting
- Refactor generate, agenerate, stream, astream to reuse gemini param setting and response parsing
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.
big fan of the gemini/non gemini split
note: tests/integration_tests/test_chat_models.py::test_chat_vertexai_gemini_function_calling_with_multiple_parts fails non-deterministically |
tools: Sequence[Union[Type[BaseModel], Callable, BaseTool]], | ||
tool_config: Optional[Dict[str, Any]] = None, | ||
tools: Sequence[Union[_FunctionDeclarationLike, VertexTool]], | ||
tool_config: Optional[_ToolConfigDict] = None, |
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.
nits: should we add to the docstring an example how to contsruct _ToolConfigDict and _ToolChoiceType?
so that people don't need to read the source code.
return False | ||
return float(self.model_name.split("-")[1]) > 1.0 | ||
except (ValueError, IndexError): | ||
return False | ||
|
||
def _generate( |
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.
+1 for splitting gemini vs non-gemini.
WDYT if we make this method to only select which function to run, and move gemini implementation to _generate_gemini
(and obviously in other methods) like this:
def _generate(
self,
messages: List[BaseMessage],
stop: Optional[List[str]] = None,
run_manager: Optional[CallbackManagerForLLMRun] = None,
stream: Optional[bool] = None,
**kwargs: Any,
) -> ChatResult:
if stream is True or (stream is None and self.streaming):
stream_iter = self._stream(
messages, stop=stop, run_manager=run_manager, **kwargs
)
return generate_from_stream(stream_iter)
if not self._is_gemini_model:
return self._generate_non_gemini(messages, stop=stop, **kwargs)
return self._generate_gemini(messages, stop=stop, **kwargs)
def _generate_gemini(
self,
messages: List[BaseMessage],
stop: Optional[List[str]] = None,
**kwargs: Any,
):
client, contents = self._gemini_client_and_contents(messages)
params = self._gemini_params(stop=stop, **kwargs)
with telemetry.tool_context_manager(self._user_agent):
response = client.generate_content(contents, **params)
return self._gemini_response_to_chat_result(response)
def _generate_non_gemini(
self,
messages: List[BaseMessage],
stop: Optional[List[str]] = None,
**kwargs: Any,
) -> ChatResult:
...
That would make testing a little bit cleaner.