add new param to langchain test

getsentry · Jul 15, 2024 · 25ca6eb · 25ca6eb
1 parent 42c7c3b
commit 25ca6eb
Showing 1 changed file with 15 additions and 1 deletion.
diff --git a/tests/integrations/langchain/test_langchain.py b/tests/integrations/langchain/test_langchain.py
@@ -46,6 +46,15 @@ def _llm_type(self) -> str:
         return llm_type
 
 
+def tiktoken_encoding_if_installed():
+    try:
+        import tiktoken  # type: ignore # noqa # pylint: disable=unused-import
+
+        return "cl100k_base"
+    except ImportError:
+        return None
+
+
 @pytest.mark.parametrize(
     "send_default_pii, include_prompts, use_unknown_llm_type",
     [
@@ -62,7 +71,12 @@ def test_langchain_agent(
     llm_type = "acme-llm" if use_unknown_llm_type else "openai-chat"
 
     sentry_init(
-        integrations=[LangchainIntegration(include_prompts=include_prompts)],
+        integrations=[
+            LangchainIntegration(
+                include_prompts=include_prompts,
+                tiktoken_encoding_name=tiktoken_encoding_if_installed(),
+            )
+        ],
         traces_sample_rate=1.0,
         send_default_pii=send_default_pii,
     )