Improve error handling in langchain.py

Author: Ickerday

splunklib/ai/engines/langchain.py

@@ -322,7 +322,7 @@ async def awrap_model_call(
                 agent_thread_ids: dict[str, set[str]] = {}
 
                 # Update the subagent schema definitions to include all thread_ids that the
-                # LLM could use to continue a corespondation with a subagent.
+                # LLM could use to continue correspondence with a subagent.
                 new_tools: list[BaseTool | dict[str, Any]] = []
                 for tool in request.tools:
                     assert isinstance(tool, StructuredTool)
@@ -342,8 +342,8 @@ async def awrap_model_call(
 
                         # Create an updated tool, with an updated input schema, that
                         # has a thread_id as an enum instead of a str, disallowing the
-                        # LLM from halucinating a thread_id, instead requiring that it is one
-                        # of the few specified thread_ids.
+                        # LLM from hallucinating a thread_id, instead requiring that
+                        # it is one of the few specified thread_ids.
                         tool = tool.model_copy(
                             update={
                                 "args_schema": create_model(
@@ -358,7 +358,7 @@ async def awrap_model_call(
                                             description=(
                                                 "Provide previous thread id to continue an existing conversation with an agent. "
                                                 + "Never issue a call to the same thread_id more than once. "
-                                                + "Do not provide to start a new corespondation"
+                                                + "Omit thread_id to start a new correspondence"
                                             ),
                                             default=None,
                                         ),
@@ -388,7 +388,7 @@ async def awrap_model_call(
                         args = SubagentLCArgs(**call["args"])
                         possible_thread_ids = agent_thread_ids.get(call["name"], set())
                         if args.thread_id and args.thread_id not in possible_thread_ids:
-                            # LLM halucinated a thread_id, start a new conversation instead.
+                            # LLM hallucinated a thread_id, start a new conversation instead.
                             # This should not happen, since we provide an enum above, but just
                             # in case.
                             args.thread_id = _thread_id_new_uuid()
@@ -418,7 +418,7 @@ class _SubagentArgumentPacker(LC_AgentMiddleware):
             # to differentiate that we wrap the resulting args in an SubagentLCArgs.
             #
             # This middleware performs the corresponding pack/unpack at the two
-            # points in the LangChain call graph where raw args are needed/retreived.
+            # points in the LangChain call graph where raw args are needed/retrieved.
             #
             # TODO: we could move this logic to  _Middleware.
             @override
@@ -586,12 +586,12 @@ async def awrap_tool_call(
             lc_middleware.append(_DEBUGMiddleware())
 
         if agent.limits.max_tokens is not None:
-            # Other limits are implemented using SDK middlewres, but this one
+            # Other limits are implemented using SDK middlewares, but this one
             # cannot be easily implemented that way, since count_tokens_approximately needs
             # access to list[BaseTool] and the langchain model. We don't expose these
             # in our SDK middleware, thus we use the langchain middlewares directly here.
             #
-            # Potentially we could implement count_tokens_approximately puerly using our SDK,
+            # Potentially we could implement count_tokens_approximately purely using our SDK,
             # that would additionally require exposing list[Tool] to AgentState, such that
             # middlewares get access to the tools that are passed to LLMs.
             #
@@ -750,7 +750,7 @@ async def invoke_agent(req: AgentRequest) -> AgentResponse[Any | None]:
         except _InvalidMessagesException as e:
             raise _InvalidMessagesException(
                 f"Agent middleware modified messages and made it invalid: {e}"
-            )
+            ) from e
 
         if self._sdk_agent.output_schema:
             if result.structured_output is None:
@@ -884,18 +884,18 @@ async def llm_handler(req: ModelRequest) -> ModelResponse:
                         raise StructuredOutputGenerationException(
                             message=msg,
                             error=StructuredOutputMultipleToolCallsError(),
-                        )
+                        ) from e
                     case LC_StructuredOutputValidationError():
                         raise StructuredOutputGenerationException(
                             message=msg,
                             error=StructuredOutputValidationError(validation_error=str(e.source)),
-                        )
+                        ) from e
                     case LC_StructuredOutputError():
                         # Langchain only returns the above handled exceptions, LC_StructuredOutputError
                         # is never returned alone (it is the base class for above exceptions).
                         raise AssertionError(
                             "internal error: LC_StructuredOutputError has been returned"
-                        )
+                        ) from e
 
         try:
             sdk_response = await self._with_model_middleware(llm_handler)(req)
@@ -937,7 +937,7 @@ async def llm_handler(req: ModelRequest) -> ModelResponse:
                 return LC_ModelResponse(result=result)
             else:
                 # Provider strategy
-                assert isinstance(e.error, StructuredOutputValidationError)
+                assert isinstance(e.error, StructuredOutputValidationError)  # noqa: PT017
 
                 request.runtime.context.retry = LC_HumanMessage(
                     content=create_structured_prompt(
@@ -1228,10 +1228,9 @@ def _convert_tool_message_from_lc(
 ) -> ToolMessage | SubagentMessage | StructuredOutputMessage:
     match message:
         case LC_ToolMessage(name=name) if name and name.startswith(AGENT_PREFIX):
-            assert (
-                isinstance(message.artifact, SubagentStructuredResult)
-                or isinstance(message.artifact, SubagentTextResult)
-                or isinstance(message.artifact, SubagentFailureResult)
+            assert isinstance(
+                message.artifact,
+                (SubagentStructuredResult, SubagentTextResult, SubagentFailureResult),
             )
             return SubagentMessage(
                 name=_denormalize_agent_name(name),
@@ -1250,9 +1249,7 @@ def _convert_tool_message_from_lc(
                     content=str(message.content),  # pyright: ignore[reportUnknownArgumentType]
                 )
 
-            assert isinstance(message.artifact, ToolResult) or isinstance(
-                message.artifact, ToolFailureResult
-            )
+            assert isinstance(message.artifact, (ToolResult, ToolFailureResult))
 
             tool_type: ToolType = (
                 ToolType.LOCAL if message.name.startswith(LOCAL_TOOL_PREFIX) else ToolType.REMOTE
@@ -1398,10 +1395,10 @@ async def _tool_call(
             result = await tool.func(**kwargs)
         except ToolException as e:
             raise LC_ToolException(*e.args) from e
-        except LC_ToolException:
-            assert False, (  # noqa: PT015
+        except LC_ToolException as lcte:
+            raise AssertionError(
                 "ToolException from LangChain should not be raised in tool.func"
-            )
+            ) from lcte
 
         artifact = ToolResult(content=result.content, structured_content=result.structured_content)
 
@@ -1801,14 +1798,14 @@ def _create_langchain_model(model: PredefinedModel) -> BaseChatModel:
                     extra_body=model.extra_body,
                     http_async_client=model.httpx_client,
                 )
-            except ImportError:
+            except ImportError as ie:
                 raise ImportError(
                     "OpenAI support is not installed.\n"
                     + "To enable OpenAI / ChatGPT models, install the optional extra:\n"
                     + 'pip install "splunk-sdk[openai]"\n'
                     + "# or if using uv:\n"
                     + "uv add splunk-sdk[openai]"
-                )
+                ) from ie
         case AnthropicModel():
             try:
                 from langchain_anthropic import ChatAnthropic
@@ -1822,14 +1819,14 @@ def _create_langchain_model(model: PredefinedModel) -> BaseChatModel:
                     kwargs["temperature"] = model.temperature
 
                 return ChatAnthropic(**kwargs)
-            except ImportError:
+            except ImportError as ie:
                 raise ImportError(
                     "Anthropic support is not installed.\n"
                     + "To enable Anthropic models, install the optional extra:\n"
                     + 'pip install "splunk-sdk[anthropic]"\n'
                     + "# or if using uv:\n"
                     + "uv add splunk-sdk[anthropic]"
-                )
+                ) from ie
         case GoogleModel():
             try:
                 from langchain_google_genai import ChatGoogleGenerativeAI
@@ -1849,14 +1846,14 @@ def _create_langchain_model(model: PredefinedModel) -> BaseChatModel:
                     google_kwargs["temperature"] = model.temperature
 
                 return ChatGoogleGenerativeAI(**google_kwargs)
-            except ImportError:
+            except ImportError as ie:
                 raise ImportError(
                     "Google GenAI support is not installed.\n"
                     + "To enable Google / Gemini models, install the optional extra:\n"
                     + 'pip install "splunk-sdk[google]"\n'
                     + "# or if using uv:\n"
                     + "uv add splunk-sdk[google]"
-                )
+                ) from ie
         case _:
             raise InvalidModelError("Cannot create langchain model - invalid SDK model provided")
 
@@ -2083,8 +2080,8 @@ async def model_middleware(
     ) -> ModelResponse:
         try:
             return await handler(request)
-        except StructuredOutputGenerationException:
+        except StructuredOutputGenerationException as soge:
             self._retries_per_thread_id[request.state.thread_id] += 1
             if self._retries_per_thread_id[request.state.thread_id] > self._limit:
-                raise StructuredOutputRetryLimitExceededException(self._limit)
+                raise StructuredOutputRetryLimitExceededException(self._limit) from soge
             raise  # re-raise, to retry structured output generation