Add support for langchain 1.0.4, also fix other bugs.

Signed-off-by: Joe Klein <joseph.klein@oracle.com>

Author: joseph-klein

libs/oci/langchain_oci/chat_models/oci_generative_ai.py

@@ -91,25 +91,73 @@ def remove_signature_from_tool_description(name: str, description: str) -> str:
         description = re.sub(r"(?s)(?:\n?\n\s*?)?Args:.*$", "", description)
         return description
 
+    @staticmethod
+    def oci_to_dict(obj):
+        if obj is None or isinstance(obj, (str, int, float, bool)):
+            return obj
+        if isinstance(obj, dict):
+            return {k: OCIUtils.oci_to_dict(v) for k, v in obj.items()}
+        if isinstance(obj, list):
+            return [OCIUtils.oci_to_dict(x) for x in obj]
+        if hasattr(obj, "__dict__"):
+            out = {k: OCIUtils.oci_to_dict(getattr(obj, k)) for k in obj.__dict__ if not k.startswith("_")}
+            if hasattr(obj, "attribute_map"):
+                attribute_map = getattr(obj, "attribute_map", {})
+                out = {attribute_map.get(k, k): v for k, v in out.items()}
+            return out
+        return str(obj)
+    
+    
     @staticmethod
     def convert_oci_tool_call_to_langchain(tool_call: Any) -> ToolCall:
         """Convert an OCI tool call to a LangChain ToolCall."""
-        parsed = json.loads(tool_call.arguments)
+        # Handle various possible provider API shapes ("arguments" or "parameters")
+        # If object has an 'arguments' attr use it, else use 'parameters'
+
+        # Cohere models return a "parameters" dict, Meta models .arguments (JSON string)
+        if hasattr(tool_call, "arguments"):
+            parsed = json.loads(tool_call.arguments)
+            # If the parsed result is a string, it means the JSON was escaped, so parse again
+            if isinstance(parsed, str):
+                try:
+                    parsed = json.loads(parsed)
+                except json.JSONDecodeError:
+                    pass
+            args = parsed
+        elif hasattr(tool_call, "parameters"):
+            args = tool_call.parameters  # Already a dict
+        elif isinstance(tool_call, dict):
+            # Defensive: handle raw dict
+            if "arguments" in tool_call:
+                value = tool_call["arguments"]
+                args = json.loads(value) if isinstance(value, str) else value
+            elif "parameters" in tool_call:
+                args = tool_call["parameters"]
+            else:
+                args = {}
+        else:
+            args = {}
 
-        # If the parsed result is a string, it means the JSON was escaped, so parse again
-        if isinstance(parsed, str):
-            try:
-                parsed = json.loads(parsed)
-            except json.JSONDecodeError:
-                # If it's not valid JSON, keep it as a string
-                pass
+        # Try to get id, fallback to random if missing
+        if hasattr(tool_call, "id"):
+            tool_call_id = tool_call.id
+        elif isinstance(tool_call, dict) and "id" in tool_call:
+            tool_call_id = tool_call["id"]
+        else:
+            tool_call_id = uuid.uuid4().hex[:]
+
+        # Try to get name
+        if hasattr(tool_call, "name"):
+            name = tool_call.name
+        elif isinstance(tool_call, dict) and "name" in tool_call:
+            name = tool_call["name"]
+        else:
+            name = "unknown_tool"
 
         return ToolCall(
-            name=tool_call.name,
-            args=parsed
-            if "arguments" in tool_call.attribute_map
-            else tool_call.parameters,
-            id=tool_call.id if "id" in tool_call.attribute_map else uuid.uuid4().hex[:],
+            name=name,
+            args=args,
+            id=tool_call_id,
         )
 
 
@@ -223,9 +271,6 @@ def __init__(self) -> None:
             "SYSTEM": models.CohereSystemMessage,
             "TOOL": models.CohereToolMessage,
         }
-
-        self.oci_response_json_schema = models.ResponseJsonSchema
-        self.oci_json_schema_response_format = models.JsonSchemaResponseFormat
         self.chat_api_format = models.BaseChatRequest.API_FORMAT_COHERE
 
     def chat_response_to_text(self, response: Any) -> str:
@@ -255,11 +300,6 @@ def chat_generation_info(self, response: Any) -> Dict[str, Any]:
             "is_search_required": response.data.chat_response.is_search_required,
             "finish_reason": response.data.chat_response.finish_reason,
         }
-
-        # Include token usage if available
-        if hasattr(response.data.chat_response, "usage") and response.data.chat_response.usage:
-            generation_info["total_tokens"] = response.data.chat_response.usage.total_tokens
-
         # Include tool calls if available
         if self.chat_tool_calls(response):
             generation_info["tool_calls"] = self.format_response_tool_calls(
@@ -302,8 +342,12 @@ def format_response_tool_calls(
                 {
                     "id": uuid.uuid4().hex[:],
                     "function": {
-                        "name": tool_call.name,
-                        "arguments": json.dumps(tool_call.parameters),
+                        "name": getattr(tool_call, "name", getattr(tool_call, "name", "unknown_tool")),
+                        "arguments": (
+                            json.dumps(json.loads(tool_call.arguments))
+                            if hasattr(tool_call, "arguments")
+                            else json.dumps(getattr(tool_call, "parameters", {}))
+                        ),
                     },
                     "type": "function",
                 }
@@ -385,7 +429,9 @@ def messages_to_oci_params(
                     self.oci_chat_message[self.get_role(msg)](
                         tool_results=[
                             self.oci_tool_result(
-                                call=self.oci_tool_call(name=msg.name, parameters={}),
+                                call=self.oci_tool_call(
+                                    name=msg.name, parameters={}
+                                ),
                                 outputs=[{"output": msg.content}],
                             )
                         ],
@@ -397,17 +443,9 @@ def messages_to_oci_params(
         for i, message in enumerate(messages[::-1]):
             current_turn.append(message)
             if isinstance(message, HumanMessage):
-                if len(messages) > i and isinstance(
-                    messages[len(messages) - i - 2], ToolMessage
-                ):
-                    # add dummy message REPEATING the tool_result to avoid
-                    # the error about ToolMessage needing to be followed
-                    # by an AI message
-                    oci_chat_history.append(
-                        self.oci_chat_message["CHATBOT"](
-                            message=messages[len(messages) - i - 2].content
-                        )
-                    )
+                if len(messages) > i and isinstance(messages[len(messages) - i - 2], ToolMessage):
+                    # add dummy message REPEATING the tool_result to avoid the error about ToolMessage needing to be followed by an AI message
+                    oci_chat_history.append(self.oci_chat_message['CHATBOT'](message=messages[len(messages) - i - 2].content))
                 break
         current_turn = list(reversed(current_turn))
 
@@ -601,10 +639,6 @@ def __init__(self) -> None:
         self.oci_tool_call = models.FunctionCall
         self.oci_tool_message = models.ToolMessage
 
-        # Response format models
-        self.oci_response_json_schema = models.ResponseJsonSchema
-        self.oci_json_schema_response_format = models.JsonSchemaResponseFormat
-
         self.chat_api_format = models.BaseChatRequest.API_FORMAT_GENERIC
 
     def chat_response_to_text(self, response: Any) -> str:
@@ -630,11 +664,6 @@ def chat_generation_info(self, response: Any) -> Dict[str, Any]:
             "finish_reason": response.data.chat_response.choices[0].finish_reason,
             "time_created": str(response.data.chat_response.time_created),
         }
-
-        # Include token usage if available
-        if hasattr(response.data.chat_response, "usage") and response.data.chat_response.usage:
-            generation_info["total_tokens"] = response.data.chat_response.usage.total_tokens
-            
         if self.chat_tool_calls(response):
             generation_info["tool_calls"] = self.format_response_tool_calls(
                 self.chat_tool_calls(response)
@@ -668,8 +697,12 @@ def format_response_tool_calls(self, tool_calls: List[Any]) -> List[Dict]:
                 {
                     "id": tool_call.id,
                     "function": {
-                        "name": tool_call.name,
-                        "arguments": json.loads(tool_call.arguments),
+                        "name": getattr(tool_call, "name", getattr(tool_call, "name", "unknown_tool")),
+                        "arguments": (
+                            json.dumps(json.loads(tool_call.arguments))
+                            if hasattr(tool_call, "arguments")
+                            else json.dumps(getattr(tool_call, "parameters", {}))
+                        ),
                     },
                     "type": "function",
                 }
@@ -695,7 +728,11 @@ def format_stream_tool_calls(
                     "id": tool_call.get("id", ""),
                     "function": {
                         "name": tool_call.get("name", ""),
-                        "arguments": tool_call.get("arguments", ""),
+                        "arguments": (
+                            json.dumps(json.loads(tool_call["arguments"]))
+                            if "arguments" in tool_call
+                            else json.dumps(tool_call.get("parameters", {}))
+                        ),
                     },
                     "type": "function",
                 }
@@ -745,8 +782,8 @@ def messages_to_oci_params(
                     )
                 else:
                     oci_message = self.oci_chat_message[role](content=tool_content)
-            elif isinstance(message, AIMessage) and (
-                message.tool_calls or message.additional_kwargs.get("tool_calls")
+            elif isinstance(message, AIMessage) and message.additional_kwargs.get(
+                "tool_calls"
             ):
                 # Process content and tool calls for assistant messages
                 content = self._process_message_content(message.content)
@@ -904,22 +941,7 @@ def convert_to_oci_tool(
         Raises:
             ValueError: If the tool type is not supported.
         """
-        if (isinstance(tool, type) and issubclass(tool, BaseModel)) or callable(tool):
-            as_json_schema_function = convert_to_openai_function(tool)
-            parameters = as_json_schema_function.get("parameters", {})
-            return self.oci_function_definition(
-                name=as_json_schema_function.get("name"),
-                description=as_json_schema_function.get(
-                    "description",
-                    as_json_schema_function.get("name"),
-                ),
-                parameters={
-                    "type": "object",
-                    "properties": parameters.get("properties", {}),
-                    "required": parameters.get("required", []),
-                },
-            )
-        elif isinstance(tool, BaseTool):
+        if isinstance(tool, BaseTool):
             return self.oci_function_definition(
                 name=tool.name,
                 description=OCIUtils.remove_signature_from_tool_description(
@@ -941,6 +963,21 @@ def convert_to_oci_tool(
                     ],
                 },
             )
+        elif (isinstance(tool, type) and issubclass(tool, BaseModel)) or callable(tool):
+            as_json_schema_function = convert_to_openai_function(tool)
+            parameters = as_json_schema_function.get("parameters", {})
+            return self.oci_function_definition(
+                name=as_json_schema_function.get("name"),
+                description=as_json_schema_function.get(
+                    "description",
+                    as_json_schema_function.get("name"),
+                ),
+                parameters={
+                    "type": "object",
+                    "properties": parameters.get("properties", {}),
+                    "required": parameters.get("required", []),
+                },
+            )
         raise ValueError(
             f"Unsupported tool type {type(tool)}. "
             "Tool must be passed in as a BaseTool "
@@ -1034,7 +1071,6 @@ def process_stream_tool_calls(
 
 class MetaProvider(GenericProvider):
     """Provider for Meta models. This provider is for backward compatibility."""
-
     pass
 
 
@@ -1151,27 +1187,14 @@ def _prepare_request(
                 "Please make sure you have the oci package installed."
             ) from ex
 
-        oci_params = self._provider.messages_to_oci_params(
-            messages,
-            max_sequential_tool_calls=self.max_sequential_tool_calls,
-            **kwargs
-        )
+        oci_params = self._provider.messages_to_oci_params(messages, **kwargs)
 
         oci_params["is_stream"] = stream
         _model_kwargs = self.model_kwargs or {}
 
         if stop is not None:
             _model_kwargs[self._provider.stop_sequence_key] = stop
 
-        # Warn if using max_tokens with OpenAI models
-        if self.model_id and self.model_id.startswith("openai.") and "max_tokens" in _model_kwargs:
-            import warnings
-            warnings.warn(
-                f"OpenAI models require 'max_completion_tokens' instead of 'max_tokens'.",
-                UserWarning,
-                stacklevel=2
-            )
-
         chat_params = {**_model_kwargs, **kwargs, **oci_params}
 
         if not self.model_id:
@@ -1247,14 +1270,14 @@ def with_structured_output(
                 `method` is "function_calling" and `schema` is a dict, then the dict
                 must match the OCI Generative AI function-calling spec.
             method:
-                The method for steering model generation, either "function_calling" (default method)
-                or "json_mode" or "json_schema". If "function_calling" then the schema
+                The method for steering model generation, either "function_calling"
+                or "json_mode" or "json_schema. If "function_calling" then the schema
                 will be converted to an OCI function and the returned model will make
                 use of the function-calling API. If "json_mode" then Cohere's JSON mode will be
                 used. Note that if using "json_mode" then you must include instructions
                 for formatting the output into the desired schema into the model call.
                 If "json_schema" then it allows the user to pass a json schema (or pydantic)
-                to the model for structured output. 
+                to the model for structured output. This is the default method.
             include_raw:
                 If False then only the parsed structured output is returned. If
                 an error occurs during model output parsing it will be raised. If True
@@ -1305,24 +1328,19 @@ def with_structured_output(
                 else JsonOutputParser()
             )
         elif method == "json_schema":
-            json_schema_dict = (
-                schema.model_json_schema()  # type: ignore[union-attr]
+            response_format = (
+                dict(
+                    schema.model_json_schema().items()  # type: ignore[union-attr]
+                )
                 if is_pydantic_schema
                 else schema
             )
-            
-            response_json_schema = self._provider.oci_response_json_schema(
-                name=json_schema_dict.get("title", "response"),
-                description=json_schema_dict.get("description", ""),
-                schema=json_schema_dict,
-                is_strict=True
-            )
-            
-            response_format_obj = self._provider.oci_json_schema_response_format(
-                json_schema=response_json_schema
-            )
-            
-            llm = self.bind(response_format=response_format_obj)
+            llm_response_format: Dict[Any, Any] = {"type": "JSON_OBJECT"}
+            llm_response_format["schema"] = {
+                k: v
+                for k, v in response_format.items()  # type: ignore[union-attr]
+            }
+            llm = self.bind(response_format=llm_response_format)
             if is_pydantic_schema:
                 output_parser = PydanticOutputParser(pydantic_object=schema)
             else:
@@ -1400,7 +1418,7 @@ def _generate(
                 for tool_call in self._provider.chat_tool_calls(response)
             ]
         message = AIMessage(
-            content=content or "",
+            content=content,
             additional_kwargs=generation_info,
             tool_calls=tool_calls,
         )