Refactor libs/oci to meet mypy stardard

Removed unnecessary type: ignore comments, improved type annotations, and updated function signatures for better type safety across chat models, embeddings, and tests. Also enhanced mypy configuration for stricter type checking and plugin support.

Author: paxiaatucsdedu

libs/oci/langchain_oci/chat_models/oci_data_science.py

@@ -31,7 +31,12 @@
     agenerate_from_stream,
     generate_from_stream,
 )
-from langchain_core.messages import AIMessageChunk, BaseMessage, BaseMessageChunk
+from langchain_core.messages import (
+    AIMessage,
+    AIMessageChunk,
+    BaseMessage,
+    BaseMessageChunk,
+)
 from langchain_core.output_parsers import (
     JsonOutputParser,
     PydanticOutputParser,
@@ -598,7 +603,7 @@ def with_structured_output(
         if method == "json_mode":
             llm = self.bind(response_format={"type": "json_object"})
             output_parser = (
-                PydanticOutputParser(pydantic_object=schema)  # type: ignore[type-var, arg-type]
+                PydanticOutputParser(pydantic_object=schema)
                 if is_pydantic_schema
                 else JsonOutputParser()
             )
@@ -767,7 +772,7 @@ def bind_tools(
         self,
         tools: Sequence[Union[Dict[str, Any], Type[BaseModel], Callable, BaseTool]],
         **kwargs: Any,
-    ) -> Runnable[LanguageModelInput, BaseMessage]:
+    ) -> Runnable[LanguageModelInput, AIMessage]:
         formatted_tools = [convert_to_openai_tool(tool) for tool in tools]
         return super().bind(tools=formatted_tools, **kwargs)
 

libs/oci/langchain_oci/chat_models/oci_generative_ai.py

@@ -31,7 +31,6 @@
     AIMessage,
     AIMessageChunk,
     BaseMessage,
-    ChatMessage,
     HumanMessage,
     SystemMessage,
     ToolCall,
@@ -350,7 +349,7 @@ def get_role(self, message: BaseMessage) -> str:
         raise ValueError(f"Unknown message type: {type(message)}")
 
     def messages_to_oci_params(
-        self, messages: Sequence[ChatMessage], **kwargs: Any
+        self, messages: Sequence[BaseMessage], **kwargs: Any
     ) -> Dict[str, Any]:
         """
         Convert LangChain messages to OCI parameters for Cohere.
@@ -417,7 +416,7 @@ def messages_to_oci_params(
         current_turn = list(reversed(current_turn))
 
         # Process tool results from the current turn
-        oci_tool_results: List[Any] = []
+        oci_tool_results: Optional[List[Any]] = []
         for message in current_turn:
             if isinstance(message, ToolMessage):
                 tool_msg = message
@@ -434,7 +433,7 @@ def messages_to_oci_params(
                                 parameters=lc_tool_call["args"],
                             )
                             tool_result.outputs = [{"output": tool_msg.content}]
-                            oci_tool_results.append(tool_result)
+                            oci_tool_results.append(tool_result)  # type: ignore[union-attr]
         if not oci_tool_results:
             oci_tool_results = None
 
@@ -552,7 +551,7 @@ def process_stream_tool_calls(
         Returns:
             List of ToolCallChunk objects
         """
-        tool_call_chunks = []
+        tool_call_chunks: List[ToolCallChunk] = []
         tool_call_response = self.chat_stream_tool_calls(event_data)
 
         if not tool_call_response:
@@ -813,7 +812,7 @@ def _should_allow_more_tool_calls(
                 return False
 
             # Detect infinite loop: same tool called with same arguments in succession
-            recent_calls = []
+            recent_calls: list = []
             for msg in reversed(messages):
                 if hasattr(msg, "tool_calls") and msg.tool_calls:
                     for tc in msg.tool_calls:
@@ -895,7 +894,7 @@ def _process_message_content(
 
     def convert_to_oci_tool(
         self,
-        tool: Union[Type[BaseModel], Callable, BaseTool],
+        tool: Union[Dict[str, Any], Type[BaseModel], Callable, BaseTool],
     ) -> Dict[str, Any]:
         """Convert a BaseTool instance, TypedDict or BaseModel type
         to a OCI tool in Meta's format.
@@ -1016,7 +1015,7 @@ def process_stream_tool_calls(
         Returns:
             List of ToolCallChunk objects
         """
-        tool_call_chunks = []
+        tool_call_chunks: List[ToolCallChunk] = []
         tool_call_response = self.chat_stream_tool_calls(event_data)
 
         if not tool_call_response:
@@ -1142,7 +1141,7 @@ def _prepare_request(
         stop: Optional[List[str]],
         stream: bool,
         **kwargs: Any,
-    ) -> Dict[str, Any]:
+    ) -> Any:
         """
         Prepare the OCI chat request from LangChain messages.
 
@@ -1206,7 +1205,7 @@ def bind_tools(
             Union[dict, str, Literal["auto", "none", "required", "any"], bool]
         ] = None,
         **kwargs: Any,
-    ) -> Runnable[LanguageModelInput, BaseMessage]:
+    ) -> Runnable[LanguageModelInput, AIMessage]:
         """Bind tool-like objects to this chat model.
 
         Assumes model is compatible with Meta's tool-calling API.
@@ -1299,8 +1298,8 @@ def with_structured_output(
             tool_name = getattr(self._provider.convert_to_oci_tool(schema), "name")
             if is_pydantic_schema:
                 output_parser: OutputParserLike = PydanticToolsParser(
-                    tools=[schema],  # type: ignore[list-item]
-                    first_tool_only=True,  # type: ignore[list-item]
+                    tools=[schema],
+                    first_tool_only=True,
                 )
             else:
                 output_parser = JsonOutputKeyToolsParser(
@@ -1309,15 +1308,15 @@ def with_structured_output(
         elif method == "json_mode":
             llm = self.bind(response_format={"type": "JSON_OBJECT"})
             output_parser = (
-                PydanticOutputParser(pydantic_object=schema)  # type: ignore[type-var, arg-type]
+                PydanticOutputParser(pydantic_object=schema)
                 if is_pydantic_schema
                 else JsonOutputParser()
             )
         elif method == "json_schema":
-            json_schema_dict = (
+            json_schema_dict: Dict[str, Any] = (
                 schema.model_json_schema()  # type: ignore[union-attr]
                 if is_pydantic_schema
-                else schema
+                else schema  # type: ignore[assignment]
             )
 
             response_json_schema = self._provider.oci_response_json_schema(

libs/oci/langchain_oci/embeddings/oci_data_science_model_deployment_endpoint.py

@@ -104,12 +104,9 @@ def _completion_with_retry(**kwargs: Any) -> Any:
                 response.raise_for_status()
                 return response
             except requests.exceptions.HTTPError as http_err:
-                if response.status_code == 401 and self._refresh_signer():
-                    raise TokenExpiredError() from http_err
-                else:
-                    raise ValueError(
-                        f"Server error: {str(http_err)}. Message: {response.text}"
-                    ) from http_err
+                raise ValueError(
+                    f"Server error: {str(http_err)}. Message: {response.text}"
+                ) from http_err
             except Exception as e:
                 raise ValueError(f"Error occurs by inference endpoint: {str(e)}") from e
 

libs/oci/langchain_oci/embeddings/oci_generative_ai.py

@@ -125,7 +125,7 @@ def validate_environment(cls, values: Dict) -> Dict:  # pylint: disable=no-self-
                 client_kwargs.pop("signer", None)
             elif values["auth_type"] == OCIAuthType(2).name:
 
-                def make_security_token_signer(oci_config):  # type: ignore[no-untyped-def]
+                def make_security_token_signer(oci_config):
                     pk = oci.signer.load_private_key_from_file(
                         oci_config.get("key_file"), None
                     )

libs/oci/langchain_oci/llms/oci_generative_ai.py

@@ -151,7 +151,7 @@ def validate_environment(cls, values: Dict) -> Dict:
                 client_kwargs.pop("signer", None)
             elif values["auth_type"] == OCIAuthType(2).name:
 
-                def make_security_token_signer(oci_config):  # type: ignore[no-untyped-def]
+                def make_security_token_signer(oci_config):
                     pk = oci.signer.load_private_key_from_file(
                         oci_config.get("key_file"), None
                     )

libs/oci/pyproject.toml

@@ -77,20 +77,26 @@ ignore = [
 ]
 
 [tool.mypy]
-ignore_missing_imports = "True"
-
-# Disable specific error codes that are causing issues
-disallow_untyped_defs = "False"
-disable_error_code = ["attr-defined", "assignment", "var-annotated", "override", "union-attr", "arg-type"]
-
-# TODO: LangChain Google settings
-# plugins = ["pydantic.mypy"]
-# strict = true
-# disallow_untyped_defs = true
-
-# # TODO: activate for 'strict' checking
-# disallow_any_generics = false
-# warn_return_any = false
+plugins = ["pydantic.mypy"]
+check_untyped_defs = true
+error_summary = false
+pretty = true
+show_column_numbers = true
+show_error_codes = true
+show_error_context = true
+warn_redundant_casts = true
+warn_unreachable = true
+warn_unused_configs = true
+warn_unused_ignores = true
+
+# Ignore missing imports only for specific untyped packages
+[[tool.mypy.overrides]]
+module = [
+    "oci.*",
+    "ads.*",
+    "langchain_openai.*",
+]
+ignore_missing_imports = true
 
 [tool.coverage.run]
 omit = ["tests/*"]

libs/oci/tests/integration_tests/chat_models/test_tool_calling.py

@@ -428,7 +428,7 @@ def should_continue(state: MessagesState):
 
     # Invoke agent with a diagnostic scenario
     result = agent.invoke(
-        {
+        {  # type: ignore[arg-type]
             "messages": [
                 SystemMessage(content=system_prompt),
                 HumanMessage(

libs/oci/tests/unit_tests/chat_models/test_oci_data_science.py

@@ -80,7 +80,7 @@ def __init__(self, json_data: Dict, status_code: int = 200):
     def raise_for_status(self) -> None:
         """Mocked raise for status."""
         if 400 <= self.status_code < 600:
-            raise HTTPError()  # type: ignore[call-arg]
+            raise HTTPError()
 
     def json(self) -> Dict:
         """Returns mocked json data."""

libs/oci/tests/unit_tests/chat_models/test_oci_generative_ai.py

@@ -13,12 +13,12 @@
 
 
 class MockResponseDict(dict):
-    def __getattr__(self, val):  # type: ignore[no-untyped-def]
+    def __getattr__(self, val):
         return self.get(val)
 
 
 class MockToolCall(dict):
-    def __getattr__(self, val):  # type: ignore[no-untyped-def]
+    def __getattr__(self, val):
         return self[val]
 
 
@@ -37,7 +37,7 @@ def test_llm_chat(monkeypatch: MonkeyPatch, test_model_id: str) -> None:
 
     provider = model_id.split(".")[0].lower()
 
-    def mocked_response(*args):  # type: ignore[no-untyped-def]
+    def mocked_response(*args):
         response_text = "Assistant chat reply."
         response = None
         if provider == "cohere":
@@ -166,7 +166,7 @@ def test_meta_tool_calling(monkeypatch: MonkeyPatch) -> None:
     oci_gen_ai_client = MagicMock()
     llm = ChatOCIGenAI(model_id="meta.llama-3-70b-instruct", client=oci_gen_ai_client)
 
-    def mocked_response(*args, **kwargs):  # type: ignore[no-untyped-def]
+    def mocked_response(*args, **kwargs):
         # Mock response with tool calls
         return MockResponseDict(
             {
@@ -231,7 +231,7 @@ def get_weather(location: str) -> str:
     messages = [HumanMessage(content="What's the weather like?")]
 
     # Test different tool choice options
-    tool_choices = [
+    tool_choices: list[str | bool | dict[str, str | dict[str, str]]] = [
         "get_weather",  # Specific tool
         "auto",  # Auto mode
         "none",  # No tools
@@ -254,7 +254,7 @@ def get_weather(location: str) -> str:
             assert tool_call["function"]["name"] == "get_weather"
 
     # Test escaped JSON arguments (issue #52)
-    def mocked_response_escaped(*args, **kwargs):  # type: ignore[no-untyped-def]
+    def mocked_response_escaped(*args, **kwargs):
         """Mock response with escaped JSON arguments."""
         return MockResponseDict(
             {
@@ -310,6 +310,7 @@ def mocked_response_escaped(*args, **kwargs):  # type: ignore[no-untyped-def]
     response_escaped = llm.bind_tools(tools=[get_weather]).invoke(messages)
 
     # Verify escaped JSON was correctly parsed to a dict
+    assert isinstance(response_escaped, AIMessage)
     assert len(response_escaped.tool_calls) == 1
     assert response_escaped.tool_calls[0]["name"] == "get_weather"
     assert response_escaped.tool_calls[0]["args"] == {"location": "San Francisco"}
@@ -337,7 +338,7 @@ def get_weather(location: str) -> str:
         ).invoke(messages)
 
     # Mock response for the case without tool choice
-    def mocked_response(*args, **kwargs):  # type: ignore[no-untyped-def]
+    def mocked_response(*args, **kwargs):
         return MockResponseDict(
             {
                 "status": 200,
@@ -378,7 +379,7 @@ def test_meta_tool_conversion(monkeypatch: MonkeyPatch) -> None:
     oci_gen_ai_client = MagicMock()
     llm = ChatOCIGenAI(model_id="meta.llama-3.3-70b-instruct", client=oci_gen_ai_client)
 
-    def mocked_response(*args, **kwargs):  # type: ignore[no-untyped-def]
+    def mocked_response(*args, **kwargs):
         request = args[0]
         # Check the conversion of tools to oci generic API spec
         # Function tool
@@ -467,6 +468,7 @@ class PydanticTool(BaseModel):
 
     # For tool calls, the response content should be empty.
     assert response.content == ""
+    assert isinstance(response, AIMessage)
     assert len(response.tool_calls) == 1
     assert response.tool_calls[0]["name"] == "function_tool"
 
@@ -483,7 +485,7 @@ class WeatherResponse(BaseModel):
     oci_gen_ai_client = MagicMock()
     llm = ChatOCIGenAI(model_id="cohere.command-r-16k", client=oci_gen_ai_client)
 
-    def mocked_response(*args, **kwargs):  # type: ignore[no-untyped-def]
+    def mocked_response(*args, **kwargs):
         return MockResponseDict(
             {
                 "status": 200,
@@ -533,7 +535,7 @@ class WeatherResponse(BaseModel):
     oci_gen_ai_client = MagicMock()
     llm = ChatOCIGenAI(model_id="cohere.command-latest", client=oci_gen_ai_client)
 
-    def mocked_response(*args, **kwargs):  # type: ignore[no-untyped-def]
+    def mocked_response(*args, **kwargs):
         # Verify that response_format is a JsonSchemaResponseFormat object
         request = args[0]
         assert hasattr(request.chat_request, "response_format")
@@ -610,7 +612,7 @@ class WeatherResponse(BaseModel):
     oci_gen_ai_client = MagicMock()
     llm = ChatOCIGenAI(model_id="cohere.command-r-16k", client=oci_gen_ai_client)
 
-    def mocked_response(*args, **kwargs):  # type: ignore[no-untyped-def]
+    def mocked_response(*args, **kwargs):
         return MockResponseDict(
             {
                 "status": 200,
@@ -663,7 +665,7 @@ def test_ai_message_tool_calls_direct_field(monkeypatch: MonkeyPatch) -> None:
     # Track if the tool_calls processing branch is executed
     tool_calls_processed = False
 
-    def mocked_response(*args, **kwargs):  # type: ignore[no-untyped-def]
+    def mocked_response(*args, **kwargs):
         nonlocal tool_calls_processed
         # Check if the request contains tool_calls in the message
         request = args[0]
@@ -746,7 +748,7 @@ def test_ai_message_tool_calls_additional_kwargs(monkeypatch: MonkeyPatch) -> No
     oci_gen_ai_client = MagicMock()
     llm = ChatOCIGenAI(model_id="meta.llama-3.3-70b-instruct", client=oci_gen_ai_client)
 
-    def mocked_response(*args, **kwargs):  # type: ignore[no-untyped-def]
+    def mocked_response(*args, **kwargs):
         return MockResponseDict(
             {
                 "status": 200,
@@ -884,8 +886,11 @@ def get_weather(city: str) -> str:
     ]
 
     # Prepare the request - need to pass tools from the bound model kwargs
-    request = llm_with_tools._prepare_request(
-        messages, stop=None, stream=False, **llm_with_tools.kwargs
+    request = llm._prepare_request(
+        messages,
+        stop=None,
+        stream=False,
+        **llm_with_tools.kwargs,  # type: ignore[attr-defined]
     )
 
     # Verify that tool_choice is set to 'none' because limit was reached