restore test

Author: alexander-alderman-webb

tests/integrations/langchain/test_langchain.py

@@ -5793,6 +5793,139 @@ def test_transform_google_file_data(self):
         }
 
 
+@pytest.mark.parametrize("span_streaming", [True, False])
+@pytest.mark.parametrize("stream_gen_ai_spans", [True, False])
+@pytest.mark.parametrize(
+    "ai_type,expected_system",
+    [
+        # Real LangChain _type values (from _llm_type properties)
+        # OpenAI
+        ("openai-chat", "openai-chat"),
+        ("openai", "openai"),
+        # Azure OpenAI
+        ("azure-openai-chat", "azure-openai-chat"),
+        ("azure", "azure"),
+        # Anthropic
+        ("anthropic-chat", "anthropic-chat"),
+        # Google
+        ("vertexai", "vertexai"),
+        ("chat-google-generative-ai", "chat-google-generative-ai"),
+        ("google_gemini", "google_gemini"),
+        # AWS Bedrock
+        ("amazon_bedrock_chat", "amazon_bedrock_chat"),
+        ("amazon_bedrock", "amazon_bedrock"),
+        # Cohere
+        ("cohere-chat", "cohere-chat"),
+        # Ollama
+        ("chat-ollama", "chat-ollama"),
+        ("ollama-llm", "ollama-llm"),
+        # Mistral
+        ("mistralai-chat", "mistralai-chat"),
+        # Fireworks
+        ("fireworks-chat", "fireworks-chat"),
+        ("fireworks", "fireworks"),
+        # HuggingFace
+        ("huggingface-chat-wrapper", "huggingface-chat-wrapper"),
+        # Groq
+        ("groq-chat", "groq-chat"),
+        # NVIDIA
+        ("chat-nvidia-ai-playground", "chat-nvidia-ai-playground"),
+        # xAI
+        ("xai-chat", "xai-chat"),
+        # DeepSeek
+        ("chat-deepseek", "chat-deepseek"),
+        # Edge cases
+        ("", None),
+        (None, None),
+    ],
+)
+def test_langchain_ai_system_detection(
+    sentry_init,
+    capture_events,
+    capture_items,
+    ai_type,
+    expected_system,
+    stream_gen_ai_spans,
+    span_streaming,
+):
+    sentry_init(
+        integrations=[LangchainIntegration()],
+        disabled_integrations=[StdlibIntegration],
+        traces_sample_rate=1.0,
+        stream_gen_ai_spans=stream_gen_ai_spans,
+        trace_lifecycle="stream" if span_streaming else "static",
+    )
+
+    callback = SentryLangchainCallback(max_span_map_size=100, include_prompts=True)
+
+    run_id = "test-ai-system-uuid"
+    serialized = {"_type": ai_type} if ai_type is not None else {}
+    prompts = ["Test prompt"]
+
+    if span_streaming or stream_gen_ai_spans:
+        items = capture_items("span")
+
+        with start_transaction():
+            callback.on_llm_start(
+                serialized=serialized,
+                prompts=prompts,
+                run_id=run_id,
+                invocation_params={"_type": ai_type, "model": "test-model"},
+            )
+
+            generation = Mock(text="Test response", message=None)
+            response = Mock(generations=[[generation]])
+            callback.on_llm_end(response=response, run_id=run_id)
+
+        sentry_sdk.flush()
+        spans = [item.payload for item in items]
+        llm_spans = [
+            span
+            for span in spans
+            if span["attributes"].get("sentry.op") == "gen_ai.text_completion"
+        ]
+
+        assert len(llm_spans) > 0
+        llm_span = llm_spans[0]
+
+        if expected_system is not None:
+            assert llm_span["attributes"][SPANDATA.GEN_AI_SYSTEM] == expected_system
+        else:
+            assert SPANDATA.GEN_AI_SYSTEM not in llm_span.get("attributes", {})
+    else:
+        events = capture_events()
+
+        with start_transaction():
+            callback.on_llm_start(
+                serialized=serialized,
+                prompts=prompts,
+                run_id=run_id,
+                invocation_params={"_type": ai_type, "model": "test-model"},
+            )
+
+            generation = Mock(text="Test response", message=None)
+            response = Mock(generations=[[generation]])
+            callback.on_llm_end(response=response, run_id=run_id)
+
+        assert len(events) > 0
+        tx = events[0]
+        assert tx["type"] == "transaction"
+
+        llm_spans = [
+            span
+            for span in tx.get("spans", [])
+            if span.get("op") == "gen_ai.text_completion"
+        ]
+
+        assert len(llm_spans) > 0
+        llm_span = llm_spans[0]
+
+        if expected_system is not None:
+            assert llm_span["data"][SPANDATA.GEN_AI_SYSTEM] == expected_system
+        else:
+            assert SPANDATA.GEN_AI_SYSTEM not in llm_span.get("data", {})
+
+
 class TestTransformLangchainMessageContent:
     """Tests for _transform_langchain_message_content function."""