chore(langchain-ibm): disallow Any generics (#147)

Author: cbornet

libs/ibm/langchain_ibm/agent_toolkits/sql/tool.py

@@ -136,7 +136,7 @@ class QuerySQLCheckerTool(BaseSQLDatabaseTool, BaseTool):
     """Use an LLM to check if a query is correct."""
 
     template: str = QUERY_CHECKER
-    llm: BaseLanguageModel
+    llm: BaseLanguageModel[Any]
     llm_chain: Any = Field(init=False)
     name: str = "sql_db_query_checker"
     description: str = """
@@ -154,7 +154,7 @@ def initialize_llm_chain(cls, values: dict[str, Any]) -> Any:
                 template=QUERY_CHECKER,
                 input_variables=["query", "schema"],
             )
-            llm = cast("BaseLanguageModel", values.get("llm"))
+            llm = cast("BaseLanguageModel[Any]", values.get("llm"))
 
             values["llm_chain"] = prompt | llm
 

libs/ibm/langchain_ibm/agent_toolkits/sql/toolkit.py

@@ -1,5 +1,7 @@
 """IBM watsonx.ai SQL Toolkit wrapper."""
 
+from typing import Any
+
 from langchain_core.language_models import BaseLanguageModel
 from langchain_core.tools import BaseTool
 from langchain_core.tools.base import BaseToolkit
@@ -20,7 +22,7 @@ class WatsonxSQLDatabaseToolkit(BaseToolkit):
     db: WatsonxSQLDatabase = Field(exclude=True)
     """Instance of the watsonx SQL database."""
 
-    llm: BaseLanguageModel = Field(exclude=True)
+    llm: BaseLanguageModel[Any] = Field(exclude=True)
     """Instance of the LLM."""
 
     model_config = ConfigDict(
@@ -70,7 +72,7 @@ def get_tools(self) -> list[BaseTool]:
             query_sql_checker_tool,
         ]
 
-    def get_context(self) -> dict:
+    def get_context(self) -> dict[str, Any]:
         """Return db context that you may want in agent prompt."""
         return self.db.get_context()
 

libs/ibm/langchain_ibm/agent_toolkits/utility/toolkit.py

@@ -43,18 +43,18 @@ class WatsonxTool(BaseTool):
     """The precise instruction to agent LLMs
     and should be treated as part of the system prompt."""
 
-    tool_input_schema: dict | None = None
+    tool_input_schema: dict[str, Any] | None = None
     """Schema of the input that is provided when running the tool if applicable."""
 
-    tool_config_schema: dict | None = None
+    tool_config_schema: dict[str, Any] | None = None
     """Schema of the config that can be provided when running the tool if applicable."""
 
-    tool_config: dict | None = None
+    tool_config: dict[str, Any] | None = None
     """Config properties to be used when running a tool if applicable."""
 
     args_schema: type[BaseModel] = BaseModel
 
-    _watsonx_tool: Tool | None = PrivateAttr(default=None)  #: :meta private:
+    _watsonx_tool: Tool  #: :meta private:
 
     watsonx_client: APIClient = Field(exclude=True)
 
@@ -84,7 +84,7 @@ def _run(
         *args: Any,
         run_manager: CallbackManagerForToolRun | None = None,
         **kwargs: Any,
-    ) -> dict:
+    ) -> Any:
         """Run the tool."""
         if self.tool_input_schema is None:
             input_data = kwargs.get("input") or args[0]
@@ -95,9 +95,9 @@ def _run(
                 if k in self.tool_input_schema["properties"]
             }
 
-        return cast("dict", self._watsonx_tool.run(input_data, self.tool_config))  # type: ignore[union-attr]
+        return self._watsonx_tool.run(input_data, self.tool_config)
 
-    def set_tool_config(self, tool_config: dict) -> None:
+    def set_tool_config(self, tool_config: dict[str, Any]) -> None:
         """Set tool config properties.
 
         ???+ example "Example"

libs/ibm/langchain_ibm/agent_toolkits/utility/utils.py

@@ -1,13 +1,13 @@
 """Utility helpers."""
 
 from copy import deepcopy
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any
 
 if TYPE_CHECKING:
     from langchain_ibm.agent_toolkits.utility.toolkit import WatsonxTool
 
 
-def convert_to_watsonx_tool(tool: "WatsonxTool") -> dict:
+def convert_to_watsonx_tool(tool: "WatsonxTool") -> dict[str, Any]:
     """Convert `WatsonxTool` to watsonx tool structure.
 
     Args:
@@ -53,7 +53,7 @@ def convert_to_watsonx_tool(tool: "WatsonxTool") -> dict:
         ```
     """
 
-    def parse_parameters(input_schema: dict | None) -> dict:
+    def parse_parameters(input_schema: dict[str, Any] | None) -> dict[str, Any]:
         if input_schema:
             parameters = deepcopy(input_schema)
         else:

libs/ibm/langchain_ibm/chat_models.py

@@ -153,7 +153,7 @@ def _convert_dict_to_message(_dict: Mapping[str, Any], call_id: str) -> BaseMess
         return HumanMessage(content=_dict.get("content", ""), id=id_, name=name)
     if role == "assistant":
         content = _dict.get("content", "") or ""
-        additional_kwargs: dict = {}
+        additional_kwargs: dict[str, Any] = {}
         if function_call := _dict.get("function_call"):
             additional_kwargs["function_call"] = dict(function_call)
         tool_calls = []
@@ -242,7 +242,7 @@ def _format_message_content(content: Any) -> Any:
     return formatted_content
 
 
-def _convert_message_to_dict(message: BaseMessage) -> dict:
+def _convert_message_to_dict(message: BaseMessage) -> dict[str, Any]:
     """Convert a LangChain message to a dictionary.
 
     Args:
@@ -327,7 +327,7 @@ def _convert_delta_to_message_chunk(
     id_ = call_id
     role = cast("str", _dict.get("role"))
     content = cast("str", _dict.get("content") or "")
-    additional_kwargs: dict = {}
+    additional_kwargs: dict[str, Any] = {}
     if _dict.get("function_call"):
         function_call = dict(_dict["function_call"])
         if "name" in function_call and function_call["name"] is None:
@@ -379,7 +379,7 @@ def _convert_delta_to_message_chunk(
 
 
 def _convert_chunk_to_generation_chunk(
-    chunk: dict,
+    chunk: dict[str, Any],
     default_chunk_class: type,
     *,
     is_first_tool_chunk: bool,
@@ -939,7 +939,7 @@ class Joke(BaseModel):
     version: SecretStr | None = None
     """Version of the CPD instance."""
 
-    params: dict | TextChatParameters | None = None
+    params: dict[str, Any] | TextChatParameters | None = None
     """Model parameters to use during request generation.
 
     !!! note
@@ -988,7 +988,7 @@ class Joke(BaseModel):
 
     We generally recommend altering this or top_p but not both."""
 
-    response_format: dict | None = None
+    response_format: dict[str, Any] | None = None
     """The chat response format parameters."""
 
     top_p: float | None = None
@@ -1003,7 +1003,7 @@ class Joke(BaseModel):
     """Time limit in milliseconds - if not completed within this time,
     generation will stop."""
 
-    logit_bias: dict | None = None
+    logit_bias: dict[str, int] | None = None
     """Increasing or decreasing probability of tokens being selected
     during generation."""
 
@@ -1015,7 +1015,7 @@ class Joke(BaseModel):
     """Stop sequences are one or more strings which will cause the text generation
     to stop if/when they are produced as part of the output."""
 
-    chat_template_kwargs: dict | None = None
+    chat_template_kwargs: dict[str, Any] | None = None
     """Additional chat template parameters."""
 
     verify: str | bool | None = None
@@ -1187,7 +1187,9 @@ def validate_environment(self) -> Self:
         return self
 
     @gateway_error_handler
-    def _call_model_gateway(self, *, model: str, messages: list, **params: Any) -> Any:
+    def _call_model_gateway(
+        self, *, model: str, messages: list[dict[str, Any]], **params: Any
+    ) -> Any:
         return self.watsonx_model_gateway.chat.completions.create(
             model=model,
             messages=messages,
@@ -1199,7 +1201,7 @@ async def _acall_model_gateway(
         self,
         *,
         model: str,
-        messages: list,
+        messages: list[dict[str, Any]],
         **params: Any,
     ) -> Any:
         return await self.watsonx_model_gateway.chat.completions.acreate(
@@ -1405,7 +1407,7 @@ async def _astream(
             yield generation_chunk
 
     @staticmethod
-    def _merge_params(params: dict, kwargs: dict) -> dict:
+    def _merge_params(params: dict[str, Any], kwargs: dict[str, Any]) -> dict[str, Any]:
         param_updates = {}
         for k in ChatWatsonx._get_supported_chat_params():
             if kwargs.get(k) is not None:
@@ -1438,8 +1440,8 @@ def _create_message_dicts(
 
     def _create_chat_result(
         self,
-        response: dict,
-        generation_info: dict | None = None,
+        response: dict[str, Any],
+        generation_info: dict[str, Any] | None = None,
     ) -> ChatResult:
         generations = []
 
@@ -1496,9 +1498,9 @@ def _get_supported_chat_params() -> list[str]:
 
     def bind_tools(
         self,
-        tools: Sequence[dict[str, Any] | type | Callable | BaseTool],
+        tools: Sequence[dict[str, Any] | type | Callable[..., Any] | BaseTool],
         *,
-        tool_choice: dict | str | bool | None = None,
+        tool_choice: dict[str, Any] | str | bool | None = None,
         strict: bool | None = None,
         **kwargs: Any,
     ) -> Runnable[LanguageModelInput, AIMessage]:
@@ -1580,7 +1582,7 @@ def bind_tools(
     @override
     def with_structured_output(
         self,
-        schema: dict | type | None = None,
+        schema: dict[str, Any] | type | None = None,
         *,
         method: Literal[
             "function_calling",
@@ -1590,7 +1592,7 @@ def with_structured_output(
         include_raw: bool = False,
         strict: bool | None = None,
         **kwargs: Any,
-    ) -> Runnable[LanguageModelInput, dict | BaseModel]:
+    ) -> Runnable[LanguageModelInput, dict[str, Any] | BaseModel]:
         r"""Model wrapper that returns outputs formatted to match the given schema.
 
         Args:
@@ -1947,7 +1949,7 @@ class AnswerWithJustification(BaseModel):
                 },
             )
             if is_pydantic_schema:
-                output_parser: Runnable = PydanticToolsParser(
+                output_parser: Runnable[Any, Any] = PydanticToolsParser(
                     tools=[schema],
                     first_tool_only=True,
                 )
@@ -2016,7 +2018,7 @@ def _is_pydantic_class(obj: Any) -> bool:
     return isinstance(obj, type) and is_basemodel_subclass(obj)
 
 
-def _lc_tool_call_to_watsonx_tool_call(tool_call: ToolCall) -> dict:
+def _lc_tool_call_to_watsonx_tool_call(tool_call: ToolCall) -> dict[str, Any]:
     return {
         "type": "function",
         "id": tool_call["id"],
@@ -2029,7 +2031,7 @@ def _lc_tool_call_to_watsonx_tool_call(tool_call: ToolCall) -> dict:
 
 def _lc_invalid_tool_call_to_watsonx_tool_call(
     invalid_tool_call: InvalidToolCall,
-) -> dict:
+) -> dict[str, Any]:
     return {
         "type": "function",
         "id": invalid_tool_call["id"],
@@ -2041,7 +2043,7 @@ def _lc_invalid_tool_call_to_watsonx_tool_call(
 
 
 def _create_usage_metadata(
-    oai_token_usage: dict,
+    oai_token_usage: dict[str, Any],
     *,
     _prompt_tokens_included: bool,
 ) -> UsageMetadata:
@@ -2059,7 +2061,7 @@ def _create_usage_metadata(
 
 def _convert_to_openai_response_format(
     schema: dict[str, Any] | type, *, strict: bool | None = None
-) -> dict | TypeBaseModel:
+) -> dict[str, Any] | TypeBaseModel:
     if isinstance(schema, type) and is_basemodel_subclass(schema):
         return schema
 

libs/ibm/langchain_ibm/embeddings.py

@@ -177,7 +177,7 @@ class WatsonxEmbeddings(BaseModel, LangChainEmbeddings):
     version: SecretStr | None = None
     """Version of the CPD instance."""
 
-    params: dict | None = None
+    params: dict[str, Any] | None = None
     """Model parameters to use during request generation."""
 
     verify: str | bool | None = None

libs/ibm/langchain_ibm/llms.py

@@ -209,7 +209,7 @@ class WatsonxLLM(BaseLLM):
     version: SecretStr | None = None
     """Version of the CPD instance."""
 
-    params: dict | None = None
+    params: dict[str, Any] | None = None
     """Model parameters to use during request generation."""
 
     verify: str | bool | None = None
@@ -343,7 +343,9 @@ def validate_environment(self) -> Self:
         return self
 
     @gateway_error_handler
-    def _call_model_gateway(self, *, model: str, prompt: list, **params: Any) -> Any:
+    def _call_model_gateway(
+        self, *, model: str, prompt: str | list[str] | list[int], **params: Any
+    ) -> Any:
         return self.watsonx_model_gateway.completions.create(
             model=model,
             prompt=prompt,
@@ -355,7 +357,7 @@ async def _acall_model_gateway(
         self,
         *,
         model: str,
-        prompt: list,
+        prompt: str | list[str] | list[int],
         **params: Any,
     ) -> Any:
         return await self.watsonx_model_gateway.completions.acreate(
@@ -455,7 +457,7 @@ def _get_chat_params(
             params = (params or {}) | {"stop_sequences": stop}
         return params, kwargs
 
-    def _create_llm_result(self, response: list[dict]) -> LLMResult:
+    def _create_llm_result(self, response: list[dict[str, Any]]) -> LLMResult:
         """Create the LLMResult from the choices and prompts."""
         generations = [
             [
@@ -480,7 +482,7 @@ def _create_llm_result(self, response: list[dict]) -> LLMResult:
         }
         return LLMResult(generations=generations, llm_output=llm_output)
 
-    def _create_llm_gateway_result(self, response: dict) -> LLMResult:
+    def _create_llm_gateway_result(self, response: dict[str, Any]) -> LLMResult:
         """Create the LLMResult from the choices and prompts."""
         choices = response["choices"]
 

libs/ibm/langchain_ibm/rerank.py

@@ -148,7 +148,7 @@ class WatsonxRerank(BaseDocumentCompressor):
     version: SecretStr | None = None
     """Version of the CPD instance."""
 
-    params: dict | RerankParameters | None = None
+    params: dict[str, Any] | RerankParameters | None = None
     """Model parameters to use during request generation."""
 
     verify: str | bool | None = None
@@ -235,7 +235,7 @@ def validate_environment(self) -> Self:
 
     def rerank(
         self,
-        documents: Sequence[str | Document | dict],
+        documents: Sequence[str | Document | dict[str, Any]],
         query: str,
         **kwargs: Any,
     ) -> list[dict[str, Any]]: