langchain-b12 0.1.6__py3-none-any.whl → 0.1.8__py3-none-any.whl

langchain_b12/genai/embeddings.py

@@ -1,6 +1,7 @@
 import os
 
 from google.genai import Client
+from google.genai.types import EmbedContentConfigOrDict
 from google.oauth2 import service_account
 from langchain_core.embeddings import Embeddings
 from pydantic import BaseModel, ConfigDict, Field
@@ -19,6 +20,7 @@ class GenAIEmbeddings(Embeddings, BaseModel):
         ),
         exclude=True,
     )
+    embed_content_config: EmbedContentConfigOrDict | None = Field(default=None)
     model_config = ConfigDict(
         arbitrary_types_allowed=True,
     )
@@ -33,19 +35,17 @@ class GenAIEmbeddings(Embeddings, BaseModel):
             list[list[float]]: The embedding vectors.
         """
         embeddings = []
-        for text in texts:
-            response = self.client.models.embed_content(
-                model=self.model_name,
-                contents=[text],
-            )
+        response = self.client.models.embed_content(
+            model=self.model_name,
+            contents=texts,
+            config=self.embed_content_config,
+        )
+        assert response.embeddings is not None, "No embeddings found in the response."
+        for embedding in response.embeddings:
             assert (
-                response.embeddings is not None
-            ), "No embeddings found in the response."
-            for embedding in response.embeddings:
-                assert (
-                    embedding.values is not None
-                ), "No embedding values found in the response."
-                embeddings.append(embedding.values)
+                embedding.values is not None
+            ), "No embedding values found in the response."
+            embeddings.append(embedding.values)
         assert len(embeddings) == len(
             texts
         ), "The number of embeddings does not match the number of texts."
@@ -75,6 +75,7 @@ class GenAIEmbeddings(Embeddings, BaseModel):
         response = await self.client.aio.models.embed_content(
             model=self.model_name,
             contents=texts,
+            config=self.embed_content_config,
         )
         assert response.embeddings is not None, "No embeddings found in the response."
         for embedding in response.embeddings:

langchain_b12/genai/genai.py

@@ -75,6 +75,8 @@ class ChatGenAI(BaseChatModel):
     """How many completions to generate for each prompt."""
     seed: int | None = None
     """Random seed for the generation."""
+    max_retries: int | None = Field(default=3)
+    """Maximum number of retries when generation fails. None disables retries."""
     safety_settings: list[types.SafetySetting] | None = None
     """The default safety settings to use for all generations.
 
@@ -173,10 +175,24 @@ class ChatGenAI(BaseChatModel):
         run_manager: CallbackManagerForLLMRun | None = None,
         **kwargs: Any,
     ) -> ChatResult:
-        stream_iter = self._stream(
-            messages, stop=stop, run_manager=run_manager, **kwargs
-        )
-        return generate_from_stream(stream_iter)
+        attempts = 0
+        while True:
+            try:
+                stream_iter = self._stream(
+                    messages, stop=stop, run_manager=run_manager, **kwargs
+                )
+                return generate_from_stream(stream_iter)
+            except Exception as e:  # noqa: BLE001
+                if self.max_retries is None or attempts >= self.max_retries:
+                    raise
+                attempts += 1
+                logger.warning(
+                    "ChatGenAI._generate failed (attempt %d/%d). "
+                    "Retrying... Error: %s",
+                    attempts,
+                    self.max_retries,
+                    e,
+                )
 
     async def _agenerate(
         self,
@@ -185,10 +201,24 @@ class ChatGenAI(BaseChatModel):
         run_manager: AsyncCallbackManagerForLLMRun | None = None,
         **kwargs: Any,
     ) -> ChatResult:
-        stream_iter = self._astream(
-            messages, stop=stop, run_manager=run_manager, **kwargs
-        )
-        return await agenerate_from_stream(stream_iter)
+        attempts = 0
+        while True:
+            try:
+                stream_iter = self._astream(
+                    messages, stop=stop, run_manager=run_manager, **kwargs
+                )
+                return await agenerate_from_stream(stream_iter)
+            except Exception as e:  # noqa: BLE001
+                if self.max_retries is None or attempts >= self.max_retries:
+                    raise
+                attempts += 1
+                logger.warning(
+                    "ChatGenAI._agenerate failed (attempt %d/%d). "
+                    "Retrying... Error: %s",
+                    attempts,
+                    self.max_retries,
+                    e,
+                )
 
     def _stream(
         self,
@@ -419,6 +449,7 @@ class ChatGenAI(BaseChatModel):
             # add model name if final chunk
             if top_candidate.finish_reason is not None:
                 message.response_metadata["model_name"] = self.model_name
+                message.response_metadata["tags"] = self.tags or []
 
         return (
             ChatGenerationChunk(

{langchain_b12-0.1.6.dist-info → langchain_b12-0.1.8.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: langchain-b12
-Version: 0.1.6
+Version: 0.1.8
 Summary: A reusable collection of tools and implementations for Langchain
 Author-email: Vincent Min <vincent.min@b12-consulting.com>
 Requires-Python: >=3.11

langchain_b12-0.1.8.dist-info/RECORD

@@ -0,0 +1,9 @@
+langchain_b12/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+langchain_b12/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+langchain_b12/citations/citations.py,sha256=ZQvYayjQXIUaRosJ0qwL3Nc7kC8sBzmaIkE-BOslaVI,12261
+langchain_b12/genai/embeddings.py,sha256=h0Z-5PltDW9q79AjSrLemsz-_QKMB-043XXDvYSRQds,3483
+langchain_b12/genai/genai.py,sha256=4Q0j2YsAPLrHhIy_pYXoncb4hqJJkEUpxB3oD3qaECI,18120
+langchain_b12/genai/genai_utils.py,sha256=tA6UiJURK25-11vtaX4768UV47jDCYwVKIIWydD4Egw,10736
+langchain_b12-0.1.8.dist-info/METADATA,sha256=0-KZr-PXjE16ar4LpQbdWHX8CrViLBxlfV9uGwE0Qw0,1204
+langchain_b12-0.1.8.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+langchain_b12-0.1.8.dist-info/RECORD,,

langchain_b12-0.1.6.dist-info/RECORD

@@ -1,9 +0,0 @@
-langchain_b12/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-langchain_b12/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-langchain_b12/citations/citations.py,sha256=ZQvYayjQXIUaRosJ0qwL3Nc7kC8sBzmaIkE-BOslaVI,12261
-langchain_b12/genai/embeddings.py,sha256=od2bVIgt7v9aNAHG0PVypVF1H_XgHto2nTd8vwfvyN8,3355
-langchain_b12/genai/genai.py,sha256=r7v_Z97N_Vd0zIR5mcQrlY3eWCPWhThvOvnXg59Ls8c,16868
-langchain_b12/genai/genai_utils.py,sha256=tA6UiJURK25-11vtaX4768UV47jDCYwVKIIWydD4Egw,10736
-langchain_b12-0.1.6.dist-info/METADATA,sha256=c2_hq-9spfRCb1nNRT6ztrCf1s1KNl8lkg-tLiB0mho,1204
-langchain_b12-0.1.6.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-langchain_b12-0.1.6.dist-info/RECORD,,