Merge branch 'main' into taegyunkim/prof-12434-uwsgi

Author: taegyunkim

ddtrace/_trace/_span_pointer.py

@@ -64,7 +64,8 @@ def __post_init__(self):
     def __repr__(self):
         return (
             f"SpanPointer(trace_id={self.trace_id}, span_id={self.span_id}, kind={self.kind}, "
-            f"direction={self.direction}, hash={self.hash}, attributes={self.attributes})"
+            f"direction={self.attributes.get('ptr.dir')}, hash={self.attributes.get('ptr.hash')}, "
+            f"attributes={self.attributes})"
         )
 
 

ddtrace/appsec/_ai_guard/_langchain.py

@@ -13,6 +13,7 @@
 from ddtrace.contrib.internal.trace_utils import unwrap
 from ddtrace.contrib.internal.trace_utils import wrap
 import ddtrace.internal.logger as ddlogger
+from ddtrace.internal.utils import get_argument_value
 
 
 logger = ddlogger.get_logger(__name__)
@@ -164,28 +165,57 @@ def _handle_agent_action_result(client: AIGuardClient, result, kwargs):
 
 
 def _langchain_chatmodel_generate_before(client: AIGuardClient, message_lists):
-    from langchain_core.messages import HumanMessage
-
     for messages in message_lists:
-        # only call evaluator when the last message is an actual user prompt
-        if len(messages) > 0 and isinstance(messages[-1], HumanMessage):
-            history = _convert_messages(messages)
-            prompt = history.pop(-1)
-            try:
-                if not client.evaluate_prompt(prompt["role"], prompt["content"], history=history):  # type: ignore[typeddict-item]
-                    return AIGuardAbortError()
-            except AIGuardAbortError as e:
-                return e
-            except Exception:
-                logger.debug("Failed to evaluate chat model prompt", exc_info=True)
+        result = _evaluate_langchain_messages(client, messages)
+        if result:
+            return result
+    return None
 
 
 def _langchain_llm_generate_before(client: AIGuardClient, prompts):
     for prompt in prompts:
+        result = _evaluate_langchain_prompt(client, prompt)
+        if result:
+            return result
+    return None
+
+
+def _langchain_chatmodel_stream_before(client: AIGuardClient, instance, args, kwargs):
+    input_arg = get_argument_value(args, kwargs, 0, "input")
+    messages = instance._convert_input(input_arg).to_messages()
+    return _evaluate_langchain_messages(client, messages)
+
+
+def _langchain_llm_stream_before(client: AIGuardClient, instance, args, kwargs):
+    input_arg = get_argument_value(args, kwargs, 0, "input")
+    prompt = instance._convert_input(input_arg).to_string()
+    return _evaluate_langchain_prompt(client, prompt)
+
+
+def _evaluate_langchain_messages(client: AIGuardClient, messages):
+    from langchain_core.messages import HumanMessage
+
+    # only call evaluator when the last message is an actual user prompt
+    if len(messages) > 0 and isinstance(messages[-1], HumanMessage):
+        history = _convert_messages(messages)
+        prompt = history.pop(-1)
         try:
-            if not client.evaluate_prompt("user", prompt):
+            role, content = (prompt["role"], prompt["content"])  # type: ignore[typeddict-item]
+            if not client.evaluate_prompt(role, content, history=history):
                 return AIGuardAbortError()
         except AIGuardAbortError as e:
             return e
         except Exception:
-            logger.debug("Failed to evaluate llm prompt", exc_info=True)
+            logger.debug("Failed to evaluate chat model prompt", exc_info=True)
+    return None
+
+
+def _evaluate_langchain_prompt(client: AIGuardClient, prompt):
+    try:
+        if not client.evaluate_prompt("user", prompt):
+            return AIGuardAbortError()
+    except AIGuardAbortError as e:
+        return e
+    except Exception:
+        logger.debug("Failed to evaluate llm prompt", exc_info=True)
+    return None

ddtrace/appsec/_ai_guard/_listener.py

@@ -1,7 +1,9 @@
 from functools import partial
 
 from ddtrace.appsec._ai_guard._langchain import _langchain_chatmodel_generate_before
+from ddtrace.appsec._ai_guard._langchain import _langchain_chatmodel_stream_before
 from ddtrace.appsec._ai_guard._langchain import _langchain_llm_generate_before
+from ddtrace.appsec._ai_guard._langchain import _langchain_llm_stream_before
 from ddtrace.appsec._ai_guard._langchain import _langchain_patch
 from ddtrace.appsec._ai_guard._langchain import _langchain_unpatch
 from ddtrace.appsec.ai_guard import AIGuardClient
@@ -20,6 +22,8 @@ def _langchain_listen(client: AIGuardClient):
 
     core.on("langchain.chatmodel.generate.before", partial(_langchain_chatmodel_generate_before, client))
     core.on("langchain.chatmodel.agenerate.before", partial(_langchain_chatmodel_generate_before, client))
+    core.on("langchain.chatmodel.stream.before", partial(_langchain_chatmodel_stream_before, client))
 
     core.on("langchain.llm.generate.before", partial(_langchain_llm_generate_before, client))
     core.on("langchain.llm.agenerate.before", partial(_langchain_llm_generate_before, client))
+    core.on("langchain.llm.stream.before", partial(_langchain_llm_stream_before, client))

ddtrace/appsec/_common_module_patches.py

@@ -222,7 +222,7 @@ def wrapped_open_ED4CF71136E15EBF(original_open_callable, instance, args, kwargs
                     # api10 response handler for regular reponses
                     if response.__class__.__name__ == "HTTPResponse":
                         addresses = {
-                            "DOWN_RES_STATUS": response.status,
+                            "DOWN_RES_STATUS": str(response.status),
                             "DOWN_RES_HEADERS": _build_headers(response.getheaders()),
                         }
                         if use_body:
@@ -242,7 +242,7 @@ def wrapped_open_ED4CF71136E15EBF(original_open_callable, instance, args, kwargs
                             response_headers = None
                         if status_code is not None or response_headers is not None:
                             call_waf_callback(
-                                {"DOWN_RES_STATUS": status_code, "DOWN_RES_HEADERS": response_headers},
+                                {"DOWN_RES_STATUS": str(status_code), "DOWN_RES_HEADERS": response_headers},
                                 rule_type=EXPLOIT_PREVENTION.TYPE.SSRF,
                             )
                     raise

ddtrace/contrib/internal/langchain/patch.py

@@ -412,6 +412,8 @@ def traced_chat_stream(langchain, pin, func, instance, args, kwargs):
     llm_provider = instance._llm_type
     model = _extract_model_name(instance)
 
+    _raising_dispatch("langchain.chatmodel.stream.before", (instance, args, kwargs))
+
     def _on_span_started(span: Span):
         integration.record_instance(instance, span)
 
@@ -443,6 +445,15 @@ def traced_llm_stream(langchain, pin, func, instance, args, kwargs):
     llm_provider = instance._llm_type
     model = _extract_model_name(instance)
 
+    _raising_dispatch(
+        "langchain.llm.stream.before",
+        (
+            instance,
+            args,
+            kwargs,
+        ),
+    )
+
     def _on_span_start(span: Span):
         integration.record_instance(instance, span)
 

ddtrace/llmobs/_integrations/google_genai.py

@@ -3,12 +3,6 @@
 from typing import List
 from typing import Optional
 
-
-try:
-    from google.genai.types import FunctionDeclaration
-except ImportError:
-    FunctionDeclaration = None
-
 from ddtrace._trace.span import Span
 from ddtrace.llmobs._constants import INPUT_DOCUMENTS
 from ddtrace.llmobs._constants import INPUT_MESSAGES
@@ -187,6 +181,10 @@ def _function_declaration_to_tool_definition(self, function_declaration) -> Tool
         )
 
     def _extract_tools(self, config) -> List[ToolDefinition]:
+        try:
+            from google.genai.types import FunctionDeclaration
+        except ImportError:
+            FunctionDeclaration = None
         tool_definitions = []
         tools = _get_attr(config, "tools", []) or []
         for tool in tools:

ddtrace/llmobs/_integrations/openai_agents.py

@@ -37,6 +37,7 @@
 from ddtrace.llmobs._utils import _get_nearest_llmobs_ancestor
 from ddtrace.llmobs._utils import _get_span_name
 from ddtrace.llmobs._utils import load_data_value
+from ddtrace.llmobs._utils import safe_json
 from ddtrace.trace import Span
 
 
@@ -232,13 +233,13 @@ def _llmobs_set_response_attributes(self, span: Span, oai_span: OaiSpanAdapter)
         if oai_span.response and oai_span.response.output:
             messages, tool_call_outputs = oai_span.llmobs_output_messages()
 
-            for tool_id, tool_name, tool_args in tool_call_outputs:
+            for tool_call_output in tool_call_outputs:
                 core.dispatch(
                     DISPATCH_ON_LLM_TOOL_CHOICE,
                     (
-                        tool_id,
-                        tool_name,
-                        tool_args,
+                        tool_call_output["tool_id"],
+                        tool_call_output["name"],
+                        safe_json(tool_call_output["arguments"]),
                         {
                             "trace_id": format_trace_id(span.trace_id),
                             "span_id": str(span.span_id),