feat: add support for vertex ai #1042
Conversation
🦋 Changeset detectedLatest commit: 877b7d5 The changes in this PR will be included in the next version bump. This PR includes changesets to release 1 package
Not sure what this means? Click here to learn what changesets are. Click here if you're a maintainer who wants to add another changeset to this PR |
livekit-plugins/livekit-plugins-google/livekit/plugins/google/models.py
Outdated
Show resolved
Hide resolved
| logger.info( | ||
| "Starting chat with temperature: %s, candidate count: %s", temperature, n | ||
| ) |
|
Do you think there is a way we reuse https://github.com/livekit/agents/blob/main/livekit-agents/livekit/agents/llm/_oai_api.py when running/checking function calls. The good think about it is that it verify the types are correct before running any code |
|
@theomonnom I think we could reuse |
|
@theomonnom I've updated code to utilize |
| if not self._project: | ||
| logger.error("Project is required for Vertex AI initialization") | ||
| raise ValueError | ||
| if not self._location: | ||
| logger.error("Location is required for Vertex AI initialization") | ||
| raise ValueError | ||
| if not self._gac: | ||
| logger.error( | ||
| "`GOOGLE_APPLICATION_CREDENTIALS` environment variable is not set. please set it to the path of the service account key file." | ||
| ) | ||
| raise ValueError | ||
|
|
There was a problem hiding this comment.
Can you add the messages directly inside the Exception instead of logging them.
| self._location = location | ||
| self._model_name = model | ||
| self._kwargs = kwargs | ||
| self._gac = os.getenv("GOOGLE_APPLICATION_CREDENTIALS") |
There was a problem hiding this comment.
Can you follow what we did for auth on the google.STT.
As well as the extra arguments for the credentials
| project: Optional[str] = None, | ||
| location: Optional[str] = None, | ||
| model: str | GeminiModels = "gemini-1.5-flash-002", | ||
| **kwargs, |
There was a problem hiding this comment.
Let's not use **kwargs inside our constructor. This is error-prone
|
|
||
| def _build_vertex_context(chat_ctx: llm.ChatContext, cache_key: Any) -> List[dict]: | ||
| contents = [] | ||
| current_entry: Optional[dict] = None |
There was a problem hiding this comment.
The typing isn't very useful here
| from typing import Any, AsyncIterator, Awaitable, Dict, List, Optional | ||
| from uuid import uuid4 as uuid | ||
|
|
||
| import vertexai # type: ignore |
There was a problem hiding this comment.
I don't think we should ignore the types of those packages
| self._function_calls_info.append(function_call_info) | ||
| choice_delta.tool_calls = self._function_calls_info | ||
| return llm.ChatChunk( | ||
| request_id=str(uuid()), |
There was a problem hiding this comment.
Can you use utils.shortuuid()
Also it seems like we're generating a new request_id for each chunk?
It should be the same for the entire request
| cmp = model.generate_content_async( | ||
| contents=contents, | ||
| generation_config=generation_config, | ||
| tools=tools or None, |
There was a problem hiding this comment.
This can be done above on line 96
| ) | ||
| raise ValueError | ||
|
|
||
| vertexai.init(project=project, location=location) |
There was a problem hiding this comment.
Seems like it is a singleton, what happens if the user already initialized it himself?
Is this function returning smthg?
| model = GenerativeModel( | ||
| model_name=self._model_name, | ||
| system_instruction=system_instruction, | ||
| **self._kwargs, |
There was a problem hiding this comment.
| **self._kwargs, |
| def _build_parameters(arguments: Dict[str, llm.FunctionArgInfo]) -> Dict[str, Any]: | ||
| properties = { | ||
| arg_name: { | ||
| "type": JSON_SCHEMA_TYPE_MAP.get(arg_info.type, "string"), |
There was a problem hiding this comment.
I don't think we should default to string?
There was a problem hiding this comment.
Let's raise an error instead if the type isn't supported
No description provided.