From c44e5a3d93e0eed884fa88e64d743025f54aa8fd Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Sat, 22 Jun 2024 13:40:34 -0700 Subject: [PATCH] testing - use in memory langfuse client cache --- litellm/tests/test_alangfuse.py | 22 +++++++++++++++++----- 1 file changed, 17 insertions(+), 5 deletions(-) diff --git a/litellm/tests/test_alangfuse.py b/litellm/tests/test_alangfuse.py index 4e3d42352f..a074b110be 100644 --- a/litellm/tests/test_alangfuse.py +++ b/litellm/tests/test_alangfuse.py @@ -4,6 +4,7 @@ import json import logging import os import sys +from typing import Any from unittest.mock import MagicMock, patch logging.basicConfig(level=logging.DEBUG) @@ -19,17 +20,28 @@ import time import pytest +in_memory_langfuse_client: dict[str, Any] = {} + @pytest.fixture def langfuse_client(): import langfuse - langfuse_client = langfuse.Langfuse( - public_key=os.environ["LANGFUSE_PUBLIC_KEY"], - secret_key=os.environ["LANGFUSE_SECRET_KEY"], - host=None, + _langfuse_cache_key = ( + f"{os.environ['LANGFUSE_PUBLIC_KEY']}-{os.environ['LANGFUSE_SECRET_KEY']}" ) - print("NEW LANGFUSE CLIENT") + # use a in memory langfuse client for testing, RAM util on ci/cd gets too high when we init many langfuse clients + if _langfuse_cache_key in in_memory_langfuse_clients: + langfuse_client = in_memory_langfuse_clients[_langfuse_cache_key] + else: + langfuse_client = langfuse.Langfuse( + public_key=os.environ["LANGFUSE_PUBLIC_KEY"], + secret_key=os.environ["LANGFUSE_SECRET_KEY"], + host=None, + ) + in_memory_langfuse_clients[_langfuse_cache_key] = langfuse_client + + print("NEW LANGFUSE CLIENT") with patch( "langfuse.Langfuse", MagicMock(return_value=langfuse_client)