Skip to content
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

Merged
merged 20 commits into from
Apr 20, 2024

Conversation

baskaryan
Copy link
Contributor

@baskaryan baskaryan commented Apr 18, 2024

  • 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

Copy link
Member

@efriis efriis left a 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

@baskaryan baskaryan changed the title WIP vertexai[patch]: refactor bind_tools vertexai[patch]: ChatVertexAI.generate, stream, bind_tools Apr 18, 2024
@baskaryan baskaryan marked this pull request as ready for review April 18, 2024 22:50
@baskaryan
Copy link
Contributor Author

note: tests/integration_tests/test_chat_models.py::test_chat_vertexai_gemini_function_calling_with_multiple_parts fails non-deterministically

@baskaryan baskaryan requested a review from eyurtsev April 19, 2024 00:48
tools: Sequence[Union[Type[BaseModel], Callable, BaseTool]],
tool_config: Optional[Dict[str, Any]] = None,
tools: Sequence[Union[_FunctionDeclarationLike, VertexTool]],
tool_config: Optional[_ToolConfigDict] = None,
Copy link
Collaborator

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(
Copy link
Collaborator

@alx13 alx13 Apr 19, 2024

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.

@lkuligin lkuligin merged commit 9c10541 into main Apr 20, 2024
15 checks passed
@lkuligin lkuligin deleted the bagatur/refac_bind_tools branch April 20, 2024 16:17
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

None yet

4 participants