test(test_proxy_server_langfuse.py): cleanup tests causing OOM issues.

This commit is contained in:
Krrish Dholakia 2024-06-22 23:53:18 -07:00
parent d857d82b0e
commit 9cff2b03f3

View file

@ -1,19 +1,24 @@
import sys, os
import os
import sys
import traceback
from dotenv import load_dotenv
load_dotenv()
import os, io
import io
import os
# this file is to test litellm/proxy
sys.path.insert(
0, os.path.abspath("../..")
) # Adds the parent directory to the system path
import pytest, logging
import logging
import pytest
import litellm
from litellm import embedding, completion, completion_cost, Timeout
from litellm import RateLimitError
from litellm import RateLimitError, Timeout, completion, completion_cost, embedding
# Configure logging
logging.basicConfig(
@ -21,14 +26,16 @@ logging.basicConfig(
format="%(asctime)s - %(levelname)s - %(message)s",
)
from fastapi import FastAPI
# test /chat/completion request to the proxy
from fastapi.testclient import TestClient
from fastapi import FastAPI
from litellm.proxy.proxy_server import (
from litellm.proxy.proxy_server import ( # Replace with the actual module where your FastAPI router is defined
router,
save_worker_config,
startup_event,
) # Replace with the actual module where your FastAPI router is defined
)
filepath = os.path.dirname(os.path.abspath(__file__))
config_fp = f"{filepath}/test_configs/test_config.yaml"
@ -67,6 +74,9 @@ def client():
yield client
@pytest.mark.skip(
reason="Init multiple Langfuse clients causing OOM issues. Reduce init clients on ci/cd. "
)
def test_chat_completion(client):
try:
# Your test data