langchain-google-genai 1.0.2__py3-none-any.whl → 1.0.4__py3-none-any.whl

langchain_google_genai/embeddings.py

@@ -1,12 +1,19 @@
 from typing import Any, Dict, List, Optional
 
 # TODO: remove ignore once the google package is published with types
-import google.generativeai as genai  # type: ignore[import]
+from google.ai.generativelanguage_v1beta.types import (
+    BatchEmbedContentsRequest,
+    EmbedContentRequest,
+)
 from langchain_core.embeddings import Embeddings
 from langchain_core.pydantic_v1 import BaseModel, Field, SecretStr, root_validator
 from langchain_core.utils import get_from_dict_or_env
 
-from langchain_google_genai._common import GoogleGenerativeAIError
+from langchain_google_genai._common import (
+    GoogleGenerativeAIError,
+    get_client_info,
+)
+from langchain_google_genai._genai_extension import build_generative_service
 
 
 class GoogleGenerativeAIEmbeddings(BaseModel, Embeddings):
@@ -27,6 +34,7 @@ class GoogleGenerativeAIEmbeddings(BaseModel, Embeddings):
             embeddings.embed_query("What's our Q1 revenue?")
     """
 
+    client: Any  #: :meta private:
     model: str = Field(
         ...,
         description="The name of the embedding model to use. "
@@ -61,47 +69,52 @@ class GoogleGenerativeAIEmbeddings(BaseModel, Embeddings):
         None,
         description="A string, one of: [`rest`, `grpc`, `grpc_asyncio`].",
     )
+    request_options: Optional[Dict] = Field(
+        None,
+        description="A dictionary of request options to pass to the Google API client."
+        "Example: `{'timeout': 10}`",
+    )
 
     @root_validator()
     def validate_environment(cls, values: Dict) -> Dict:
         """Validates params and passes them to google-generativeai package."""
-        if values.get("credentials"):
-            genai.configure(
-                credentials=values.get("credentials"),
-                transport=values.get("transport"),
-                client_options=values.get("client_options"),
-            )
-        else:
-            google_api_key = get_from_dict_or_env(
-                values, "google_api_key", "GOOGLE_API_KEY"
-            )
-            if isinstance(google_api_key, SecretStr):
-                google_api_key = google_api_key.get_secret_value()
-
-            genai.configure(
-                api_key=google_api_key,
-                transport=values.get("transport"),
-                client_options=values.get("client_options"),
-            )
+        google_api_key = get_from_dict_or_env(
+            values, "google_api_key", "GOOGLE_API_KEY"
+        )
+        client_info = get_client_info("GoogleGenerativeAIEmbeddings")
+
+        values["client"] = build_generative_service(
+            credentials=values.get("credentials"),
+            api_key=google_api_key,
+            client_info=client_info,
+            client_options=values.get("client_options"),
+        )
         return values
 
-    def _embed(
-        self, texts: List[str], task_type: str, title: Optional[str] = None
-    ) -> List[List[float]]:
-        task_type = self.task_type or "retrieval_document"
-        try:
-            result = genai.embed_content(
-                model=self.model,
-                content=texts,
-                task_type=task_type,
-                title=title,
-            )
-        except Exception as e:
-            raise GoogleGenerativeAIError(f"Error embedding content: {e}") from e
-        return result["embedding"]
+    def _prepare_request(
+        self,
+        text: str,
+        task_type: Optional[str] = None,
+        title: Optional[str] = None,
+        output_dimensionality: Optional[int] = None,
+    ) -> EmbedContentRequest:
+        task_type = self.task_type or task_type or "RETRIEVAL_DOCUMENT"
+        # https://ai.google.dev/api/rest/v1/models/batchEmbedContents#EmbedContentRequest
+        request = EmbedContentRequest(
+            content={"parts": [{"text": text}]},
+            model=self.model,
+            task_type=task_type.upper(),
+            title=title,
+            output_dimensionality=output_dimensionality,
+        )
+        return request
 
     def embed_documents(
-        self, texts: List[str], batch_size: int = 5
+        self,
+        texts: List[str],
+        task_type: Optional[str] = None,
+        titles: Optional[List[str]] = None,
+        output_dimensionality: Optional[int] = None,
     ) -> List[List[float]]:
         """Embed a list of strings. Vertex AI currently
         sets a max batch size of 5 strings.
@@ -109,21 +122,58 @@ class GoogleGenerativeAIEmbeddings(BaseModel, Embeddings):
         Args:
             texts: List[str] The list of strings to embed.
             batch_size: [int] The batch size of embeddings to send to the model
+            task_type: task_type (https://ai.google.dev/api/rest/v1/TaskType)
+            titles: An optional list of titles for texts provided.
+            Only applicable when TaskType is RETRIEVAL_DOCUMENT.
+            output_dimensionality: Optional reduced dimension for the output embedding.
+            https://ai.google.dev/api/rest/v1/models/batchEmbedContents#EmbedContentRequest
 
         Returns:
             List of embeddings, one for each text.
         """
-        task_type = self.task_type or "retrieval_document"
-        return self._embed(texts, task_type=task_type)
+        titles = titles if titles else [None] * len(texts)  # type: ignore[list-item]
+        requests = [
+            self._prepare_request(
+                text=text,
+                task_type=task_type,
+                title=title,
+                output_dimensionality=output_dimensionality,
+            )
+            for text, title in zip(texts, titles)
+        ]
 
-    def embed_query(self, text: str) -> List[float]:
+        try:
+            result = self.client.batch_embed_contents(
+                BatchEmbedContentsRequest(requests=requests, model=self.model)
+            )
+        except Exception as e:
+            raise GoogleGenerativeAIError(f"Error embedding content: {e}") from e
+        return [e.values for e in result.embeddings]
+
+    def embed_query(
+        self,
+        text: str,
+        task_type: Optional[str] = None,
+        title: Optional[str] = None,
+        output_dimensionality: Optional[int] = None,
+    ) -> List[float]:
         """Embed a text.
 
         Args:
             text: The text to embed.
+            task_type: task_type (https://ai.google.dev/api/rest/v1/TaskType)
+            title: An optional title for the text.
+            Only applicable when TaskType is RETRIEVAL_DOCUMENT.
+            output_dimensionality: Optional reduced dimension for the output embedding.
+            https://ai.google.dev/api/rest/v1/models/batchEmbedContents#EmbedContentRequest
 
         Returns:
             Embedding for the text.
         """
-        task_type = self.task_type or "retrieval_query"
-        return self._embed([text], task_type=task_type)[0]
+        task_type = self.task_type or "RETRIEVAL_QUERY"
+        return self.embed_documents(
+            [text],
+            task_type=task_type,
+            titles=[title] if title else None,
+            output_dimensionality=output_dimensionality,
+        )[0]

langchain_google_genai/llms.py

@@ -86,6 +86,7 @@ def _completion_with_retry(
                     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:
@@ -143,6 +144,10 @@ Supported examples:
        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(
         None,
         description=(
@@ -169,7 +174,6 @@ Supported examples:
             from google.generativeai.types.safety_types import HarmBlockThreshold, HarmCategory
 
             safety_settings = {
-                HarmCategory.HARM_CATEGORY_UNSPECIFIED: HarmBlockThreshold.BLOCK_NONE,
                 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,
@@ -259,6 +263,9 @@ class GoogleGenerativeAI(_BaseGoogleGenerativeAI, BaseLLM):
         if values["max_output_tokens"] is not None and values["max_output_tokens"] <= 0:
             raise ValueError("max_output_tokens must be greater than zero")
 
+        if values["timeout"] is not None and values["timeout"] <= 0:
+            raise ValueError("timeout must be greater than zero")
+
         return values
 
     def _generate(

{langchain_google_genai-1.0.2.dist-info → langchain_google_genai-1.0.4.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: langchain-google-genai
-Version: 1.0.2
+Version: 1.0.4
 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
 Provides-Extra: images
-Requires-Dist: google-generativeai (>=0.5.0,<0.6.0)
-Requires-Dist: langchain-core (>=0.1.27,<0.2)
+Requires-Dist: google-generativeai (>=0.5.2,<0.6.0)
+Requires-Dist: langchain-core (>=0.1.45,<0.3)
 Requires-Dist: pillow (>=10.1.0,<11.0.0) ; extra == "images"
 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-1.0.4.dist-info/RECORD

@@ -0,0 +1,16 @@
+langchain_google_genai/__init__.py,sha256=Oji-S2KYWrku1wyQEskY84IOfY8MfRhujjJ4d7hbsk4,2758
+langchain_google_genai/_common.py,sha256=ASlwE8hEbvOm55BVF_D4rf2nl7RYsnpsi5xbM6DW3Cc,1576
+langchain_google_genai/_enums.py,sha256=KLPmxS1K83K4HjBIXFaXoL_sFEOv8Hq-2B2PDMKyDgo,197
+langchain_google_genai/_function_utils.py,sha256=d0ApSCjoV9Em1CteBaGznilxrw-PDXqQ4sQa5p7cJfM,8232
+langchain_google_genai/_genai_extension.py,sha256=ZwNwLV22RSf9LB7FOCLsoHzLlQDF-EQmRNYM1an2uSw,20769
+langchain_google_genai/_image_utils.py,sha256=-0XgCMdYkvrIktFvUpy-2GPbFgfSVKZICawB2hiJzus,4999
+langchain_google_genai/chat_models.py,sha256=3ubY5qCaZjFSHKHiiP5XCOPDusUxO8kJvpj9DwLtUG4,33014
+langchain_google_genai/embeddings.py,sha256=kQW6pl1TUGKKSxiUjc7rptp0iEu_Rer1m1LLHKVaW14,6578
+langchain_google_genai/genai_aqa.py,sha256=zcC5cdFYtqLK7DGPhYGvWNeHHeU-CQKA9KhewmsA5lw,4303
+langchain_google_genai/google_vector_store.py,sha256=PPIk-4FmD5UUdmYA2u7VcEhGsiztvRVN59QoGLXdfoA,16139
+langchain_google_genai/llms.py,sha256=S7tOy-c37DElcHtkGl8rwvvg1zOzCxb9PEyJ4E-j7qU,13431
+langchain_google_genai/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+langchain_google_genai-1.0.4.dist-info/LICENSE,sha256=DppmdYJVSc1jd0aio6ptnMUn5tIHrdAhQ12SclEBfBg,1072
+langchain_google_genai-1.0.4.dist-info/METADATA,sha256=nHJMXwe6iuI99J_VO1GQrVgv1BLvVGW5x0mSNswDWq0,3818
+langchain_google_genai-1.0.4.dist-info/WHEEL,sha256=FMvqSimYX_P7y0a7UY-_Mc83r5zkBZsCYPm7Lr0Bsq4,88
+langchain_google_genai-1.0.4.dist-info/RECORD,,

langchain_google_genai-1.0.2.dist-info/RECORD

@@ -1,15 +0,0 @@
-langchain_google_genai/__init__.py,sha256=NiLfU4IyEQSlSl8hbN_fTdX-hrO5tuOflxFlEK0Sy4c,2762
-langchain_google_genai/_common.py,sha256=1r0VrrBSTZfGprmICZ5OV-W5SK31jKRFFCNE3vJ3jmk,136
-langchain_google_genai/_enums.py,sha256=q8IYAqufV-_yZ98FDnsZ3x-1w4804J_e8PrTKT0sdhY,163
-langchain_google_genai/_function_utils.py,sha256=9IVMPQq5lQB8F_whG3mrGOM_tjmP-TMEY3URHfJnjgI,3640
-langchain_google_genai/_genai_extension.py,sha256=2Uqg7vSF0vu1J4AhAyIPzadtpM5JJwZsBXvpItO2TY4,18736
-langchain_google_genai/chat_models.py,sha256=0a_bjJUSuSDVPo4AKk4bcdVhdliKchcO48y4uwYv67U,23557
-langchain_google_genai/embeddings.py,sha256=VbrVGX5f_aKo-XA_2lATB4GblomXYWoj9J1Cz20Pr4E,4562
-langchain_google_genai/genai_aqa.py,sha256=zcC5cdFYtqLK7DGPhYGvWNeHHeU-CQKA9KhewmsA5lw,4303
-langchain_google_genai/google_vector_store.py,sha256=PPIk-4FmD5UUdmYA2u7VcEhGsiztvRVN59QoGLXdfoA,16139
-langchain_google_genai/llms.py,sha256=DEz_KxEjmJafVUXtvlHccfGjMhZWnkcBei9Ey0evNgw,13193
-langchain_google_genai/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-langchain_google_genai-1.0.2.dist-info/LICENSE,sha256=DppmdYJVSc1jd0aio6ptnMUn5tIHrdAhQ12SclEBfBg,1072
-langchain_google_genai-1.0.2.dist-info/METADATA,sha256=GZNA06KlCUNaISeR0xYiZrltekAa7-8sg8XIt_GwFwE,3818
-langchain_google_genai-1.0.2.dist-info/WHEEL,sha256=FMvqSimYX_P7y0a7UY-_Mc83r5zkBZsCYPm7Lr0Bsq4,88
-langchain_google_genai-1.0.2.dist-info/RECORD,,