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

langchain_google_genai/_function_utils.py

@@ -7,6 +7,7 @@ from typing import (
     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
@@ -14,55 +15,36 @@ from langchain_core.utils.json_schema import dereference_refs
 FunctionCallType = Union[BaseTool, Type[BaseModel], Dict]
 
 TYPE_ENUM = {
-    "string": 1,
-    "number": 2,
-    "integer": 3,
-    "boolean": 4,
-    "array": 5,
-    "object": 6,
+    "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],
-) -> Dict:
-    function_declarations = []
-    for fc in function_calls:
-        function_declarations.append(_convert_to_genai_function(fc))
-    return {
-        "function_declarations": function_declarations,
-    }
+) -> 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) -> Dict:
-    """
-        Produce
-
-        {
-      "name": "get_weather",
-      "description": "Determine weather in my location",
-      "parameters": {
-        "properties": {
-          "location": {
-            "description": "The city and state e.g. San Francisco, CA",
-            "type_": 1
-          },
-          "unit": { "enum": ["c", "f"], "type_": 1 }
-        },
-        "required": ["location"],
-        "type_": 6
-      }
-    }
-
-    """
+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 {
-            **fc,
-            "parameters": {
+        return glm.FunctionDeclaration(
+            name=fc["name"],
+            description=fc.get("description"),
+            parameters={
                 "properties": {
                     k: {
                         "type_": TYPE_ENUM[v["type"]],
@@ -73,20 +55,20 @@ def _convert_to_genai_function(fc: FunctionCallType) -> Dict:
                 "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) -> Dict:
+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 {
-            "name": tool.name or schema["title"],
-            "description": tool.description or schema["description"],
-            "parameters": {
+        return glm.FunctionDeclaration(
+            name=tool.name or schema["title"],
+            description=tool.description or schema["description"],
+            parameters={
                 "properties": {
                     k: {
                         "type_": TYPE_ENUM[v["type"]],
@@ -97,31 +79,30 @@ def _convert_tool_to_genai_function(tool: BaseTool) -> Dict:
                 "required": schema["required"],
                 "type_": TYPE_ENUM[schema["type"]],
             },
-        }
+        )
     else:
-        return {
-            "name": tool.name,
-            "description": tool.description,
-            "parameters": {
+        return glm.FunctionDeclaration(
+            name=tool.name,
+            description=tool.description,
+            parameters={
                 "properties": {
-                    "__arg1": {"type": "string"},
+                    "__arg1": {"type_": TYPE_ENUM["string"]},
                 },
                 "required": ["__arg1"],
                 "type_": TYPE_ENUM["object"],
             },
-        }
+        )
 
 
 def _convert_pydantic_to_genai_function(
     pydantic_model: Type[BaseModel],
-) -> Dict:
+) -> glm.FunctionDeclaration:
     schema = dereference_refs(pydantic_model.schema())
     schema.pop("definitions", None)
-
-    return {
-        "name": schema["title"],
-        "description": schema.get("description", ""),
-        "parameters": {
+    return glm.FunctionDeclaration(
+        name=schema["title"],
+        description=schema.get("description", ""),
+        parameters={
             "properties": {
                 k: {
                     "type_": TYPE_ENUM[v["type"]],
@@ -132,4 +113,4 @@ def _convert_pydantic_to_genai_function(
             "required": schema["required"],
             "type_": TYPE_ENUM[schema["type"]],
         },
-    }
+    )

langchain_google_genai/chat_models.py

@@ -16,18 +16,18 @@ 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 google.ai.generativelanguage_v1beta import FunctionCall
 from langchain_core.callbacks.manager import (
     AsyncCallbackManagerForLLMRun,
     CallbackManagerForLLMRun,
@@ -37,11 +37,8 @@ from langchain_core.messages import (
     AIMessage,
     AIMessageChunk,
     BaseMessage,
-    ChatMessage,
-    ChatMessageChunk,
     FunctionMessage,
     HumanMessage,
-    HumanMessageChunk,
     SystemMessage,
 )
 from langchain_core.outputs import ChatGeneration, ChatGenerationChunk, ChatResult
@@ -327,15 +324,42 @@ llm = ChatGoogleGenerativeAI(model="gemini-pro", convert_system_message_to_human
             continue
         elif isinstance(message, AIMessage):
             role = "model"
-            # TODO: Handle AImessage with function call
-            parts = _convert_to_parts(message.content)
+            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"
-            # TODO: Handle FunctionMessage
-            parts = _convert_to_parts(message.content)
+            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}."
@@ -353,100 +377,51 @@ llm = ChatGoogleGenerativeAI(model="gemini-pro", convert_system_message_to_human
     return messages
 
 
-def _retrieve_function_call_response(
-    parts: List[genai.types.PartType],
-) -> Optional[Dict]:
-    for idx, part in enumerate(parts):
-        if part.function_call and part.function_call.name:
-            fc: FunctionCall = part.function_call
-            return {
-                "function_call": {
-                    "name": fc.name,
-                    "arguments": json.dumps(
-                        dict(fc.args.items())
-                    ),  # dump to match other function calling llms for now
-                }
-            }
-    return None
-
-
-def _parts_to_content(
-    parts: List[genai.types.PartType],
-) -> Tuple[Union[str, List[Union[Dict[Any, Any], str]]], Optional[Dict]]:
-    """Converts a list of Gemini API Part objects into a list of LangChain messages."""
-    function_call_resp = _retrieve_function_call_response(parts)
-
-    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, function_call_resp
-    elif not parts:
-        logger.warning("Gemini produced an empty response.")
-        return "", function_call_resp
-    messages: List[Union[Dict[Any, Any], str]] = []
-    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, function_call_resp
+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, additional_kwargs = _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,
-                additional_kwargs=additional_kwargs or {},
-            )
-        else:
-            msg = role_map[content.role](
-                content=parts_content,
-                additional_kwargs=additional_kwargs or {},
-            )
         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,
@@ -455,7 +430,12 @@ 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)
 
 
@@ -616,17 +596,11 @@ class ChatGoogleGenerativeAI(_BaseGoogleGenerativeAI, BaseChatModel):
             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,
@@ -646,17 +620,11 @@ class ChatGoogleGenerativeAI(_BaseGoogleGenerativeAI, BaseChatModel):
             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,

langchain_google_genai/llms.py

@@ -151,6 +151,18 @@ Supported examples:
     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.

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: langchain-google-genai
-Version: 0.0.7
+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

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.7.dist-info/RECORD

@@ -1,11 +0,0 @@
-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=aCLchIQE2SJCQu9lv5qSsSRJiFeQPsCRZHfibwg4_No,3810
-langchain_google_genai/chat_models.py,sha256=fvXsT2wLuaQkwxZ0EcprMkYkp7aEXNMcdDQJwUimjkk,24263
-langchain_google_genai/embeddings.py,sha256=EMa-sDGXUpAPMSyjA2-YXF_TGrlSlqljNeqysAh574s,3951
-langchain_google_genai/llms.py,sha256=BjP9hTBVMu5rb6FQRzK9qnDoBeRqMpDUT1sBSOyJMhA,10609
-langchain_google_genai/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-langchain_google_genai-0.0.7.dist-info/LICENSE,sha256=DppmdYJVSc1jd0aio6ptnMUn5tIHrdAhQ12SclEBfBg,1072
-langchain_google_genai-0.0.7.dist-info/METADATA,sha256=ki5KGByrSZsLg3ZMqoVPLQXhcEbKksxioZk1AJyw4YQ,2851
-langchain_google_genai-0.0.7.dist-info/WHEEL,sha256=FMvqSimYX_P7y0a7UY-_Mc83r5zkBZsCYPm7Lr0Bsq4,88
-langchain_google_genai-0.0.7.dist-info/RECORD,,