langchain-google-genai 0.0.6__py3-none-any.whl → 0.0.8__py3-none-any.whl

langchain_google_genai/_function_utils.py

@@ -0,0 +1,116 @@
+from __future__ import annotations
+
+from typing import (
+    Dict,
+    List,
+    Type,
+    Union,
+)
+
+import google.ai.generativelanguage as glm
+from langchain_core.pydantic_v1 import BaseModel
+from langchain_core.tools import BaseTool
+from langchain_core.utils.json_schema import dereference_refs
+
+FunctionCallType = Union[BaseTool, Type[BaseModel], Dict]
+
+TYPE_ENUM = {
+    "string": glm.Type.STRING,
+    "number": glm.Type.NUMBER,
+    "integer": glm.Type.INTEGER,
+    "boolean": glm.Type.BOOLEAN,
+    "array": glm.Type.ARRAY,
+    "object": glm.Type.OBJECT,
+}
+
+
+def convert_to_genai_function_declarations(
+    function_calls: List[FunctionCallType],
+) -> List[glm.Tool]:
+    return [
+        glm.Tool(
+            function_declarations=[_convert_to_genai_function(fc)],
+        )
+        for fc in function_calls
+    ]
+
+
+def _convert_to_genai_function(fc: FunctionCallType) -> glm.FunctionDeclaration:
+    if isinstance(fc, BaseTool):
+        return _convert_tool_to_genai_function(fc)
+    elif isinstance(fc, type) and issubclass(fc, BaseModel):
+        return _convert_pydantic_to_genai_function(fc)
+    elif isinstance(fc, dict):
+        return glm.FunctionDeclaration(
+            name=fc["name"],
+            description=fc.get("description"),
+            parameters={
+                "properties": {
+                    k: {
+                        "type_": TYPE_ENUM[v["type"]],
+                        "description": v.get("description"),
+                    }
+                    for k, v in fc["parameters"]["properties"].items()
+                },
+                "required": fc["parameters"].get("required", []),
+                "type_": TYPE_ENUM[fc["parameters"]["type"]],
+            },
+        )
+    else:
+        raise ValueError(f"Unsupported function call type {fc}")
+
+
+def _convert_tool_to_genai_function(tool: BaseTool) -> glm.FunctionDeclaration:
+    if tool.args_schema:
+        schema = dereference_refs(tool.args_schema.schema())
+        schema.pop("definitions", None)
+
+        return glm.FunctionDeclaration(
+            name=tool.name or schema["title"],
+            description=tool.description or schema["description"],
+            parameters={
+                "properties": {
+                    k: {
+                        "type_": TYPE_ENUM[v["type"]],
+                        "description": v.get("description"),
+                    }
+                    for k, v in schema["properties"].items()
+                },
+                "required": schema["required"],
+                "type_": TYPE_ENUM[schema["type"]],
+            },
+        )
+    else:
+        return glm.FunctionDeclaration(
+            name=tool.name,
+            description=tool.description,
+            parameters={
+                "properties": {
+                    "__arg1": {"type_": TYPE_ENUM["string"]},
+                },
+                "required": ["__arg1"],
+                "type_": TYPE_ENUM["object"],
+            },
+        )
+
+
+def _convert_pydantic_to_genai_function(
+    pydantic_model: Type[BaseModel],
+) -> glm.FunctionDeclaration:
+    schema = dereference_refs(pydantic_model.schema())
+    schema.pop("definitions", None)
+    return glm.FunctionDeclaration(
+        name=schema["title"],
+        description=schema.get("description", ""),
+        parameters={
+            "properties": {
+                k: {
+                    "type_": TYPE_ENUM[v["type"]],
+                    "description": v.get("description"),
+                }
+                for k, v in schema["properties"].items()
+            },
+            "required": schema["required"],
+            "type_": TYPE_ENUM[schema["type"]],
+        },
+    )

langchain_google_genai/chat_models.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import base64
+import json
 import logging
 import os
 from io import BytesIO
@@ -15,14 +16,17 @@ from typing import (
     Optional,
     Sequence,
     Tuple,
-    Type,
     Union,
     cast,
 )
 from urllib.parse import urlparse
 
+import google.ai.generativelanguage as glm
+import google.api_core
+
 # TODO: remove ignore once the google package is published with types
 import google.generativeai as genai  # type: ignore[import]
+import proto  # type: ignore[import]
 import requests
 from langchain_core.callbacks.manager import (
     AsyncCallbackManagerForLLMRun,
@@ -33,14 +37,12 @@ from langchain_core.messages import (
     AIMessage,
     AIMessageChunk,
     BaseMessage,
-    ChatMessage,
-    ChatMessageChunk,
+    FunctionMessage,
     HumanMessage,
-    HumanMessageChunk,
     SystemMessage,
 )
 from langchain_core.outputs import ChatGeneration, ChatGenerationChunk, ChatResult
-from langchain_core.pydantic_v1 import Field, SecretStr, root_validator
+from langchain_core.pydantic_v1 import SecretStr, root_validator
 from langchain_core.utils import get_from_dict_or_env
 from tenacity import (
     before_sleep_log,
@@ -51,6 +53,10 @@ from tenacity import (
 )
 
 from langchain_google_genai._common import GoogleGenerativeAIError
+from langchain_google_genai._function_utils import (
+    convert_to_genai_function_declarations,
+)
+from langchain_google_genai.llms import GoogleModelFamily, _BaseGoogleGenerativeAI
 
 IMAGE_TYPES: Tuple = ()
 try:
@@ -87,8 +93,6 @@ def _create_retry_decorator() -> Callable[[Any], Any]:
         Callable[[Any], Any]: A retry decorator configured for handling specific
         Google API exceptions.
     """
-    import google.api_core.exceptions
-
     multiplier = 2
     min_seconds = 1
     max_seconds = 60
@@ -123,14 +127,22 @@ def _chat_with_retry(generation_method: Callable, **kwargs: Any) -> Any:
         Any: The result from the chat generation method.
     """
     retry_decorator = _create_retry_decorator()
-    from google.api_core.exceptions import InvalidArgument  # type: ignore
 
     @retry_decorator
     def _chat_with_retry(**kwargs: Any) -> Any:
         try:
             return generation_method(**kwargs)
-        except InvalidArgument as e:
-            # Do not retry for these errors.
+        # Do not retry for these errors.
+        except google.api_core.exceptions.FailedPrecondition as exc:
+            if "location is not supported" in exc.message:
+                error_msg = (
+                    "Your location is not supported by google-generativeai "
+                    "at the moment. Try to use ChatVertexAI LLM from "
+                    "langchain_google_vertexai."
+                )
+                raise ValueError(error_msg)
+
+        except google.api_core.exceptions.InvalidArgument as e:
             raise ChatGoogleGenerativeAIError(
                 f"Invalid argument provided to Gemini: {e}"
             ) from e
@@ -312,14 +324,47 @@ llm = ChatGoogleGenerativeAI(model="gemini-pro", convert_system_message_to_human
             continue
         elif isinstance(message, AIMessage):
             role = "model"
+            raw_function_call = message.additional_kwargs.get("function_call")
+            if raw_function_call:
+                function_call = glm.FunctionCall(
+                    {
+                        "name": raw_function_call["name"],
+                        "args": json.loads(raw_function_call["arguments"]),
+                    }
+                )
+                parts = [glm.Part(function_call=function_call)]
+            else:
+                parts = _convert_to_parts(message.content)
         elif isinstance(message, HumanMessage):
             role = "user"
+            parts = _convert_to_parts(message.content)
+        elif isinstance(message, FunctionMessage):
+            role = "user"
+            response: Any
+            if not isinstance(message.content, str):
+                response = message.content
+            else:
+                try:
+                    response = json.loads(message.content)
+                except json.JSONDecodeError:
+                    response = message.content  # leave as str representation
+            parts = [
+                glm.Part(
+                    function_response=glm.FunctionResponse(
+                        name=message.name,
+                        response=(
+                            {"output": response}
+                            if not isinstance(response, dict)
+                            else response
+                        ),
+                    )
+                )
+            ]
         else:
             raise ValueError(
                 f"Unexpected message with type {type(message)} at the position {i}."
             )
 
-        parts = _convert_to_parts(message.content)
         if raw_system_message:
             if role == "model":
                 raise ValueError(
@@ -332,71 +377,51 @@ llm = ChatGoogleGenerativeAI(model="gemini-pro", convert_system_message_to_human
     return messages
 
 
-def _parts_to_content(parts: List[genai.types.PartType]) -> Union[List[dict], str]:
-    """Converts a list of Gemini API Part objects into a list of LangChain messages."""
-    if len(parts) == 1 and parts[0].text is not None and not parts[0].inline_data:
-        # Simple text response. The typical response
-        return parts[0].text
-    elif not parts:
-        logger.warning("Gemini produced an empty response.")
-        return ""
-    messages = []
-    for part in parts:
-        if part.text is not None:
-            messages.append(
-                {
-                    "type": "text",
-                    "text": part.text,
-                }
-            )
-        else:
-            # TODO: Handle inline_data if that's a thing?
-            raise ChatGoogleGenerativeAIError(f"Unexpected part type. {part}")
-    return messages
+def _parse_response_candidate(
+    response_candidate: glm.Candidate, stream: bool
+) -> AIMessage:
+    first_part = response_candidate.content.parts[0]
+    if first_part.function_call:
+        function_call = proto.Message.to_dict(first_part.function_call)
+        function_call["arguments"] = json.dumps(function_call.pop("args", {}))
+        return (AIMessageChunk if stream else AIMessage)(
+            content="", additional_kwargs={"function_call": function_call}
+        )
+    else:
+        parts = response_candidate.content.parts
+
+    if len(parts) == 1 and parts[0].text:
+        content: Union[str, List[Union[str, Dict]]] = parts[0].text
+    else:
+        content = [proto.Message.to_dict(part) for part in parts]
+    return (AIMessageChunk if stream else AIMessage)(
+        content=content, additional_kwargs={}
+    )
 
 
 def _response_to_result(
-    response: genai.types.GenerateContentResponse,
-    ai_msg_t: Type[BaseMessage] = AIMessage,
-    human_msg_t: Type[BaseMessage] = HumanMessage,
-    chat_msg_t: Type[BaseMessage] = ChatMessage,
-    generation_t: Type[ChatGeneration] = ChatGeneration,
+    response: glm.GenerateContentResponse,
+    stream: bool = False,
 ) -> ChatResult:
     """Converts a PaLM API response into a LangChain ChatResult."""
-    llm_output = {}
-    if response.prompt_feedback:
-        try:
-            prompt_feedback = type(response.prompt_feedback).to_dict(
-                response.prompt_feedback, use_integers_for_enums=False
-            )
-            llm_output["prompt_feedback"] = prompt_feedback
-        except Exception as e:
-            logger.debug(f"Unable to convert prompt_feedback to dict: {e}")
+    llm_output = {"prompt_feedback": proto.Message.to_dict(response.prompt_feedback)}
 
     generations: List[ChatGeneration] = []
 
-    role_map = {
-        "model": ai_msg_t,
-        "user": human_msg_t,
-    }
     for candidate in response.candidates:
-        content = candidate.content
-        parts_content = _parts_to_content(content.parts)
-        if content.role not in role_map:
-            logger.warning(
-                f"Unrecognized role: {content.role}. Treating as a ChatMessage."
-            )
-            msg = chat_msg_t(content=parts_content, role=content.role)
-        else:
-            msg = role_map[content.role](content=parts_content)
         generation_info = {}
         if candidate.finish_reason:
             generation_info["finish_reason"] = candidate.finish_reason.name
-        if candidate.safety_ratings:
-            generation_info["safety_ratings"] = [
-                type(rating).to_dict(rating) for rating in candidate.safety_ratings
-            ]
-        generations.append(generation_t(message=msg, generation_info=generation_info))
+        generation_info["safety_ratings"] = [
+            proto.Message.to_dict(safety_rating, use_integers_for_enums=False)
+            for safety_rating in candidate.safety_ratings
+        ]
+        generations.append(
+            (ChatGenerationChunk if stream else ChatGeneration)(
+                message=_parse_response_candidate(candidate, stream=stream),
+                generation_info=generation_info,
+            )
+        )
     if not response.candidates:
         # Likely a "prompt feedback" violation (e.g., toxic input)
         # Raising an error would be different than how OpenAI handles it,
@@ -405,11 +430,16 @@ def _response_to_result(
             "Gemini produced an empty response. Continuing with empty message\n"
             f"Feedback: {response.prompt_feedback}"
         )
-        generations = [generation_t(message=ai_msg_t(content=""), generation_info={})]
+        generations = [
+            (ChatGenerationChunk if stream else ChatGeneration)(
+                message=(AIMessageChunk if stream else AIMessage)(content=""),
+                generation_info={},
+            )
+        ]
     return ChatResult(generations=generations, llm_output=llm_output)
 
 
-class ChatGoogleGenerativeAI(BaseChatModel):
+class ChatGoogleGenerativeAI(_BaseGoogleGenerativeAI, BaseChatModel):
     """`Google Generative AI` Chat models API.
 
     To use, you must have either:
@@ -427,53 +457,13 @@ class ChatGoogleGenerativeAI(BaseChatModel):
 
     """
 
-    model: str = Field(
-        ...,
-        description="""The name of the model to use.
-Supported examples:
-    - gemini-pro""",
-    )
-    max_output_tokens: int = Field(default=None, description="Max output tokens")
-
     client: Any  #: :meta private:
-    google_api_key: Optional[SecretStr] = None
-    temperature: Optional[float] = None
-    """Run inference with this temperature. Must by in the closed
-       interval [0.0, 1.0]."""
-    top_k: Optional[int] = None
-    """Decode using top-k sampling: consider the set of top_k most probable tokens.
-       Must be positive."""
-    top_p: Optional[int] = None
-    """The maximum cumulative probability of tokens to consider when sampling.
-
-        The model uses combined Top-k and nucleus sampling.
-
-        Tokens are sorted based on their assigned probabilities so
-        that only the most likely tokens are considered. Top-k
-        sampling directly limits the maximum number of tokens to
-        consider, while Nucleus sampling limits number of tokens
-        based on the cumulative probability.
-
-        Note: The default value varies by model, see the
-        `Model.top_p` attribute of the `Model` returned the
-        `genai.get_model` function.
-    """
-    n: int = Field(default=1, alias="candidate_count")
-    """Number of chat completions to generate for each prompt. Note that the API may
-       not return the full n completions if duplicates are generated."""
+
     convert_system_message_to_human: bool = False
     """Whether to merge any leading SystemMessage into the following HumanMessage.
     
     Gemini does not support system messages; any unsupported messages will 
     raise an error."""
-    client_options: Optional[Dict] = Field(
-        None,
-        description="Client options to pass to the Google API client.",
-    )
-    transport: Optional[str] = Field(
-        None,
-        description="A string, one of: [`rest`, `grpc`, `grpc_asyncio`].",
-    )
 
     class Config:
         allow_population_by_field_name = True
@@ -486,10 +476,6 @@ Supported examples:
     def _llm_type(self) -> str:
         return "chat-google-generative-ai"
 
-    @property
-    def _is_geminiai(self) -> bool:
-        return self.model is not None and "gemini" in self.model
-
     @classmethod
     def is_lc_serializable(self) -> bool:
         return True
@@ -560,7 +546,11 @@ Supported examples:
         run_manager: Optional[CallbackManagerForLLMRun] = None,
         **kwargs: Any,
     ) -> ChatResult:
-        params, chat, message = self._prepare_chat(messages, stop=stop)
+        params, chat, message = self._prepare_chat(
+            messages,
+            stop=stop,
+            functions=kwargs.get("functions"),
+        )
         response: genai.types.GenerateContentResponse = _chat_with_retry(
             content=message,
             **params,
@@ -575,7 +565,11 @@ Supported examples:
         run_manager: Optional[AsyncCallbackManagerForLLMRun] = None,
         **kwargs: Any,
     ) -> ChatResult:
-        params, chat, message = self._prepare_chat(messages, stop=stop)
+        params, chat, message = self._prepare_chat(
+            messages,
+            stop=stop,
+            functions=kwargs.get("functions"),
+        )
         response: genai.types.GenerateContentResponse = await _achat_with_retry(
             content=message,
             **params,
@@ -590,7 +584,11 @@ Supported examples:
         run_manager: Optional[CallbackManagerForLLMRun] = None,
         **kwargs: Any,
     ) -> Iterator[ChatGenerationChunk]:
-        params, chat, message = self._prepare_chat(messages, stop=stop)
+        params, chat, message = self._prepare_chat(
+            messages,
+            stop=stop,
+            functions=kwargs.get("functions"),
+        )
         response: genai.types.GenerateContentResponse = _chat_with_retry(
             content=message,
             **params,
@@ -598,17 +596,11 @@ Supported examples:
             stream=True,
         )
         for chunk in response:
-            _chat_result = _response_to_result(
-                chunk,
-                ai_msg_t=AIMessageChunk,
-                human_msg_t=HumanMessageChunk,
-                chat_msg_t=ChatMessageChunk,
-                generation_t=ChatGenerationChunk,
-            )
+            _chat_result = _response_to_result(chunk, stream=True)
             gen = cast(ChatGenerationChunk, _chat_result.generations[0])
-            yield gen
             if run_manager:
                 run_manager.on_llm_new_token(gen.text)
+            yield gen
 
     async def _astream(
         self,
@@ -617,24 +609,22 @@ Supported examples:
         run_manager: Optional[AsyncCallbackManagerForLLMRun] = None,
         **kwargs: Any,
     ) -> AsyncIterator[ChatGenerationChunk]:
-        params, chat, message = self._prepare_chat(messages, stop=stop)
+        params, chat, message = self._prepare_chat(
+            messages,
+            stop=stop,
+            functions=kwargs.get("functions"),
+        )
         async for chunk in await _achat_with_retry(
             content=message,
             **params,
             generation_method=chat.send_message_async,
             stream=True,
         ):
-            _chat_result = _response_to_result(
-                chunk,
-                ai_msg_t=AIMessageChunk,
-                human_msg_t=HumanMessageChunk,
-                chat_msg_t=ChatMessageChunk,
-                generation_t=ChatGenerationChunk,
-            )
+            _chat_result = _response_to_result(chunk, stream=True)
             gen = cast(ChatGenerationChunk, _chat_result.generations[0])
-            yield gen
             if run_manager:
                 await run_manager.on_llm_new_token(gen.text)
+            yield gen
 
     def _prepare_chat(
         self,
@@ -642,11 +632,37 @@ Supported examples:
         stop: Optional[List[str]] = None,
         **kwargs: Any,
     ) -> Tuple[Dict[str, Any], genai.ChatSession, genai.types.ContentDict]:
+        client = self.client
+        functions = kwargs.pop("functions", None)
+        if functions:
+            tools = convert_to_genai_function_declarations(functions)
+            client = genai.GenerativeModel(model_name=self.model, tools=tools)
+
         params = self._prepare_params(stop, **kwargs)
         history = _parse_chat_history(
             messages,
             convert_system_message_to_human=self.convert_system_message_to_human,
         )
         message = history.pop()
-        chat = self.client.start_chat(history=history)
+        chat = client.start_chat(history=history)
         return params, chat, message
+
+    def get_num_tokens(self, text: str) -> int:
+        """Get the number of tokens present in the text.
+
+        Useful for checking if an input will fit in a model's context window.
+
+        Args:
+            text: The string input to tokenize.
+
+        Returns:
+            The integer number of tokens in the text.
+        """
+        if self._model_family == GoogleModelFamily.GEMINI:
+            result = self.client.count_tokens(text)
+            token_count = result.total_tokens
+        else:
+            result = self.client.count_text_tokens(model=self.model, prompt=text)
+            token_count = result["token_count"]
+
+        return token_count

langchain_google_genai/llms.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+from enum import Enum, auto
 from typing import Any, Callable, Dict, Iterator, List, Optional, Union
 
 import google.api_core
@@ -15,6 +16,19 @@ from langchain_core.pydantic_v1 import BaseModel, Field, SecretStr, root_validat
 from langchain_core.utils import get_from_dict_or_env
 
 
+class GoogleModelFamily(str, Enum):
+    GEMINI = auto()
+    PALM = auto()
+
+    @classmethod
+    def _missing_(cls, value: Any) -> Optional["GoogleModelFamily"]:
+        if "gemini" in value.lower():
+            return GoogleModelFamily.GEMINI
+        elif "text-bison" in value.lower():
+            return GoogleModelFamily.PALM
+        return None
+
+
 def _create_retry_decorator(
     llm: BaseLLM,
     *,
@@ -56,21 +70,25 @@ def _completion_with_retry(
         prompt: LanguageModelInput, is_gemini: bool, stream: bool, **kwargs: Any
     ) -> Any:
         generation_config = kwargs.get("generation_config", {})
-        if is_gemini:
-            return llm.client.generate_content(
-                contents=prompt, stream=stream, generation_config=generation_config
-            )
-        return llm.client.generate_text(prompt=prompt, **kwargs)
+        error_msg = (
+            "Your location is not supported by google-generativeai at the moment. "
+            "Try to use VertexAI LLM from langchain_google_vertexai"
+        )
+        try:
+            if is_gemini:
+                return llm.client.generate_content(
+                    contents=prompt, stream=stream, generation_config=generation_config
+                )
+            return llm.client.generate_text(prompt=prompt, **kwargs)
+        except google.api_core.exceptions.FailedPrecondition as exc:
+            if "location is not supported" in exc.message:
+                raise ValueError(error_msg)
 
     return _completion_with_retry(
         prompt=prompt, is_gemini=is_gemini, stream=stream, **kwargs
     )
 
 
-def _is_gemini_model(model_name: str) -> bool:
-    return "gemini" in model_name
-
-
 def _strip_erroneous_leading_spaces(text: str) -> str:
     """Strip erroneous leading spaces from text.
 
@@ -84,17 +102,9 @@ def _strip_erroneous_leading_spaces(text: str) -> str:
         return text
 
 
-class GoogleGenerativeAI(BaseLLM, BaseModel):
-    """Google GenerativeAI models.
-
-    Example:
-        .. code-block:: python
+class _BaseGoogleGenerativeAI(BaseModel):
+    """Base class for Google Generative AI LLMs"""
 
-            from langchain_google_genai import GoogleGenerativeAI
-            llm = GoogleGenerativeAI(model="gemini-pro")
-    """
-
-    client: Any  #: :meta private:
     model: str = Field(
         ...,
         description="""The name of the model to use.
@@ -133,15 +143,39 @@ Supported examples:
         description="A string, one of: [`rest`, `grpc`, `grpc_asyncio`].",
     )
 
-    @property
-    def is_gemini(self) -> bool:
-        """Returns whether a model is belongs to a Gemini family or not."""
-        return _is_gemini_model(self.model)
-
     @property
     def lc_secrets(self) -> Dict[str, str]:
         return {"google_api_key": "GOOGLE_API_KEY"}
 
+    @property
+    def _model_family(self) -> str:
+        return GoogleModelFamily(self.model)
+
+    @property
+    def _identifying_params(self) -> Dict[str, Any]:
+        """Get the identifying parameters."""
+        return {
+            "model": self.model,
+            "temperature": self.temperature,
+            "top_p": self.top_p,
+            "top_k": self.top_k,
+            "max_output_tokens": self.max_output_tokens,
+            "candidate_count": self.n,
+        }
+
+
+class GoogleGenerativeAI(_BaseGoogleGenerativeAI, BaseLLM):
+    """Google GenerativeAI models.
+
+    Example:
+        .. code-block:: python
+
+            from langchain_google_genai import GoogleGenerativeAI
+            llm = GoogleGenerativeAI(model="gemini-pro")
+    """
+
+    client: Any  #: :meta private:
+
     @root_validator()
     def validate_environment(cls, values: Dict) -> Dict:
         """Validates params and passes them to google-generativeai package."""
@@ -159,7 +193,7 @@ Supported examples:
             client_options=values.get("client_options"),
         )
 
-        if _is_gemini_model(model_name):
+        if GoogleModelFamily(model_name) == GoogleModelFamily.GEMINI:
             values["client"] = genai.GenerativeModel(model_name=model_name)
         else:
             values["client"] = genai
@@ -195,7 +229,7 @@ Supported examples:
             "candidate_count": self.n,
         }
         for prompt in prompts:
-            if self.is_gemini:
+            if self._model_family == GoogleModelFamily.GEMINI:
                 res = _completion_with_retry(
                     self,
                     prompt=prompt,
@@ -271,7 +305,11 @@ Supported examples:
         Returns:
             The integer number of tokens in the text.
         """
-        if self.is_gemini:
-            raise ValueError("Counting tokens is not yet supported!")
-        result = self.client.count_text_tokens(model=self.model, prompt=text)
-        return result["token_count"]
+        if self._model_family == GoogleModelFamily.GEMINI:
+            result = self.client.count_tokens(text)
+            token_count = result.total_tokens
+        else:
+            result = self.client.count_text_tokens(model=self.model, prompt=text)
+            token_count = result["token_count"]
+
+        return token_count

{langchain_google_genai-0.0.6.dist-info → langchain_google_genai-0.0.8.dist-info}/METADATA

@@ -1,13 +1,16 @@
 Metadata-Version: 2.1
 Name: langchain-google-genai
-Version: 0.0.6
+Version: 0.0.8
 Summary: An integration package connecting Google's genai package and LangChain
 Home-page: https://github.com/langchain-ai/langchain
+License: MIT
 Requires-Python: >=3.9,<4.0
+Classifier: License :: OSI Approved :: MIT License
 Classifier: Programming Language :: Python :: 3
 Classifier: Programming Language :: Python :: 3.9
 Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
 Provides-Extra: images
 Requires-Dist: google-generativeai (>=0.3.1,<0.4.0)
 Requires-Dist: langchain-core (>=0.1,<0.2)

langchain_google_genai-0.0.8.dist-info/RECORD

@@ -0,0 +1,11 @@
+langchain_google_genai/__init__.py,sha256=cDMb1xbsenQtYBACNP0dYPwA7Rt015MT7HC_XP3X-4Y,2304
+langchain_google_genai/_common.py,sha256=1r0VrrBSTZfGprmICZ5OV-W5SK31jKRFFCNE3vJ3jmk,136
+langchain_google_genai/_function_utils.py,sha256=9IVMPQq5lQB8F_whG3mrGOM_tjmP-TMEY3URHfJnjgI,3640
+langchain_google_genai/chat_models.py,sha256=flq1xYC2OYoijOtU9qJy12CZwJ8sdY3va3x0004pl1M,23208
+langchain_google_genai/embeddings.py,sha256=EMa-sDGXUpAPMSyjA2-YXF_TGrlSlqljNeqysAh574s,3951
+langchain_google_genai/llms.py,sha256=Kk7fCrWbbfR1tpiFRYBX6fgEhgAYq7HQVSNvR6UvWqY,10990
+langchain_google_genai/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+langchain_google_genai-0.0.8.dist-info/LICENSE,sha256=DppmdYJVSc1jd0aio6ptnMUn5tIHrdAhQ12SclEBfBg,1072
+langchain_google_genai-0.0.8.dist-info/METADATA,sha256=mwV8mJuT-jVGcq-HQa7IiWjd6HDvVOIVa01lSTe1R-A,2851
+langchain_google_genai-0.0.8.dist-info/WHEEL,sha256=FMvqSimYX_P7y0a7UY-_Mc83r5zkBZsCYPm7Lr0Bsq4,88
+langchain_google_genai-0.0.8.dist-info/RECORD,,

{langchain_google_genai-0.0.6.dist-info → langchain_google_genai-0.0.8.dist-info}/WHEEL

@@ -1,4 +1,4 @@
 Wheel-Version: 1.0
-Generator: poetry-core 1.7.0
+Generator: poetry-core 1.8.1
 Root-Is-Purelib: true
 Tag: py3-none-any

langchain_google_genai-0.0.6.dist-info/RECORD

@@ -1,10 +0,0 @@
-langchain_google_genai/__init__.py,sha256=cDMb1xbsenQtYBACNP0dYPwA7Rt015MT7HC_XP3X-4Y,2304
-langchain_google_genai/_common.py,sha256=1r0VrrBSTZfGprmICZ5OV-W5SK31jKRFFCNE3vJ3jmk,136
-langchain_google_genai/chat_models.py,sha256=451I60uXY24h7bpTxkF3HNccdY0vIRm6wly1PL9URXk,22894
-langchain_google_genai/embeddings.py,sha256=EMa-sDGXUpAPMSyjA2-YXF_TGrlSlqljNeqysAh574s,3951
-langchain_google_genai/llms.py,sha256=xoypoTqhw-p3_1Htk8yURoTpRtjVjFvsqP1WPWwk8eg,9707
-langchain_google_genai/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-langchain_google_genai-0.0.6.dist-info/LICENSE,sha256=DppmdYJVSc1jd0aio6ptnMUn5tIHrdAhQ12SclEBfBg,1072
-langchain_google_genai-0.0.6.dist-info/METADATA,sha256=9IOHkgYT2g3XbCU0kdxETvOJaK5z4FjckdV6I1hmn2w,2736
-langchain_google_genai-0.0.6.dist-info/WHEEL,sha256=d2fvjOD7sXsVzChCqf0Ty0JbHKBaLYwDbGQDwQTnJ50,88
-langchain_google_genai-0.0.6.dist-info/RECORD,,