testing - use in memory langfuse client cache

This commit is contained in:
Ishaan Jaff 2024-06-22 13:40:34 -07:00
parent b201402b0b
commit c44e5a3d93

View file

@ -4,6 +4,7 @@ import json
import logging import logging
import os import os
import sys import sys
from typing import Any
from unittest.mock import MagicMock, patch from unittest.mock import MagicMock, patch
logging.basicConfig(level=logging.DEBUG) logging.basicConfig(level=logging.DEBUG)
@ -19,16 +20,27 @@ import time
import pytest import pytest
in_memory_langfuse_client: dict[str, Any] = {}
@pytest.fixture @pytest.fixture
def langfuse_client(): def langfuse_client():
import langfuse import langfuse
_langfuse_cache_key = (
f"{os.environ['LANGFUSE_PUBLIC_KEY']}-{os.environ['LANGFUSE_SECRET_KEY']}"
)
# 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( langfuse_client = langfuse.Langfuse(
public_key=os.environ["LANGFUSE_PUBLIC_KEY"], public_key=os.environ["LANGFUSE_PUBLIC_KEY"],
secret_key=os.environ["LANGFUSE_SECRET_KEY"], secret_key=os.environ["LANGFUSE_SECRET_KEY"],
host=None, host=None,
) )
in_memory_langfuse_clients[_langfuse_cache_key] = langfuse_client
print("NEW LANGFUSE CLIENT") print("NEW LANGFUSE CLIENT")
with patch( with patch(