langchain-google-genai 2.0.9__py3-none-any.whl → 2.0.11__py3-none-any.whl

langchain_google_genai/_common.py

@@ -1,7 +1,9 @@
 from importlib import metadata
-from typing import Optional, Tuple, TypedDict
+from typing import Any, Dict, Optional, Tuple, TypedDict
 
 from google.api_core.gapic_v1.client_info import ClientInfo
+from langchain_core.utils import secret_from_env
+from pydantic import BaseModel, Field, SecretStr
 
 from langchain_google_genai._enums import HarmBlockThreshold, HarmCategory
 
@@ -12,6 +14,97 @@ class GoogleGenerativeAIError(Exception):
     """
 
 
+class _BaseGoogleGenerativeAI(BaseModel):
+    """Base class for Google Generative AI LLMs"""
+
+    model: str = Field(
+        ...,
+        description="""The name of the model to use.
+Supported examples:
+    - gemini-pro
+    - models/text-bison-001""",
+    )
+    """Model name to use."""
+    google_api_key: Optional[SecretStr] = Field(
+        alias="api_key", default_factory=secret_from_env("GOOGLE_API_KEY", default=None)
+    )
+    """Google AI API key.
+    If not specified will be read from env var ``GOOGLE_API_KEY``."""
+    credentials: Any = None
+    "The default custom credentials (google.auth.credentials.Credentials) to use "
+    "when making API calls. If not provided, credentials will be ascertained from "
+    "the GOOGLE_API_KEY envvar"
+    temperature: float = 0.7
+    """Run inference with this temperature. Must by in the closed interval
+       [0.0, 1.0]."""
+    top_p: Optional[float] = None
+    """Decode using nucleus sampling: consider the smallest set of tokens whose
+       probability sum is at least top_p. Must be 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."""
+    max_output_tokens: Optional[int] = Field(default=None, alias="max_tokens")
+    """Maximum number of tokens to include in a candidate. Must be greater than zero.
+       If unset, will default to 64."""
+    n: int = 1
+    """Number of chat completions to generate for each prompt. Note that the API may
+       not return the full n completions if duplicates are generated."""
+    max_retries: int = 6
+    """The maximum number of retries to make when generating."""
+
+    timeout: Optional[float] = None
+    """The maximum number of seconds to wait for a response."""
+
+    client_options: Optional[Dict] = Field(
+        default=None,
+        description=(
+            "A dictionary of client options to pass to the Google API client, "
+            "such as `api_endpoint`."
+        ),
+    )
+    transport: Optional[str] = Field(
+        default=None,
+        description="A string, one of: [`rest`, `grpc`, `grpc_asyncio`].",
+    )
+    additional_headers: Optional[Dict[str, str]] = Field(
+        default=None,
+        description=(
+            "A key-value dictionary representing additional headers for the model call"
+        ),
+    )
+
+    safety_settings: Optional[Dict[HarmCategory, HarmBlockThreshold]] = None
+    """The default safety settings to use for all generations. 
+    
+        For example: 
+
+            from google.generativeai.types.safety_types import HarmBlockThreshold, HarmCategory
+
+            safety_settings = {
+                HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT: HarmBlockThreshold.BLOCK_MEDIUM_AND_ABOVE,
+                HarmCategory.HARM_CATEGORY_HATE_SPEECH: HarmBlockThreshold.BLOCK_ONLY_HIGH,
+                HarmCategory.HARM_CATEGORY_HARASSMENT: HarmBlockThreshold.BLOCK_LOW_AND_ABOVE,
+                HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT: HarmBlockThreshold.BLOCK_NONE,
+            }
+            """  # noqa: E501
+
+    @property
+    def lc_secrets(self) -> Dict[str, str]:
+        return {"google_api_key": "GOOGLE_API_KEY"}
+
+    @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,
+        }
+
+
 def get_user_agent(module: Optional[str] = None) -> Tuple[str, str]:
     r"""Returns a custom user agent header.
 

langchain_google_genai/_function_utils.py

@@ -7,7 +7,6 @@ import logging
 from typing import (
     Any,
     Callable,
-    Collection,
     Dict,
     List,
     Literal,
@@ -22,7 +21,6 @@ from typing import (
 import google.ai.generativelanguage as glm
 import google.ai.generativelanguage_v1beta.types as gapic
 import proto  # type: ignore[import]
-from google.generativeai.types.content_types import ToolDict  # type: ignore[import]
 from langchain_core.tools import BaseTool
 from langchain_core.tools import tool as callable_as_lc_tool
 from langchain_core.utils.function_calling import (
@@ -59,38 +57,21 @@ _ALLOWED_SCHEMA_FIELDS.extend(
 _ALLOWED_SCHEMA_FIELDS_SET = set(_ALLOWED_SCHEMA_FIELDS)
 
 
-class _ToolDictLike(TypedDict):
-    function_declarations: _FunctionDeclarationLikeList
-
-
-class _FunctionDeclarationDict(TypedDict):
-    name: str
-    description: str
-    parameters: Dict[str, Collection[str]]
-
-
-class _ToolDict(TypedDict):
-    function_declarations: Sequence[_FunctionDeclarationDict]
-
-
 # Info: This is a FunctionDeclaration(=fc).
 _FunctionDeclarationLike = Union[
     BaseTool, Type[BaseModel], gapic.FunctionDeclaration, Callable, Dict[str, Any]
 ]
 
-# Info: This mean one tool.
-_FunctionDeclarationLikeList = Sequence[_FunctionDeclarationLike]
+
+class _ToolDict(TypedDict):
+    function_declarations: Sequence[_FunctionDeclarationLike]
 
 
 # Info: This means one tool=Sequence of FunctionDeclaration
 # The dict should be gapic.Tool like. {"function_declarations": [ { "name": ...}.
 # OpenAI like dict is not be accepted. {{'type': 'function', 'function': {'name': ...}
 _ToolsType = Union[
-    gapic.Tool,
-    ToolDict,
-    _ToolDictLike,
-    _FunctionDeclarationLikeList,
-    _FunctionDeclarationLike,
+    gapic.Tool, _ToolDict, _FunctionDeclarationLike, Sequence[_FunctionDeclarationLike]
 ]
 
 
@@ -152,12 +133,12 @@ def convert_to_genai_function_declarations(
     gapic_tool = gapic.Tool()
     for tool in tools:
         if isinstance(tool, gapic.Tool):
-            gapic_tool.function_declarations.extend(tool.function_declarations)
+            gapic_tool.function_declarations.extend(tool.function_declarations)  # type: ignore[union-attr]
         elif isinstance(tool, dict) and "function_declarations" not in tool:
             fd = _format_to_gapic_function_declaration(tool)
             gapic_tool.function_declarations.append(fd)
         elif isinstance(tool, dict):
-            function_declarations = cast(_ToolDictLike, tool)["function_declarations"]
+            function_declarations = cast(_ToolDict, tool)["function_declarations"]
             if not isinstance(function_declarations, collections.abc.Sequence):
                 raise ValueError(
                     "function_declarations should be a list"
@@ -170,7 +151,7 @@ def convert_to_genai_function_declarations(
                 ]
                 gapic_tool.function_declarations.extend(fds)
         else:
-            fd = _format_to_gapic_function_declaration(tool)
+            fd = _format_to_gapic_function_declaration(tool)  # type: ignore[arg-type]
             gapic_tool.function_declarations.append(fd)
     return gapic_tool
 
@@ -235,13 +216,16 @@ def _format_base_tool_to_function_declaration(
             ),
         )
 
-    if issubclass(tool.args_schema, BaseModel):
+    if isinstance(tool.args_schema, dict):
+        schema = tool.args_schema
+    elif issubclass(tool.args_schema, BaseModel):
         schema = tool.args_schema.model_json_schema()
     elif issubclass(tool.args_schema, BaseModelV1):
         schema = tool.args_schema.schema()
     else:
         raise NotImplementedError(
-            f"args_schema must be a Pydantic BaseModel, got {tool.args_schema}."
+            "args_schema must be a Pydantic BaseModel or JSON schema, "
+            f"got {tool.args_schema}."
         )
     parameters = _dict_to_gapic_schema(schema)
 
@@ -301,10 +285,18 @@ def _get_properties_from_schema(schema: Dict) -> Dict[str, Any]:
             continue
         properties_item: Dict[str, Union[str, int, Dict, List]] = {}
         if v.get("type") or v.get("anyOf") or v.get("type_"):
-            properties_item["type_"] = _get_type_from_schema(v)
+            item_type_ = _get_type_from_schema(v)
+            properties_item["type_"] = item_type_
             if _is_nullable_schema(v):
                 properties_item["nullable"] = True
 
+            # Replace `v` with chosen definition for array / object json types
+            any_of_types = v.get("anyOf")
+            if any_of_types and item_type_ in [glm.Type.ARRAY, glm.Type.OBJECT]:
+                json_type_ = "array" if item_type_ == glm.Type.ARRAY else "object"
+                # Use Index -1 for consistency with `_get_nullable_type_from_schema`
+                v = [val for val in any_of_types if val.get("type") == json_type_][-1]
+
         if v.get("enum"):
             properties_item["enum"] = v["enum"]
 
@@ -364,6 +356,8 @@ def _get_items_from_schema(schema: Union[Dict, List, str]) -> Dict[str, Any]:
             )
         if _is_nullable_schema(schema):
             items["nullable"] = True
+        if "required" in schema:
+            items["required"] = schema["required"]
     else:
         # str
         items["type_"] = _get_type_from_schema({"type": schema})

langchain_google_genai/chat_models.py

@@ -35,6 +35,7 @@ from google.ai.generativelanguage_v1beta.types import (
     Content,
     FileData,
     FunctionCall,
+    FunctionDeclaration,
     FunctionResponse,
     GenerateContentRequest,
     GenerateContentResponse,
@@ -44,12 +45,8 @@ from google.ai.generativelanguage_v1beta.types import (
     ToolConfig,
     VideoMetadata,
 )
-from google.generativeai.caching import CachedContent  # type: ignore[import]
-from google.generativeai.types import Tool as GoogleTool  # type: ignore[import]
-from google.generativeai.types import caching_types, content_types
-from google.generativeai.types.content_types import (  # type: ignore[import]
-    FunctionDeclarationType,
-    ToolDict,
+from google.ai.generativelanguage_v1beta.types import (
+    Tool as GoogleTool,
 )
 from langchain_core.callbacks.manager import (
     AsyncCallbackManagerForLLMRun,
@@ -76,7 +73,7 @@ from langchain_core.output_parsers.openai_tools import (
 )
 from langchain_core.outputs import ChatGeneration, ChatGenerationChunk, ChatResult
 from langchain_core.runnables import Runnable, RunnablePassthrough
-from langchain_core.utils import secret_from_env
+from langchain_core.tools import BaseTool
 from langchain_core.utils.function_calling import convert_to_openai_tool
 from pydantic import (
     BaseModel,
@@ -97,24 +94,32 @@ from typing_extensions import Self
 from langchain_google_genai._common import (
     GoogleGenerativeAIError,
     SafetySettingDict,
+    _BaseGoogleGenerativeAI,
     get_client_info,
 )
 from langchain_google_genai._function_utils import (
     _tool_choice_to_tool_config,
     _ToolChoiceType,
     _ToolConfigDict,
+    _ToolDict,
     convert_to_genai_function_declarations,
     is_basemodel_subclass_safe,
     tool_to_dict,
 )
 from langchain_google_genai._image_utils import ImageBytesLoader
-from langchain_google_genai.llms import _BaseGoogleGenerativeAI
 
 from . import _genai_extension as genaix
 
 logger = logging.getLogger(__name__)
 
 
+_FunctionDeclarationType = Union[
+    FunctionDeclaration,
+    dict[str, Any],
+    Callable[..., Any],
+]
+
+
 class ChatGoogleGenerativeAIError(GoogleGenerativeAIError):
     """
     Custom exception class for errors associated with the `Google GenAI` API.
@@ -301,9 +306,12 @@ def _convert_to_parts(
     return parts
 
 
-def _convert_tool_message_to_part(message: ToolMessage | FunctionMessage) -> Part:
+def _convert_tool_message_to_part(
+    message: ToolMessage | FunctionMessage, name: Optional[str] = None
+) -> Part:
     """Converts a tool or function message to a google part."""
-    name = message.name
+    # Legacy agent stores tool name in message.additional_kwargs instead of message.name
+    name = message.name or name or message.additional_kwargs.get("name")
     response: Any
     if not isinstance(message.content, str):
         response = message.content
@@ -331,16 +339,17 @@ def _get_ai_message_tool_messages_parts(
     list of Parts.
     """
     # We are interested only in the tool messages that are part of the AI message
-    tool_calls_ids = [tool_call["id"] for tool_call in ai_message.tool_calls]
+    tool_calls_ids = {tool_call["id"]: tool_call for tool_call in ai_message.tool_calls}
     parts = []
     for i, message in enumerate(tool_messages):
         if not tool_calls_ids:
             break
         if message.tool_call_id in tool_calls_ids:
-            # remove the id from the list, so that we do not iterate over it again
-            tool_calls_ids.remove(message.tool_call_id)
-            part = _convert_tool_message_to_part(message)
+            tool_call = tool_calls_ids[message.tool_call_id]
+            part = _convert_tool_message_to_part(message, name=tool_call.get("name"))
             parts.append(part)
+            # remove the id from the dict, so that we do not iterate over it again
+            tool_calls_ids.pop(message.tool_call_id)
     return parts
 
 
@@ -360,8 +369,14 @@ def _parse_chat_history(
         message for message in input_messages if isinstance(message, ToolMessage)
     ]
     for i, message in enumerate(messages_without_tool_messages):
-        if i == 0 and isinstance(message, SystemMessage):
-            system_instruction = Content(parts=_convert_to_parts(message.content))
+        if isinstance(message, SystemMessage):
+            system_parts = _convert_to_parts(message.content)
+            if i == 0:
+                system_instruction = Content(parts=system_parts)
+            elif system_instruction is not None:
+                system_instruction.parts.extend(system_parts)
+            else:
+                pass
             continue
         elif isinstance(message, AIMessage):
             role = "model"
@@ -776,11 +791,6 @@ class ChatGoogleGenerativeAI(_BaseGoogleGenerativeAI, BaseChatModel):
 
     client: Any = Field(default=None, exclude=True)  #: :meta private:
     async_client_running: Any = Field(default=None, exclude=True)  #: :meta private:
-    google_api_key: Optional[SecretStr] = Field(
-        alias="api_key", default_factory=secret_from_env("GOOGLE_API_KEY", default=None)
-    )
-    """Google AI API key.
-    If not specified will be read from env var ``GOOGLE_API_KEY``."""
     default_metadata: Sequence[Tuple[str, str]] = Field(
         default_factory=list
     )  #: :meta private:
@@ -928,8 +938,8 @@ class ChatGoogleGenerativeAI(_BaseGoogleGenerativeAI, BaseChatModel):
         stop: Optional[List[str]] = None,
         run_manager: Optional[CallbackManagerForLLMRun] = None,
         *,
-        tools: Optional[Sequence[Union[ToolDict, GoogleTool]]] = None,
-        functions: Optional[Sequence[FunctionDeclarationType]] = None,
+        tools: Optional[Sequence[Union[_ToolDict, GoogleTool]]] = None,
+        functions: Optional[Sequence[_FunctionDeclarationType]] = None,
         safety_settings: Optional[SafetySettingDict] = None,
         tool_config: Optional[Union[Dict, _ToolConfigDict]] = None,
         generation_config: Optional[Dict[str, Any]] = None,
@@ -962,8 +972,8 @@ class ChatGoogleGenerativeAI(_BaseGoogleGenerativeAI, BaseChatModel):
         stop: Optional[List[str]] = None,
         run_manager: Optional[AsyncCallbackManagerForLLMRun] = None,
         *,
-        tools: Optional[Sequence[Union[ToolDict, GoogleTool]]] = None,
-        functions: Optional[Sequence[FunctionDeclarationType]] = None,
+        tools: Optional[Sequence[Union[_ToolDict, GoogleTool]]] = None,
+        functions: Optional[Sequence[_FunctionDeclarationType]] = None,
         safety_settings: Optional[SafetySettingDict] = None,
         tool_config: Optional[Union[Dict, _ToolConfigDict]] = None,
         generation_config: Optional[Dict[str, Any]] = None,
@@ -1011,8 +1021,8 @@ class ChatGoogleGenerativeAI(_BaseGoogleGenerativeAI, BaseChatModel):
         stop: Optional[List[str]] = None,
         run_manager: Optional[CallbackManagerForLLMRun] = None,
         *,
-        tools: Optional[Sequence[Union[ToolDict, GoogleTool]]] = None,
-        functions: Optional[Sequence[FunctionDeclarationType]] = None,
+        tools: Optional[Sequence[Union[_ToolDict, GoogleTool]]] = None,
+        functions: Optional[Sequence[_FunctionDeclarationType]] = None,
         safety_settings: Optional[SafetySettingDict] = None,
         tool_config: Optional[Union[Dict, _ToolConfigDict]] = None,
         generation_config: Optional[Dict[str, Any]] = None,
@@ -1073,8 +1083,8 @@ class ChatGoogleGenerativeAI(_BaseGoogleGenerativeAI, BaseChatModel):
         stop: Optional[List[str]] = None,
         run_manager: Optional[AsyncCallbackManagerForLLMRun] = None,
         *,
-        tools: Optional[Sequence[Union[ToolDict, GoogleTool]]] = None,
-        functions: Optional[Sequence[FunctionDeclarationType]] = None,
+        tools: Optional[Sequence[Union[_ToolDict, GoogleTool]]] = None,
+        functions: Optional[Sequence[_FunctionDeclarationType]] = None,
         safety_settings: Optional[SafetySettingDict] = None,
         tool_config: Optional[Union[Dict, _ToolConfigDict]] = None,
         generation_config: Optional[Dict[str, Any]] = None,
@@ -1148,8 +1158,8 @@ class ChatGoogleGenerativeAI(_BaseGoogleGenerativeAI, BaseChatModel):
         messages: List[BaseMessage],
         *,
         stop: Optional[List[str]] = None,
-        tools: Optional[Sequence[Union[ToolDict, GoogleTool]]] = None,
-        functions: Optional[Sequence[FunctionDeclarationType]] = None,
+        tools: Optional[Sequence[Union[_ToolDict, GoogleTool]]] = None,
+        functions: Optional[Sequence[_FunctionDeclarationType]] = None,
         safety_settings: Optional[SafetySettingDict] = None,
         tool_config: Optional[Union[Dict, _ToolConfigDict]] = None,
         tool_choice: Optional[Union[_ToolChoiceType, bool]] = None,
@@ -1241,8 +1251,8 @@ class ChatGoogleGenerativeAI(_BaseGoogleGenerativeAI, BaseChatModel):
             )
         else:
             parser = JsonOutputToolsParser()
-        tool_choice = _get_tool_name(schema) if self._supports_tool_choice else None
-        llm = self.bind_tools([schema], tool_choice=tool_choice)
+        tool_choice = _get_tool_name(schema) if self._supports_tool_choice else None  # type: ignore[arg-type]
+        llm = self.bind_tools([schema], tool_choice=tool_choice)  # type: ignore[list-item]
         if include_raw:
             parser_with_fallback = RunnablePassthrough.assign(
                 parsed=itemgetter("raw") | parser, parsing_error=lambda _: None
@@ -1256,7 +1266,9 @@ class ChatGoogleGenerativeAI(_BaseGoogleGenerativeAI, BaseChatModel):
 
     def bind_tools(
         self,
-        tools: Sequence[Union[ToolDict, GoogleTool]],
+        tools: Sequence[
+            dict[str, Any] | type | Callable[..., Any] | BaseTool | GoogleTool
+        ],
         tool_config: Optional[Union[Dict, _ToolConfigDict]] = None,
         *,
         tool_choice: Optional[Union[_ToolChoiceType, bool]] = None,
@@ -1293,90 +1305,6 @@ class ChatGoogleGenerativeAI(_BaseGoogleGenerativeAI, BaseChatModel):
             pass
         return self.bind(tools=formatted_tools, **kwargs)
 
-    def create_cached_content(
-        self,
-        contents: Union[List[BaseMessage], content_types.ContentsType],
-        *,
-        display_name: str | None = None,
-        tools: Union[ToolDict, GoogleTool, None] = None,
-        tool_choice: Optional[Union[_ToolChoiceType, bool]] = None,
-        ttl: Optional[caching_types.TTLTypes] = None,
-        expire_time: Optional[caching_types.ExpireTimeTypes] = None,
-    ) -> str:
-        """
-
-        Args:
-            display_name: The user-generated meaningful display name
-                of the cached content. `display_name` must be no
-                more than 128 unicode characters.
-            contents: Contents to cache.
-            tools: A list of `Tools` the model may use to generate response.
-            tool_choice: Which tool to require the model to call.
-            ttl: TTL for cached resource (in seconds). Defaults to 1 hour.
-                 `ttl` and `expire_time` are exclusive arguments.
-            expire_time: Expiration time for cached resource.
-                `ttl` and `expire_time` are exclusive arguments.
-        """
-        system: Optional[content_types.ContentType] = None
-        genai_contents: list = []
-        if all(isinstance(c, BaseMessage) for c in contents):
-            system, genai_contents = _parse_chat_history(
-                contents,
-                convert_system_message_to_human=self.convert_system_message_to_human,
-            )
-        elif any(isinstance(c, BaseMessage) for c in contents):
-            raise ValueError(
-                f"'contents' must either be a list of "
-                f"langchain_core.messages.BaseMessage or a list "
-                f"google.generativeai.types.content_types.ContentType, but not a mix "
-                f"of the two. Received {contents}"
-            )
-        else:
-            for content in contents:
-                if hasattr(content, "role") and content.role == "system":
-                    if system is not None:
-                        warnings.warn(
-                            "Received multiple pieces of content with role 'system'. "
-                            "Should only be one set of system instructions. Ignoring "
-                            "all but the first 'system' content."
-                        )
-                    else:
-                        system = content
-                elif isinstance(content, dict) and content.get("role") == "system":
-                    if system is not None:
-                        warnings.warn(
-                            "Received multiple pieces of content with role 'system'. "
-                            "Should only be one set of system instructions. Ignoring "
-                            "all but the first 'system' content."
-                        )
-                    else:
-                        system = content
-                else:
-                    genai_contents.append(content)
-        if tools:
-            genai_tools = [convert_to_genai_function_declarations(tools)]
-        else:
-            genai_tools = None
-        if tool_choice and genai_tools:
-            all_names = [f.name for t in genai_tools for f in t.function_declarations]
-            tool_config = _tool_choice_to_tool_config(tool_choice, all_names)
-            genai_tool_config = ToolConfig(
-                function_calling_config=tool_config["function_calling_config"]
-            )
-        else:
-            genai_tool_config = None
-        cached_content = CachedContent.create(
-            model=self.model,
-            system_instruction=system,
-            contents=genai_contents,
-            display_name=display_name,
-            tools=genai_tools,
-            tool_config=genai_tool_config,
-            ttl=ttl,
-            expire_time=expire_time,
-        )
-        return cached_content.name
-
     @property
     def _supports_tool_choice(self) -> bool:
         return (
@@ -1387,7 +1315,7 @@ class ChatGoogleGenerativeAI(_BaseGoogleGenerativeAI, BaseChatModel):
 
 
 def _get_tool_name(
-    tool: Union[ToolDict, GoogleTool],
+    tool: Union[_ToolDict, GoogleTool],
 ) -> str:
     genai_tool = tool_to_dict(convert_to_genai_function_declarations([tool]))
     return [f["name"] for f in genai_tool["function_declarations"]][0]  # type: ignore[index]

langchain_google_genai/llms.py

@@ -1,208 +1,21 @@
 from __future__ import annotations
 
-from enum import Enum, auto
-from typing import Any, Callable, Dict, Iterator, List, Optional, Union
+from typing import Any, Iterator, List, Optional
 
-import google.api_core
-import google.generativeai as genai  # type: ignore[import]
 from langchain_core.callbacks import (
-    AsyncCallbackManagerForLLMRun,
     CallbackManagerForLLMRun,
 )
-from langchain_core.language_models import LangSmithParams, LanguageModelInput
-from langchain_core.language_models.llms import BaseLLM, create_base_retry_decorator
+from langchain_core.language_models import LangSmithParams
+from langchain_core.language_models.llms import BaseLLM
+from langchain_core.messages import HumanMessage
 from langchain_core.outputs import Generation, GenerationChunk, LLMResult
-from langchain_core.utils import secret_from_env
-from pydantic import BaseModel, ConfigDict, Field, SecretStr, model_validator
+from pydantic import ConfigDict, model_validator
 from typing_extensions import Self
 
-from langchain_google_genai._enums import (
-    HarmBlockThreshold,
-    HarmCategory,
+from langchain_google_genai._common import (
+    _BaseGoogleGenerativeAI,
 )
-
-
-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,
-    *,
-    max_retries: int = 1,
-    run_manager: Optional[
-        Union[AsyncCallbackManagerForLLMRun, CallbackManagerForLLMRun]
-    ] = None,
-) -> Callable[[Any], Any]:
-    """Creates a retry decorator for Vertex / Palm LLMs."""
-
-    errors = [
-        google.api_core.exceptions.ResourceExhausted,
-        google.api_core.exceptions.ServiceUnavailable,
-        google.api_core.exceptions.Aborted,
-        google.api_core.exceptions.DeadlineExceeded,
-        google.api_core.exceptions.GoogleAPIError,
-    ]
-    decorator = create_base_retry_decorator(
-        error_types=errors, max_retries=max_retries, run_manager=run_manager
-    )
-    return decorator
-
-
-def _completion_with_retry(
-    llm: GoogleGenerativeAI,
-    prompt: LanguageModelInput,
-    is_gemini: bool = False,
-    stream: bool = False,
-    run_manager: Optional[CallbackManagerForLLMRun] = None,
-    **kwargs: Any,
-) -> Any:
-    """Use tenacity to retry the completion call."""
-    retry_decorator = _create_retry_decorator(
-        llm, max_retries=llm.max_retries, run_manager=run_manager
-    )
-
-    @retry_decorator
-    def _completion_with_retry(
-        prompt: LanguageModelInput, is_gemini: bool, stream: bool, **kwargs: Any
-    ) -> Any:
-        generation_config = kwargs.get("generation_config", {})
-        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,
-                    safety_settings=kwargs.pop("safety_settings", None),
-                    request_options={"timeout": llm.timeout} if llm.timeout else None,
-                )
-            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 _strip_erroneous_leading_spaces(text: str) -> str:
-    """Strip erroneous leading spaces from text.
-
-    The PaLM API will sometimes erroneously return a single leading space in all
-    lines > 1. This function strips that space.
-    """
-    has_leading_space = all(not line or line[0] == " " for line in text.split("\n")[1:])
-    if has_leading_space:
-        return text.replace("\n ", "\n")
-    else:
-        return text
-
-
-class _BaseGoogleGenerativeAI(BaseModel):
-    """Base class for Google Generative AI LLMs"""
-
-    model: str = Field(
-        ...,
-        description="""The name of the model to use.
-Supported examples:
-    - gemini-pro
-    - models/text-bison-001""",
-    )
-    """Model name to use."""
-    google_api_key: Optional[SecretStr] = Field(
-        alias="api_key", default_factory=secret_from_env("GOOGLE_API_KEY", default=None)
-    )
-    credentials: Any = None
-    "The default custom credentials (google.auth.credentials.Credentials) to use "
-    "when making API calls. If not provided, credentials will be ascertained from "
-    "the GOOGLE_API_KEY envvar"
-    temperature: float = 0.7
-    """Run inference with this temperature. Must by in the closed interval
-       [0.0, 1.0]."""
-    top_p: Optional[float] = None
-    """Decode using nucleus sampling: consider the smallest set of tokens whose
-       probability sum is at least top_p. Must be 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."""
-    max_output_tokens: Optional[int] = Field(default=None, alias="max_tokens")
-    """Maximum number of tokens to include in a candidate. Must be greater than zero.
-       If unset, will default to 64."""
-    n: int = 1
-    """Number of chat completions to generate for each prompt. Note that the API may
-       not return the full n completions if duplicates are generated."""
-    max_retries: int = 6
-    """The maximum number of retries to make when generating."""
-
-    timeout: Optional[float] = None
-    """The maximum number of seconds to wait for a response."""
-
-    client_options: Optional[Dict] = Field(
-        default=None,
-        description=(
-            "A dictionary of client options to pass to the Google API client, "
-            "such as `api_endpoint`."
-        ),
-    )
-    transport: Optional[str] = Field(
-        default=None,
-        description="A string, one of: [`rest`, `grpc`, `grpc_asyncio`].",
-    )
-    additional_headers: Optional[Dict[str, str]] = Field(
-        default=None,
-        description=(
-            "A key-value dictionary representing additional headers for the model call"
-        ),
-    )
-
-    safety_settings: Optional[Dict[HarmCategory, HarmBlockThreshold]] = None
-    """The default safety settings to use for all generations. 
-    
-        For example: 
-
-            from google.generativeai.types.safety_types import HarmBlockThreshold, HarmCategory
-
-            safety_settings = {
-                HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT: HarmBlockThreshold.BLOCK_MEDIUM_AND_ABOVE,
-                HarmCategory.HARM_CATEGORY_HATE_SPEECH: HarmBlockThreshold.BLOCK_ONLY_HIGH,
-                HarmCategory.HARM_CATEGORY_HARASSMENT: HarmBlockThreshold.BLOCK_LOW_AND_ABOVE,
-                HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT: HarmBlockThreshold.BLOCK_NONE,
-            }
-            """  # noqa: E501
-
-    @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,
-        }
+from langchain_google_genai.chat_models import ChatGoogleGenerativeAI
 
 
 class GoogleGenerativeAI(_BaseGoogleGenerativeAI, BaseLLM):
@@ -223,53 +36,21 @@ class GoogleGenerativeAI(_BaseGoogleGenerativeAI, BaseLLM):
     @model_validator(mode="after")
     def validate_environment(self) -> Self:
         """Validates params and passes them to google-generativeai package."""
-        if self.credentials:
-            genai.configure(
-                credentials=self.credentials,
-                transport=self.transport,
-                client_options=self.client_options,
-            )
-        else:
-            if isinstance(self.google_api_key, SecretStr):
-                google_api_key: Optional[str] = self.google_api_key.get_secret_value()
-            else:
-                google_api_key = self.google_api_key
-            genai.configure(
-                api_key=google_api_key,
-                transport=self.transport,
-                client_options=self.client_options,
-            )
-
-        model_name = self.model
-
-        safety_settings = self.safety_settings
-
-        if safety_settings and (
-            not GoogleModelFamily(model_name) == GoogleModelFamily.GEMINI
-        ):
-            raise ValueError("Safety settings are only supported for Gemini models")
-
-        if GoogleModelFamily(model_name) == GoogleModelFamily.GEMINI:
-            self.client = genai.GenerativeModel(
-                model_name=model_name, safety_settings=safety_settings
-            )
-        else:
-            self.client = genai
-
-        if self.temperature is not None and not 0 <= self.temperature <= 1:
-            raise ValueError("temperature must be in the range [0.0, 1.0]")
-
-        if self.top_p is not None and not 0 <= self.top_p <= 1:
-            raise ValueError("top_p must be in the range [0.0, 1.0]")
-
-        if self.top_k is not None and self.top_k <= 0:
-            raise ValueError("top_k must be positive")
 
-        if self.max_output_tokens is not None and self.max_output_tokens <= 0:
-            raise ValueError("max_output_tokens must be greater than zero")
-
-        if self.timeout is not None and self.timeout <= 0:
-            raise ValueError("timeout must be greater than zero")
+        self.client = ChatGoogleGenerativeAI(
+            api_key=self.google_api_key,
+            credentials=self.credentials,
+            temperature=self.temperature,
+            top_p=self.top_p,
+            top_k=self.top_k,
+            max_tokens=self.max_output_tokens,
+            timeout=self.timeout,
+            model=self.model,
+            client_options=self.client_options,
+            transport=self.transport,
+            additional_headers=self.additional_headers,
+            safety_settings=self.safety_settings,
+        )
 
         return self
 
@@ -290,58 +71,26 @@ class GoogleGenerativeAI(_BaseGoogleGenerativeAI, BaseLLM):
         run_manager: Optional[CallbackManagerForLLMRun] = None,
         **kwargs: Any,
     ) -> LLMResult:
-        generations: List[List[Generation]] = []
-        generation_config = {
-            "stop_sequences": stop,
-            "temperature": self.temperature,
-            "top_p": self.top_p,
-            "top_k": self.top_k,
-            "max_output_tokens": self.max_output_tokens,
-            "candidate_count": self.n,
-        }
+        generations = []
         for prompt in prompts:
-            if self._model_family == GoogleModelFamily.GEMINI:
-                res = _completion_with_retry(
-                    self,
-                    prompt=prompt,
-                    stream=False,
-                    is_gemini=True,
-                    run_manager=run_manager,
-                    generation_config=generation_config,
-                    safety_settings=kwargs.pop("safety_settings", None),
-                )
-                generation_info = None
-                if res.usage_metadata is not None:
-                    generation_info = {
-                        "usage_metadata": res.to_dict().get("usage_metadata")
-                    }
-
-                candidates = [
-                    "".join([p.text for p in c.content.parts]) for c in res.candidates
+            chat_result = self.client._generate(
+                [HumanMessage(content=prompt)],
+                stop=stop,
+                run_manager=run_manager,
+                **kwargs,
+            )
+            generations.append(
+                [
+                    Generation(
+                        text=g.message.content,
+                        generation_info={
+                            **g.generation_info,
+                            **{"usage_metadata": g.message.usage_metadata},
+                        },
+                    )
+                    for g in chat_result.generations
                 ]
-                generations.append(
-                    [
-                        Generation(text=c, generation_info=generation_info)
-                        for c in candidates
-                    ]
-                )
-            else:
-                res = _completion_with_retry(
-                    self,
-                    model=self.model,
-                    prompt=prompt,
-                    stream=False,
-                    is_gemini=False,
-                    run_manager=run_manager,
-                    **generation_config,
-                )
-                prompt_generations = []
-                for candidate in res.candidates:
-                    raw_text = candidate["output"]
-                    stripped_text = _strip_erroneous_leading_spaces(raw_text)
-                    prompt_generations.append(Generation(text=stripped_text))
-                generations.append(prompt_generations)
-
+            )
         return LLMResult(generations=generations)
 
     def _stream(
@@ -351,31 +100,17 @@ class GoogleGenerativeAI(_BaseGoogleGenerativeAI, BaseLLM):
         run_manager: Optional[CallbackManagerForLLMRun] = None,
         **kwargs: Any,
     ) -> Iterator[GenerationChunk]:
-        generation_config = {
-            "stop_sequences": stop,
-            "temperature": self.temperature,
-            "top_p": self.top_p,
-            "top_k": self.top_k,
-            "max_output_tokens": self.max_output_tokens,
-            "candidate_count": self.n,
-        }
-        generation_config = generation_config | kwargs.get("generation_config", {})
-
-        for stream_resp in _completion_with_retry(
-            self,
-            prompt,
-            stream=True,
-            is_gemini=True,
+        for stream_chunk in self.client._stream(
+            [HumanMessage(content=prompt)],
+            stop=stop,
             run_manager=run_manager,
-            generation_config=generation_config,
-            safety_settings=kwargs.pop("safety_settings", None),
             **kwargs,
         ):
-            chunk = GenerationChunk(text=stream_resp.text)
+            chunk = GenerationChunk(text=stream_chunk.message.content)
             yield chunk
             if run_manager:
                 run_manager.on_llm_new_token(
-                    stream_resp.text,
+                    chunk.text,
                     chunk=chunk,
                     verbose=self.verbose,
                 )
@@ -383,7 +118,7 @@ class GoogleGenerativeAI(_BaseGoogleGenerativeAI, BaseLLM):
     @property
     def _llm_type(self) -> str:
         """Return type of llm."""
-        return "google_palm"
+        return "google_gemini"
 
     def get_num_tokens(self, text: str) -> int:
         """Get the number of tokens present in the text.
@@ -396,11 +131,4 @@ class GoogleGenerativeAI(_BaseGoogleGenerativeAI, BaseLLM):
         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
+        return self.client.get_num_tokens(text)

{langchain_google_genai-2.0.9.dist-info → langchain_google_genai-2.0.11.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: langchain-google-genai
-Version: 2.0.9
+Version: 2.0.11
 Summary: An integration package connecting Google's genai package and LangChain
 Home-page: https://github.com/langchain-ai/langchain-google
 License: MIT
@@ -12,8 +12,8 @@ Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
 Requires-Dist: filetype (>=1.2.0,<2.0.0)
-Requires-Dist: google-generativeai (>=0.8.0,<0.9.0)
-Requires-Dist: langchain-core (>=0.3.27,<0.4.0)
+Requires-Dist: google-ai-generativelanguage (>=0.6.16,<0.7.0)
+Requires-Dist: langchain-core (>=0.3.37,<0.4.0)
 Requires-Dist: pydantic (>=2,<3)
 Project-URL: Repository, https://github.com/langchain-ai/langchain-google
 Project-URL: Source Code, https://github.com/langchain-ai/langchain-google/tree/main/libs/genai

{langchain_google_genai-2.0.9.dist-info → langchain_google_genai-2.0.11.dist-info}/RECORD

@@ -1,16 +1,16 @@
 langchain_google_genai/__init__.py,sha256=Oji-S2KYWrku1wyQEskY84IOfY8MfRhujjJ4d7hbsk4,2758
-langchain_google_genai/_common.py,sha256=ASlwE8hEbvOm55BVF_D4rf2nl7RYsnpsi5xbM6DW3Cc,1576
+langchain_google_genai/_common.py,sha256=KaNewDLkaQtDEgTMc631kjOKqKl1QoUKqD9iOIXNf-0,5275
 langchain_google_genai/_enums.py,sha256=KLPmxS1K83K4HjBIXFaXoL_sFEOv8Hq-2B2PDMKyDgo,197
-langchain_google_genai/_function_utils.py,sha256=xcUwX2DmGM4UwH7bhBC6W9E5oVAE8k_l9lMYzJUJwA0,17433
+langchain_google_genai/_function_utils.py,sha256=FPZ4CxI4iTFco_W6oWNCNV_lgNbCLRbj7Gf-D0zeTCY,17736
 langchain_google_genai/_genai_extension.py,sha256=81a4ly5ZHlqMf37uJfdB8K41qE6J5ujLnbUypIfFf2o,20775
 langchain_google_genai/_image_utils.py,sha256=tPrQyMvVmO8xkuow1SvA91omxUEv9ZUy1EMHNGjMAKY,5202
-langchain_google_genai/chat_models.py,sha256=DpTIfMdO_HvuM7s6T3BMuVLfVH0n_ndAq-eiGHKncug,54547
+langchain_google_genai/chat_models.py,sha256=hRCxo2eagOnz7suKPiaQO6W1XyBxKr1UCy9AeK71xRs,50920
 langchain_google_genai/embeddings.py,sha256=jQRWPXD9twXoVBlXJQG7Duz0fb8UC0kgRzzwAmW3Dic,10146
 langchain_google_genai/genai_aqa.py,sha256=qB6h3-BSXqe0YLR3eeVllYzmNKK6ofI6xJLdBahUVZo,4300
 langchain_google_genai/google_vector_store.py,sha256=4wvhIiOmc3Fo046FyafPmT9NBCLek-9bgluvuTfrbpQ,16148
-langchain_google_genai/llms.py,sha256=EPUgkz5aqKOyKbztT7br8w60Uo5D_X_bF5qP-zd6iLs,14593
+langchain_google_genai/llms.py,sha256=QNPitkORf86w8WQpTbjuPFCQFkB-qKRMW2phhRBwAEA,4318
 langchain_google_genai/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-langchain_google_genai-2.0.9.dist-info/LICENSE,sha256=DppmdYJVSc1jd0aio6ptnMUn5tIHrdAhQ12SclEBfBg,1072
-langchain_google_genai-2.0.9.dist-info/METADATA,sha256=lYcZ7V1qdeG84ekvofzoFACmjjMolZ-Hevgigsdw7dQ,3594
-langchain_google_genai-2.0.9.dist-info/WHEEL,sha256=FMvqSimYX_P7y0a7UY-_Mc83r5zkBZsCYPm7Lr0Bsq4,88
-langchain_google_genai-2.0.9.dist-info/RECORD,,
+langchain_google_genai-2.0.11.dist-info/LICENSE,sha256=DppmdYJVSc1jd0aio6ptnMUn5tIHrdAhQ12SclEBfBg,1072
+langchain_google_genai-2.0.11.dist-info/METADATA,sha256=QQSJFXoI4IeDHVyAFT5Xtic6aJ5cd_kY1qCP-B6NX1c,3605
+langchain_google_genai-2.0.11.dist-info/WHEEL,sha256=FMvqSimYX_P7y0a7UY-_Mc83r5zkBZsCYPm7Lr0Bsq4,88
+langchain_google_genai-2.0.11.dist-info/RECORD,,