| import itertools
|
|
|
| import orjson as json
|
|
|
| from ..constants import GRPC
|
| from ..exceptions import APIError
|
| from ..types import Gem, GemJar, RPCData
|
| from ..utils import extract_json_from_response, get_nested_value, logger
|
|
|
|
|
| class GemMixin:
|
| """
|
| Mixin class providing gem-related functionality for GeminiClient.
|
| """
|
|
|
| def __init__(self, *args, **kwargs):
|
| super().__init__(*args, **kwargs)
|
| self._gems: GemJar | None = None
|
|
|
| @property
|
| def gems(self) -> GemJar:
|
| """
|
| Returns a `GemJar` object containing cached gems.
|
| Only available after calling `GeminiClient.fetch_gems()`.
|
|
|
| Returns
|
| -------
|
| :class:`GemJar`
|
| Refer to `gemini_webapi.types.GemJar`.
|
|
|
| Raises
|
| ------
|
| `RuntimeError`
|
| If `GeminiClient.fetch_gems()` has not been called before accessing this property.
|
| """
|
|
|
| if self._gems is None:
|
| raise RuntimeError(
|
| "Gems not fetched yet. Call `GeminiClient.fetch_gems()` method to fetch gems from gemini.google.com."
|
| )
|
|
|
| return self._gems
|
|
|
| async def fetch_gems(
|
| self, include_hidden: bool = False, language: str = "en", **kwargs
|
| ) -> GemJar:
|
| """
|
| Get a list of available gems from gemini, including system predefined gems and user-created custom gems.
|
|
|
| Note that network request will be sent every time this method is called.
|
| Once the gems are fetched, they will be cached and accessible via `GeminiClient.gems` property.
|
|
|
| Parameters
|
| ----------
|
| include_hidden: `bool`, optional
|
| There are some predefined gems that by default are not shown to users (and therefore may not work properly).
|
| Set this parameter to `True` to include them in the fetched gem list.
|
| language: `str`, optional
|
| Language code for the gems to fetch. Default is 'en'.
|
|
|
| Returns
|
| -------
|
| :class:`GemJar`
|
| Refer to `gemini_webapi.types.GemJar`.
|
| """
|
|
|
| response = await self._batch_execute(
|
| [
|
| RPCData(
|
| rpcid=GRPC.LIST_GEMS,
|
| payload=(
|
| f"[4,['{language}'],0]"
|
| if include_hidden
|
| else f"[3,['{language}'],0]"
|
| ),
|
| identifier="system",
|
| ),
|
| RPCData(
|
| rpcid=GRPC.LIST_GEMS,
|
| payload=f"[2,['{language}'],0]",
|
| identifier="custom",
|
| ),
|
| ],
|
| **kwargs,
|
| )
|
|
|
| try:
|
| response_json = extract_json_from_response(response.text)
|
|
|
| predefined_gems, custom_gems = [], []
|
|
|
| for part in response_json:
|
| try:
|
| identifier = get_nested_value(part, [-1])
|
| part_body_str = get_nested_value(part, [2])
|
| if not part_body_str:
|
| continue
|
|
|
| part_body = json.loads(part_body_str)
|
| if identifier == "system":
|
| predefined_gems = get_nested_value(part_body, [2], [])
|
| elif identifier == "custom":
|
| custom_gems = get_nested_value(part_body, [2], [])
|
| except json.JSONDecodeError:
|
| continue
|
|
|
| if not predefined_gems and not custom_gems:
|
| raise Exception
|
| except Exception:
|
| await self.close()
|
| logger.debug(f"Unexpected response data structure: {response.text}")
|
| raise APIError(
|
| "Failed to fetch gems. Unexpected response data structure. Client will try to re-initialize on next request."
|
| )
|
|
|
| self._gems = GemJar(
|
| itertools.chain(
|
| (
|
| (
|
| gem[0],
|
| Gem(
|
| id=gem[0],
|
| name=gem[1][0],
|
| description=gem[1][1],
|
| prompt=gem[2] and gem[2][0] or None,
|
| predefined=True,
|
| ),
|
| )
|
| for gem in predefined_gems
|
| ),
|
| (
|
| (
|
| gem[0],
|
| Gem(
|
| id=gem[0],
|
| name=gem[1][0],
|
| description=gem[1][1],
|
| prompt=gem[2] and gem[2][0] or None,
|
| predefined=False,
|
| ),
|
| )
|
| for gem in custom_gems
|
| ),
|
| )
|
| )
|
|
|
| return self._gems
|
|
|
| async def create_gem(self, name: str, prompt: str, description: str = "") -> Gem:
|
| """
|
| Create a new custom gem.
|
|
|
| Parameters
|
| ----------
|
| name: `str`
|
| Name of the custom gem.
|
| prompt: `str`
|
| System instructions for the custom gem.
|
| description: `str`, optional
|
| Description of the custom gem (has no effect on the model's behavior).
|
|
|
| Returns
|
| -------
|
| :class:`Gem`
|
| The created gem.
|
| """
|
|
|
| response = await self._batch_execute(
|
| [
|
| RPCData(
|
| rpcid=GRPC.CREATE_GEM,
|
| payload=json.dumps(
|
| [
|
| [
|
| name,
|
| description,
|
| prompt,
|
| None,
|
| None,
|
| None,
|
| None,
|
| None,
|
| 0,
|
| None,
|
| 1,
|
| None,
|
| None,
|
| None,
|
| [],
|
| ]
|
| ]
|
| ).decode("utf-8"),
|
| )
|
| ]
|
| )
|
|
|
| try:
|
| response_json = extract_json_from_response(response.text)
|
| part_body_str = get_nested_value(response_json, [0, 2], verbose=True)
|
| if not part_body_str:
|
| raise Exception
|
|
|
| part_body = json.loads(part_body_str)
|
| gem_id = get_nested_value(part_body, [0], verbose=True)
|
| if not gem_id:
|
| raise Exception
|
| except Exception:
|
| await self.close()
|
| logger.debug(f"Unexpected response data structure: {response.text}")
|
| raise APIError(
|
| "Failed to create gem. Unexpected response data structure. Client will try to re-initialize on next request."
|
| )
|
|
|
| return Gem(
|
| id=gem_id,
|
| name=name,
|
| description=description,
|
| prompt=prompt,
|
| predefined=False,
|
| )
|
|
|
| async def update_gem(
|
| self, gem: Gem | str, name: str, prompt: str, description: str = ""
|
| ) -> Gem:
|
| """
|
| Update an existing custom gem.
|
|
|
| Parameters
|
| ----------
|
| gem: `Gem | str`
|
| Gem to update, can be either a `gemini_webapi.types.Gem` object or a gem id string.
|
| name: `str`
|
| New name for the custom gem.
|
| prompt: `str`
|
| New system instructions for the custom gem.
|
| description: `str`, optional
|
| New description of the custom gem (has no effect on the model's behavior).
|
|
|
| Returns
|
| -------
|
| :class:`Gem`
|
| The updated gem.
|
| """
|
|
|
| if isinstance(gem, Gem):
|
| gem_id = gem.id
|
| else:
|
| gem_id = gem
|
|
|
| await self._batch_execute(
|
| [
|
| RPCData(
|
| rpcid=GRPC.UPDATE_GEM,
|
| payload=json.dumps(
|
| [
|
| gem_id,
|
| [
|
| name,
|
| description,
|
| prompt,
|
| None,
|
| None,
|
| None,
|
| None,
|
| None,
|
| 0,
|
| None,
|
| 1,
|
| None,
|
| None,
|
| None,
|
| [],
|
| 0,
|
| ],
|
| ]
|
| ).decode("utf-8"),
|
| )
|
| ]
|
| )
|
|
|
| return Gem(
|
| id=gem_id,
|
| name=name,
|
| description=description,
|
| prompt=prompt,
|
| predefined=False,
|
| )
|
|
|
| async def delete_gem(self, gem: Gem | str, **kwargs) -> None:
|
| """
|
| Delete a custom gem.
|
|
|
| Parameters
|
| ----------
|
| gem: `Gem | str`
|
| Gem to delete, can be either a `gemini_webapi.types.Gem` object or a gem id string.
|
| """
|
|
|
| if isinstance(gem, Gem):
|
| gem_id = gem.id
|
| else:
|
| gem_id = gem
|
|
|
| await self._batch_execute(
|
| [
|
| RPCData(
|
| rpcid=GRPC.DELETE_GEM, payload=json.dumps([gem_id]).decode("utf-8")
|
| )
|
| ],
|
| **kwargs,
|
| )
|
|
|
