diff --git a/.circleci/config.yml b/.circleci/config.yml index 88e83fa7fd3b..7961cfddb60d 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -686,6 +686,7 @@ jobs: pip install "pytest-retry==1.6.3" pip install "pytest-cov==5.0.0" pip install "pytest-asyncio==0.21.1" + pip install pytest-mock pip install "respx==0.21.1" pip install "google-generativeai==0.3.2" pip install "google-cloud-aiplatform==1.43.0" diff --git a/deploy/charts/litellm-helm/templates/migrations-job.yaml b/deploy/charts/litellm-helm/templates/migrations-job.yaml new file mode 100644 index 000000000000..fc1aacf16e6c --- /dev/null +++ b/deploy/charts/litellm-helm/templates/migrations-job.yaml @@ -0,0 +1,30 @@ +# This job runs the prisma migrations for the LiteLLM DB. + +apiVersion: batch/v1 +kind: Job +metadata: + name: {{ include "litellm.fullname" . }}-migrations + annotations: + argocd.argoproj.io/hook: PreSync + argocd.argoproj.io/hook-delete-policy: Never # keep this resource so we can debug status on ArgoCD + checksum/config: {{ toYaml .Values | sha256sum }} +spec: + template: + spec: + containers: + - name: prisma-migrations + image: "ghcr.io/berriai/litellm:main-stable" + command: ["python", "litellm/proxy/prisma_migration.py"] + workingDir: "/app" + env: + {{- if .Values.db.deployStandalone }} + - name: DATABASE_URL + value: postgresql://{{ .Values.postgresql.auth.username }}:{{ .Values.postgresql.auth.password }}@{{ .Release.Name }}-postgresql/{{ .Values.postgresql.auth.database }} + {{- else if .Values.db.useExisting }} + - name: DATABASE_URL + value: {{ .Values.db.url | quote }} + {{- end }} + - name: DISABLE_SCHEMA_UPDATE + value: "{{ .Values.migrationJob.disableSchemaUpdate }}" + restartPolicy: OnFailure + backoffLimit: {{ .Values.migrationJob.backoffLimit }} diff --git a/deploy/charts/litellm-helm/values.yaml b/deploy/charts/litellm-helm/values.yaml index a2c55f2faa8d..c8e4aa1f2ec7 100644 --- a/deploy/charts/litellm-helm/values.yaml +++ b/deploy/charts/litellm-helm/values.yaml @@ -179,3 +179,12 @@ postgresql: redis: enabled: false architecture: standalone + +# Prisma migration job settings +migrationJob: + enabled: true # Enable or disable the schema migration Job + retries: 3 # Number of retries for the Job in case of failure + backoffLimit: 4 # Backoff limit for Job restarts + disableSchemaUpdate: false # Skip schema migrations for specific environments. When True, the job will exit with code 0. + + diff --git a/docs/my-website/.gitignore b/docs/my-website/.gitignore index b2d6de30624f..4d8604572301 100644 --- a/docs/my-website/.gitignore +++ b/docs/my-website/.gitignore @@ -18,3 +18,4 @@ npm-debug.log* yarn-debug.log* yarn-error.log* +yarn.lock diff --git a/docs/my-website/docs/benchmarks.md b/docs/my-website/docs/benchmarks.md new file mode 100644 index 000000000000..86699008bdd6 --- /dev/null +++ b/docs/my-website/docs/benchmarks.md @@ -0,0 +1,41 @@ +# Benchmarks + +Benchmarks for LiteLLM Gateway (Proxy Server) + +Locust Settings: +- 2500 Users +- 100 user Ramp Up + + +## Basic Benchmarks + +Overhead when using a Deployed Proxy vs Direct to LLM +- Latency overhead added by LiteLLM Proxy: 107ms + +| Metric | Direct to Fake Endpoint | Basic Litellm Proxy | +|--------|------------------------|---------------------| +| RPS | 1196 | 1133.2 | +| Median Latency (ms) | 33 | 140 | + + +## Logging Callbacks + +### [GCS Bucket Logging](https://docs.litellm.ai/docs/proxy/bucket) + +Using GCS Bucket has **no impact on latency, RPS compared to Basic Litellm Proxy** + +| Metric | Basic Litellm Proxy | LiteLLM Proxy with GCS Bucket Logging | +|--------|------------------------|---------------------| +| RPS | 1133.2 | 1137.3 | +| Median Latency (ms) | 140 | 138 | + + +### [LangSmith logging](https://docs.litellm.ai/docs/proxy/logging) + +Using LangSmith has **no impact on latency, RPS compared to Basic Litellm Proxy** + +| Metric | Basic Litellm Proxy | LiteLLM Proxy with LangSmith | +|--------|------------------------|---------------------| +| RPS | 1133.2 | 1135 | +| Median Latency (ms) | 140 | 132 | + diff --git a/docs/my-website/docs/observability/opentelemetry_integration.md b/docs/my-website/docs/observability/opentelemetry_integration.md index ba5ef2ff8589..218064b3d1e8 100644 --- a/docs/my-website/docs/observability/opentelemetry_integration.md +++ b/docs/my-website/docs/observability/opentelemetry_integration.md @@ -49,9 +49,19 @@ OTEL_ENDPOINT="http://0.0.0.0:4317" + + +```shell +OTEL_EXPORTER="otlp_grpc" +OTEL_ENDPOINT="https://api.lmnr.ai:8443" +OTEL_HEADERS="authorization=Bearer " +``` + + + -Use just 2 lines of code, to instantly log your LLM responses **across all providers** with OpenTelemetry: +Use just 1 line of code, to instantly log your LLM responses **across all providers** with OpenTelemetry: ```python litellm.callbacks = ["otel"] diff --git a/docs/my-website/docs/providers/anthropic.md b/docs/my-website/docs/providers/anthropic.md index 0c7b2a442d85..290e094d0906 100644 --- a/docs/my-website/docs/providers/anthropic.md +++ b/docs/my-website/docs/providers/anthropic.md @@ -864,3 +864,96 @@ Human: How do I boil water? Assistant: ``` + +## Usage - PDF + +Pass base64 encoded PDF files to Anthropic models using the `image_url` field. + + + + +### **using base64** +```python +from litellm import completion, supports_pdf_input +import base64 +import requests + +# URL of the file +url = "https://storage.googleapis.com/cloud-samples-data/generative-ai/pdf/2403.05530.pdf" + +# Download the file +response = requests.get(url) +file_data = response.content + +encoded_file = base64.b64encode(file_data).decode("utf-8") + +## check if model supports pdf input - (2024/11/11) only claude-3-5-haiku-20241022 supports it +supports_pdf_input("anthropic/claude-3-5-haiku-20241022") # True + +response = completion( + model="anthropic/claude-3-5-haiku-20241022", + messages=[ + { + "role": "user", + "content": [ + {"type": "text", "text": "You are a very professional document summarization specialist. Please summarize the given document."}, + { + "type": "image_url", + "image_url": f"data:application/pdf;base64,{encoded_file}", # 👈 PDF + }, + ], + } + ], + max_tokens=300, +) + +print(response.choices[0]) +``` + + + +1. Add model to config + +```yaml +- model_name: claude-3-5-haiku-20241022 + litellm_params: + model: anthropic/claude-3-5-haiku-20241022 + api_key: os.environ/ANTHROPIC_API_KEY +``` + +2. Start Proxy + +``` +litellm --config /path/to/config.yaml +``` + +3. Test it! + +```bash +curl http://0.0.0.0:4000/v1/chat/completions \ + -H "Content-Type: application/json" \ + -H "Authorization: Bearer " \ + -d '{ + "model": "claude-3-5-haiku-20241022", + "messages": [ + { + "role": "user", + "content": [ + { + "type": "text", + "text": "You are a very professional document summarization specialist. Please summarize the given document" + }, + { + "type": "image_url", + "image_url": "data:application/pdf;base64,{encoded_file}" # 👈 PDF + } + } + ] + } + ], + "max_tokens": 300 + }' + +``` + + diff --git a/docs/my-website/docs/proxy/bucket.md b/docs/my-website/docs/proxy/bucket.md index 3422d0371f8c..d1b9e607694e 100644 --- a/docs/my-website/docs/proxy/bucket.md +++ b/docs/my-website/docs/proxy/bucket.md @@ -9,7 +9,7 @@ LiteLLM Supports Logging to the following Cloud Buckets - (Enterprise) ✨ [Google Cloud Storage Buckets](#logging-proxy-inputoutput-to-google-cloud-storage-buckets) - (Free OSS) [Amazon s3 Buckets](#logging-proxy-inputoutput---s3-buckets) -## Logging Proxy Input/Output to Google Cloud Storage Buckets +## Google Cloud Storage Buckets Log LLM Logs to [Google Cloud Storage Buckets](https://cloud.google.com/storage?hl=en) @@ -20,6 +20,14 @@ Log LLM Logs to [Google Cloud Storage Buckets](https://cloud.google.com/storage? ::: +| Property | Details | +|----------|---------| +| Description | Log LLM Input/Output to cloud storage buckets | +| Load Test Benchmarks | [Benchmarks](https://docs.litellm.ai/docs/benchmarks) | +| Google Docs on Cloud Storage | [Google Cloud Storage](https://cloud.google.com/storage?hl=en) | + + + ### Usage 1. Add `gcs_bucket` to LiteLLM Config.yaml @@ -85,7 +93,7 @@ curl --location 'http://0.0.0.0:4000/chat/completions' \ 6. Save the JSON file and add the path to `GCS_PATH_SERVICE_ACCOUNT` -## Logging Proxy Input/Output - s3 Buckets +## s3 Buckets We will use the `--config` to set diff --git a/docs/my-website/docs/proxy/caching.md b/docs/my-website/docs/proxy/caching.md index d81db5b93623..3f5342c7e60a 100644 --- a/docs/my-website/docs/proxy/caching.md +++ b/docs/my-website/docs/proxy/caching.md @@ -136,6 +136,7 @@ litellm_settings: type: "redis" service_name: "mymaster" sentinel_nodes: [["localhost", 26379]] + sentinel_password: "password" # [OPTIONAL] ``` @@ -149,6 +150,7 @@ You can configure redis sentinel in your .env by setting `REDIS_SENTINEL_NODES` ```env REDIS_SENTINEL_NODES='[["localhost", 26379]]' REDIS_SERVICE_NAME = "mymaster" +REDIS_SENTINEL_PASSWORD = "password" ``` :::note diff --git a/docs/my-website/docs/proxy/configs.md b/docs/my-website/docs/proxy/configs.md index b4d70a4e73cc..c6b9f2d451c5 100644 --- a/docs/my-website/docs/proxy/configs.md +++ b/docs/my-website/docs/proxy/configs.md @@ -934,8 +934,8 @@ router_settings: | EMAIL_SUPPORT_CONTACT | Support contact email address | GCS_BUCKET_NAME | Name of the Google Cloud Storage bucket | GCS_PATH_SERVICE_ACCOUNT | Path to the Google Cloud service account JSON file -| GCS_FLUSH_INTERVAL | Flush interval for GCS logging (in seconds). Specify how often you want a log to be sent to GCS. -| GCS_BATCH_SIZE | Batch size for GCS logging. Specify after how many logs you want to flush to GCS. If `BATCH_SIZE` is set to 10, logs are flushed every 10 logs. +| GCS_FLUSH_INTERVAL | Flush interval for GCS logging (in seconds). Specify how often you want a log to be sent to GCS. **Default is 20 seconds** +| GCS_BATCH_SIZE | Batch size for GCS logging. Specify after how many logs you want to flush to GCS. If `BATCH_SIZE` is set to 10, logs are flushed every 10 logs. **Default is 2048** | GENERIC_AUTHORIZATION_ENDPOINT | Authorization endpoint for generic OAuth providers | GENERIC_CLIENT_ID | Client ID for generic OAuth providers | GENERIC_CLIENT_SECRET | Client secret for generic OAuth providers diff --git a/docs/my-website/docs/proxy/deploy.md b/docs/my-website/docs/proxy/deploy.md index 0287af2a28ae..20e108abfa25 100644 --- a/docs/my-website/docs/proxy/deploy.md +++ b/docs/my-website/docs/proxy/deploy.md @@ -688,8 +688,35 @@ Provide an ssl certificate when starting litellm proxy server Use this if you want to run the proxy with hypercorn to support http/2 -**Usage** -Pass the `--run_hypercorn` flag when starting the proxy +Step 1. Build your custom docker image with hypercorn + +```shell +# Use the provided base image +FROM ghcr.io/berriai/litellm:main-latest + +# Set the working directory to /app +WORKDIR /app + +# Copy the configuration file into the container at /app +COPY config.yaml . + +# Make sure your docker/entrypoint.sh is executable +RUN chmod +x ./docker/entrypoint.sh + +# Expose the necessary port +EXPOSE 4000/tcp + +# 👉 Key Change: Install hypercorn +RUN pip install hypercorn + +# Override the CMD instruction with your desired command and arguments +# WARNING: FOR PROD DO NOT USE `--detailed_debug` it slows down response times, instead use the following CMD +# CMD ["--port", "4000", "--config", "config.yaml"] + +CMD ["--port", "4000", "--config", "config.yaml", "--detailed_debug"] +``` + +Step 2. Pass the `--run_hypercorn` flag when starting the proxy ```shell docker run \ @@ -699,7 +726,7 @@ docker run \ -e SERVER_ROOT_PATH="/api/v1"\ -e DATABASE_URL=postgresql://:@:/ \ -e LITELLM_MASTER_KEY="sk-1234"\ - ghcr.io/berriai/litellm:main-latest \ + your_custom_docker_image \ --config /app/config.yaml --run_hypercorn ``` diff --git a/docs/my-website/docs/proxy/logging.md b/docs/my-website/docs/proxy/logging.md index 94faa7734a35..5867a8f23885 100644 --- a/docs/my-website/docs/proxy/logging.md +++ b/docs/my-website/docs/proxy/logging.md @@ -107,7 +107,7 @@ class StandardLoggingModelInformation(TypedDict): model_map_value: Optional[ModelInfo] ``` -## Logging Proxy Input/Output - Langfuse +## Langfuse We will use the `--config` to set `litellm.success_callback = ["langfuse"]` this will log all successfull LLM calls to langfuse. Make sure to set `LANGFUSE_PUBLIC_KEY` and `LANGFUSE_SECRET_KEY` in your environment @@ -463,7 +463,7 @@ You will see `raw_request` in your Langfuse Metadata. This is the RAW CURL comma -## Logging Proxy Input/Output in OpenTelemetry format +## OpenTelemetry format :::info @@ -1216,7 +1216,7 @@ litellm_settings: Start the LiteLLM Proxy and make a test request to verify the logs reached your callback API -## Logging LLM IO to Langsmith +## Langsmith 1. Set `success_callback: ["langsmith"]` on litellm config.yaml @@ -1261,7 +1261,7 @@ Expect to see your log on Langfuse -## Logging LLM IO to Arize AI +## Arize AI 1. Set `success_callback: ["arize"]` on litellm config.yaml @@ -1309,7 +1309,7 @@ Expect to see your log on Langfuse -## Logging LLM IO to Langtrace +## Langtrace 1. Set `success_callback: ["langtrace"]` on litellm config.yaml @@ -1351,7 +1351,7 @@ curl --location 'http://0.0.0.0:4000/chat/completions' \ ' ``` -## Logging LLM IO to Galileo +## Galileo [BETA] @@ -1466,7 +1466,7 @@ curl --location 'http://0.0.0.0:4000/chat/completions' \ -## Logging Proxy Input/Output - DataDog +## DataDog LiteLLM Supports logging to the following Datdog Integrations: - `datadog` [Datadog Logs](https://docs.datadoghq.com/logs/) @@ -1543,7 +1543,7 @@ Expected output on Datadog -## Logging Proxy Input/Output - DynamoDB +## DynamoDB We will use the `--config` to set @@ -1669,7 +1669,7 @@ Your logs should be available on DynamoDB } ``` -## Logging Proxy Input/Output - Sentry +## Sentry If api calls fail (llm/database) you can log those to Sentry: @@ -1711,7 +1711,7 @@ Test Request litellm --test ``` -## Logging Proxy Input/Output Athina +## Athina [Athina](https://athina.ai/) allows you to log LLM Input/Output for monitoring, analytics, and observability. diff --git a/docs/my-website/docs/proxy/team_logging.md b/docs/my-website/docs/proxy/team_logging.md index e2fcfa4b5e8f..8286ac449add 100644 --- a/docs/my-website/docs/proxy/team_logging.md +++ b/docs/my-website/docs/proxy/team_logging.md @@ -281,6 +281,51 @@ curl -X POST 'http://0.0.0.0:4000/key/generate' \ }' ``` + + + + +1. Create Virtual Key to log to a specific Langsmith Project + + ```bash + curl -X POST 'http://0.0.0.0:4000/key/generate' \ + -H 'Authorization: Bearer sk-1234' \ + -H 'Content-Type: application/json' \ + -d '{ + "metadata": { + "logging": [{ + "callback_name": "langsmith", # "otel", "gcs_bucket" + "callback_type": "success", # "success", "failure", "success_and_failure" + "callback_vars": { + "langsmith_api_key": "os.environ/LANGSMITH_API_KEY", # API Key for Langsmith logging + "langsmith_project": "pr-brief-resemblance-72", # project name on langsmith + "langsmith_base_url": "https://api.smith.langchain.com" + } + }] + } + }' + + ``` + +2. Test it - `/chat/completions` request + + Use the virtual key from step 3 to make a `/chat/completions` request + + You should see your logs on your Langsmith project on a successful request + + ```shell + curl -i http://localhost:4000/v1/chat/completions \ + -H "Content-Type: application/json" \ + -H "Authorization: Bearer sk-Fxq5XSyWKeXDKfPdqXZhPg" \ + -d '{ + "model": "fake-openai-endpoint", + "messages": [ + {"role": "user", "content": "Hello, Claude"} + ], + "user": "hello", + }' + ``` + diff --git a/docs/my-website/sidebars.js b/docs/my-website/sidebars.js index 18ad940f8aad..1dc33f554e7f 100644 --- a/docs/my-website/sidebars.js +++ b/docs/my-website/sidebars.js @@ -266,6 +266,7 @@ const sidebars = { type: "category", label: "Load Testing", items: [ + "benchmarks", "load_test", "load_test_advanced", "load_test_sdk", diff --git a/litellm/__init__.py b/litellm/__init__.py index b739afb9315d..9812de1d8098 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -375,6 +375,7 @@ def identify(event_details): cohere_models: List = [] cohere_chat_models: List = [] mistral_chat_models: List = [] +text_completion_codestral_models: List = [] anthropic_models: List = [] empower_models: List = [] openrouter_models: List = [] @@ -401,6 +402,19 @@ def identify(event_details): perplexity_models: List = [] watsonx_models: List = [] gemini_models: List = [] +xai_models: List = [] +deepseek_models: List = [] +azure_ai_models: List = [] +voyage_models: List = [] +databricks_models: List = [] +cloudflare_models: List = [] +codestral_models: List = [] +friendliai_models: List = [] +palm_models: List = [] +groq_models: List = [] +azure_models: List = [] +anyscale_models: List = [] +cerebras_models: List = [] def add_known_models(): @@ -477,6 +491,34 @@ def add_known_models(): # ignore the 'up-to', '-to-' model names -> not real models. just for cost tracking based on model params. if "-to-" not in key: fireworks_ai_embedding_models.append(key) + elif value.get("litellm_provider") == "text-completion-codestral": + text_completion_codestral_models.append(key) + elif value.get("litellm_provider") == "xai": + xai_models.append(key) + elif value.get("litellm_provider") == "deepseek": + deepseek_models.append(key) + elif value.get("litellm_provider") == "azure_ai": + azure_ai_models.append(key) + elif value.get("litellm_provider") == "voyage": + voyage_models.append(key) + elif value.get("litellm_provider") == "databricks": + databricks_models.append(key) + elif value.get("litellm_provider") == "cloudflare": + cloudflare_models.append(key) + elif value.get("litellm_provider") == "codestral": + codestral_models.append(key) + elif value.get("litellm_provider") == "friendliai": + friendliai_models.append(key) + elif value.get("litellm_provider") == "palm": + palm_models.append(key) + elif value.get("litellm_provider") == "groq": + groq_models.append(key) + elif value.get("litellm_provider") == "azure": + azure_models.append(key) + elif value.get("litellm_provider") == "anyscale": + anyscale_models.append(key) + elif value.get("litellm_provider") == "cerebras": + cerebras_models.append(key) add_known_models() @@ -722,6 +764,20 @@ def add_known_models(): + vertex_language_models + watsonx_models + gemini_models + + text_completion_codestral_models + + xai_models + + deepseek_models + + azure_ai_models + + voyage_models + + databricks_models + + cloudflare_models + + codestral_models + + friendliai_models + + palm_models + + groq_models + + azure_models + + anyscale_models + + cerebras_models ) @@ -778,6 +834,7 @@ class LlmProviders(str, Enum): FIREWORKS_AI = "fireworks_ai" FRIENDLIAI = "friendliai" WATSONX = "watsonx" + WATSONX_TEXT = "watsonx_text" TRITON = "triton" PREDIBASE = "predibase" DATABRICKS = "databricks" @@ -794,6 +851,7 @@ class LlmProviders(str, Enum): models_by_provider: dict = { "openai": open_ai_chat_completion_models + open_ai_text_completion_models, + "text-completion-openai": open_ai_text_completion_models, "cohere": cohere_models + cohere_chat_models, "cohere_chat": cohere_chat_models, "anthropic": anthropic_models, @@ -817,6 +875,23 @@ class LlmProviders(str, Enum): "watsonx": watsonx_models, "gemini": gemini_models, "fireworks_ai": fireworks_ai_models + fireworks_ai_embedding_models, + "aleph_alpha": aleph_alpha_models, + "text-completion-codestral": text_completion_codestral_models, + "xai": xai_models, + "deepseek": deepseek_models, + "mistral": mistral_chat_models, + "azure_ai": azure_ai_models, + "voyage": voyage_models, + "databricks": databricks_models, + "cloudflare": cloudflare_models, + "codestral": codestral_models, + "nlp_cloud": nlp_cloud_models, + "friendliai": friendliai_models, + "palm": palm_models, + "groq": groq_models, + "azure": azure_models, + "anyscale": anyscale_models, + "cerebras": cerebras_models, } # mapping for those models which have larger equivalents @@ -889,7 +964,6 @@ class LlmProviders(str, Enum): supports_system_messages, get_litellm_params, acreate, - get_model_list, get_max_tokens, get_model_info, register_prompt_template, diff --git a/litellm/_redis.py b/litellm/_redis.py index c058a0d3a879..2fba9d146952 100644 --- a/litellm/_redis.py +++ b/litellm/_redis.py @@ -12,13 +12,13 @@ # s/o [@Frank Colson](https://www.linkedin.com/in/frank-colson-422b9b183/) for this redis implementation import os -from typing import List, Optional, Union +from typing import Dict, List, Optional, Union import redis # type: ignore import redis.asyncio as async_redis # type: ignore import litellm -from litellm import get_secret +from litellm import get_secret, get_secret_str from ._logging import verbose_logger @@ -141,6 +141,13 @@ def _get_redis_client_logic(**env_overrides): if _sentinel_nodes is not None and isinstance(_sentinel_nodes, str): redis_kwargs["sentinel_nodes"] = json.loads(_sentinel_nodes) + _sentinel_password: Optional[str] = redis_kwargs.get( + "sentinel_password", None + ) or get_secret_str("REDIS_SENTINEL_PASSWORD") + + if _sentinel_password is not None: + redis_kwargs["sentinel_password"] = _sentinel_password + _service_name: Optional[str] = redis_kwargs.get("service_name", None) or get_secret( # type: ignore "REDIS_SERVICE_NAME" ) @@ -217,6 +224,7 @@ def _init_redis_sentinel(redis_kwargs) -> redis.Redis: def _init_async_redis_sentinel(redis_kwargs) -> async_redis.Redis: sentinel_nodes = redis_kwargs.get("sentinel_nodes") + sentinel_password = redis_kwargs.get("sentinel_password") service_name = redis_kwargs.get("service_name") if not sentinel_nodes or not service_name: @@ -227,7 +235,11 @@ def _init_async_redis_sentinel(redis_kwargs) -> async_redis.Redis: verbose_logger.debug("init_redis_sentinel: sentinel nodes are being initialized.") # Set up the Sentinel client - sentinel = async_redis.Sentinel(sentinel_nodes, socket_timeout=0.1) + sentinel = async_redis.Sentinel( + sentinel_nodes, + socket_timeout=0.1, + password=sentinel_password, + ) # Return the master instance for the given service diff --git a/litellm/caching/base_cache.py b/litellm/caching/base_cache.py index a50e09bf97f5..7109951d1599 100644 --- a/litellm/caching/base_cache.py +++ b/litellm/caching/base_cache.py @@ -8,6 +8,7 @@ - async_get_cache """ +from abc import ABC, abstractmethod from typing import TYPE_CHECKING, Any, Optional if TYPE_CHECKING: @@ -18,7 +19,7 @@ Span = Any -class BaseCache: +class BaseCache(ABC): def __init__(self, default_ttl: int = 60): self.default_ttl = default_ttl @@ -37,6 +38,10 @@ def set_cache(self, key, value, **kwargs): async def async_set_cache(self, key, value, **kwargs): raise NotImplementedError + @abstractmethod + async def async_set_cache_pipeline(self, cache_list, **kwargs): + pass + def get_cache(self, key, **kwargs): raise NotImplementedError diff --git a/litellm/caching/caching.py b/litellm/caching/caching.py index 5fd972a76ffd..17c09b997763 100644 --- a/litellm/caching/caching.py +++ b/litellm/caching/caching.py @@ -233,19 +233,18 @@ def __init__( if self.namespace is not None and isinstance(self.cache, RedisCache): self.cache.namespace = self.namespace - def get_cache_key(self, *args, **kwargs) -> str: + def get_cache_key(self, **kwargs) -> str: """ Get the cache key for the given arguments. Args: - *args: args to litellm.completion() or embedding() **kwargs: kwargs to litellm.completion() or embedding() Returns: str: The cache key generated from the arguments, or None if no cache key could be generated. """ cache_key = "" - verbose_logger.debug("\nGetting Cache key. Kwargs: %s", kwargs) + # verbose_logger.debug("\nGetting Cache key. Kwargs: %s", kwargs) preset_cache_key = self._get_preset_cache_key_from_kwargs(**kwargs) if preset_cache_key is not None: @@ -521,7 +520,7 @@ def _get_cache_logic( return cached_response return cached_result - def get_cache(self, *args, **kwargs): + def get_cache(self, **kwargs): """ Retrieves the cached result for the given arguments. @@ -533,13 +532,13 @@ def get_cache(self, *args, **kwargs): The cached result if it exists, otherwise None. """ try: # never block execution - if self.should_use_cache(*args, **kwargs) is not True: + if self.should_use_cache(**kwargs) is not True: return messages = kwargs.get("messages", []) if "cache_key" in kwargs: cache_key = kwargs["cache_key"] else: - cache_key = self.get_cache_key(*args, **kwargs) + cache_key = self.get_cache_key(**kwargs) if cache_key is not None: cache_control_args = kwargs.get("cache", {}) max_age = cache_control_args.get( @@ -553,29 +552,28 @@ def get_cache(self, *args, **kwargs): print_verbose(f"An exception occurred: {traceback.format_exc()}") return None - async def async_get_cache(self, *args, **kwargs): + async def async_get_cache(self, **kwargs): """ Async get cache implementation. Used for embedding calls in async wrapper """ + try: # never block execution - if self.should_use_cache(*args, **kwargs) is not True: + if self.should_use_cache(**kwargs) is not True: return kwargs.get("messages", []) if "cache_key" in kwargs: cache_key = kwargs["cache_key"] else: - cache_key = self.get_cache_key(*args, **kwargs) + cache_key = self.get_cache_key(**kwargs) if cache_key is not None: cache_control_args = kwargs.get("cache", {}) max_age = cache_control_args.get( "s-max-age", cache_control_args.get("s-maxage", float("inf")) ) - cached_result = await self.cache.async_get_cache( - cache_key, *args, **kwargs - ) + cached_result = await self.cache.async_get_cache(cache_key, **kwargs) return self._get_cache_logic( cached_result=cached_result, max_age=max_age ) @@ -583,7 +581,7 @@ async def async_get_cache(self, *args, **kwargs): print_verbose(f"An exception occurred: {traceback.format_exc()}") return None - def _add_cache_logic(self, result, *args, **kwargs): + def _add_cache_logic(self, result, **kwargs): """ Common implementation across sync + async add_cache functions """ @@ -591,7 +589,7 @@ def _add_cache_logic(self, result, *args, **kwargs): if "cache_key" in kwargs: cache_key = kwargs["cache_key"] else: - cache_key = self.get_cache_key(*args, **kwargs) + cache_key = self.get_cache_key(**kwargs) if cache_key is not None: if isinstance(result, BaseModel): result = result.model_dump_json() @@ -613,7 +611,7 @@ def _add_cache_logic(self, result, *args, **kwargs): except Exception as e: raise e - def add_cache(self, result, *args, **kwargs): + def add_cache(self, result, **kwargs): """ Adds a result to the cache. @@ -625,41 +623,42 @@ def add_cache(self, result, *args, **kwargs): None """ try: - if self.should_use_cache(*args, **kwargs) is not True: + if self.should_use_cache(**kwargs) is not True: return cache_key, cached_data, kwargs = self._add_cache_logic( - result=result, *args, **kwargs + result=result, **kwargs ) self.cache.set_cache(cache_key, cached_data, **kwargs) except Exception as e: verbose_logger.exception(f"LiteLLM Cache: Excepton add_cache: {str(e)}") - async def async_add_cache(self, result, *args, **kwargs): + async def async_add_cache(self, result, **kwargs): """ Async implementation of add_cache """ try: - if self.should_use_cache(*args, **kwargs) is not True: + if self.should_use_cache(**kwargs) is not True: return if self.type == "redis" and self.redis_flush_size is not None: # high traffic - fill in results in memory and then flush - await self.batch_cache_write(result, *args, **kwargs) + await self.batch_cache_write(result, **kwargs) else: cache_key, cached_data, kwargs = self._add_cache_logic( - result=result, *args, **kwargs + result=result, **kwargs ) + await self.cache.async_set_cache(cache_key, cached_data, **kwargs) except Exception as e: verbose_logger.exception(f"LiteLLM Cache: Excepton add_cache: {str(e)}") - async def async_add_cache_pipeline(self, result, *args, **kwargs): + async def async_add_cache_pipeline(self, result, **kwargs): """ Async implementation of add_cache for Embedding calls Does a bulk write, to prevent using too many clients """ try: - if self.should_use_cache(*args, **kwargs) is not True: + if self.should_use_cache(**kwargs) is not True: return # set default ttl if not set @@ -668,29 +667,27 @@ async def async_add_cache_pipeline(self, result, *args, **kwargs): cache_list = [] for idx, i in enumerate(kwargs["input"]): - preset_cache_key = self.get_cache_key(*args, **{**kwargs, "input": i}) + preset_cache_key = self.get_cache_key(**{**kwargs, "input": i}) kwargs["cache_key"] = preset_cache_key embedding_response = result.data[idx] cache_key, cached_data, kwargs = self._add_cache_logic( result=embedding_response, - *args, **kwargs, ) cache_list.append((cache_key, cached_data)) - async_set_cache_pipeline = getattr( - self.cache, "async_set_cache_pipeline", None - ) - if async_set_cache_pipeline: - await async_set_cache_pipeline(cache_list=cache_list, **kwargs) - else: - tasks = [] - for val in cache_list: - tasks.append(self.cache.async_set_cache(val[0], val[1], **kwargs)) - await asyncio.gather(*tasks) + + await self.cache.async_set_cache_pipeline(cache_list=cache_list, **kwargs) + # if async_set_cache_pipeline: + # await async_set_cache_pipeline(cache_list=cache_list, **kwargs) + # else: + # tasks = [] + # for val in cache_list: + # tasks.append(self.cache.async_set_cache(val[0], val[1], **kwargs)) + # await asyncio.gather(*tasks) except Exception as e: verbose_logger.exception(f"LiteLLM Cache: Excepton add_cache: {str(e)}") - def should_use_cache(self, *args, **kwargs): + def should_use_cache(self, **kwargs): """ Returns true if we should use the cache for LLM API calls @@ -708,10 +705,8 @@ def should_use_cache(self, *args, **kwargs): return True return False - async def batch_cache_write(self, result, *args, **kwargs): - cache_key, cached_data, kwargs = self._add_cache_logic( - result=result, *args, **kwargs - ) + async def batch_cache_write(self, result, **kwargs): + cache_key, cached_data, kwargs = self._add_cache_logic(result=result, **kwargs) await self.cache.batch_cache_write(cache_key, cached_data, **kwargs) async def ping(self): diff --git a/litellm/caching/caching_handler.py b/litellm/caching/caching_handler.py index f4e7d8476dde..11ae600b74a2 100644 --- a/litellm/caching/caching_handler.py +++ b/litellm/caching/caching_handler.py @@ -137,7 +137,7 @@ async def _async_get_cache( if litellm.cache is not None and self._is_call_type_supported_by_cache( original_function=original_function ): - print_verbose("Checking Cache") + verbose_logger.debug("Checking Cache") cached_result = await self._retrieve_from_cache( call_type=call_type, kwargs=kwargs, @@ -145,7 +145,7 @@ async def _async_get_cache( ) if cached_result is not None and not isinstance(cached_result, list): - print_verbose("Cache Hit!") + verbose_logger.debug("Cache Hit!") cache_hit = True end_time = datetime.datetime.now() model, _, _, _ = litellm.get_llm_provider( @@ -215,6 +215,7 @@ async def _async_get_cache( final_embedding_cached_response=final_embedding_cached_response, embedding_all_elements_cache_hit=embedding_all_elements_cache_hit, ) + verbose_logger.debug(f"CACHE RESULT: {cached_result}") return CachingHandlerResponse( cached_result=cached_result, final_embedding_cached_response=final_embedding_cached_response, @@ -233,12 +234,19 @@ def _sync_get_cache( from litellm.utils import CustomStreamWrapper args = args or () + new_kwargs = kwargs.copy() + new_kwargs.update( + convert_args_to_kwargs( + self.original_function, + args, + ) + ) cached_result: Optional[Any] = None if litellm.cache is not None and self._is_call_type_supported_by_cache( original_function=original_function ): print_verbose("Checking Cache") - cached_result = litellm.cache.get_cache(*args, **kwargs) + cached_result = litellm.cache.get_cache(**new_kwargs) if cached_result is not None: if "detail" in cached_result: # implies an error occurred @@ -475,14 +483,21 @@ async def _retrieve_from_cache( if litellm.cache is None: return None + new_kwargs = kwargs.copy() + new_kwargs.update( + convert_args_to_kwargs( + self.original_function, + args, + ) + ) cached_result: Optional[Any] = None if call_type == CallTypes.aembedding.value and isinstance( - kwargs["input"], list + new_kwargs["input"], list ): tasks = [] - for idx, i in enumerate(kwargs["input"]): + for idx, i in enumerate(new_kwargs["input"]): preset_cache_key = litellm.cache.get_cache_key( - *args, **{**kwargs, "input": i} + **{**new_kwargs, "input": i} ) tasks.append(litellm.cache.async_get_cache(cache_key=preset_cache_key)) cached_result = await asyncio.gather(*tasks) @@ -493,9 +508,9 @@ async def _retrieve_from_cache( cached_result = None else: if litellm.cache._supports_async() is True: - cached_result = await litellm.cache.async_get_cache(*args, **kwargs) + cached_result = await litellm.cache.async_get_cache(**new_kwargs) else: # for s3 caching. [NOT RECOMMENDED IN PROD - this will slow down responses since boto3 is sync] - cached_result = litellm.cache.get_cache(*args, **kwargs) + cached_result = litellm.cache.get_cache(**new_kwargs) return cached_result def _convert_cached_result_to_model_response( @@ -580,6 +595,7 @@ def _convert_cached_result_to_model_response( model_response_object=EmbeddingResponse(), response_type="embedding", ) + elif ( call_type == CallTypes.arerank.value or call_type == CallTypes.rerank.value ) and isinstance(cached_result, dict): @@ -603,6 +619,13 @@ def _convert_cached_result_to_model_response( response_type="audio_transcription", hidden_params=hidden_params, ) + + if ( + hasattr(cached_result, "_hidden_params") + and cached_result._hidden_params is not None + and isinstance(cached_result._hidden_params, dict) + ): + cached_result._hidden_params["cache_hit"] = True return cached_result def _convert_cached_stream_response( @@ -658,12 +681,19 @@ async def async_set_cache( Raises: None """ - kwargs.update(convert_args_to_kwargs(result, original_function, kwargs, args)) + + new_kwargs = kwargs.copy() + new_kwargs.update( + convert_args_to_kwargs( + original_function, + args, + ) + ) if litellm.cache is None: return # [OPTIONAL] ADD TO CACHE if self._should_store_result_in_cache( - original_function=original_function, kwargs=kwargs + original_function=original_function, kwargs=new_kwargs ): if ( isinstance(result, litellm.ModelResponse) @@ -673,29 +703,29 @@ async def async_set_cache( ): if ( isinstance(result, EmbeddingResponse) - and isinstance(kwargs["input"], list) + and isinstance(new_kwargs["input"], list) and litellm.cache is not None and not isinstance( litellm.cache.cache, S3Cache ) # s3 doesn't support bulk writing. Exclude. ): asyncio.create_task( - litellm.cache.async_add_cache_pipeline(result, **kwargs) + litellm.cache.async_add_cache_pipeline(result, **new_kwargs) ) elif isinstance(litellm.cache.cache, S3Cache): threading.Thread( target=litellm.cache.add_cache, args=(result,), - kwargs=kwargs, + kwargs=new_kwargs, ).start() else: asyncio.create_task( litellm.cache.async_add_cache( - result.model_dump_json(), **kwargs + result.model_dump_json(), **new_kwargs ) ) else: - asyncio.create_task(litellm.cache.async_add_cache(result, **kwargs)) + asyncio.create_task(litellm.cache.async_add_cache(result, **new_kwargs)) def sync_set_cache( self, @@ -706,16 +736,20 @@ def sync_set_cache( """ Sync internal method to add the result to the cache """ - kwargs.update( - convert_args_to_kwargs(result, self.original_function, kwargs, args) + new_kwargs = kwargs.copy() + new_kwargs.update( + convert_args_to_kwargs( + self.original_function, + args, + ) ) if litellm.cache is None: return if self._should_store_result_in_cache( - original_function=self.original_function, kwargs=kwargs + original_function=self.original_function, kwargs=new_kwargs ): - litellm.cache.add_cache(result, **kwargs) + litellm.cache.add_cache(result, **new_kwargs) return @@ -865,9 +899,7 @@ def _update_litellm_logging_obj_environment( def convert_args_to_kwargs( - result: Any, original_function: Callable, - kwargs: Dict[str, Any], args: Optional[Tuple[Any, ...]] = None, ) -> Dict[str, Any]: # Get the signature of the original function diff --git a/litellm/caching/disk_cache.py b/litellm/caching/disk_cache.py index 2c086ed50c32..94f82926d33e 100644 --- a/litellm/caching/disk_cache.py +++ b/litellm/caching/disk_cache.py @@ -24,7 +24,6 @@ def __init__(self, disk_cache_dir: Optional[str] = None): self.disk_cache = dc.Cache(disk_cache_dir) def set_cache(self, key, value, **kwargs): - print_verbose("DiskCache: set_cache") if "ttl" in kwargs: self.disk_cache.set(key, value, expire=kwargs["ttl"]) else: @@ -33,10 +32,10 @@ def set_cache(self, key, value, **kwargs): async def async_set_cache(self, key, value, **kwargs): self.set_cache(key=key, value=value, **kwargs) - async def async_set_cache_pipeline(self, cache_list, ttl=None): + async def async_set_cache_pipeline(self, cache_list, **kwargs): for cache_key, cache_value in cache_list: - if ttl is not None: - self.set_cache(key=cache_key, value=cache_value, ttl=ttl) + if "ttl" in kwargs: + self.set_cache(key=cache_key, value=cache_value, ttl=kwargs["ttl"]) else: self.set_cache(key=cache_key, value=cache_value) diff --git a/litellm/caching/dual_cache.py b/litellm/caching/dual_cache.py index ddcd02abe5f6..a6c218c0148a 100644 --- a/litellm/caching/dual_cache.py +++ b/litellm/caching/dual_cache.py @@ -314,7 +314,8 @@ async def async_set_cache(self, key, value, local_only: bool = False, **kwargs): f"LiteLLM Cache: Excepton async add_cache: {str(e)}" ) - async def async_batch_set_cache( + # async_batch_set_cache + async def async_set_cache_pipeline( self, cache_list: list, local_only: bool = False, **kwargs ): """ diff --git a/litellm/caching/qdrant_semantic_cache.py b/litellm/caching/qdrant_semantic_cache.py index be67001f6917..acaa8e918928 100644 --- a/litellm/caching/qdrant_semantic_cache.py +++ b/litellm/caching/qdrant_semantic_cache.py @@ -9,6 +9,7 @@ """ import ast +import asyncio import json from typing import Any @@ -422,3 +423,9 @@ async def async_get_cache(self, key, **kwargs): async def _collection_info(self): return self.collection_info + + async def async_set_cache_pipeline(self, cache_list, **kwargs): + tasks = [] + for val in cache_list: + tasks.append(self.async_set_cache(val[0], val[1], **kwargs)) + await asyncio.gather(*tasks) diff --git a/litellm/caching/redis_cache.py b/litellm/caching/redis_cache.py index 40bb49f448bf..e15a3f83d683 100644 --- a/litellm/caching/redis_cache.py +++ b/litellm/caching/redis_cache.py @@ -404,7 +404,7 @@ async def async_set_cache_pipeline( parent_otel_span=_get_parent_otel_span_from_kwargs(kwargs), ) ) - return results + return None except Exception as e: ## LOGGING ## end_time = time.time() diff --git a/litellm/caching/redis_semantic_cache.py b/litellm/caching/redis_semantic_cache.py index 444a3259f256..e3098f085625 100644 --- a/litellm/caching/redis_semantic_cache.py +++ b/litellm/caching/redis_semantic_cache.py @@ -9,6 +9,7 @@ """ import ast +import asyncio import json from typing import Any @@ -331,3 +332,9 @@ async def async_get_cache(self, key, **kwargs): async def _index_info(self): return await self.index.ainfo() + + async def async_set_cache_pipeline(self, cache_list, **kwargs): + tasks = [] + for val in cache_list: + tasks.append(self.async_set_cache(val[0], val[1], **kwargs)) + await asyncio.gather(*tasks) diff --git a/litellm/caching/s3_cache.py b/litellm/caching/s3_cache.py index c22347a7f2db..6be16e289a4f 100644 --- a/litellm/caching/s3_cache.py +++ b/litellm/caching/s3_cache.py @@ -10,6 +10,7 @@ """ import ast +import asyncio import json from typing import Any, Optional @@ -153,3 +154,9 @@ def flush_cache(self): async def disconnect(self): pass + + async def async_set_cache_pipeline(self, cache_list, **kwargs): + tasks = [] + for val in cache_list: + tasks.append(self.async_set_cache(val[0], val[1], **kwargs)) + await asyncio.gather(*tasks) diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index 0be7f1d38bf1..0aa8a8e36b6f 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -28,6 +28,9 @@ from litellm.llms.AzureOpenAI.cost_calculation import ( cost_per_token as azure_openai_cost_per_token, ) +from litellm.llms.bedrock.image.cost_calculator import ( + cost_calculator as bedrock_image_cost_calculator, +) from litellm.llms.cohere.cost_calculator import ( cost_per_query as cohere_rerank_cost_per_query, ) @@ -168,7 +171,6 @@ def cost_per_token( # noqa: PLR0915 model_with_provider = model_with_provider_and_region else: _, custom_llm_provider, _, _ = litellm.get_llm_provider(model=model) - model_without_prefix = model model_parts = model.split("/", 1) if len(model_parts) > 1: @@ -451,7 +453,6 @@ def _select_model_name_for_cost_calc( if base_model is not None: return base_model - return_model = model if isinstance(completion_response, str): return return_model @@ -521,12 +522,13 @@ def completion_cost( # noqa: PLR0915 custom_llm_provider=None, region_name=None, # used for bedrock pricing ### IMAGE GEN ### - size=None, + size: Optional[str] = None, quality=None, n=None, # number of images ### CUSTOM PRICING ### custom_cost_per_token: Optional[CostPerToken] = None, custom_cost_per_second: Optional[float] = None, + optional_params: Optional[dict] = None, ) -> float: """ Calculate the cost of a given completion call fot GPT-3.5-turbo, llama2, any litellm supported llm. @@ -616,7 +618,8 @@ def completion_cost( # noqa: PLR0915 f"completion_response response ms: {getattr(completion_response, '_response_ms', None)} " ) model = _select_model_name_for_cost_calc( - model=model, completion_response=completion_response + model=model, + completion_response=completion_response, ) hidden_params = getattr(completion_response, "_hidden_params", None) if hidden_params is not None: @@ -667,7 +670,17 @@ def completion_cost( # noqa: PLR0915 # https://cloud.google.com/vertex-ai/generative-ai/pricing # Vertex Charges Flat $0.20 per image return 0.020 - + elif custom_llm_provider == "bedrock": + if isinstance(completion_response, ImageResponse): + return bedrock_image_cost_calculator( + model=model, + size=size, + image_response=completion_response, + optional_params=optional_params, + ) + raise TypeError( + "completion_response must be of type ImageResponse for bedrock image cost calculation" + ) if size is None: size = "1024-x-1024" # openai default # fix size to match naming convention @@ -677,9 +690,9 @@ def completion_cost( # noqa: PLR0915 image_gen_model_name_with_quality = image_gen_model_name if quality is not None: image_gen_model_name_with_quality = f"{quality}/{image_gen_model_name}" - size = size.split("-x-") - height = int(size[0]) # if it's 1024-x-1024 vs. 1024x1024 - width = int(size[1]) + size_parts = size.split("-x-") + height = int(size_parts[0]) # if it's 1024-x-1024 vs. 1024x1024 + width = int(size_parts[1]) verbose_logger.debug(f"image_gen_model_name: {image_gen_model_name}") verbose_logger.debug( f"image_gen_model_name_with_quality: {image_gen_model_name_with_quality}" @@ -839,11 +852,14 @@ def response_cost_calculator( if isinstance(response_object, BaseModel): response_object._hidden_params["optional_params"] = optional_params if isinstance(response_object, ImageResponse): + if base_model is not None: + model = base_model response_cost = completion_cost( completion_response=response_object, model=model, call_type=call_type, custom_llm_provider=custom_llm_provider, + optional_params=optional_params, ) else: if custom_pricing is True: # override defaults if custom pricing is set diff --git a/litellm/integrations/SlackAlerting/slack_alerting.py b/litellm/integrations/SlackAlerting/slack_alerting.py index 85d54a337d16..d585e235b799 100644 --- a/litellm/integrations/SlackAlerting/slack_alerting.py +++ b/litellm/integrations/SlackAlerting/slack_alerting.py @@ -423,7 +423,7 @@ async def send_daily_reports(self, router) -> bool: # noqa: PLR0915 latency_cache_keys = [(key, 0) for key in latency_keys] failed_request_cache_keys = [(key, 0) for key in failed_request_keys] combined_metrics_cache_keys = latency_cache_keys + failed_request_cache_keys - await self.internal_usage_cache.async_batch_set_cache( + await self.internal_usage_cache.async_set_cache_pipeline( cache_list=combined_metrics_cache_keys ) diff --git a/litellm/integrations/gcs_bucket/gcs_bucket.py b/litellm/integrations/gcs_bucket/gcs_bucket.py index 0b637f9b6ec4..83b831904995 100644 --- a/litellm/integrations/gcs_bucket/gcs_bucket.py +++ b/litellm/integrations/gcs_bucket/gcs_bucket.py @@ -115,7 +115,17 @@ async def async_log_failure_event(self, kwargs, response_obj, start_time, end_ti verbose_logger.exception(f"GCS Bucket logging error: {str(e)}") async def async_send_batch(self): - """Process queued logs in batch - sends logs to GCS Bucket""" + """ + Process queued logs in batch - sends logs to GCS Bucket + + + GCS Bucket does not have a Batch endpoint to batch upload logs + + Instead, we + - collect the logs to flush every `GCS_FLUSH_INTERVAL` seconds + - during async_send_batch, we make 1 POST request per log to GCS Bucket + + """ if not self.log_queue: return diff --git a/litellm/integrations/langfuse/langfuse.py b/litellm/integrations/langfuse/langfuse.py index 18892871e46a..73485a0bdbfb 100644 --- a/litellm/integrations/langfuse/langfuse.py +++ b/litellm/integrations/langfuse/langfuse.py @@ -3,8 +3,9 @@ import copy import os import traceback +import types from collections.abc import MutableMapping, MutableSequence, MutableSet -from typing import TYPE_CHECKING, Any, Dict, Optional +from typing import TYPE_CHECKING, Any, Dict, Optional, cast from packaging.version import Version from pydantic import BaseModel @@ -355,17 +356,28 @@ def _log_langfuse_v1( ) ) - def _prepare_metadata(self, metadata) -> Any: + def is_base_type(self, value: Any) -> bool: + # Check if the value is of a base type + base_types = (int, float, str, bool, list, dict, tuple) + return isinstance(value, base_types) + + def _prepare_metadata(self, metadata: Optional[dict]) -> Any: try: - return copy.deepcopy(metadata) # Avoid modifying the original metadata - except (TypeError, copy.Error) as e: - verbose_logger.warning(f"Langfuse Layer Error - {e}") + if metadata is None: + return None + + # Filter out function types from the metadata + sanitized_metadata = {k: v for k, v in metadata.items() if not callable(v)} + + return copy.deepcopy(sanitized_metadata) + except Exception as e: + verbose_logger.debug(f"Langfuse Layer Error - {e}, metadata: {metadata}") new_metadata: Dict[str, Any] = {} # if metadata is not a MutableMapping, return an empty dict since we can't call items() on it if not isinstance(metadata, MutableMapping): - verbose_logger.warning( + verbose_logger.debug( "Langfuse Layer Logging - metadata is not a MutableMapping, returning empty dict" ) return new_metadata @@ -373,25 +385,40 @@ def _prepare_metadata(self, metadata) -> Any: for key, value in metadata.items(): try: if isinstance(value, MutableMapping): - new_metadata[key] = self._prepare_metadata(value) - elif isinstance(value, (MutableSequence, MutableSet)): - new_metadata[key] = type(value)( - *( - ( - self._prepare_metadata(v) - if isinstance(v, MutableMapping) - else copy.deepcopy(v) - ) - for v in value + new_metadata[key] = self._prepare_metadata(cast(dict, value)) + elif isinstance(value, MutableSequence): + # For lists or other mutable sequences + new_metadata[key] = list( + ( + self._prepare_metadata(cast(dict, v)) + if isinstance(v, MutableMapping) + else copy.deepcopy(v) ) + for v in value + ) + elif isinstance(value, MutableSet): + # For sets specifically, create a new set by passing an iterable + new_metadata[key] = set( + ( + self._prepare_metadata(cast(dict, v)) + if isinstance(v, MutableMapping) + else copy.deepcopy(v) + ) + for v in value ) elif isinstance(value, BaseModel): new_metadata[key] = value.model_dump() + elif self.is_base_type(value): + new_metadata[key] = value else: - new_metadata[key] = copy.deepcopy(value) + verbose_logger.debug( + f"Langfuse Layer Error - Unsupported metadata type: {type(value)} for key: {key}" + ) + continue + except (TypeError, copy.Error): - verbose_logger.warning( - f"Langfuse Layer Error - Couldn't copy metadata key: {key} - {traceback.format_exc()}" + verbose_logger.debug( + f"Langfuse Layer Error - Couldn't copy metadata key: {key}, type of key: {type(key)}, type of value: {type(value)} - {traceback.format_exc()}" ) return new_metadata diff --git a/litellm/integrations/langsmith.py b/litellm/integrations/langsmith.py index 9513934457f5..4abd2a2c3365 100644 --- a/litellm/integrations/langsmith.py +++ b/litellm/integrations/langsmith.py @@ -23,34 +23,8 @@ get_async_httpx_client, httpxSpecialProvider, ) -from litellm.types.utils import StandardLoggingPayload - - -class LangsmithInputs(BaseModel): - model: Optional[str] = None - messages: Optional[List[Any]] = None - stream: Optional[bool] = None - call_type: Optional[str] = None - litellm_call_id: Optional[str] = None - completion_start_time: Optional[datetime] = None - temperature: Optional[float] = None - max_tokens: Optional[int] = None - custom_llm_provider: Optional[str] = None - input: Optional[List[Any]] = None - log_event_type: Optional[str] = None - original_response: Optional[Any] = None - response_cost: Optional[float] = None - - # LiteLLM Virtual Key specific fields - user_api_key: Optional[str] = None - user_api_key_user_id: Optional[str] = None - user_api_key_team_alias: Optional[str] = None - - -class LangsmithCredentialsObject(TypedDict): - LANGSMITH_API_KEY: str - LANGSMITH_PROJECT: str - LANGSMITH_BASE_URL: str +from litellm.types.integrations.langsmith import * +from litellm.types.utils import StandardCallbackDynamicParams, StandardLoggingPayload def is_serializable(value): @@ -93,15 +67,16 @@ def __init__( ) if _batch_size: self.batch_size = int(_batch_size) + self.log_queue: List[LangsmithQueueObject] = [] asyncio.create_task(self.periodic_flush()) self.flush_lock = asyncio.Lock() super().__init__(**kwargs, flush_lock=self.flush_lock) def get_credentials_from_env( self, - langsmith_api_key: Optional[str], - langsmith_project: Optional[str], - langsmith_base_url: Optional[str], + langsmith_api_key: Optional[str] = None, + langsmith_project: Optional[str] = None, + langsmith_base_url: Optional[str] = None, ) -> LangsmithCredentialsObject: _credentials_api_key = langsmith_api_key or os.getenv("LANGSMITH_API_KEY") @@ -132,42 +107,19 @@ def get_credentials_from_env( LANGSMITH_PROJECT=_credentials_project, ) - def _prepare_log_data( # noqa: PLR0915 - self, kwargs, response_obj, start_time, end_time + def _prepare_log_data( + self, + kwargs, + response_obj, + start_time, + end_time, + credentials: LangsmithCredentialsObject, ): - import json - from datetime import datetime as dt - try: _litellm_params = kwargs.get("litellm_params", {}) or {} metadata = _litellm_params.get("metadata", {}) or {} - new_metadata = {} - for key, value in metadata.items(): - if ( - isinstance(value, list) - or isinstance(value, str) - or isinstance(value, int) - or isinstance(value, float) - ): - new_metadata[key] = value - elif isinstance(value, BaseModel): - new_metadata[key] = value.model_dump_json() - elif isinstance(value, dict): - for k, v in value.items(): - if isinstance(v, dt): - value[k] = v.isoformat() - new_metadata[key] = value - - metadata = new_metadata - - kwargs["user_api_key"] = metadata.get("user_api_key", None) - kwargs["user_api_key_user_id"] = metadata.get("user_api_key_user_id", None) - kwargs["user_api_key_team_alias"] = metadata.get( - "user_api_key_team_alias", None - ) - project_name = metadata.get( - "project_name", self.default_credentials["LANGSMITH_PROJECT"] + "project_name", credentials["LANGSMITH_PROJECT"] ) run_name = metadata.get("run_name", self.langsmith_default_run_name) run_id = metadata.get("id", None) @@ -175,16 +127,10 @@ def _prepare_log_data( # noqa: PLR0915 trace_id = metadata.get("trace_id", None) session_id = metadata.get("session_id", None) dotted_order = metadata.get("dotted_order", None) - tags = metadata.get("tags", []) or [] verbose_logger.debug( f"Langsmith Logging - project_name: {project_name}, run_name {run_name}" ) - # filter out kwargs to not include any dicts, langsmith throws an erros when trying to log kwargs - # logged_kwargs = LangsmithInputs(**kwargs) - # kwargs = logged_kwargs.model_dump() - - # new_kwargs = {} # Ensure everything in the payload is converted to str payload: Optional[StandardLoggingPayload] = kwargs.get( "standard_logging_object", None @@ -193,7 +139,6 @@ def _prepare_log_data( # noqa: PLR0915 if payload is None: raise Exception("Error logging request payload. Payload=none.") - new_kwargs = payload metadata = payload[ "metadata" ] # ensure logged metadata is json serializable @@ -201,12 +146,12 @@ def _prepare_log_data( # noqa: PLR0915 data = { "name": run_name, "run_type": "llm", # this should always be llm, since litellm always logs llm calls. Langsmith allow us to log "chain" - "inputs": new_kwargs, - "outputs": new_kwargs["response"], + "inputs": payload, + "outputs": payload["response"], "session_name": project_name, - "start_time": new_kwargs["startTime"], - "end_time": new_kwargs["endTime"], - "tags": tags, + "start_time": payload["startTime"], + "end_time": payload["endTime"], + "tags": payload["request_tags"], "extra": metadata, } @@ -243,37 +188,6 @@ def _prepare_log_data( # noqa: PLR0915 except Exception: raise - def _send_batch(self): - if not self.log_queue: - return - - langsmith_api_key = self.default_credentials["LANGSMITH_API_KEY"] - langsmith_api_base = self.default_credentials["LANGSMITH_BASE_URL"] - - url = f"{langsmith_api_base}/runs/batch" - - headers = {"x-api-key": langsmith_api_key} - - try: - response = requests.post( - url=url, - json=self.log_queue, - headers=headers, - ) - - if response.status_code >= 300: - verbose_logger.error( - f"Langsmith Error: {response.status_code} - {response.text}" - ) - else: - verbose_logger.debug( - f"Batch of {len(self.log_queue)} runs successfully created" - ) - - self.log_queue.clear() - except Exception: - verbose_logger.exception("Langsmith Layer Error - Error sending batch.") - def log_success_event(self, kwargs, response_obj, start_time, end_time): try: sampling_rate = ( @@ -295,8 +209,20 @@ def log_success_event(self, kwargs, response_obj, start_time, end_time): kwargs, response_obj, ) - data = self._prepare_log_data(kwargs, response_obj, start_time, end_time) - self.log_queue.append(data) + credentials = self._get_credentials_to_use_for_request(kwargs=kwargs) + data = self._prepare_log_data( + kwargs=kwargs, + response_obj=response_obj, + start_time=start_time, + end_time=end_time, + credentials=credentials, + ) + self.log_queue.append( + LangsmithQueueObject( + data=data, + credentials=credentials, + ) + ) verbose_logger.debug( f"Langsmith, event added to queue. Will flush in {self.flush_interval} seconds..." ) @@ -323,8 +249,20 @@ async def async_log_success_event(self, kwargs, response_obj, start_time, end_ti kwargs, response_obj, ) - data = self._prepare_log_data(kwargs, response_obj, start_time, end_time) - self.log_queue.append(data) + credentials = self._get_credentials_to_use_for_request(kwargs=kwargs) + data = self._prepare_log_data( + kwargs=kwargs, + response_obj=response_obj, + start_time=start_time, + end_time=end_time, + credentials=credentials, + ) + self.log_queue.append( + LangsmithQueueObject( + data=data, + credentials=credentials, + ) + ) verbose_logger.debug( "Langsmith logging: queue length %s, batch size %s", len(self.log_queue), @@ -349,8 +287,20 @@ async def async_log_failure_event(self, kwargs, response_obj, start_time, end_ti return # Skip logging verbose_logger.info("Langsmith Failure Event Logging!") try: - data = self._prepare_log_data(kwargs, response_obj, start_time, end_time) - self.log_queue.append(data) + credentials = self._get_credentials_to_use_for_request(kwargs=kwargs) + data = self._prepare_log_data( + kwargs=kwargs, + response_obj=response_obj, + start_time=start_time, + end_time=end_time, + credentials=credentials, + ) + self.log_queue.append( + LangsmithQueueObject( + data=data, + credentials=credentials, + ) + ) verbose_logger.debug( "Langsmith logging: queue length %s, batch size %s", len(self.log_queue), @@ -365,31 +315,58 @@ async def async_log_failure_event(self, kwargs, response_obj, start_time, end_ti async def async_send_batch(self): """ - sends runs to /batch endpoint + Handles sending batches of runs to Langsmith - Sends runs from self.log_queue + self.log_queue contains LangsmithQueueObjects + Each LangsmithQueueObject has the following: + - "credentials" - credentials to use for the request (langsmith_api_key, langsmith_project, langsmith_base_url) + - "data" - data to log on to langsmith for the request - Returns: None - Raises: Does not raise an exception, will only verbose_logger.exception() + This function + - groups the queue objects by credentials + - loops through each unique credentials and sends batches to Langsmith + + + This was added to support key/team based logging on langsmith """ if not self.log_queue: return - langsmith_api_base = self.default_credentials["LANGSMITH_BASE_URL"] + batch_groups = self._group_batches_by_credentials() + for batch_group in batch_groups.values(): + await self._log_batch_on_langsmith( + credentials=batch_group.credentials, + queue_objects=batch_group.queue_objects, + ) - url = f"{langsmith_api_base}/runs/batch" + async def _log_batch_on_langsmith( + self, + credentials: LangsmithCredentialsObject, + queue_objects: List[LangsmithQueueObject], + ): + """ + Logs a batch of runs to Langsmith + sends runs to /batch endpoint for the given credentials - langsmith_api_key = self.default_credentials["LANGSMITH_API_KEY"] + Args: + credentials: LangsmithCredentialsObject + queue_objects: List[LangsmithQueueObject] + Returns: None + + Raises: Does not raise an exception, will only verbose_logger.exception() + """ + langsmith_api_base = credentials["LANGSMITH_BASE_URL"] + langsmith_api_key = credentials["LANGSMITH_API_KEY"] + url = f"{langsmith_api_base}/runs/batch" headers = {"x-api-key": langsmith_api_key} + elements_to_log = [queue_object["data"] for queue_object in queue_objects] try: response = await self.async_httpx_client.post( url=url, - json={ - "post": self.log_queue, - }, + json={"post": elements_to_log}, headers=headers, ) response.raise_for_status() @@ -411,6 +388,74 @@ async def async_send_batch(self): f"Langsmith Layer Error - {traceback.format_exc()}" ) + def _group_batches_by_credentials(self) -> Dict[CredentialsKey, BatchGroup]: + """Groups queue objects by credentials using a proper key structure""" + log_queue_by_credentials: Dict[CredentialsKey, BatchGroup] = {} + + for queue_object in self.log_queue: + credentials = queue_object["credentials"] + key = CredentialsKey( + api_key=credentials["LANGSMITH_API_KEY"], + project=credentials["LANGSMITH_PROJECT"], + base_url=credentials["LANGSMITH_BASE_URL"], + ) + + if key not in log_queue_by_credentials: + log_queue_by_credentials[key] = BatchGroup( + credentials=credentials, queue_objects=[] + ) + + log_queue_by_credentials[key].queue_objects.append(queue_object) + + return log_queue_by_credentials + + def _get_credentials_to_use_for_request( + self, kwargs: Dict[str, Any] + ) -> LangsmithCredentialsObject: + """ + Handles key/team based logging + + If standard_callback_dynamic_params are provided, use those credentials. + + Otherwise, use the default credentials. + """ + standard_callback_dynamic_params: Optional[StandardCallbackDynamicParams] = ( + kwargs.get("standard_callback_dynamic_params", None) + ) + if standard_callback_dynamic_params is not None: + credentials = self.get_credentials_from_env( + langsmith_api_key=standard_callback_dynamic_params.get( + "langsmith_api_key", None + ), + langsmith_project=standard_callback_dynamic_params.get( + "langsmith_project", None + ), + langsmith_base_url=standard_callback_dynamic_params.get( + "langsmith_base_url", None + ), + ) + else: + credentials = self.default_credentials + return credentials + + def _send_batch(self): + """Calls async_send_batch in an event loop""" + if not self.log_queue: + return + + try: + # Try to get the existing event loop + loop = asyncio.get_event_loop() + if loop.is_running(): + # If we're already in an event loop, create a task + asyncio.create_task(self.async_send_batch()) + else: + # If no event loop is running, run the coroutine directly + loop.run_until_complete(self.async_send_batch()) + except RuntimeError: + # If we can't get an event loop, create a new one + asyncio.run(self.async_send_batch()) + def get_run_by_id(self, run_id): langsmith_api_key = self.default_credentials["LANGSMITH_API_KEY"] diff --git a/litellm/integrations/opentelemetry.py b/litellm/integrations/opentelemetry.py index 8102f2c60328..30a280e5734e 100644 --- a/litellm/integrations/opentelemetry.py +++ b/litellm/integrations/opentelemetry.py @@ -2,14 +2,16 @@ from dataclasses import dataclass from datetime import datetime from functools import wraps -from typing import TYPE_CHECKING, Any, Dict, Optional, Union +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union import litellm from litellm._logging import verbose_logger from litellm.integrations.custom_logger import CustomLogger from litellm.types.services import ServiceLoggerPayload from litellm.types.utils import ( + ChatCompletionMessageToolCall, EmbeddingResponse, + Function, ImageResponse, ModelResponse, StandardLoggingPayload, @@ -403,6 +405,28 @@ def cast_as_primitive_value_type(self, value) -> Union[str, bool, int, float]: except Exception: return "" + @staticmethod + def _tool_calls_kv_pair( + tool_calls: List[ChatCompletionMessageToolCall], + ) -> Dict[str, Any]: + from litellm.proxy._types import SpanAttributes + + kv_pairs: Dict[str, Any] = {} + for idx, tool_call in enumerate(tool_calls): + _function = tool_call.get("function") + if not _function: + continue + + keys = Function.__annotations__.keys() + for key in keys: + _value = _function.get(key) + if _value: + kv_pairs[ + f"{SpanAttributes.LLM_COMPLETIONS}.{idx}.function_call.{key}" + ] = _value + + return kv_pairs + def set_attributes( # noqa: PLR0915 self, span: Span, kwargs, response_obj: Optional[Any] ): @@ -597,18 +621,13 @@ def set_attributes( # noqa: PLR0915 message = choice.get("message") tool_calls = message.get("tool_calls") if tool_calls: - self.safe_set_attribute( - span=span, - key=f"{SpanAttributes.LLM_COMPLETIONS}.{idx}.function_call.name", - value=tool_calls[0].get("function").get("name"), - ) - self.safe_set_attribute( - span=span, - key=f"{SpanAttributes.LLM_COMPLETIONS}.{idx}.function_call.arguments", - value=tool_calls[0] - .get("function") - .get("arguments"), - ) + kv_pairs = OpenTelemetry._tool_calls_kv_pair(tool_calls) # type: ignore + for key, value in kv_pairs.items(): + self.safe_set_attribute( + span=span, + key=key, + value=value, + ) except Exception as e: verbose_logger.exception( diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 2ab905e850eb..d2e65742c8a0 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -2774,11 +2774,6 @@ def get_standard_logging_object_payload( metadata=metadata ) - if litellm.cache is not None: - cache_key = litellm.cache.get_cache_key(**kwargs) - else: - cache_key = None - saved_cache_cost: float = 0.0 if cache_hit is True: @@ -2820,7 +2815,7 @@ def get_standard_logging_object_payload( completionStartTime=completion_start_time_float, model=kwargs.get("model", "") or "", metadata=clean_metadata, - cache_key=cache_key, + cache_key=clean_hidden_params["cache_key"], response_cost=response_cost, total_tokens=usage.total_tokens, prompt_tokens=usage.prompt_tokens, diff --git a/litellm/llms/anthropic/chat/handler.py b/litellm/llms/anthropic/chat/handler.py index da95ac075f04..2d119a28f278 100644 --- a/litellm/llms/anthropic/chat/handler.py +++ b/litellm/llms/anthropic/chat/handler.py @@ -71,11 +71,12 @@ def validate_environment( prompt_caching_set = AnthropicConfig().is_cache_control_set(messages=messages) computer_tool_used = AnthropicConfig().is_computer_tool_used(tools=tools) - + pdf_used = AnthropicConfig().is_pdf_used(messages=messages) headers = AnthropicConfig().get_anthropic_headers( anthropic_version=anthropic_version, computer_tool_used=computer_tool_used, prompt_caching_set=prompt_caching_set, + pdf_used=pdf_used, api_key=api_key, ) diff --git a/litellm/llms/anthropic/chat/transformation.py b/litellm/llms/anthropic/chat/transformation.py index ec32854735e8..e222d8721b76 100644 --- a/litellm/llms/anthropic/chat/transformation.py +++ b/litellm/llms/anthropic/chat/transformation.py @@ -7,6 +7,7 @@ AllAnthropicToolsValues, AnthropicComputerTool, AnthropicHostedTools, + AnthropicInputSchema, AnthropicMessageRequestBase, AnthropicMessagesRequest, AnthropicMessagesTool, @@ -104,6 +105,7 @@ def get_anthropic_headers( anthropic_version: Optional[str] = None, computer_tool_used: bool = False, prompt_caching_set: bool = False, + pdf_used: bool = False, ) -> dict: import json @@ -112,6 +114,8 @@ def get_anthropic_headers( betas.append("prompt-caching-2024-07-31") if computer_tool_used: betas.append("computer-use-2024-10-22") + if pdf_used: + betas.append("pdfs-2024-09-25") headers = { "anthropic-version": anthropic_version or "2023-06-01", "x-api-key": api_key, @@ -156,15 +160,17 @@ def _map_tool_helper( returned_tool: Optional[AllAnthropicToolsValues] = None if tool["type"] == "function" or tool["type"] == "custom": + _input_schema: dict = tool["function"].get( + "parameters", + { + "type": "object", + "properties": {}, + }, + ) + input_schema: AnthropicInputSchema = AnthropicInputSchema(**_input_schema) _tool = AnthropicMessagesTool( name=tool["function"]["name"], - input_schema=tool["function"].get( - "parameters", - { - "type": "object", - "properties": {}, - }, - ), + input_schema=input_schema, ) _description = tool["function"].get("description") @@ -301,17 +307,10 @@ def map_openai_params( - You should set tool_choice (see Forcing tool use) to instruct the model to explicitly use that tool - Remember that the model will pass the input to the tool, so the name of the tool and description should be from the model’s perspective. """ - _tool_choice = None _tool_choice = {"name": "json_tool_call", "type": "tool"} - - _tool = AnthropicMessagesTool( - name="json_tool_call", - input_schema={ - "type": "object", - "properties": {"values": json_schema}, # type: ignore - }, + _tool = self._create_json_tool_call_for_response_format( + json_schema=json_schema, ) - optional_params["tools"] = [_tool] optional_params["tool_choice"] = _tool_choice optional_params["json_mode"] = True @@ -338,6 +337,34 @@ def map_openai_params( return optional_params + def _create_json_tool_call_for_response_format( + self, + json_schema: Optional[dict] = None, + ) -> AnthropicMessagesTool: + """ + Handles creating a tool call for getting responses in JSON format. + + Args: + json_schema (Optional[dict]): The JSON schema the response should be in + + Returns: + AnthropicMessagesTool: The tool call to send to Anthropic API to get responses in JSON format + """ + _input_schema: AnthropicInputSchema = AnthropicInputSchema( + type="object", + ) + + if json_schema is None: + # Anthropic raises a 400 BadRequest error if properties is passed as None + # see usage with additionalProperties (Example 5) https://github.com/anthropics/anthropic-cookbook/blob/main/tool_use/extracting_structured_json.ipynb + _input_schema["additionalProperties"] = True + _input_schema["properties"] = {} + else: + _input_schema["properties"] = json_schema + + _tool = AnthropicMessagesTool(name="json_tool_call", input_schema=_input_schema) + return _tool + def is_cache_control_set(self, messages: List[AllMessageValues]) -> bool: """ Return if {"cache_control": ..} in message content block @@ -365,6 +392,21 @@ def is_computer_tool_used( return True return False + def is_pdf_used(self, messages: List[AllMessageValues]) -> bool: + """ + Set to true if media passed into messages. + """ + for message in messages: + if ( + "content" in message + and message["content"] is not None + and isinstance(message["content"], list) + ): + for content in message["content"]: + if "type" in content: + return True + return False + def translate_system_message( self, messages: List[AllMessageValues] ) -> List[AnthropicSystemMessageContent]: diff --git a/litellm/llms/bedrock/image/cost_calculator.py b/litellm/llms/bedrock/image/cost_calculator.py new file mode 100644 index 000000000000..0a20b44cb388 --- /dev/null +++ b/litellm/llms/bedrock/image/cost_calculator.py @@ -0,0 +1,41 @@ +from typing import Optional + +import litellm +from litellm.types.utils import ImageResponse + + +def cost_calculator( + model: str, + image_response: ImageResponse, + size: Optional[str] = None, + optional_params: Optional[dict] = None, +) -> float: + """ + Bedrock image generation cost calculator + + Handles both Stability 1 and Stability 3 models + """ + if litellm.AmazonStability3Config()._is_stability_3_model(model=model): + pass + else: + # Stability 1 models + optional_params = optional_params or {} + + # see model_prices_and_context_window.json for details on how steps is used + # Reference pricing by steps for stability 1: https://aws.amazon.com/bedrock/pricing/ + _steps = optional_params.get("steps", 50) + steps = "max-steps" if _steps > 50 else "50-steps" + + # size is stored in model_prices_and_context_window.json as 1024-x-1024 + # current size has 1024x1024 + size = size or "1024-x-1024" + model = f"{size}/{steps}/{model}" + + _model_info = litellm.get_model_info( + model=model, + custom_llm_provider="bedrock", + ) + + output_cost_per_image: float = _model_info.get("output_cost_per_image") or 0.0 + num_images: int = len(image_response.data) + return output_cost_per_image * num_images diff --git a/litellm/llms/mistral/mistral_chat_transformation.py b/litellm/llms/mistral/mistral_chat_transformation.py index 5d1a54c3a88d..aeb1a90fdb16 100644 --- a/litellm/llms/mistral/mistral_chat_transformation.py +++ b/litellm/llms/mistral/mistral_chat_transformation.py @@ -10,6 +10,7 @@ from typing import List, Literal, Optional, Tuple, Union from litellm.secret_managers.main import get_secret_str +from litellm.types.llms.openai import AllMessageValues class MistralConfig: @@ -148,3 +149,59 @@ def _get_openai_compatible_provider_info( or get_secret_str("MISTRAL_API_KEY") ) return api_base, dynamic_api_key + + @classmethod + def _transform_messages(cls, messages: List[AllMessageValues]): + """ + - handles scenario where content is list and not string + - content list is just text, and no images + - if image passed in, then just return as is (user-intended) + - if `name` is passed, then drop it for mistral API: https://github.com/BerriAI/litellm/issues/6696 + + Motivation: mistral api doesn't support content as a list + """ + new_messages = [] + for m in messages: + special_keys = ["role", "content", "tool_calls", "function_call"] + extra_args = {} + if isinstance(m, dict): + for k, v in m.items(): + if k not in special_keys: + extra_args[k] = v + texts = "" + _content = m.get("content") + if _content is not None and isinstance(_content, list): + for c in _content: + _text: Optional[str] = c.get("text") + if c["type"] == "image_url": + return messages + elif c["type"] == "text" and isinstance(_text, str): + texts += _text + elif _content is not None and isinstance(_content, str): + texts = _content + + new_m = {"role": m["role"], "content": texts, **extra_args} + + if m.get("tool_calls"): + new_m["tool_calls"] = m.get("tool_calls") + + new_m = cls._handle_name_in_message(new_m) + + new_messages.append(new_m) + return new_messages + + @classmethod + def _handle_name_in_message(cls, message: dict) -> dict: + """ + Mistral API only supports `name` in tool messages + + If role == tool, then we keep `name` + Otherwise, we drop `name` + """ + if message.get("name") is not None: + if message["role"] == "tool": + message["name"] = message.get("name") + else: + message.pop("name", None) + + return message diff --git a/litellm/llms/prompt_templates/factory.py b/litellm/llms/prompt_templates/factory.py index aee304760a68..29028e0530ce 100644 --- a/litellm/llms/prompt_templates/factory.py +++ b/litellm/llms/prompt_templates/factory.py @@ -259,43 +259,6 @@ def mistral_instruct_pt(messages): return prompt -def mistral_api_pt(messages): - """ - - handles scenario where content is list and not string - - content list is just text, and no images - - if image passed in, then just return as is (user-intended) - - Motivation: mistral api doesn't support content as a list - """ - new_messages = [] - for m in messages: - special_keys = ["role", "content", "tool_calls", "function_call"] - extra_args = {} - if isinstance(m, dict): - for k, v in m.items(): - if k not in special_keys: - extra_args[k] = v - texts = "" - if m.get("content", None) is not None and isinstance(m["content"], list): - for c in m["content"]: - if c["type"] == "image_url": - return messages - elif c["type"] == "text" and isinstance(c["text"], str): - texts += c["text"] - elif m.get("content", None) is not None and isinstance(m["content"], str): - texts = m["content"] - - new_m = {"role": m["role"], "content": texts, **extra_args} - - if new_m["role"] == "tool" and m.get("name"): - new_m["name"] = m["name"] - if m.get("tool_calls"): - new_m["tool_calls"] = m["tool_calls"] - - new_messages.append(new_m) - return new_messages - - # Falcon prompt template - from https://github.com/lm-sys/FastChat/blob/main/fastchat/conversation.py#L110 def falcon_instruct_pt(messages): prompt = "" @@ -1330,7 +1293,10 @@ def convert_to_anthropic_tool_invoke( def add_cache_control_to_content( anthropic_content_element: Union[ - dict, AnthropicMessagesImageParam, AnthropicMessagesTextParam + dict, + AnthropicMessagesImageParam, + AnthropicMessagesTextParam, + AnthropicMessagesDocumentParam, ], orignal_content_element: Union[dict, AllMessageValues], ): @@ -1343,6 +1309,32 @@ def add_cache_control_to_content( return anthropic_content_element +def _anthropic_content_element_factory( + image_chunk: GenericImageParsingChunk, +) -> Union[AnthropicMessagesImageParam, AnthropicMessagesDocumentParam]: + if image_chunk["media_type"] == "application/pdf": + _anthropic_content_element: Union[ + AnthropicMessagesDocumentParam, AnthropicMessagesImageParam + ] = AnthropicMessagesDocumentParam( + type="document", + source=AnthropicContentParamSource( + type="base64", + media_type=image_chunk["media_type"], + data=image_chunk["data"], + ), + ) + else: + _anthropic_content_element = AnthropicMessagesImageParam( + type="image", + source=AnthropicContentParamSource( + type="base64", + media_type=image_chunk["media_type"], + data=image_chunk["data"], + ), + ) + return _anthropic_content_element + + def anthropic_messages_pt( # noqa: PLR0915 messages: List[AllMessageValues], model: str, @@ -1400,15 +1392,9 @@ def anthropic_messages_pt( # noqa: PLR0915 openai_image_url=m["image_url"]["url"] ) - _anthropic_content_element = AnthropicMessagesImageParam( - type="image", - source=AnthropicImageParamSource( - type="base64", - media_type=image_chunk["media_type"], - data=image_chunk["data"], - ), + _anthropic_content_element = ( + _anthropic_content_element_factory(image_chunk) ) - _content_element = add_cache_control_to_content( anthropic_content_element=_anthropic_content_element, orignal_content_element=dict(m), @@ -2830,7 +2816,7 @@ def prompt_factory( else: return gemini_text_image_pt(messages=messages) elif custom_llm_provider == "mistral": - return mistral_api_pt(messages=messages) + return litellm.MistralConfig._transform_messages(messages=messages) elif custom_llm_provider == "bedrock": if "amazon.titan-text" in model: return amazon_titan_pt(messages=messages) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index e8aeac2cb116..fb8fb105c705 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -26,16 +26,17 @@ "supports_prompt_caching": true }, "gpt-4o": { - "max_tokens": 4096, + "max_tokens": 16384, "max_input_tokens": 128000, - "max_output_tokens": 4096, - "input_cost_per_token": 0.000005, - "output_cost_per_token": 0.000015, + "max_output_tokens": 16384, + "input_cost_per_token": 0.0000025, + "output_cost_per_token": 0.000010, "cache_read_input_token_cost": 0.00000125, "litellm_provider": "openai", "mode": "chat", "supports_function_calling": true, "supports_parallel_function_calling": true, + "supports_response_schema": true, "supports_vision": true, "supports_prompt_caching": true }, @@ -1898,7 +1899,8 @@ "supports_function_calling": true, "tool_use_system_prompt_tokens": 264, "supports_assistant_prefill": true, - "supports_prompt_caching": true + "supports_prompt_caching": true, + "supports_pdf_input": true }, "claude-3-opus-20240229": { "max_tokens": 4096, diff --git a/litellm/proxy/_new_secret_config.yaml b/litellm/proxy/_new_secret_config.yaml index c44a46a6720c..911f15b86373 100644 --- a/litellm/proxy/_new_secret_config.yaml +++ b/litellm/proxy/_new_secret_config.yaml @@ -1,5 +1,5 @@ model_list: - - model_name: claude-3-5-sonnet-20240620 + - model_name: "*" litellm_params: model: claude-3-5-sonnet-20240620 api_key: os.environ/ANTHROPIC_API_KEY @@ -48,6 +48,13 @@ model_list: aws_access_key_id: os.environ/BEDROCK_AWS_ACCESS_KEY_ID aws_secret_access_key: os.environ/BEDROCK_AWS_SECRET_ACCESS_KEY aws_region_name: os.environ/AWS_REGION_NAME + + - model_name: "bedrock/*" + litellm_params: + model: bedrock/* + aws_access_key_id: os.environ/BEDROCK_AWS_ACCESS_KEY_ID + aws_secret_access_key: os.environ/BEDROCK_AWS_SECRET_ACCESS_KEY + aws_region_name: os.environ/AWS_REGION_NAME - model_name: gpt-4 litellm_params: @@ -61,8 +68,13 @@ model_list: litellm_settings: fallbacks: [{ "claude-3-5-sonnet-20240620": ["claude-3-5-sonnet-aihubmix"] }] - callbacks: ["otel", "prometheus"] + # callbacks: ["otel", "prometheus"] default_redis_batch_cache_expiry: 10 + # default_team_settings: + # - team_id: "dbe2f686-a686-4896-864a-4c3924458709" + # success_callback: ["langfuse"] + # langfuse_public_key: os.environ/LANGFUSE_PUB_KEY_1 # Project 1 + # langfuse_secret: os.environ/LANGFUSE_PRIVATE_KEY_1 # Project 1 # litellm_settings: # cache: True diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index fd9ef8556c46..2d869af85120 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -1894,6 +1894,7 @@ class ProxyErrorTypes(str, enum.Enum): auth_error = "auth_error" internal_server_error = "internal_server_error" bad_request_error = "bad_request_error" + not_found_error = "not_found_error" class SSOUserDefinedValues(TypedDict): diff --git a/litellm/proxy/auth/route_checks.py b/litellm/proxy/auth/route_checks.py index a237b0bdd60e..1b593162c8de 100644 --- a/litellm/proxy/auth/route_checks.py +++ b/litellm/proxy/auth/route_checks.py @@ -44,14 +44,8 @@ def non_proxy_admin_allowed_routes_check( route in LiteLLMRoutes.info_routes.value ): # check if user allowed to call an info route if route == "/key/info": - # check if user can access this route - query_params = request.query_params - key = query_params.get("key") - if key is not None and hash_token(token=key) != api_key: - raise HTTPException( - status_code=status.HTTP_403_FORBIDDEN, - detail="user not allowed to access this key's info", - ) + # handled by function itself + pass elif route == "/user/info": # check if user can access this route query_params = request.query_params diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index ff1acc3c92ee..6032a72af166 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -1236,7 +1236,6 @@ def _return_user_api_key_auth_obj( start_time: datetime, user_role: Optional[LitellmUserRoles] = None, ) -> UserAPIKeyAuth: - traceback.print_stack() end_time = datetime.now() user_api_key_service_logger_obj.service_success_hook( service=ServiceTypes.AUTH, diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 2c240a17f8f3..01baa5a43903 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -32,7 +32,7 @@ ) from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.management_helpers.utils import management_endpoint_wrapper -from litellm.proxy.utils import _duration_in_seconds +from litellm.proxy.utils import _duration_in_seconds, _hash_token_if_needed from litellm.secret_managers.main import get_secret router = APIRouter() @@ -303,21 +303,17 @@ async def generate_key_fn( # noqa: PLR0915 ) -async def prepare_key_update_data( +def prepare_key_update_data( data: Union[UpdateKeyRequest, RegenerateKeyRequest], existing_key_row ): - data_json: dict = data.dict(exclude_unset=True) + data_json: dict = data.model_dump(exclude_unset=True) data_json.pop("key", None) _metadata_fields = ["model_rpm_limit", "model_tpm_limit", "guardrails"] non_default_values = {} for k, v in data_json.items(): if k in _metadata_fields: continue - if v is not None: - if not isinstance(v, bool) and v in ([], {}, 0): - pass - else: - non_default_values[k] = v + non_default_values[k] = v if "duration" in non_default_values: duration = non_default_values.pop("duration") @@ -379,7 +375,7 @@ async def update_key_fn( ) try: - data_json: dict = data.json() + data_json: dict = data.model_dump(exclude_unset=True) key = data_json.pop("key") # get the row from db if prisma_client is None: @@ -395,7 +391,7 @@ async def update_key_fn( detail={"error": f"Team not found, passed team_id={data.team_id}"}, ) - non_default_values = await prepare_key_update_data( + non_default_values = prepare_key_update_data( data=data, existing_key_row=existing_key_row ) @@ -734,13 +730,37 @@ async def info_key_fn( raise Exception( "Database not connected. Connect a database to your proxy - https://docs.litellm.ai/docs/simple_proxy#managing-auth---virtual-keys" ) - if key is None: - key = user_api_key_dict.api_key - key_info = await prisma_client.get_data(token=key) + + # default to using Auth token if no key is passed in + key = key or user_api_key_dict.api_key + hashed_key: Optional[str] = key + if key is not None: + hashed_key = _hash_token_if_needed(token=key) + key_info = await prisma_client.db.litellm_verificationtoken.find_unique( + where={"token": hashed_key}, # type: ignore + include={"litellm_budget_table": True}, + ) if key_info is None: + raise ProxyException( + message="Key not found in database", + type=ProxyErrorTypes.not_found_error, + param="key", + code=status.HTTP_404_NOT_FOUND, + ) + + if ( + _can_user_query_key_info( + user_api_key_dict=user_api_key_dict, + key=key, + key_info=key_info, + ) + is not True + ): raise HTTPException( - status_code=status.HTTP_404_NOT_FOUND, - detail={"message": "No keys found"}, + status_code=status.HTTP_403_FORBIDDEN, + detail="You are not allowed to access this key's info. Your role={}".format( + user_api_key_dict.user_role + ), ) ## REMOVE HASHED TOKEN INFO BEFORE RETURNING ## try: @@ -1120,7 +1140,7 @@ async def regenerate_key_fn( non_default_values = {} if data is not None: # Update with any provided parameters from GenerateKeyRequest - non_default_values = await prepare_key_update_data( + non_default_values = prepare_key_update_data( data=data, existing_key_row=_key_in_db ) @@ -1540,6 +1560,27 @@ async def key_health( ) +def _can_user_query_key_info( + user_api_key_dict: UserAPIKeyAuth, + key: Optional[str], + key_info: LiteLLM_VerificationToken, +) -> bool: + """ + Helper to check if the user has access to the key's info + """ + if ( + user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value + or user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY.value + ): + return True + elif user_api_key_dict.api_key == key: + return True + # user can query their own key info + elif key_info.user_id == user_api_key_dict.user_id: + return True + return False + + async def test_key_logging( user_api_key_dict: UserAPIKeyAuth, request: Request, diff --git a/litellm/proxy/proxy_config.yaml b/litellm/proxy/proxy_config.yaml index b4a18baa4a27..29d14c910c52 100644 --- a/litellm/proxy/proxy_config.yaml +++ b/litellm/proxy/proxy_config.yaml @@ -6,5 +6,7 @@ model_list: api_base: https://exampleopenaiendpoint-production.up.railway.app/ + litellm_settings: callbacks: ["gcs_bucket"] + diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 12e80876c368..c9c6af77f1cf 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -1308,7 +1308,7 @@ async def _update_team_cache(): await _update_team_cache() asyncio.create_task( - user_api_key_cache.async_batch_set_cache( + user_api_key_cache.async_set_cache_pipeline( cache_list=values_to_update_in_cache, ttl=60, litellm_parent_otel_span=parent_otel_span, @@ -2978,7 +2978,7 @@ async def initialize_scheduled_background_jobs( if ( proxy_logging_obj is not None - and proxy_logging_obj.slack_alerting_instance is not None + and proxy_logging_obj.slack_alerting_instance.alerting is not None and prisma_client is not None ): print("Alerting: Initializing Weekly/Monthly Spend Reports") # noqa diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 9d33244a0074..c143d30e4801 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -175,7 +175,7 @@ async def async_batch_set_cache( local_only: bool = False, **kwargs, ) -> None: - return await self.dual_cache.async_batch_set_cache( + return await self.dual_cache.async_set_cache_pipeline( cache_list=cache_list, local_only=local_only, litellm_parent_otel_span=litellm_parent_otel_span, @@ -1424,9 +1424,7 @@ async def get_data( # noqa: PLR0915 # check if plain text or hash if token is not None: if isinstance(token, str): - hashed_token = token - if token.startswith("sk-"): - hashed_token = self.hash_token(token=token) + hashed_token = _hash_token_if_needed(token=token) verbose_proxy_logger.debug( f"PrismaClient: find_unique for token: {hashed_token}" ) @@ -1493,8 +1491,7 @@ async def get_data( # noqa: PLR0915 if token is not None: where_filter["token"] = {} if isinstance(token, str): - if token.startswith("sk-"): - token = self.hash_token(token=token) + token = _hash_token_if_needed(token=token) where_filter["token"]["in"] = [token] elif isinstance(token, list): hashed_tokens = [] @@ -1630,9 +1627,7 @@ async def get_data( # noqa: PLR0915 # check if plain text or hash if token is not None: if isinstance(token, str): - hashed_token = token - if token.startswith("sk-"): - hashed_token = self.hash_token(token=token) + hashed_token = _hash_token_if_needed(token=token) verbose_proxy_logger.debug( f"PrismaClient: find_unique for token: {hashed_token}" ) @@ -1912,8 +1907,7 @@ async def update_data( # noqa: PLR0915 if token is not None: print_verbose(f"token: {token}") # check if plain text or hash - if token.startswith("sk-"): - token = self.hash_token(token=token) + token = _hash_token_if_needed(token=token) db_data["token"] = token response = await self.db.litellm_verificationtoken.update( where={"token": token}, # type: ignore @@ -2424,6 +2418,18 @@ def hash_token(token: str): return hashed_token +def _hash_token_if_needed(token: str) -> str: + """ + Hash the token if it's a string and starts with "sk-" + + Else return the token as is + """ + if token.startswith("sk-"): + return hash_token(token=token) + else: + return token + + def _extract_from_regex(duration: str) -> Tuple[int, str]: match = re.match(r"(\d+)(mo|[smhd]?)", duration) diff --git a/litellm/router.py b/litellm/router.py index 0bdd1d1e0b46..4735d422b95c 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -339,11 +339,7 @@ def __init__( # noqa: PLR0915 cache_config: Dict[str, Any] = {} self.client_ttl = client_ttl - if redis_url is not None or ( - redis_host is not None - and redis_port is not None - and redis_password is not None - ): + if redis_url is not None or (redis_host is not None and redis_port is not None): cache_type = "redis" if redis_url is not None: diff --git a/litellm/types/integrations/langsmith.py b/litellm/types/integrations/langsmith.py new file mode 100644 index 000000000000..48c8e2e0a29b --- /dev/null +++ b/litellm/types/integrations/langsmith.py @@ -0,0 +1,61 @@ +from dataclasses import dataclass +from datetime import datetime +from typing import Any, Dict, List, NamedTuple, Optional, TypedDict + +from pydantic import BaseModel + + +class LangsmithInputs(BaseModel): + model: Optional[str] = None + messages: Optional[List[Any]] = None + stream: Optional[bool] = None + call_type: Optional[str] = None + litellm_call_id: Optional[str] = None + completion_start_time: Optional[datetime] = None + temperature: Optional[float] = None + max_tokens: Optional[int] = None + custom_llm_provider: Optional[str] = None + input: Optional[List[Any]] = None + log_event_type: Optional[str] = None + original_response: Optional[Any] = None + response_cost: Optional[float] = None + + # LiteLLM Virtual Key specific fields + user_api_key: Optional[str] = None + user_api_key_user_id: Optional[str] = None + user_api_key_team_alias: Optional[str] = None + + +class LangsmithCredentialsObject(TypedDict): + LANGSMITH_API_KEY: str + LANGSMITH_PROJECT: str + LANGSMITH_BASE_URL: str + + +class LangsmithQueueObject(TypedDict): + """ + Langsmith Queue Object - this is what gets stored in the internal system queue before flushing to Langsmith + + We need to store: + - data[Dict] - data that should get logged on langsmith + - credentials[LangsmithCredentialsObject] - credentials to use for logging to langsmith + """ + + data: Dict + credentials: LangsmithCredentialsObject + + +class CredentialsKey(NamedTuple): + """Immutable key for grouping credentials""" + + api_key: str + project: str + base_url: str + + +@dataclass +class BatchGroup: + """Groups credentials with their associated queue objects""" + + credentials: LangsmithCredentialsObject + queue_objects: List[LangsmithQueueObject] diff --git a/litellm/types/llms/anthropic.py b/litellm/types/llms/anthropic.py index bb65a372d23b..55e37ad97109 100644 --- a/litellm/types/llms/anthropic.py +++ b/litellm/types/llms/anthropic.py @@ -12,10 +12,16 @@ class AnthropicMessagesToolChoice(TypedDict, total=False): disable_parallel_tool_use: bool # default is false +class AnthropicInputSchema(TypedDict, total=False): + type: Optional[str] + properties: Optional[dict] + additionalProperties: Optional[bool] + + class AnthropicMessagesTool(TypedDict, total=False): name: Required[str] description: str - input_schema: Required[dict] + input_schema: Optional[AnthropicInputSchema] type: Literal["custom"] cache_control: Optional[Union[dict, ChatCompletionCachedContent]] @@ -74,7 +80,7 @@ class AnthopicMessagesAssistantMessageParam(TypedDict, total=False): """ -class AnthropicImageParamSource(TypedDict): +class AnthropicContentParamSource(TypedDict): type: Literal["base64"] media_type: str data: str @@ -82,7 +88,13 @@ class AnthropicImageParamSource(TypedDict): class AnthropicMessagesImageParam(TypedDict, total=False): type: Required[Literal["image"]] - source: Required[AnthropicImageParamSource] + source: Required[AnthropicContentParamSource] + cache_control: Optional[Union[dict, ChatCompletionCachedContent]] + + +class AnthropicMessagesDocumentParam(TypedDict, total=False): + type: Required[Literal["document"]] + source: Required[AnthropicContentParamSource] cache_control: Optional[Union[dict, ChatCompletionCachedContent]] @@ -108,6 +120,7 @@ class AnthropicMessagesToolResultParam(TypedDict, total=False): AnthropicMessagesTextParam, AnthropicMessagesImageParam, AnthropicMessagesToolResultParam, + AnthropicMessagesDocumentParam, ] diff --git a/litellm/types/utils.py b/litellm/types/utils.py index c0a9764e85b7..e3df357bea26 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -1322,11 +1322,6 @@ def json(self, **kwargs): # type: ignore class GenericImageParsingChunk(TypedDict): - # { - # "type": "base64", - # "media_type": f"image/{image_format}", - # "data": base64_data, - # } type: str media_type: str data: str @@ -1600,3 +1595,8 @@ class StandardCallbackDynamicParams(TypedDict, total=False): # GCS dynamic params gcs_bucket_name: Optional[str] gcs_path_service_account: Optional[str] + + # Langsmith dynamic params + langsmith_api_key: Optional[str] + langsmith_project: Optional[str] + langsmith_base_url: Optional[str] diff --git a/litellm/utils.py b/litellm/utils.py index d07d86f7dbbd..802bcfc04d2e 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -796,7 +796,7 @@ def wrapper(*args, **kwargs): # noqa: PLR0915 and kwargs.get("_arealtime", False) is not True ): # allow users to control returning cached responses from the completion function # checking cache - print_verbose("INSIDE CHECKING CACHE") + verbose_logger.debug("INSIDE CHECKING SYNC CACHE") caching_handler_response: CachingHandlerResponse = ( _llm_caching_handler._sync_get_cache( model=model or "", @@ -808,6 +808,7 @@ def wrapper(*args, **kwargs): # noqa: PLR0915 args=args, ) ) + if caching_handler_response.cached_result is not None: return caching_handler_response.cached_result @@ -1835,6 +1836,13 @@ def supports_audio_input(model: str, custom_llm_provider: Optional[str] = None) ) +def supports_pdf_input(model: str, custom_llm_provider: Optional[str] = None) -> bool: + """Check if a given model supports pdf input in a chat completion call""" + return _supports_factory( + model=model, custom_llm_provider=custom_llm_provider, key="supports_pdf_input" + ) + + def supports_audio_output( model: str, custom_llm_provider: Optional[str] = None ) -> bool: @@ -4629,6 +4637,7 @@ def _get_max_position_embeddings(model_name): "output_cost_per_character_above_128k_tokens", None ), output_cost_per_second=_model_info.get("output_cost_per_second", None), + output_cost_per_image=_model_info.get("output_cost_per_image", None), output_vector_size=_model_info.get("output_vector_size", None), litellm_provider=_model_info.get( "litellm_provider", custom_llm_provider @@ -5420,2121 +5429,6 @@ def register_prompt_template( return litellm.custom_prompt_dict -####### DEPRECATED ################ - - -def get_all_keys(llm_provider=None): - try: - global last_fetched_at_keys - # if user is using hosted product -> instantiate their env with their hosted api keys - refresh every 5 minutes - print_verbose(f"Reaches get all keys, llm_provider: {llm_provider}") - user_email = ( - os.getenv("LITELLM_EMAIL") - or litellm.email - or litellm.token - or os.getenv("LITELLM_TOKEN") - ) - if user_email: - time_delta = 0 - if last_fetched_at_keys is not None: - current_time = time.time() - time_delta = current_time - last_fetched_at_keys - if ( - time_delta > 300 or last_fetched_at_keys is None or llm_provider - ): # if the llm provider is passed in , assume this happening due to an AuthError for that provider - # make the api call - last_fetched_at = time.time() - print_verbose(f"last_fetched_at: {last_fetched_at}") - response = requests.post( - url="http://api.litellm.ai/get_all_keys", - headers={"content-type": "application/json"}, - data=json.dumps({"user_email": user_email}), - ) - print_verbose(f"get model key response: {response.text}") - data = response.json() - # update model list - for key, value in data[ - "model_keys" - ].items(): # follows the LITELLM API KEY format - _API_KEY - e.g. HUGGINGFACE_API_KEY - os.environ[key] = value - # set model alias map - for model_alias, value in data["model_alias_map"].items(): - litellm.model_alias_map[model_alias] = value - return "it worked!" - return None - return None - except Exception: - print_verbose( - f"[Non-Blocking Error] get_all_keys error - {traceback.format_exc()}" - ) - pass - - -def get_model_list(): - global last_fetched_at, print_verbose - try: - # if user is using hosted product -> get their updated model list - user_email = ( - os.getenv("LITELLM_EMAIL") - or litellm.email - or litellm.token - or os.getenv("LITELLM_TOKEN") - ) - if user_email: - # make the api call - last_fetched_at = time.time() - print_verbose(f"last_fetched_at: {last_fetched_at}") - response = requests.post( - url="http://api.litellm.ai/get_model_list", - headers={"content-type": "application/json"}, - data=json.dumps({"user_email": user_email}), - ) - print_verbose(f"get_model_list response: {response.text}") - data = response.json() - # update model list - model_list = data["model_list"] - # # check if all model providers are in environment - # model_providers = data["model_providers"] - # missing_llm_provider = None - # for item in model_providers: - # if f"{item.upper()}_API_KEY" not in os.environ: - # missing_llm_provider = item - # break - # # update environment - if required - # threading.Thread(target=get_all_keys, args=(missing_llm_provider)).start() - return model_list - return [] # return empty list by default - except Exception: - print_verbose( - f"[Non-Blocking Error] get_model_list error - {traceback.format_exc()}" - ) - - -######## Streaming Class ############################ -# wraps the completion stream to return the correct format for the model -# replicate/anthropic/cohere - -# class CustomStreamWrapper: -# def __init__( -# self, -# completion_stream, -# model, -# logging_obj: Any, -# custom_llm_provider: Optional[str] = None, -# stream_options=None, -# make_call: Optional[Callable] = None, -# _response_headers: Optional[dict] = None, -# ): -# self.model = model -# self.make_call = make_call -# self.custom_llm_provider = custom_llm_provider -# self.logging_obj: LiteLLMLoggingObject = logging_obj -# self.completion_stream = completion_stream -# self.sent_first_chunk = False -# self.sent_last_chunk = False -# self.system_fingerprint: Optional[str] = None -# self.received_finish_reason: Optional[str] = None -# self.special_tokens = [ -# "<|assistant|>", -# "<|system|>", -# "<|user|>", -# "", -# "", -# "<|im_end|>", -# "<|im_start|>", -# ] -# self.holding_chunk = "" -# self.complete_response = "" -# self.response_uptil_now = "" -# _model_info = ( -# self.logging_obj.model_call_details.get("litellm_params", {}).get( -# "model_info", {} -# ) -# or {} -# ) -# self._hidden_params = { -# "model_id": (_model_info.get("id", None)), -# } # returned as x-litellm-model-id response header in proxy - -# self._hidden_params["additional_headers"] = process_response_headers( -# _response_headers or {} -# ) # GUARANTEE OPENAI HEADERS IN RESPONSE - -# self._response_headers = _response_headers -# self.response_id = None -# self.logging_loop = None -# self.rules = Rules() -# self.stream_options = stream_options or getattr( -# logging_obj, "stream_options", None -# ) -# self.messages = getattr(logging_obj, "messages", None) -# self.sent_stream_usage = False -# self.send_stream_usage = ( -# True if self.check_send_stream_usage(self.stream_options) else False -# ) -# self.tool_call = False -# self.chunks: List = ( -# [] -# ) # keep track of the returned chunks - used for calculating the input/output tokens for stream options -# self.is_function_call = self.check_is_function_call(logging_obj=logging_obj) - -# def __iter__(self): -# return self - -# def __aiter__(self): -# return self - -# def check_send_stream_usage(self, stream_options: Optional[dict]): -# return ( -# stream_options is not None -# and stream_options.get("include_usage", False) is True -# ) - -# def check_is_function_call(self, logging_obj) -> bool: -# if hasattr(logging_obj, "optional_params") and isinstance( -# logging_obj.optional_params, dict -# ): -# if ( -# "litellm_param_is_function_call" in logging_obj.optional_params -# and logging_obj.optional_params["litellm_param_is_function_call"] -# is True -# ): -# return True - -# return False - -# def process_chunk(self, chunk: str): -# """ -# NLP Cloud streaming returns the entire response, for each chunk. Process this, to only return the delta. -# """ -# try: -# chunk = chunk.strip() -# self.complete_response = self.complete_response.strip() - -# if chunk.startswith(self.complete_response): -# # Remove last_sent_chunk only if it appears at the start of the new chunk -# chunk = chunk[len(self.complete_response) :] - -# self.complete_response += chunk -# return chunk -# except Exception as e: -# raise e - -# def safety_checker(self) -> None: -# """ -# Fixes - https://github.com/BerriAI/litellm/issues/5158 - -# if the model enters a loop and starts repeating the same chunk again, break out of loop and raise an internalservererror - allows for retries. - -# Raises - InternalServerError, if LLM enters infinite loop while streaming -# """ -# if len(self.chunks) >= litellm.REPEATED_STREAMING_CHUNK_LIMIT: -# # Get the last n chunks -# last_chunks = self.chunks[-litellm.REPEATED_STREAMING_CHUNK_LIMIT :] - -# # Extract the relevant content from the chunks -# last_contents = [chunk.choices[0].delta.content for chunk in last_chunks] - -# # Check if all extracted contents are identical -# if all(content == last_contents[0] for content in last_contents): -# if ( -# last_contents[0] is not None -# and isinstance(last_contents[0], str) -# and len(last_contents[0]) > 2 -# ): # ignore empty content - https://github.com/BerriAI/litellm/issues/5158#issuecomment-2287156946 -# # All last n chunks are identical -# raise litellm.InternalServerError( -# message="The model is repeating the same chunk = {}.".format( -# last_contents[0] -# ), -# model="", -# llm_provider="", -# ) - -# def check_special_tokens(self, chunk: str, finish_reason: Optional[str]): -# """ -# Output parse / special tokens for sagemaker + hf streaming. -# """ -# hold = False -# if ( -# self.custom_llm_provider != "huggingface" -# and self.custom_llm_provider != "sagemaker" -# ): -# return hold, chunk - -# if finish_reason: -# for token in self.special_tokens: -# if token in chunk: -# chunk = chunk.replace(token, "") -# return hold, chunk - -# if self.sent_first_chunk is True: -# return hold, chunk - -# curr_chunk = self.holding_chunk + chunk -# curr_chunk = curr_chunk.strip() - -# for token in self.special_tokens: -# if len(curr_chunk) < len(token) and curr_chunk in token: -# hold = True -# self.holding_chunk = curr_chunk -# elif len(curr_chunk) >= len(token): -# if token in curr_chunk: -# self.holding_chunk = curr_chunk.replace(token, "") -# hold = True -# else: -# pass - -# if hold is False: # reset -# self.holding_chunk = "" -# return hold, curr_chunk - -# def handle_anthropic_text_chunk(self, chunk): -# """ -# For old anthropic models - claude-1, claude-2. - -# Claude-3 is handled from within Anthropic.py VIA ModelResponseIterator() -# """ -# str_line = chunk -# if isinstance(chunk, bytes): # Handle binary data -# str_line = chunk.decode("utf-8") # Convert bytes to string -# text = "" -# is_finished = False -# finish_reason = None -# if str_line.startswith("data:"): -# data_json = json.loads(str_line[5:]) -# type_chunk = data_json.get("type", None) -# if type_chunk == "completion": -# text = data_json.get("completion") -# finish_reason = data_json.get("stop_reason") -# if finish_reason is not None: -# is_finished = True -# return { -# "text": text, -# "is_finished": is_finished, -# "finish_reason": finish_reason, -# } -# elif "error" in str_line: -# raise ValueError(f"Unable to parse response. Original response: {str_line}") -# else: -# return { -# "text": text, -# "is_finished": is_finished, -# "finish_reason": finish_reason, -# } - -# def handle_vertexai_anthropic_chunk(self, chunk): -# """ -# - MessageStartEvent(message=Message(id='msg_01LeRRgvX4gwkX3ryBVgtuYZ', content=[], model='claude-3-sonnet-20240229', role='assistant', stop_reason=None, stop_sequence=None, type='message', usage=Usage(input_tokens=8, output_tokens=1)), type='message_start'); custom_llm_provider: vertex_ai -# - ContentBlockStartEvent(content_block=ContentBlock(text='', type='text'), index=0, type='content_block_start'); custom_llm_provider: vertex_ai -# - ContentBlockDeltaEvent(delta=TextDelta(text='Hello', type='text_delta'), index=0, type='content_block_delta'); custom_llm_provider: vertex_ai -# """ -# text = "" -# prompt_tokens = None -# completion_tokens = None -# is_finished = False -# finish_reason = None -# type_chunk = getattr(chunk, "type", None) -# if type_chunk == "message_start": -# message = getattr(chunk, "message", None) -# text = "" # lets us return a chunk with usage to user -# _usage = getattr(message, "usage", None) -# if _usage is not None: -# prompt_tokens = getattr(_usage, "input_tokens", None) -# completion_tokens = getattr(_usage, "output_tokens", None) -# elif type_chunk == "content_block_delta": -# """ -# Anthropic content chunk -# chunk = {'type': 'content_block_delta', 'index': 0, 'delta': {'type': 'text_delta', 'text': 'Hello'}} -# """ -# delta = getattr(chunk, "delta", None) -# if delta is not None: -# text = getattr(delta, "text", "") -# else: -# text = "" -# elif type_chunk == "message_delta": -# """ -# Anthropic -# chunk = {'type': 'message_delta', 'delta': {'stop_reason': 'max_tokens', 'stop_sequence': None}, 'usage': {'output_tokens': 10}} -# """ -# # TODO - get usage from this chunk, set in response -# delta = getattr(chunk, "delta", None) -# if delta is not None: -# finish_reason = getattr(delta, "stop_reason", "stop") -# is_finished = True -# _usage = getattr(chunk, "usage", None) -# if _usage is not None: -# prompt_tokens = getattr(_usage, "input_tokens", None) -# completion_tokens = getattr(_usage, "output_tokens", None) - -# return { -# "text": text, -# "is_finished": is_finished, -# "finish_reason": finish_reason, -# "prompt_tokens": prompt_tokens, -# "completion_tokens": completion_tokens, -# } - -# def handle_predibase_chunk(self, chunk): -# try: -# if not isinstance(chunk, str): -# chunk = chunk.decode( -# "utf-8" -# ) # DO NOT REMOVE this: This is required for HF inference API + Streaming -# text = "" -# is_finished = False -# finish_reason = "" -# print_verbose(f"chunk: {chunk}") -# if chunk.startswith("data:"): -# data_json = json.loads(chunk[5:]) -# print_verbose(f"data json: {data_json}") -# if "token" in data_json and "text" in data_json["token"]: -# text = data_json["token"]["text"] -# if data_json.get("details", False) and data_json["details"].get( -# "finish_reason", False -# ): -# is_finished = True -# finish_reason = data_json["details"]["finish_reason"] -# elif data_json.get( -# "generated_text", False -# ): # if full generated text exists, then stream is complete -# text = "" # don't return the final bos token -# is_finished = True -# finish_reason = "stop" -# elif data_json.get("error", False): -# raise Exception(data_json.get("error")) -# return { -# "text": text, -# "is_finished": is_finished, -# "finish_reason": finish_reason, -# } -# elif "error" in chunk: -# raise ValueError(chunk) -# return { -# "text": text, -# "is_finished": is_finished, -# "finish_reason": finish_reason, -# } -# except Exception as e: -# raise e - -# def handle_huggingface_chunk(self, chunk): -# try: -# if not isinstance(chunk, str): -# chunk = chunk.decode( -# "utf-8" -# ) # DO NOT REMOVE this: This is required for HF inference API + Streaming -# text = "" -# is_finished = False -# finish_reason = "" -# print_verbose(f"chunk: {chunk}") -# if chunk.startswith("data:"): -# data_json = json.loads(chunk[5:]) -# print_verbose(f"data json: {data_json}") -# if "token" in data_json and "text" in data_json["token"]: -# text = data_json["token"]["text"] -# if data_json.get("details", False) and data_json["details"].get( -# "finish_reason", False -# ): -# is_finished = True -# finish_reason = data_json["details"]["finish_reason"] -# elif data_json.get( -# "generated_text", False -# ): # if full generated text exists, then stream is complete -# text = "" # don't return the final bos token -# is_finished = True -# finish_reason = "stop" -# elif data_json.get("error", False): -# raise Exception(data_json.get("error")) -# return { -# "text": text, -# "is_finished": is_finished, -# "finish_reason": finish_reason, -# } -# elif "error" in chunk: -# raise ValueError(chunk) -# return { -# "text": text, -# "is_finished": is_finished, -# "finish_reason": finish_reason, -# } -# except Exception as e: -# raise e - -# def handle_ai21_chunk(self, chunk): # fake streaming -# chunk = chunk.decode("utf-8") -# data_json = json.loads(chunk) -# try: -# text = data_json["completions"][0]["data"]["text"] -# is_finished = True -# finish_reason = "stop" -# return { -# "text": text, -# "is_finished": is_finished, -# "finish_reason": finish_reason, -# } -# except Exception: -# raise ValueError(f"Unable to parse response. Original response: {chunk}") - -# def handle_maritalk_chunk(self, chunk): # fake streaming -# chunk = chunk.decode("utf-8") -# data_json = json.loads(chunk) -# try: -# text = data_json["answer"] -# is_finished = True -# finish_reason = "stop" -# return { -# "text": text, -# "is_finished": is_finished, -# "finish_reason": finish_reason, -# } -# except Exception: -# raise ValueError(f"Unable to parse response. Original response: {chunk}") - -# def handle_nlp_cloud_chunk(self, chunk): -# text = "" -# is_finished = False -# finish_reason = "" -# try: -# if "dolphin" in self.model: -# chunk = self.process_chunk(chunk=chunk) -# else: -# data_json = json.loads(chunk) -# chunk = data_json["generated_text"] -# text = chunk -# if "[DONE]" in text: -# text = text.replace("[DONE]", "") -# is_finished = True -# finish_reason = "stop" -# return { -# "text": text, -# "is_finished": is_finished, -# "finish_reason": finish_reason, -# } -# except Exception: -# raise ValueError(f"Unable to parse response. Original response: {chunk}") - -# def handle_aleph_alpha_chunk(self, chunk): -# chunk = chunk.decode("utf-8") -# data_json = json.loads(chunk) -# try: -# text = data_json["completions"][0]["completion"] -# is_finished = True -# finish_reason = "stop" -# return { -# "text": text, -# "is_finished": is_finished, -# "finish_reason": finish_reason, -# } -# except Exception: -# raise ValueError(f"Unable to parse response. Original response: {chunk}") - -# def handle_cohere_chunk(self, chunk): -# chunk = chunk.decode("utf-8") -# data_json = json.loads(chunk) -# try: -# text = "" -# is_finished = False -# finish_reason = "" -# index: Optional[int] = None -# if "index" in data_json: -# index = data_json.get("index") -# if "text" in data_json: -# text = data_json["text"] -# elif "is_finished" in data_json: -# is_finished = data_json["is_finished"] -# finish_reason = data_json["finish_reason"] -# else: -# raise Exception(data_json) -# return { -# "index": index, -# "text": text, -# "is_finished": is_finished, -# "finish_reason": finish_reason, -# } -# except Exception: -# raise ValueError(f"Unable to parse response. Original response: {chunk}") - -# def handle_cohere_chat_chunk(self, chunk): -# chunk = chunk.decode("utf-8") -# data_json = json.loads(chunk) -# print_verbose(f"chunk: {chunk}") -# try: -# text = "" -# is_finished = False -# finish_reason = "" -# if "text" in data_json: -# text = data_json["text"] -# elif "is_finished" in data_json and data_json["is_finished"] is True: -# is_finished = data_json["is_finished"] -# finish_reason = data_json["finish_reason"] -# else: -# return -# return { -# "text": text, -# "is_finished": is_finished, -# "finish_reason": finish_reason, -# } -# except Exception: -# raise ValueError(f"Unable to parse response. Original response: {chunk}") - -# def handle_azure_chunk(self, chunk): -# is_finished = False -# finish_reason = "" -# text = "" -# print_verbose(f"chunk: {chunk}") -# if "data: [DONE]" in chunk: -# text = "" -# is_finished = True -# finish_reason = "stop" -# return { -# "text": text, -# "is_finished": is_finished, -# "finish_reason": finish_reason, -# } -# elif chunk.startswith("data:"): -# data_json = json.loads(chunk[5:]) # chunk.startswith("data:"): -# try: -# if len(data_json["choices"]) > 0: -# delta = data_json["choices"][0]["delta"] -# text = "" if delta is None else delta.get("content", "") -# if data_json["choices"][0].get("finish_reason", None): -# is_finished = True -# finish_reason = data_json["choices"][0]["finish_reason"] -# print_verbose( -# f"text: {text}; is_finished: {is_finished}; finish_reason: {finish_reason}" -# ) -# return { -# "text": text, -# "is_finished": is_finished, -# "finish_reason": finish_reason, -# } -# except Exception: -# raise ValueError( -# f"Unable to parse response. Original response: {chunk}" -# ) -# elif "error" in chunk: -# raise ValueError(f"Unable to parse response. Original response: {chunk}") -# else: -# return { -# "text": text, -# "is_finished": is_finished, -# "finish_reason": finish_reason, -# } - -# def handle_replicate_chunk(self, chunk): -# try: -# text = "" -# is_finished = False -# finish_reason = "" -# if "output" in chunk: -# text = chunk["output"] -# if "status" in chunk: -# if chunk["status"] == "succeeded": -# is_finished = True -# finish_reason = "stop" -# elif chunk.get("error", None): -# raise Exception(chunk["error"]) -# return { -# "text": text, -# "is_finished": is_finished, -# "finish_reason": finish_reason, -# } -# except Exception: -# raise ValueError(f"Unable to parse response. Original response: {chunk}") - -# def handle_openai_chat_completion_chunk(self, chunk): -# try: -# print_verbose(f"\nRaw OpenAI Chunk\n{chunk}\n") -# str_line = chunk -# text = "" -# is_finished = False -# finish_reason = None -# logprobs = None -# usage = None -# if str_line and str_line.choices and len(str_line.choices) > 0: -# if ( -# str_line.choices[0].delta is not None -# and str_line.choices[0].delta.content is not None -# ): -# text = str_line.choices[0].delta.content -# else: # function/tool calling chunk - when content is None. in this case we just return the original chunk from openai -# pass -# if str_line.choices[0].finish_reason: -# is_finished = True -# finish_reason = str_line.choices[0].finish_reason - -# # checking for logprobs -# if ( -# hasattr(str_line.choices[0], "logprobs") -# and str_line.choices[0].logprobs is not None -# ): -# logprobs = str_line.choices[0].logprobs -# else: -# logprobs = None - -# usage = getattr(str_line, "usage", None) - -# return { -# "text": text, -# "is_finished": is_finished, -# "finish_reason": finish_reason, -# "logprobs": logprobs, -# "original_chunk": str_line, -# "usage": usage, -# } -# except Exception as e: -# raise e - -# def handle_azure_text_completion_chunk(self, chunk): -# try: -# print_verbose(f"\nRaw OpenAI Chunk\n{chunk}\n") -# text = "" -# is_finished = False -# finish_reason = None -# choices = getattr(chunk, "choices", []) -# if len(choices) > 0: -# text = choices[0].text -# if choices[0].finish_reason is not None: -# is_finished = True -# finish_reason = choices[0].finish_reason -# return { -# "text": text, -# "is_finished": is_finished, -# "finish_reason": finish_reason, -# } - -# except Exception as e: -# raise e - -# def handle_openai_text_completion_chunk(self, chunk): -# try: -# print_verbose(f"\nRaw OpenAI Chunk\n{chunk}\n") -# text = "" -# is_finished = False -# finish_reason = None -# usage = None -# choices = getattr(chunk, "choices", []) -# if len(choices) > 0: -# text = choices[0].text -# if choices[0].finish_reason is not None: -# is_finished = True -# finish_reason = choices[0].finish_reason -# usage = getattr(chunk, "usage", None) -# return { -# "text": text, -# "is_finished": is_finished, -# "finish_reason": finish_reason, -# "usage": usage, -# } - -# except Exception as e: -# raise e - -# def handle_baseten_chunk(self, chunk): -# try: -# chunk = chunk.decode("utf-8") -# if len(chunk) > 0: -# if chunk.startswith("data:"): -# data_json = json.loads(chunk[5:]) -# if "token" in data_json and "text" in data_json["token"]: -# return data_json["token"]["text"] -# else: -# return "" -# data_json = json.loads(chunk) -# if "model_output" in data_json: -# if ( -# isinstance(data_json["model_output"], dict) -# and "data" in data_json["model_output"] -# and isinstance(data_json["model_output"]["data"], list) -# ): -# return data_json["model_output"]["data"][0] -# elif isinstance(data_json["model_output"], str): -# return data_json["model_output"] -# elif "completion" in data_json and isinstance( -# data_json["completion"], str -# ): -# return data_json["completion"] -# else: -# raise ValueError( -# f"Unable to parse response. Original response: {chunk}" -# ) -# else: -# return "" -# else: -# return "" -# except Exception as e: -# verbose_logger.exception( -# "litellm.CustomStreamWrapper.handle_baseten_chunk(): Exception occured - {}".format( -# str(e) -# ) -# ) -# return "" - -# def handle_cloudlfare_stream(self, chunk): -# try: -# print_verbose(f"\nRaw OpenAI Chunk\n{chunk}\n") -# chunk = chunk.decode("utf-8") -# str_line = chunk -# text = "" -# is_finished = False -# finish_reason = None - -# if "[DONE]" in chunk: -# return {"text": text, "is_finished": True, "finish_reason": "stop"} -# elif str_line.startswith("data:"): -# data_json = json.loads(str_line[5:]) -# print_verbose(f"delta content: {data_json}") -# text = data_json["response"] -# return { -# "text": text, -# "is_finished": is_finished, -# "finish_reason": finish_reason, -# } -# else: -# return { -# "text": text, -# "is_finished": is_finished, -# "finish_reason": finish_reason, -# } - -# except Exception as e: -# raise e - -# def handle_ollama_stream(self, chunk): -# try: -# if isinstance(chunk, dict): -# json_chunk = chunk -# else: -# json_chunk = json.loads(chunk) -# if "error" in json_chunk: -# raise Exception(f"Ollama Error - {json_chunk}") - -# text = "" -# is_finished = False -# finish_reason = None -# if json_chunk["done"] is True: -# text = "" -# is_finished = True -# finish_reason = "stop" -# return { -# "text": text, -# "is_finished": is_finished, -# "finish_reason": finish_reason, -# } -# elif json_chunk["response"]: -# print_verbose(f"delta content: {json_chunk}") -# text = json_chunk["response"] -# return { -# "text": text, -# "is_finished": is_finished, -# "finish_reason": finish_reason, -# } -# else: -# raise Exception(f"Ollama Error - {json_chunk}") -# except Exception as e: -# raise e - -# def handle_ollama_chat_stream(self, chunk): -# # for ollama_chat/ provider -# try: -# if isinstance(chunk, dict): -# json_chunk = chunk -# else: -# json_chunk = json.loads(chunk) -# if "error" in json_chunk: -# raise Exception(f"Ollama Error - {json_chunk}") - -# text = "" -# is_finished = False -# finish_reason = None -# if json_chunk["done"] is True: -# text = "" -# is_finished = True -# finish_reason = "stop" -# return { -# "text": text, -# "is_finished": is_finished, -# "finish_reason": finish_reason, -# } -# elif "message" in json_chunk: -# print_verbose(f"delta content: {json_chunk}") -# text = json_chunk["message"]["content"] -# return { -# "text": text, -# "is_finished": is_finished, -# "finish_reason": finish_reason, -# } -# else: -# raise Exception(f"Ollama Error - {json_chunk}") -# except Exception as e: -# raise e - -# def handle_watsonx_stream(self, chunk): -# try: -# if isinstance(chunk, dict): -# parsed_response = chunk -# elif isinstance(chunk, (str, bytes)): -# if isinstance(chunk, bytes): -# chunk = chunk.decode("utf-8") -# if "generated_text" in chunk: -# response = chunk.replace("data: ", "").strip() -# parsed_response = json.loads(response) -# else: -# return { -# "text": "", -# "is_finished": False, -# "prompt_tokens": 0, -# "completion_tokens": 0, -# } -# else: -# print_verbose(f"chunk: {chunk} (Type: {type(chunk)})") -# raise ValueError( -# f"Unable to parse response. Original response: {chunk}" -# ) -# results = parsed_response.get("results", []) -# if len(results) > 0: -# text = results[0].get("generated_text", "") -# finish_reason = results[0].get("stop_reason") -# is_finished = finish_reason != "not_finished" -# return { -# "text": text, -# "is_finished": is_finished, -# "finish_reason": finish_reason, -# "prompt_tokens": results[0].get("input_token_count", 0), -# "completion_tokens": results[0].get("generated_token_count", 0), -# } -# return {"text": "", "is_finished": False} -# except Exception as e: -# raise e - -# def handle_triton_stream(self, chunk): -# try: -# if isinstance(chunk, dict): -# parsed_response = chunk -# elif isinstance(chunk, (str, bytes)): -# if isinstance(chunk, bytes): -# chunk = chunk.decode("utf-8") -# if "text_output" in chunk: -# response = chunk.replace("data: ", "").strip() -# parsed_response = json.loads(response) -# else: -# return { -# "text": "", -# "is_finished": False, -# "prompt_tokens": 0, -# "completion_tokens": 0, -# } -# else: -# print_verbose(f"chunk: {chunk} (Type: {type(chunk)})") -# raise ValueError( -# f"Unable to parse response. Original response: {chunk}" -# ) -# text = parsed_response.get("text_output", "") -# finish_reason = parsed_response.get("stop_reason") -# is_finished = parsed_response.get("is_finished", False) -# return { -# "text": text, -# "is_finished": is_finished, -# "finish_reason": finish_reason, -# "prompt_tokens": parsed_response.get("input_token_count", 0), -# "completion_tokens": parsed_response.get("generated_token_count", 0), -# } -# return {"text": "", "is_finished": False} -# except Exception as e: -# raise e - -# def handle_clarifai_completion_chunk(self, chunk): -# try: -# if isinstance(chunk, dict): -# parsed_response = chunk -# elif isinstance(chunk, (str, bytes)): -# if isinstance(chunk, bytes): -# parsed_response = chunk.decode("utf-8") -# else: -# parsed_response = chunk -# else: -# raise ValueError("Unable to parse streaming chunk") -# if isinstance(parsed_response, dict): -# data_json = parsed_response -# else: -# data_json = json.loads(parsed_response) -# text = ( -# data_json.get("outputs", "")[0] -# .get("data", "") -# .get("text", "") -# .get("raw", "") -# ) -# len( -# encoding.encode( -# data_json.get("outputs", "")[0] -# .get("input", "") -# .get("data", "") -# .get("text", "") -# .get("raw", "") -# ) -# ) -# len(encoding.encode(text)) -# return { -# "text": text, -# "is_finished": True, -# } -# except Exception as e: -# verbose_logger.exception( -# "litellm.CustomStreamWrapper.handle_clarifai_chunk(): Exception occured - {}".format( -# str(e) -# ) -# ) -# return "" - -# def model_response_creator( -# self, chunk: Optional[dict] = None, hidden_params: Optional[dict] = None -# ): -# _model = self.model -# _received_llm_provider = self.custom_llm_provider -# _logging_obj_llm_provider = self.logging_obj.model_call_details.get("custom_llm_provider", None) # type: ignore -# if ( -# _received_llm_provider == "openai" -# and _received_llm_provider != _logging_obj_llm_provider -# ): -# _model = "{}/{}".format(_logging_obj_llm_provider, _model) -# if chunk is None: -# chunk = {} -# else: -# # pop model keyword -# chunk.pop("model", None) - -# model_response = ModelResponse( -# stream=True, model=_model, stream_options=self.stream_options, **chunk -# ) -# if self.response_id is not None: -# model_response.id = self.response_id -# else: -# self.response_id = model_response.id # type: ignore -# if self.system_fingerprint is not None: -# model_response.system_fingerprint = self.system_fingerprint -# if hidden_params is not None: -# model_response._hidden_params = hidden_params -# model_response._hidden_params["custom_llm_provider"] = _logging_obj_llm_provider -# model_response._hidden_params["created_at"] = time.time() -# model_response._hidden_params = { -# **model_response._hidden_params, -# **self._hidden_params, -# } - -# if ( -# len(model_response.choices) > 0 -# and getattr(model_response.choices[0], "delta") is not None -# ): -# # do nothing, if object instantiated -# pass -# else: -# model_response.choices = [StreamingChoices(finish_reason=None)] -# return model_response - -# def is_delta_empty(self, delta: Delta) -> bool: -# is_empty = True -# if delta.content is not None: -# is_empty = False -# elif delta.tool_calls is not None: -# is_empty = False -# elif delta.function_call is not None: -# is_empty = False -# return is_empty - -# def return_processed_chunk_logic( # noqa -# self, -# completion_obj: dict, -# model_response: ModelResponseStream, -# response_obj: dict, -# ): - -# print_verbose( -# f"completion_obj: {completion_obj}, model_response.choices[0]: {model_response.choices[0]}, response_obj: {response_obj}" -# ) -# if ( -# "content" in completion_obj -# and ( -# isinstance(completion_obj["content"], str) -# and len(completion_obj["content"]) > 0 -# ) -# or ( -# "tool_calls" in completion_obj -# and completion_obj["tool_calls"] is not None -# and len(completion_obj["tool_calls"]) > 0 -# ) -# or ( -# "function_call" in completion_obj -# and completion_obj["function_call"] is not None -# ) -# ): # cannot set content of an OpenAI Object to be an empty string -# self.safety_checker() -# hold, model_response_str = self.check_special_tokens( -# chunk=completion_obj["content"], -# finish_reason=model_response.choices[0].finish_reason, -# ) # filter out bos/eos tokens from openai-compatible hf endpoints -# print_verbose(f"hold - {hold}, model_response_str - {model_response_str}") -# if hold is False: -# ## check if openai/azure chunk -# original_chunk = response_obj.get("original_chunk", None) -# if original_chunk: -# model_response.id = original_chunk.id -# self.response_id = original_chunk.id -# if len(original_chunk.choices) > 0: -# choices = [] -# for choice in original_chunk.choices: -# try: -# if isinstance(choice, BaseModel): -# choice_json = choice.model_dump() -# choice_json.pop( -# "finish_reason", None -# ) # for mistral etc. which return a value in their last chunk (not-openai compatible). -# print_verbose(f"choice_json: {choice_json}") -# choices.append(StreamingChoices(**choice_json)) -# except Exception: -# choices.append(StreamingChoices()) -# print_verbose(f"choices in streaming: {choices}") -# setattr(model_response, "choices", choices) -# else: -# return -# model_response.system_fingerprint = ( -# original_chunk.system_fingerprint -# ) -# setattr( -# model_response, -# "citations", -# getattr(original_chunk, "citations", None), -# ) -# print_verbose(f"self.sent_first_chunk: {self.sent_first_chunk}") -# if self.sent_first_chunk is False: -# model_response.choices[0].delta["role"] = "assistant" -# self.sent_first_chunk = True -# elif self.sent_first_chunk is True and hasattr( -# model_response.choices[0].delta, "role" -# ): -# _initial_delta = model_response.choices[0].delta.model_dump() -# _initial_delta.pop("role", None) -# model_response.choices[0].delta = Delta(**_initial_delta) -# print_verbose( -# f"model_response.choices[0].delta: {model_response.choices[0].delta}" -# ) -# else: -# ## else -# completion_obj["content"] = model_response_str -# if self.sent_first_chunk is False: -# completion_obj["role"] = "assistant" -# self.sent_first_chunk = True - -# model_response.choices[0].delta = Delta(**completion_obj) -# _index: Optional[int] = completion_obj.get("index") -# if _index is not None: -# model_response.choices[0].index = _index -# print_verbose(f"returning model_response: {model_response}") -# return model_response -# else: -# return -# elif self.received_finish_reason is not None: -# if self.sent_last_chunk is True: -# # Bedrock returns the guardrail trace in the last chunk - we want to return this here -# if self.custom_llm_provider == "bedrock" and "trace" in model_response: -# return model_response - -# # Default - return StopIteration -# raise StopIteration -# # flush any remaining holding chunk -# if len(self.holding_chunk) > 0: -# if model_response.choices[0].delta.content is None: -# model_response.choices[0].delta.content = self.holding_chunk -# else: -# model_response.choices[0].delta.content = ( -# self.holding_chunk + model_response.choices[0].delta.content -# ) -# self.holding_chunk = "" -# # if delta is None -# _is_delta_empty = self.is_delta_empty(delta=model_response.choices[0].delta) - -# if _is_delta_empty: -# # get any function call arguments -# model_response.choices[0].finish_reason = map_finish_reason( -# finish_reason=self.received_finish_reason -# ) # ensure consistent output to openai - -# self.sent_last_chunk = True - -# return model_response -# elif ( -# model_response.choices[0].delta.tool_calls is not None -# or model_response.choices[0].delta.function_call is not None -# ): -# if self.sent_first_chunk is False: -# model_response.choices[0].delta["role"] = "assistant" -# self.sent_first_chunk = True -# return model_response -# elif ( -# len(model_response.choices) > 0 -# and hasattr(model_response.choices[0].delta, "audio") -# and model_response.choices[0].delta.audio is not None -# ): -# return model_response -# else: -# if hasattr(model_response, "usage"): -# self.chunks.append(model_response) -# return - -# def chunk_creator(self, chunk): # type: ignore # noqa: PLR0915 -# model_response = self.model_response_creator() -# response_obj: dict = {} -# try: -# # return this for all models -# completion_obj = {"content": ""} -# from litellm.litellm_core_utils.streaming_utils import ( -# generic_chunk_has_all_required_fields, -# ) -# from litellm.types.utils import GenericStreamingChunk as GChunk - -# if ( -# isinstance(chunk, dict) -# and generic_chunk_has_all_required_fields( -# chunk=chunk -# ) # check if chunk is a generic streaming chunk -# ) or ( -# self.custom_llm_provider -# and ( -# self.custom_llm_provider == "anthropic" -# or self.custom_llm_provider in litellm._custom_providers -# ) -# ): - -# if self.received_finish_reason is not None: -# if "provider_specific_fields" not in chunk: -# raise StopIteration -# anthropic_response_obj: GChunk = chunk -# completion_obj["content"] = anthropic_response_obj["text"] -# if anthropic_response_obj["is_finished"]: -# self.received_finish_reason = anthropic_response_obj[ -# "finish_reason" -# ] - -# if anthropic_response_obj["usage"] is not None: -# model_response.usage = litellm.Usage( -# **anthropic_response_obj["usage"] -# ) - -# if ( -# "tool_use" in anthropic_response_obj -# and anthropic_response_obj["tool_use"] is not None -# ): -# completion_obj["tool_calls"] = [anthropic_response_obj["tool_use"]] - -# if ( -# "provider_specific_fields" in anthropic_response_obj -# and anthropic_response_obj["provider_specific_fields"] is not None -# ): -# for key, value in anthropic_response_obj[ -# "provider_specific_fields" -# ].items(): -# setattr(model_response, key, value) - -# response_obj = anthropic_response_obj -# elif ( -# self.custom_llm_provider -# and self.custom_llm_provider == "anthropic_text" -# ): -# response_obj = self.handle_anthropic_text_chunk(chunk) -# completion_obj["content"] = response_obj["text"] -# if response_obj["is_finished"]: -# self.received_finish_reason = response_obj["finish_reason"] -# elif self.custom_llm_provider and self.custom_llm_provider == "clarifai": -# response_obj = self.handle_clarifai_completion_chunk(chunk) -# completion_obj["content"] = response_obj["text"] -# if response_obj["is_finished"]: -# self.received_finish_reason = response_obj["finish_reason"] -# elif self.model == "replicate" or self.custom_llm_provider == "replicate": -# response_obj = self.handle_replicate_chunk(chunk) -# completion_obj["content"] = response_obj["text"] -# if response_obj["is_finished"]: -# self.received_finish_reason = response_obj["finish_reason"] -# elif self.custom_llm_provider and self.custom_llm_provider == "huggingface": -# response_obj = self.handle_huggingface_chunk(chunk) -# completion_obj["content"] = response_obj["text"] -# if response_obj["is_finished"]: -# self.received_finish_reason = response_obj["finish_reason"] -# elif self.custom_llm_provider and self.custom_llm_provider == "predibase": -# response_obj = self.handle_predibase_chunk(chunk) -# completion_obj["content"] = response_obj["text"] -# if response_obj["is_finished"]: -# self.received_finish_reason = response_obj["finish_reason"] -# elif ( -# self.custom_llm_provider and self.custom_llm_provider == "baseten" -# ): # baseten doesn't provide streaming -# completion_obj["content"] = self.handle_baseten_chunk(chunk) -# elif ( -# self.custom_llm_provider and self.custom_llm_provider == "ai21" -# ): # ai21 doesn't provide streaming -# response_obj = self.handle_ai21_chunk(chunk) -# completion_obj["content"] = response_obj["text"] -# if response_obj["is_finished"]: -# self.received_finish_reason = response_obj["finish_reason"] -# elif self.custom_llm_provider and self.custom_llm_provider == "maritalk": -# response_obj = self.handle_maritalk_chunk(chunk) -# completion_obj["content"] = response_obj["text"] -# if response_obj["is_finished"]: -# self.received_finish_reason = response_obj["finish_reason"] -# elif self.custom_llm_provider and self.custom_llm_provider == "vllm": -# completion_obj["content"] = chunk[0].outputs[0].text -# elif ( -# self.custom_llm_provider and self.custom_llm_provider == "aleph_alpha" -# ): # aleph alpha doesn't provide streaming -# response_obj = self.handle_aleph_alpha_chunk(chunk) -# completion_obj["content"] = response_obj["text"] -# if response_obj["is_finished"]: -# self.received_finish_reason = response_obj["finish_reason"] -# elif self.custom_llm_provider == "nlp_cloud": -# try: -# response_obj = self.handle_nlp_cloud_chunk(chunk) -# completion_obj["content"] = response_obj["text"] -# if response_obj["is_finished"]: -# self.received_finish_reason = response_obj["finish_reason"] -# except Exception as e: -# if self.received_finish_reason: -# raise e -# else: -# if self.sent_first_chunk is False: -# raise Exception("An unknown error occurred with the stream") -# self.received_finish_reason = "stop" -# elif self.custom_llm_provider == "vertex_ai": -# import proto # type: ignore - -# if self.model.startswith("claude-3"): -# response_obj = self.handle_vertexai_anthropic_chunk(chunk=chunk) -# if response_obj is None: -# return -# completion_obj["content"] = response_obj["text"] -# setattr(model_response, "usage", Usage()) -# if response_obj.get("prompt_tokens", None) is not None: -# model_response.usage.prompt_tokens = response_obj[ -# "prompt_tokens" -# ] -# if response_obj.get("completion_tokens", None) is not None: -# model_response.usage.completion_tokens = response_obj[ -# "completion_tokens" -# ] -# if hasattr(model_response.usage, "prompt_tokens"): -# model_response.usage.total_tokens = ( -# getattr(model_response.usage, "total_tokens", 0) -# + model_response.usage.prompt_tokens -# ) -# if hasattr(model_response.usage, "completion_tokens"): -# model_response.usage.total_tokens = ( -# getattr(model_response.usage, "total_tokens", 0) -# + model_response.usage.completion_tokens -# ) - -# if response_obj["is_finished"]: -# self.received_finish_reason = response_obj["finish_reason"] -# elif hasattr(chunk, "candidates") is True: -# try: -# try: -# completion_obj["content"] = chunk.text -# except Exception as e: -# if "Part has no text." in str(e): -# ## check for function calling -# function_call = ( -# chunk.candidates[0].content.parts[0].function_call -# ) - -# args_dict = {} - -# # Check if it's a RepeatedComposite instance -# for key, val in function_call.args.items(): -# if isinstance( -# val, -# proto.marshal.collections.repeated.RepeatedComposite, -# ): -# # If so, convert to list -# args_dict[key] = [v for v in val] -# else: -# args_dict[key] = val - -# try: -# args_str = json.dumps(args_dict) -# except Exception as e: -# raise e -# _delta_obj = litellm.utils.Delta( -# content=None, -# tool_calls=[ -# { -# "id": f"call_{str(uuid.uuid4())}", -# "function": { -# "arguments": args_str, -# "name": function_call.name, -# }, -# "type": "function", -# } -# ], -# ) -# _streaming_response = StreamingChoices(delta=_delta_obj) -# _model_response = ModelResponse(stream=True) -# _model_response.choices = [_streaming_response] -# response_obj = {"original_chunk": _model_response} -# else: -# raise e -# if ( -# hasattr(chunk.candidates[0], "finish_reason") -# and chunk.candidates[0].finish_reason.name -# != "FINISH_REASON_UNSPECIFIED" -# ): # every non-final chunk in vertex ai has this -# self.received_finish_reason = chunk.candidates[ -# 0 -# ].finish_reason.name -# except Exception: -# if chunk.candidates[0].finish_reason.name == "SAFETY": -# raise Exception( -# f"The response was blocked by VertexAI. {str(chunk)}" -# ) -# else: -# completion_obj["content"] = str(chunk) -# elif self.custom_llm_provider == "cohere": -# response_obj = self.handle_cohere_chunk(chunk) -# completion_obj["content"] = response_obj["text"] -# if response_obj["is_finished"]: -# self.received_finish_reason = response_obj["finish_reason"] -# elif self.custom_llm_provider == "cohere_chat": -# response_obj = self.handle_cohere_chat_chunk(chunk) -# if response_obj is None: -# return -# completion_obj["content"] = response_obj["text"] -# if response_obj["is_finished"]: -# self.received_finish_reason = response_obj["finish_reason"] - -# elif self.custom_llm_provider == "petals": -# if len(self.completion_stream) == 0: -# if self.received_finish_reason is not None: -# raise StopIteration -# else: -# self.received_finish_reason = "stop" -# chunk_size = 30 -# new_chunk = self.completion_stream[:chunk_size] -# completion_obj["content"] = new_chunk -# self.completion_stream = self.completion_stream[chunk_size:] -# elif self.custom_llm_provider == "palm": -# # fake streaming -# response_obj = {} -# if len(self.completion_stream) == 0: -# if self.received_finish_reason is not None: -# raise StopIteration -# else: -# self.received_finish_reason = "stop" -# chunk_size = 30 -# new_chunk = self.completion_stream[:chunk_size] -# completion_obj["content"] = new_chunk -# self.completion_stream = self.completion_stream[chunk_size:] -# elif self.custom_llm_provider == "ollama": -# response_obj = self.handle_ollama_stream(chunk) -# completion_obj["content"] = response_obj["text"] -# print_verbose(f"completion obj content: {completion_obj['content']}") -# if response_obj["is_finished"]: -# self.received_finish_reason = response_obj["finish_reason"] -# elif self.custom_llm_provider == "ollama_chat": -# response_obj = self.handle_ollama_chat_stream(chunk) -# completion_obj["content"] = response_obj["text"] -# print_verbose(f"completion obj content: {completion_obj['content']}") -# if response_obj["is_finished"]: -# self.received_finish_reason = response_obj["finish_reason"] -# elif self.custom_llm_provider == "cloudflare": -# response_obj = self.handle_cloudlfare_stream(chunk) -# completion_obj["content"] = response_obj["text"] -# print_verbose(f"completion obj content: {completion_obj['content']}") -# if response_obj["is_finished"]: -# self.received_finish_reason = response_obj["finish_reason"] -# elif self.custom_llm_provider == "watsonx": -# response_obj = self.handle_watsonx_stream(chunk) -# completion_obj["content"] = response_obj["text"] -# if response_obj["is_finished"]: -# self.received_finish_reason = response_obj["finish_reason"] -# elif self.custom_llm_provider == "triton": -# response_obj = self.handle_triton_stream(chunk) -# completion_obj["content"] = response_obj["text"] -# print_verbose(f"completion obj content: {completion_obj['content']}") -# if response_obj["is_finished"]: -# self.received_finish_reason = response_obj["finish_reason"] -# elif self.custom_llm_provider == "text-completion-openai": -# response_obj = self.handle_openai_text_completion_chunk(chunk) -# completion_obj["content"] = response_obj["text"] -# print_verbose(f"completion obj content: {completion_obj['content']}") -# if response_obj["is_finished"]: -# self.received_finish_reason = response_obj["finish_reason"] -# if response_obj["usage"] is not None: -# model_response.usage = litellm.Usage( -# prompt_tokens=response_obj["usage"].prompt_tokens, -# completion_tokens=response_obj["usage"].completion_tokens, -# total_tokens=response_obj["usage"].total_tokens, -# ) -# elif self.custom_llm_provider == "text-completion-codestral": -# response_obj = litellm.MistralTextCompletionConfig()._chunk_parser( -# chunk -# ) -# completion_obj["content"] = response_obj["text"] -# print_verbose(f"completion obj content: {completion_obj['content']}") -# if response_obj["is_finished"]: -# self.received_finish_reason = response_obj["finish_reason"] -# if "usage" in response_obj is not None: -# model_response.usage = litellm.Usage( -# prompt_tokens=response_obj["usage"].prompt_tokens, -# completion_tokens=response_obj["usage"].completion_tokens, -# total_tokens=response_obj["usage"].total_tokens, -# ) -# elif self.custom_llm_provider == "azure_text": -# response_obj = self.handle_azure_text_completion_chunk(chunk) -# completion_obj["content"] = response_obj["text"] -# print_verbose(f"completion obj content: {completion_obj['content']}") -# if response_obj["is_finished"]: -# self.received_finish_reason = response_obj["finish_reason"] -# elif self.custom_llm_provider == "cached_response": -# response_obj = { -# "text": chunk.choices[0].delta.content, -# "is_finished": True, -# "finish_reason": chunk.choices[0].finish_reason, -# "original_chunk": chunk, -# "tool_calls": ( -# chunk.choices[0].delta.tool_calls -# if hasattr(chunk.choices[0].delta, "tool_calls") -# else None -# ), -# } - -# completion_obj["content"] = response_obj["text"] -# if response_obj["tool_calls"] is not None: -# completion_obj["tool_calls"] = response_obj["tool_calls"] -# print_verbose(f"completion obj content: {completion_obj['content']}") -# if hasattr(chunk, "id"): -# model_response.id = chunk.id -# self.response_id = chunk.id -# if hasattr(chunk, "system_fingerprint"): -# self.system_fingerprint = chunk.system_fingerprint -# if response_obj["is_finished"]: -# self.received_finish_reason = response_obj["finish_reason"] -# else: # openai / azure chat model -# if self.custom_llm_provider == "azure": -# if hasattr(chunk, "model"): -# # for azure, we need to pass the model from the orignal chunk -# self.model = chunk.model -# response_obj = self.handle_openai_chat_completion_chunk(chunk) -# if response_obj is None: -# return -# completion_obj["content"] = response_obj["text"] -# print_verbose(f"completion obj content: {completion_obj['content']}") -# if response_obj["is_finished"]: -# if response_obj["finish_reason"] == "error": -# raise Exception( -# "{} raised a streaming error - finish_reason: error, no content string given. Received Chunk={}".format( -# self.custom_llm_provider, response_obj -# ) -# ) -# self.received_finish_reason = response_obj["finish_reason"] -# if response_obj.get("original_chunk", None) is not None: -# if hasattr(response_obj["original_chunk"], "id"): -# model_response.id = response_obj["original_chunk"].id -# self.response_id = model_response.id -# if hasattr(response_obj["original_chunk"], "system_fingerprint"): -# model_response.system_fingerprint = response_obj[ -# "original_chunk" -# ].system_fingerprint -# self.system_fingerprint = response_obj[ -# "original_chunk" -# ].system_fingerprint -# if response_obj["logprobs"] is not None: -# model_response.choices[0].logprobs = response_obj["logprobs"] - -# if response_obj["usage"] is not None: -# if isinstance(response_obj["usage"], dict): -# model_response.usage = litellm.Usage( -# prompt_tokens=response_obj["usage"].get( -# "prompt_tokens", None -# ) -# or None, -# completion_tokens=response_obj["usage"].get( -# "completion_tokens", None -# ) -# or None, -# total_tokens=response_obj["usage"].get("total_tokens", None) -# or None, -# ) -# elif isinstance(response_obj["usage"], BaseModel): -# model_response.usage = litellm.Usage( -# **response_obj["usage"].model_dump() -# ) - -# model_response.model = self.model -# print_verbose( -# f"model_response finish reason 3: {self.received_finish_reason}; response_obj={response_obj}" -# ) -# ## FUNCTION CALL PARSING -# if ( -# response_obj is not None -# and response_obj.get("original_chunk", None) is not None -# ): # function / tool calling branch - only set for openai/azure compatible endpoints -# # enter this branch when no content has been passed in response -# original_chunk = response_obj.get("original_chunk", None) -# model_response.id = original_chunk.id -# self.response_id = original_chunk.id -# if original_chunk.choices and len(original_chunk.choices) > 0: -# delta = original_chunk.choices[0].delta -# if delta is not None and ( -# delta.function_call is not None or delta.tool_calls is not None -# ): -# try: -# model_response.system_fingerprint = ( -# original_chunk.system_fingerprint -# ) -# ## AZURE - check if arguments is not None -# if ( -# original_chunk.choices[0].delta.function_call -# is not None -# ): -# if ( -# getattr( -# original_chunk.choices[0].delta.function_call, -# "arguments", -# ) -# is None -# ): -# original_chunk.choices[ -# 0 -# ].delta.function_call.arguments = "" -# elif original_chunk.choices[0].delta.tool_calls is not None: -# if isinstance( -# original_chunk.choices[0].delta.tool_calls, list -# ): -# for t in original_chunk.choices[0].delta.tool_calls: -# if hasattr(t, "functions") and hasattr( -# t.functions, "arguments" -# ): -# if ( -# getattr( -# t.function, -# "arguments", -# ) -# is None -# ): -# t.function.arguments = "" -# _json_delta = delta.model_dump() -# print_verbose(f"_json_delta: {_json_delta}") -# if "role" not in _json_delta or _json_delta["role"] is None: -# _json_delta["role"] = ( -# "assistant" # mistral's api returns role as None -# ) -# if "tool_calls" in _json_delta and isinstance( -# _json_delta["tool_calls"], list -# ): -# for tool in _json_delta["tool_calls"]: -# if ( -# isinstance(tool, dict) -# and "function" in tool -# and isinstance(tool["function"], dict) -# and ("type" not in tool or tool["type"] is None) -# ): -# # if function returned but type set to None - mistral's api returns type: None -# tool["type"] = "function" -# model_response.choices[0].delta = Delta(**_json_delta) -# except Exception as e: -# verbose_logger.exception( -# "litellm.CustomStreamWrapper.chunk_creator(): Exception occured - {}".format( -# str(e) -# ) -# ) -# model_response.choices[0].delta = Delta() -# elif ( -# delta is not None and getattr(delta, "audio", None) is not None -# ): -# model_response.choices[0].delta.audio = delta.audio -# else: -# try: -# delta = ( -# dict() -# if original_chunk.choices[0].delta is None -# else dict(original_chunk.choices[0].delta) -# ) -# print_verbose(f"original delta: {delta}") -# model_response.choices[0].delta = Delta(**delta) -# print_verbose( -# f"new delta: {model_response.choices[0].delta}" -# ) -# except Exception: -# model_response.choices[0].delta = Delta() -# else: -# if ( -# self.stream_options is not None -# and self.stream_options["include_usage"] is True -# ): -# return model_response -# return -# print_verbose( -# f"model_response.choices[0].delta: {model_response.choices[0].delta}; completion_obj: {completion_obj}" -# ) -# print_verbose(f"self.sent_first_chunk: {self.sent_first_chunk}") - -# ## CHECK FOR TOOL USE -# if "tool_calls" in completion_obj and len(completion_obj["tool_calls"]) > 0: -# if self.is_function_call is True: # user passed in 'functions' param -# completion_obj["function_call"] = completion_obj["tool_calls"][0][ -# "function" -# ] -# completion_obj["tool_calls"] = None - -# self.tool_call = True - -# ## RETURN ARG -# return self.return_processed_chunk_logic( -# completion_obj=completion_obj, -# model_response=model_response, # type: ignore -# response_obj=response_obj, -# ) - -# except StopIteration: -# raise StopIteration -# except Exception as e: -# traceback.format_exc() -# e.message = str(e) -# raise exception_type( -# model=self.model, -# custom_llm_provider=self.custom_llm_provider, -# original_exception=e, -# ) - -# def set_logging_event_loop(self, loop): -# """ -# import litellm, asyncio - -# loop = asyncio.get_event_loop() # 👈 gets the current event loop - -# response = litellm.completion(.., stream=True) - -# response.set_logging_event_loop(loop=loop) # 👈 enables async_success callbacks for sync logging - -# for chunk in response: -# ... -# """ -# self.logging_loop = loop - -# def run_success_logging_and_cache_storage(self, processed_chunk, cache_hit: bool): -# """ -# Runs success logging in a thread and adds the response to the cache -# """ -# if litellm.disable_streaming_logging is True: -# """ -# [NOT RECOMMENDED] -# Set this via `litellm.disable_streaming_logging = True`. - -# Disables streaming logging. -# """ -# return -# ## ASYNC LOGGING -# # Create an event loop for the new thread -# if self.logging_loop is not None: -# future = asyncio.run_coroutine_threadsafe( -# self.logging_obj.async_success_handler( -# processed_chunk, None, None, cache_hit -# ), -# loop=self.logging_loop, -# ) -# future.result() -# else: -# asyncio.run( -# self.logging_obj.async_success_handler( -# processed_chunk, None, None, cache_hit -# ) -# ) -# ## SYNC LOGGING -# self.logging_obj.success_handler(processed_chunk, None, None, cache_hit) - -# ## Sync store in cache -# if self.logging_obj._llm_caching_handler is not None: -# self.logging_obj._llm_caching_handler._sync_add_streaming_response_to_cache( -# processed_chunk -# ) - -# def finish_reason_handler(self): -# model_response = self.model_response_creator() -# complete_streaming_response = litellm.stream_chunk_builder( -# chunks=self.chunks -# ) -# _finish_reason = complete_streaming_response.choices[0].finish_reason - -# print(f"_finish_reason: {_finish_reason}") -# if _finish_reason is not None: -# model_response.choices[0].finish_reason = _finish_reason -# else: -# model_response.choices[0].finish_reason = "stop" - -# ## if tool use -# if ( -# model_response.choices[0].finish_reason == "stop" and self.tool_call -# ): # don't overwrite for other - potential error finish reasons -# model_response.choices[0].finish_reason = "tool_calls" -# return model_response - -# def __next__(self): # noqa: PLR0915 -# cache_hit = False -# if ( -# self.custom_llm_provider is not None -# and self.custom_llm_provider == "cached_response" -# ): -# cache_hit = True -# try: -# if self.completion_stream is None: -# self.fetch_sync_stream() -# while True: -# if ( -# isinstance(self.completion_stream, str) -# or isinstance(self.completion_stream, bytes) -# or isinstance(self.completion_stream, ModelResponse) -# ): -# chunk = self.completion_stream -# else: -# chunk = next(self.completion_stream) -# if chunk is not None and chunk != b"": -# print_verbose( -# f"PROCESSED CHUNK PRE CHUNK CREATOR: {chunk}; custom_llm_provider: {self.custom_llm_provider}" -# ) -# response: Optional[ModelResponse] = self.chunk_creator(chunk=chunk) -# print_verbose(f"PROCESSED CHUNK POST CHUNK CREATOR: {response}") - -# if response is None: -# continue -# ## LOGGING -# threading.Thread( -# target=self.run_success_logging_and_cache_storage, -# args=(response, cache_hit), -# ).start() # log response -# choice = response.choices[0] -# if isinstance(choice, StreamingChoices): -# self.response_uptil_now += choice.delta.get("content", "") or "" -# else: -# self.response_uptil_now += "" -# self.rules.post_call_rules( -# input=self.response_uptil_now, model=self.model -# ) -# # HANDLE STREAM OPTIONS -# self.chunks.append(response) -# if hasattr( -# response, "usage" -# ): # remove usage from chunk, only send on final chunk -# # Convert the object to a dictionary -# obj_dict = response.dict() - -# # Remove an attribute (e.g., 'attr2') -# if "usage" in obj_dict: -# del obj_dict["usage"] - -# # Create a new object without the removed attribute -# response = self.model_response_creator( -# chunk=obj_dict, hidden_params=response._hidden_params -# ) -# # add usage as hidden param -# if self.sent_last_chunk is True and self.stream_options is None: -# usage = calculate_total_usage(chunks=self.chunks) -# response._hidden_params["usage"] = usage -# # RETURN RESULT -# return response - -# except StopIteration: -# if self.sent_last_chunk is True: -# complete_streaming_response = litellm.stream_chunk_builder( -# chunks=self.chunks, messages=self.messages -# ) -# response = self.model_response_creator() -# if complete_streaming_response is not None: -# setattr( -# response, -# "usage", -# getattr(complete_streaming_response, "usage"), -# ) - -# ## LOGGING -# threading.Thread( -# target=self.logging_obj.success_handler, -# args=(response, None, None, cache_hit), -# ).start() # log response - -# if self.sent_stream_usage is False and self.send_stream_usage is True: -# self.sent_stream_usage = True -# return response -# raise # Re-raise StopIteration -# else: -# self.sent_last_chunk = True -# processed_chunk = self.finish_reason_handler() -# if self.stream_options is None: # add usage as hidden param -# usage = calculate_total_usage(chunks=self.chunks) -# processed_chunk._hidden_params["usage"] = usage -# ## LOGGING -# threading.Thread( -# target=self.run_success_logging_and_cache_storage, -# args=(processed_chunk, cache_hit), -# ).start() # log response -# return processed_chunk -# except Exception as e: -# traceback_exception = traceback.format_exc() -# # LOG FAILURE - handle streaming failure logging in the _next_ object, remove `handle_failure` once it's deprecated -# threading.Thread( -# target=self.logging_obj.failure_handler, args=(e, traceback_exception) -# ).start() -# if isinstance(e, OpenAIError): -# raise e -# else: -# raise exception_type( -# model=self.model, -# original_exception=e, -# custom_llm_provider=self.custom_llm_provider, -# ) - -# def fetch_sync_stream(self): -# if self.completion_stream is None and self.make_call is not None: -# # Call make_call to get the completion stream -# self.completion_stream = self.make_call(client=litellm.module_level_client) -# self._stream_iter = self.completion_stream.__iter__() - -# return self.completion_stream - -# async def fetch_stream(self): -# if self.completion_stream is None and self.make_call is not None: -# # Call make_call to get the completion stream -# self.completion_stream = await self.make_call( -# client=litellm.module_level_aclient -# ) -# self._stream_iter = self.completion_stream.__aiter__() - -# return self.completion_stream - -# async def __anext__(self): # noqa: PLR0915 -# cache_hit = False -# if ( -# self.custom_llm_provider is not None -# and self.custom_llm_provider == "cached_response" -# ): -# cache_hit = True -# try: -# if self.completion_stream is None: -# await self.fetch_stream() - -# if ( -# self.custom_llm_provider == "openai" -# or self.custom_llm_provider == "azure" -# or self.custom_llm_provider == "custom_openai" -# or self.custom_llm_provider == "text-completion-openai" -# or self.custom_llm_provider == "text-completion-codestral" -# or self.custom_llm_provider == "azure_text" -# or self.custom_llm_provider == "anthropic" -# or self.custom_llm_provider == "anthropic_text" -# or self.custom_llm_provider == "huggingface" -# or self.custom_llm_provider == "ollama" -# or self.custom_llm_provider == "ollama_chat" -# or self.custom_llm_provider == "vertex_ai" -# or self.custom_llm_provider == "vertex_ai_beta" -# or self.custom_llm_provider == "sagemaker" -# or self.custom_llm_provider == "sagemaker_chat" -# or self.custom_llm_provider == "gemini" -# or self.custom_llm_provider == "replicate" -# or self.custom_llm_provider == "cached_response" -# or self.custom_llm_provider == "predibase" -# or self.custom_llm_provider == "databricks" -# or self.custom_llm_provider == "bedrock" -# or self.custom_llm_provider == "triton" -# or self.custom_llm_provider == "watsonx" -# or self.custom_llm_provider in litellm.openai_compatible_endpoints -# or self.custom_llm_provider in litellm._custom_providers -# ): -# async for chunk in self.completion_stream: -# if chunk == "None" or chunk is None: -# raise Exception -# elif ( -# self.custom_llm_provider == "gemini" -# and hasattr(chunk, "parts") -# and len(chunk.parts) == 0 -# ): -# continue -# # chunk_creator() does logging/stream chunk building. We need to let it know its being called in_async_func, so we don't double add chunks. -# # __anext__ also calls async_success_handler, which does logging -# print_verbose(f"PROCESSED ASYNC CHUNK PRE CHUNK CREATOR: {chunk}") - -# processed_chunk: Optional[ModelResponse] = self.chunk_creator( -# chunk=chunk -# ) -# print_verbose( -# f"PROCESSED ASYNC CHUNK POST CHUNK CREATOR: {processed_chunk}" -# ) -# if processed_chunk is None: -# continue -# ## LOGGING -# ## LOGGING -# executor.submit( -# self.logging_obj.success_handler, -# result=processed_chunk, -# start_time=None, -# end_time=None, -# cache_hit=cache_hit, -# ) - -# asyncio.create_task( -# self.logging_obj.async_success_handler( -# processed_chunk, cache_hit=cache_hit -# ) -# ) - -# if self.logging_obj._llm_caching_handler is not None: -# asyncio.create_task( -# self.logging_obj._llm_caching_handler._add_streaming_response_to_cache( -# processed_chunk=processed_chunk, -# ) -# ) - -# choice = processed_chunk.choices[0] -# if isinstance(choice, StreamingChoices): -# self.response_uptil_now += choice.delta.get("content", "") or "" -# else: -# self.response_uptil_now += "" -# self.rules.post_call_rules( -# input=self.response_uptil_now, model=self.model -# ) -# self.chunks.append(processed_chunk) -# if hasattr( -# processed_chunk, "usage" -# ): # remove usage from chunk, only send on final chunk -# # Convert the object to a dictionary -# obj_dict = processed_chunk.dict() - -# # Remove an attribute (e.g., 'attr2') -# if "usage" in obj_dict: -# del obj_dict["usage"] - -# # Create a new object without the removed attribute -# processed_chunk = self.model_response_creator(chunk=obj_dict) -# print_verbose(f"final returned processed chunk: {processed_chunk}") -# return processed_chunk -# raise StopAsyncIteration -# else: # temporary patch for non-aiohttp async calls -# # example - boto3 bedrock llms -# while True: -# if isinstance(self.completion_stream, str) or isinstance( -# self.completion_stream, bytes -# ): -# chunk = self.completion_stream -# else: -# chunk = next(self.completion_stream) -# if chunk is not None and chunk != b"": -# print_verbose(f"PROCESSED CHUNK PRE CHUNK CREATOR: {chunk}") -# processed_chunk: Optional[ModelResponse] = self.chunk_creator( -# chunk=chunk -# ) -# print_verbose( -# f"PROCESSED CHUNK POST CHUNK CREATOR: {processed_chunk}" -# ) -# if processed_chunk is None: -# continue -# ## LOGGING -# threading.Thread( -# target=self.logging_obj.success_handler, -# args=(processed_chunk, None, None, cache_hit), -# ).start() # log processed_chunk -# asyncio.create_task( -# self.logging_obj.async_success_handler( -# processed_chunk, cache_hit=cache_hit -# ) -# ) - -# choice = processed_chunk.choices[0] -# if isinstance(choice, StreamingChoices): -# self.response_uptil_now += ( -# choice.delta.get("content", "") or "" -# ) -# else: -# self.response_uptil_now += "" -# self.rules.post_call_rules( -# input=self.response_uptil_now, model=self.model -# ) -# # RETURN RESULT -# self.chunks.append(processed_chunk) -# return processed_chunk -# except (StopAsyncIteration, StopIteration): -# if self.sent_last_chunk is True: -# # log the final chunk with accurate streaming values -# complete_streaming_response = litellm.stream_chunk_builder( -# chunks=self.chunks, messages=self.messages -# ) -# response = self.model_response_creator() -# if complete_streaming_response is not None: -# setattr( -# response, -# "usage", -# getattr(complete_streaming_response, "usage"), -# ) -# ## LOGGING -# threading.Thread( -# target=self.logging_obj.success_handler, -# args=(response, None, None, cache_hit), -# ).start() # log response -# asyncio.create_task( -# self.logging_obj.async_success_handler( -# response, cache_hit=cache_hit -# ) -# ) -# if self.sent_stream_usage is False and self.send_stream_usage is True: -# self.sent_stream_usage = True -# return response -# raise StopAsyncIteration # Re-raise StopIteration -# else: -# self.sent_last_chunk = True -# processed_chunk = self.finish_reason_handler() -# ## LOGGING -# threading.Thread( -# target=self.logging_obj.success_handler, -# args=(processed_chunk, None, None, cache_hit), -# ).start() # log response -# asyncio.create_task( -# self.logging_obj.async_success_handler( -# processed_chunk, cache_hit=cache_hit -# ) -# ) -# return processed_chunk -# except httpx.TimeoutException as e: # if httpx read timeout error occues -# traceback_exception = traceback.format_exc() -# ## ADD DEBUG INFORMATION - E.G. LITELLM REQUEST TIMEOUT -# traceback_exception += "\nLiteLLM Default Request Timeout - {}".format( -# litellm.request_timeout -# ) -# if self.logging_obj is not None: -# ## LOGGING -# threading.Thread( -# target=self.logging_obj.failure_handler, -# args=(e, traceback_exception), -# ).start() # log response -# # Handle any exceptions that might occur during streaming -# asyncio.create_task( -# self.logging_obj.async_failure_handler(e, traceback_exception) -# ) -# raise e -# except Exception as e: -# traceback_exception = traceback.format_exc() -# if self.logging_obj is not None: -# ## LOGGING -# threading.Thread( -# target=self.logging_obj.failure_handler, -# args=(e, traceback_exception), -# ).start() # log response -# # Handle any exceptions that might occur during streaming -# asyncio.create_task( -# self.logging_obj.async_failure_handler(e, traceback_exception) # type: ignore -# ) -# ## Map to OpenAI Exception -# raise exception_type( -# model=self.model, -# custom_llm_provider=self.custom_llm_provider, -# original_exception=e, -# completion_kwargs={}, -# extra_kwargs={}, -# ) - - class TextCompletionStreamWrapper: def __init__( self, @@ -7977,7 +5871,6 @@ def get_valid_models() -> List[str]: if expected_provider_key in environ_keys: # key is set valid_providers.append(provider) - for provider in valid_providers: if provider == "azure": valid_models.append("Azure-LLM") @@ -8253,10 +6146,13 @@ def validate_chat_completion_user_messages(messages: List[AllMessageValues]): if isinstance(item, dict): if item.get("type") not in ValidUserMessageContentTypes: raise Exception("invalid content type") - except Exception: - raise Exception( - f"Invalid user message={m} at index {idx}. Please ensure all user messages are valid OpenAI chat completion messages." - ) + except Exception as e: + if "invalid content type" in str(e): + raise Exception( + f"Invalid user message={m} at index {idx}. Please ensure all user messages are valid OpenAI chat completion messages." + ) + else: + raise e return messages diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index e8aeac2cb116..fb8fb105c705 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -26,16 +26,17 @@ "supports_prompt_caching": true }, "gpt-4o": { - "max_tokens": 4096, + "max_tokens": 16384, "max_input_tokens": 128000, - "max_output_tokens": 4096, - "input_cost_per_token": 0.000005, - "output_cost_per_token": 0.000015, + "max_output_tokens": 16384, + "input_cost_per_token": 0.0000025, + "output_cost_per_token": 0.000010, "cache_read_input_token_cost": 0.00000125, "litellm_provider": "openai", "mode": "chat", "supports_function_calling": true, "supports_parallel_function_calling": true, + "supports_response_schema": true, "supports_vision": true, "supports_prompt_caching": true }, @@ -1898,7 +1899,8 @@ "supports_function_calling": true, "tool_use_system_prompt_tokens": 264, "supports_assistant_prefill": true, - "supports_prompt_caching": true + "supports_prompt_caching": true, + "supports_pdf_input": true }, "claude-3-opus-20240229": { "max_tokens": 4096, diff --git a/pyproject.toml b/pyproject.toml index 099f33bd8bfd..17d37c0ce14f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "litellm" -version = "1.52.4" +version = "1.52.6" description = "Library to easily interface with LLM API providers" authors = ["BerriAI"] license = "MIT" @@ -91,7 +91,7 @@ requires = ["poetry-core", "wheel"] build-backend = "poetry.core.masonry.api" [tool.commitizen] -version = "1.52.4" +version = "1.52.6" version_files = [ "pyproject.toml:^version" ] diff --git a/tests/image_gen_tests/base_image_generation_test.py b/tests/image_gen_tests/base_image_generation_test.py new file mode 100644 index 000000000000..e0652114db70 --- /dev/null +++ b/tests/image_gen_tests/base_image_generation_test.py @@ -0,0 +1,87 @@ +import asyncio +import httpx +import json +import pytest +import sys +from typing import Any, Dict, List, Optional +from unittest.mock import MagicMock, Mock, patch +import os + +sys.path.insert( + 0, os.path.abspath("../..") +) # Adds the parent directory to the system path +import litellm +from litellm.exceptions import BadRequestError +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler +from litellm.utils import CustomStreamWrapper +from openai.types.image import Image +from litellm.integrations.custom_logger import CustomLogger +from litellm.types.utils import StandardLoggingPayload + + +class TestCustomLogger(CustomLogger): + def __init__(self): + super().__init__() + self.standard_logging_payload: Optional[StandardLoggingPayload] = None + + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + self.standard_logging_payload = kwargs.get("standard_logging_object") + pass + + +# test_example.py +from abc import ABC, abstractmethod + + +class BaseImageGenTest(ABC): + """ + Abstract base test class that enforces a common test across all test classes. + """ + + @abstractmethod + def get_base_image_generation_call_args(self) -> dict: + """Must return the base image generation call args""" + pass + + @pytest.mark.asyncio(scope="module") + async def test_basic_image_generation(self): + """Test basic image generation""" + try: + custom_logger = TestCustomLogger() + litellm.callbacks = [custom_logger] + base_image_generation_call_args = self.get_base_image_generation_call_args() + litellm.set_verbose = True + response = await litellm.aimage_generation( + **base_image_generation_call_args, prompt="A image of a otter" + ) + print(response) + + await asyncio.sleep(1) + + assert response._hidden_params["response_cost"] is not None + assert response._hidden_params["response_cost"] > 0 + print("response_cost", response._hidden_params["response_cost"]) + + logged_standard_logging_payload = custom_logger.standard_logging_payload + print("logged_standard_logging_payload", logged_standard_logging_payload) + assert logged_standard_logging_payload is not None + assert logged_standard_logging_payload["response_cost"] is not None + assert logged_standard_logging_payload["response_cost"] > 0 + + from openai.types.images_response import ImagesResponse + + ImagesResponse.model_validate(response.model_dump()) + + for d in response.data: + assert isinstance(d, Image) + print("data in response.data", d) + assert d.b64_json is not None or d.url is not None + except litellm.RateLimitError as e: + pass + except litellm.ContentPolicyViolationError: + pass # Azure randomly raises these errors - skip when they occur + except Exception as e: + if "Your task failed as a result of our safety system." in str(e): + pass + else: + pytest.fail(f"An exception occurred - {str(e)}") diff --git a/tests/image_gen_tests/test_bedrock_image_gen_unit_tests.py b/tests/image_gen_tests/test_bedrock_image_gen_unit_tests.py index e04eb2a1aef1..10845a895dcc 100644 --- a/tests/image_gen_tests/test_bedrock_image_gen_unit_tests.py +++ b/tests/image_gen_tests/test_bedrock_image_gen_unit_tests.py @@ -9,12 +9,14 @@ logging.basicConfig(level=logging.DEBUG) load_dotenv() import asyncio -import os sys.path.insert( 0, os.path.abspath("../..") ) # Adds the parent directory to the system path import pytest +from litellm.llms.bedrock.image.cost_calculator import cost_calculator +from litellm.types.utils import ImageResponse, ImageObject +import os import litellm from litellm.llms.bedrock.image.amazon_stability3_transformation import ( @@ -27,7 +29,6 @@ AmazonStability3TextToImageRequest, AmazonStability3TextToImageResponse, ) -from litellm.types.utils import ImageResponse from unittest.mock import MagicMock, patch from litellm.llms.bedrock.image.image_handler import ( BedrockImageGeneration, @@ -149,7 +150,7 @@ def test_get_request_body_stability(): handler = BedrockImageGeneration() prompt = "A beautiful sunset" optional_params = {"cfg_scale": 7} - model = "stability.stable-diffusion-xl" + model = "stability.stable-diffusion-xl-v1" result = handler._get_request_body( model=model, prompt=prompt, optional_params=optional_params @@ -185,3 +186,80 @@ def test_transform_response_dict_to_openai_response_stability3(): assert len(result.data) == 2 assert all(hasattr(img, "b64_json") for img in result.data) assert [img.b64_json for img in result.data] == ["base64_image_1", "base64_image_2"] + + +def test_cost_calculator_stability3(): + # Mock image response + image_response = ImageResponse( + data=[ + ImageObject(b64_json="base64_image_1"), + ImageObject(b64_json="base64_image_2"), + ] + ) + + cost = cost_calculator( + model="stability.sd3-large-v1:0", + size="1024-x-1024", + image_response=image_response, + ) + + print("cost", cost) + + # Assert cost is calculated correctly for 2 images + assert isinstance(cost, float) + assert cost > 0 + + +def test_cost_calculator_stability1(): + # Mock image response + image_response = ImageResponse(data=[ImageObject(b64_json="base64_image_1")]) + + # Test with different step configurations + cost_default_steps = cost_calculator( + model="stability.stable-diffusion-xl-v1", + size="1024-x-1024", + image_response=image_response, + optional_params={"steps": 50}, + ) + + cost_max_steps = cost_calculator( + model="stability.stable-diffusion-xl-v1", + size="1024-x-1024", + image_response=image_response, + optional_params={"steps": 51}, + ) + + # Assert costs are calculated correctly + assert isinstance(cost_default_steps, float) + assert isinstance(cost_max_steps, float) + assert cost_default_steps > 0 + assert cost_max_steps > 0 + # Max steps should be more expensive + assert cost_max_steps > cost_default_steps + + +def test_cost_calculator_with_no_optional_params(): + image_response = ImageResponse(data=[ImageObject(b64_json="base64_image_1")]) + + cost = cost_calculator( + model="stability.stable-diffusion-xl-v0", + size="512-x-512", + image_response=image_response, + optional_params=None, + ) + + assert isinstance(cost, float) + assert cost > 0 + + +def test_cost_calculator_basic(): + image_response = ImageResponse(data=[ImageObject(b64_json="base64_image_1")]) + + cost = cost_calculator( + model="stability.stable-diffusion-xl-v1", + image_response=image_response, + optional_params=None, + ) + + assert isinstance(cost, float) + assert cost > 0 diff --git a/tests/image_gen_tests/test_image_generation.py b/tests/image_gen_tests/test_image_generation.py index cf46f90bb825..692a0e4e9b61 100644 --- a/tests/image_gen_tests/test_image_generation.py +++ b/tests/image_gen_tests/test_image_generation.py @@ -22,6 +22,11 @@ import litellm import json import tempfile +from base_image_generation_test import BaseImageGenTest +import logging +from litellm._logging import verbose_logger + +verbose_logger.setLevel(logging.DEBUG) def get_vertex_ai_creds_json() -> dict: @@ -97,67 +102,49 @@ def load_vertex_ai_credentials(): os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = os.path.abspath(temp_file.name) -def test_image_generation_openai(): - try: - litellm.set_verbose = True - response = litellm.image_generation( - prompt="A cute baby sea otter", model="dall-e-3" - ) - print(f"response: {response}") - assert len(response.data) > 0 - except litellm.RateLimitError as e: - pass - except litellm.ContentPolicyViolationError: - pass # OpenAI randomly raises these errors - skip when they occur - except Exception as e: - if "Connection error" in str(e): - pass - pytest.fail(f"An exception occurred - {str(e)}") +class TestVertexImageGeneration(BaseImageGenTest): + def get_base_image_generation_call_args(self) -> dict: + # comment this when running locally + load_vertex_ai_credentials() + litellm.in_memory_llm_clients_cache = {} + return { + "model": "vertex_ai/imagegeneration@006", + "vertex_ai_project": "adroit-crow-413218", + "vertex_ai_location": "us-central1", + "n": 1, + } -# test_image_generation_openai() +class TestBedrockSd3(BaseImageGenTest): + def get_base_image_generation_call_args(self) -> dict: + litellm.in_memory_llm_clients_cache = {} + return {"model": "bedrock/stability.sd3-large-v1:0"} -@pytest.mark.parametrize( - "sync_mode", - [ - True, - ], # False -) # -@pytest.mark.asyncio -@pytest.mark.flaky(retries=3, delay=1) -async def test_image_generation_azure(sync_mode): - try: - if sync_mode: - response = litellm.image_generation( - prompt="A cute baby sea otter", - model="azure/", - api_version="2023-06-01-preview", - ) - else: - response = await litellm.aimage_generation( - prompt="A cute baby sea otter", - model="azure/", - api_version="2023-06-01-preview", - ) - print(f"response: {response}") - assert len(response.data) > 0 - except litellm.RateLimitError as e: - pass - except litellm.ContentPolicyViolationError: - pass # Azure randomly raises these errors - skip when they occur - except litellm.InternalServerError: - pass - except Exception as e: - if "Your task failed as a result of our safety system." in str(e): - pass - if "Connection error" in str(e): - pass - else: - pytest.fail(f"An exception occurred - {str(e)}") +class TestBedrockSd1(BaseImageGenTest): + def get_base_image_generation_call_args(self) -> dict: + litellm.in_memory_llm_clients_cache = {} + return {"model": "bedrock/stability.sd3-large-v1:0"} -# test_image_generation_azure() + +class TestOpenAIDalle3(BaseImageGenTest): + def get_base_image_generation_call_args(self) -> dict: + return {"model": "dall-e-3"} + + +class TestAzureOpenAIDalle3(BaseImageGenTest): + def get_base_image_generation_call_args(self) -> dict: + litellm.set_verbose = True + return { + "model": "azure/dall-e-3-test", + "api_version": "2023-09-01-preview", + "metadata": { + "model_info": { + "base_model": "dall-e-3", + } + }, + } @pytest.mark.flaky(retries=3, delay=1) @@ -188,88 +175,13 @@ def test_image_generation_azure_dall_e_3(): pytest.fail(f"An exception occurred - {str(e)}") -# test_image_generation_azure_dall_e_3() -@pytest.mark.asyncio -async def test_async_image_generation_openai(): - try: - response = litellm.image_generation( - prompt="A cute baby sea otter", model="dall-e-3" - ) - print(f"response: {response}") - assert len(response.data) > 0 - except litellm.APIError: - pass - except litellm.RateLimitError as e: - pass - except litellm.ContentPolicyViolationError: - pass # openai randomly raises these errors - skip when they occur - except litellm.InternalServerError: - pass - except Exception as e: - if "Connection error" in str(e): - pass - pytest.fail(f"An exception occurred - {str(e)}") - - # asyncio.run(test_async_image_generation_openai()) -@pytest.mark.asyncio -async def test_async_image_generation_azure(): - try: - response = await litellm.aimage_generation( - prompt="A cute baby sea otter", - model="azure/dall-e-3-test", - api_version="2023-09-01-preview", - ) - print(f"response: {response}") - except litellm.RateLimitError as e: - pass - except litellm.ContentPolicyViolationError: - pass # Azure randomly raises these errors - skip when they occur - except litellm.InternalServerError: - pass - except Exception as e: - if "Your task failed as a result of our safety system." in str(e): - pass - if "Connection error" in str(e): - pass - else: - pytest.fail(f"An exception occurred - {str(e)}") - - -@pytest.mark.asyncio -@pytest.mark.parametrize( - "model", - ["bedrock/stability.sd3-large-v1:0", "bedrock/stability.stable-diffusion-xl-v1"], -) -def test_image_generation_bedrock(model): - try: - litellm.set_verbose = True - response = litellm.image_generation( - prompt="A cute baby sea otter", - model=model, - aws_region_name="us-west-2", - ) - - print(f"response: {response}") - from openai.types.images_response import ImagesResponse - - ImagesResponse.model_validate(response.model_dump()) - except litellm.RateLimitError as e: - pass - except litellm.ContentPolicyViolationError: - pass # Azure randomly raises these errors - skip when they occur - except Exception as e: - if "Your task failed as a result of our safety system." in str(e): - pass - else: - pytest.fail(f"An exception occurred - {str(e)}") - - @pytest.mark.asyncio async def test_aimage_generation_bedrock_with_optional_params(): try: + litellm.in_memory_llm_clients_cache = {} response = await litellm.aimage_generation( prompt="A cute baby sea otter", model="bedrock/stability.stable-diffusion-xl-v1", @@ -285,47 +197,3 @@ async def test_aimage_generation_bedrock_with_optional_params(): pass else: pytest.fail(f"An exception occurred - {str(e)}") - - -from openai.types.image import Image - - -@pytest.mark.parametrize("sync_mode", [True, False]) -@pytest.mark.asyncio -async def test_aimage_generation_vertex_ai(sync_mode): - - litellm.set_verbose = True - - load_vertex_ai_credentials() - data = { - "prompt": "An olympic size swimming pool", - "model": "vertex_ai/imagegeneration@006", - "vertex_ai_project": "adroit-crow-413218", - "vertex_ai_location": "us-central1", - "n": 1, - } - try: - if sync_mode: - response = litellm.image_generation(**data) - else: - response = await litellm.aimage_generation(**data) - assert response.data is not None - assert len(response.data) > 0 - - for d in response.data: - assert isinstance(d, Image) - print("data in response.data", d) - assert d.b64_json is not None - except litellm.ServiceUnavailableError as e: - pass - except litellm.RateLimitError as e: - pass - except litellm.InternalServerError as e: - pass - except litellm.ContentPolicyViolationError: - pass # Azure randomly raises these errors - skip when they occur - except Exception as e: - if "Your task failed as a result of our safety system." in str(e): - pass - else: - pytest.fail(f"An exception occurred - {str(e)}") diff --git a/tests/llm_translation/base_llm_unit_tests.py b/tests/llm_translation/base_llm_unit_tests.py index 4f9cd9c2517a..acb764ba14b6 100644 --- a/tests/llm_translation/base_llm_unit_tests.py +++ b/tests/llm_translation/base_llm_unit_tests.py @@ -44,3 +44,64 @@ def test_content_list_handling(self): messages=messages, ) assert response is not None + + def test_message_with_name(self): + base_completion_call_args = self.get_base_completion_call_args() + messages = [ + {"role": "user", "content": "Hello", "name": "test_name"}, + ] + response = litellm.completion(**base_completion_call_args, messages=messages) + assert response is not None + + def test_json_response_format(self): + """ + Test that the JSON response format is supported by the LLM API + """ + base_completion_call_args = self.get_base_completion_call_args() + litellm.set_verbose = True + + messages = [ + { + "role": "system", + "content": "Your output should be a JSON object with no additional properties. ", + }, + { + "role": "user", + "content": "Respond with this in json. city=San Francisco, state=CA, weather=sunny, temp=60", + }, + ] + + response = litellm.completion( + **base_completion_call_args, + messages=messages, + response_format={"type": "json_object"}, + ) + + print(response) + + @pytest.fixture + def pdf_messages(self): + import base64 + + import requests + + # URL of the file + url = "https://storage.googleapis.com/cloud-samples-data/generative-ai/pdf/2403.05530.pdf" + + response = requests.get(url) + file_data = response.content + + encoded_file = base64.b64encode(file_data).decode("utf-8") + url = f"data:application/pdf;base64,{encoded_file}" + + image_content = [ + {"type": "text", "text": "What's this file about?"}, + { + "type": "image_url", + "image_url": {"url": url}, + }, + ] + + image_messages = [{"role": "user", "content": image_content}] + + return image_messages diff --git a/tests/llm_translation/test_anthropic_completion.py b/tests/llm_translation/test_anthropic_completion.py index 46f01e0ec875..c399c3a470da 100644 --- a/tests/llm_translation/test_anthropic_completion.py +++ b/tests/llm_translation/test_anthropic_completion.py @@ -36,6 +36,7 @@ from litellm.llms.anthropic.common_utils import process_anthropic_headers from httpx import Headers +from base_llm_unit_tests import BaseLLMChatTest def test_anthropic_completion_messages_translation(): @@ -624,3 +625,72 @@ def test_anthropic_tool_helper(cache_control_location): tool = AnthropicConfig()._map_tool_helper(tool=tool) assert tool["cache_control"] == {"type": "ephemeral"} + + +def test_create_json_tool_call_for_response_format(): + """ + tests using response_format=json with anthropic + + A tool call to anthropic is made when response_format=json is used. + + """ + # Initialize AnthropicConfig + config = AnthropicConfig() + + # Test case 1: No schema provided + # See Anthropics Example 5 on how to handle cases when no schema is provided https://github.com/anthropics/anthropic-cookbook/blob/main/tool_use/extracting_structured_json.ipynb + tool = config._create_json_tool_call_for_response_format() + assert tool["name"] == "json_tool_call" + _input_schema = tool.get("input_schema") + assert _input_schema is not None + assert _input_schema.get("type") == "object" + assert _input_schema.get("additionalProperties") is True + assert _input_schema.get("properties") == {} + + # Test case 2: With custom schema + # reference: https://github.com/anthropics/anthropic-cookbook/blob/main/tool_use/extracting_structured_json.ipynb + custom_schema = {"name": {"type": "string"}, "age": {"type": "integer"}} + tool = config._create_json_tool_call_for_response_format(json_schema=custom_schema) + assert tool["name"] == "json_tool_call" + _input_schema = tool.get("input_schema") + assert _input_schema is not None + assert _input_schema.get("type") == "object" + assert _input_schema.get("properties") == custom_schema + assert "additionalProperties" not in _input_schema + + +from litellm import completion + + +class TestAnthropicCompletion(BaseLLMChatTest): + def get_base_completion_call_args(self) -> dict: + return {"model": "claude-3-haiku-20240307"} + + def test_pdf_handling(self, pdf_messages): + from litellm.llms.custom_httpx.http_handler import HTTPHandler + from litellm.types.llms.anthropic import AnthropicMessagesDocumentParam + import json + + client = HTTPHandler() + + with patch.object(client, "post", new=MagicMock()) as mock_client: + response = completion( + model="claude-3-5-sonnet-20241022", + messages=pdf_messages, + client=client, + ) + + mock_client.assert_called_once() + + json_data = json.loads(mock_client.call_args.kwargs["data"]) + headers = mock_client.call_args.kwargs["headers"] + + assert headers["anthropic-beta"] == "pdfs-2024-09-25" + + json_data["messages"][0]["role"] == "user" + _document_validation = AnthropicMessagesDocumentParam( + **json_data["messages"][0]["content"][1] + ) + assert _document_validation["type"] == "document" + assert _document_validation["source"]["media_type"] == "application/pdf" + assert _document_validation["source"]["type"] == "base64" diff --git a/tests/llm_translation/test_mistral_api.py b/tests/llm_translation/test_mistral_api.py new file mode 100644 index 000000000000..b2cb36541598 --- /dev/null +++ b/tests/llm_translation/test_mistral_api.py @@ -0,0 +1,34 @@ +import asyncio +import os +import sys +import traceback + +from dotenv import load_dotenv + +import litellm.types +import litellm.types.utils +from litellm.llms.anthropic.chat import ModelResponseIterator + +load_dotenv() +import io +import os + +sys.path.insert( + 0, os.path.abspath("../..") +) # Adds the parent directory to the system path +from typing import Optional +from unittest.mock import MagicMock, patch + +import pytest + +import litellm + +from litellm.llms.anthropic.common_utils import process_anthropic_headers +from httpx import Headers +from base_llm_unit_tests import BaseLLMChatTest + + +class TestMistralCompletion(BaseLLMChatTest): + def get_base_completion_call_args(self) -> dict: + litellm.set_verbose = True + return {"model": "mistral/mistral-small-latest"} diff --git a/tests/local_testing/cache_unit_tests.py b/tests/local_testing/cache_unit_tests.py new file mode 100644 index 000000000000..da56c773f360 --- /dev/null +++ b/tests/local_testing/cache_unit_tests.py @@ -0,0 +1,223 @@ +from abc import ABC, abstractmethod +from litellm.caching import LiteLLMCacheType +import os +import sys +import time +import traceback +import uuid + +from dotenv import load_dotenv +from test_rerank import assert_response_shape + +load_dotenv() +import os + +sys.path.insert( + 0, os.path.abspath("../..") +) # Adds the parent directory to the system path +import asyncio +import hashlib +import random + +import pytest + +import litellm +from litellm.caching import Cache +from litellm import completion, embedding + + +class LLMCachingUnitTests(ABC): + + @abstractmethod + def get_cache_type(self) -> LiteLLMCacheType: + pass + + @pytest.mark.parametrize("sync_mode", [True, False]) + @pytest.mark.asyncio + async def test_cache_completion(self, sync_mode): + litellm._turn_on_debug() + + random_number = random.randint( + 1, 100000 + ) # add a random number to ensure it's always adding / reading from cache + messages = [ + { + "role": "user", + "content": f"write a one sentence poem about: {random_number}", + } + ] + + cache_type = self.get_cache_type() + litellm.cache = Cache( + type=cache_type, + ) + + if sync_mode: + response1 = completion( + "gpt-3.5-turbo", + messages=messages, + caching=True, + max_tokens=20, + mock_response="This number is so great!", + ) + else: + response1 = await litellm.acompletion( + "gpt-3.5-turbo", + messages=messages, + caching=True, + max_tokens=20, + mock_response="This number is so great!", + ) + # response2 is mocked to a different response from response1, + # but the completion from the cache should be used instead of the mock + # response since the input is the same as response1 + await asyncio.sleep(0.5) + if sync_mode: + response2 = completion( + "gpt-3.5-turbo", + messages=messages, + caching=True, + max_tokens=20, + mock_response="This number is great!", + ) + else: + response2 = await litellm.acompletion( + "gpt-3.5-turbo", + messages=messages, + caching=True, + max_tokens=20, + mock_response="This number is great!", + ) + if ( + response1["choices"][0]["message"]["content"] + != response2["choices"][0]["message"]["content"] + ): # 1 and 2 should be the same + # 1&2 have the exact same input params. This MUST Be a CACHE HIT + print(f"response1: {response1}") + print(f"response2: {response2}") + pytest.fail( + f"Error occurred: response1 - {response1['choices'][0]['message']['content']} != response2 - {response2['choices'][0]['message']['content']}" + ) + # Since the parameters are not the same as response1, response3 should actually + # be the mock response + if sync_mode: + response3 = completion( + "gpt-3.5-turbo", + messages=messages, + caching=True, + temperature=0.5, + mock_response="This number is awful!", + ) + else: + response3 = await litellm.acompletion( + "gpt-3.5-turbo", + messages=messages, + caching=True, + temperature=0.5, + mock_response="This number is awful!", + ) + + print("\nresponse 1", response1) + print("\nresponse 2", response2) + print("\nresponse 3", response3) + # print("\nresponse 4", response4) + litellm.cache = None + litellm.success_callback = [] + litellm._async_success_callback = [] + + # 1 & 2 should be exactly the same + # 1 & 3 should be different, since input params are diff + + if ( + response1["choices"][0]["message"]["content"] + == response3["choices"][0]["message"]["content"] + ): + # if input params like max_tokens, temperature are diff it should NOT be a cache hit + print(f"response1: {response1}") + print(f"response3: {response3}") + pytest.fail( + f"Response 1 == response 3. Same model, diff params shoudl not cache Error" + f" occurred:" + ) + + assert response1.id == response2.id + assert response1.created == response2.created + assert ( + response1.choices[0].message.content == response2.choices[0].message.content + ) + + @pytest.mark.parametrize("sync_mode", [True, False]) + @pytest.mark.asyncio + async def test_disk_cache_embedding(self, sync_mode): + litellm._turn_on_debug() + + random_number = random.randint( + 1, 100000 + ) # add a random number to ensure it's always adding / reading from cache + input = [f"hello {random_number}"] + litellm.cache = Cache( + type="disk", + ) + + if sync_mode: + response1 = embedding( + "openai/text-embedding-ada-002", + input=input, + caching=True, + ) + else: + response1 = await litellm.aembedding( + "openai/text-embedding-ada-002", + input=input, + caching=True, + ) + # response2 is mocked to a different response from response1, + # but the completion from the cache should be used instead of the mock + # response since the input is the same as response1 + await asyncio.sleep(0.5) + if sync_mode: + response2 = embedding( + "openai/text-embedding-ada-002", + input=input, + caching=True, + ) + else: + response2 = await litellm.aembedding( + "openai/text-embedding-ada-002", + input=input, + caching=True, + ) + + if response2._hidden_params["cache_hit"] is not True: + pytest.fail("Cache hit should be True") + + # Since the parameters are not the same as response1, response3 should actually + # be the mock response + if sync_mode: + response3 = embedding( + "openai/text-embedding-ada-002", + input=input, + user="charlie", + caching=True, + ) + else: + response3 = await litellm.aembedding( + "openai/text-embedding-ada-002", + input=input, + caching=True, + user="charlie", + ) + + print("\nresponse 1", response1) + print("\nresponse 2", response2) + print("\nresponse 3", response3) + # print("\nresponse 4", response4) + litellm.cache = None + litellm.success_callback = [] + litellm._async_success_callback = [] + + # 1 & 2 should be exactly the same + # 1 & 3 should be different, since input params are diff + + if response3._hidden_params.get("cache_hit") is True: + pytest.fail("Cache hit should not be True") diff --git a/tests/local_testing/test_alerting.py b/tests/local_testing/test_alerting.py index b79438ffc5a7..cc668801f0d7 100644 --- a/tests/local_testing/test_alerting.py +++ b/tests/local_testing/test_alerting.py @@ -438,7 +438,7 @@ async def test_send_daily_reports_ignores_zero_values(): slack_alerting.internal_usage_cache.async_batch_get_cache = AsyncMock( return_value=[None, 0, 10, 0, 0, None] ) - slack_alerting.internal_usage_cache.async_batch_set_cache = AsyncMock() + slack_alerting.internal_usage_cache.async_set_cache_pipeline = AsyncMock() router.get_model_info.side_effect = lambda x: {"litellm_params": {"model": x}} diff --git a/tests/local_testing/test_caching.py b/tests/local_testing/test_caching.py index 479c1204e995..222013a86e80 100644 --- a/tests/local_testing/test_caching.py +++ b/tests/local_testing/test_caching.py @@ -1103,81 +1103,6 @@ async def test_redis_cache_acompletion_stream_bedrock(): raise e -def test_disk_cache_completion(): - litellm.set_verbose = False - - random_number = random.randint( - 1, 100000 - ) # add a random number to ensure it's always adding / reading from cache - messages = [ - {"role": "user", "content": f"write a one sentence poem about: {random_number}"} - ] - litellm.cache = Cache( - type="disk", - ) - - response1 = completion( - model="gpt-3.5-turbo", - messages=messages, - caching=True, - max_tokens=20, - mock_response="This number is so great!", - ) - # response2 is mocked to a different response from response1, - # but the completion from the cache should be used instead of the mock - # response since the input is the same as response1 - response2 = completion( - model="gpt-3.5-turbo", - messages=messages, - caching=True, - max_tokens=20, - mock_response="This number is awful!", - ) - # Since the parameters are not the same as response1, response3 should actually - # be the mock response - response3 = completion( - model="gpt-3.5-turbo", - messages=messages, - caching=True, - temperature=0.5, - mock_response="This number is awful!", - ) - - print("\nresponse 1", response1) - print("\nresponse 2", response2) - print("\nresponse 3", response3) - # print("\nresponse 4", response4) - litellm.cache = None - litellm.success_callback = [] - litellm._async_success_callback = [] - - # 1 & 2 should be exactly the same - # 1 & 3 should be different, since input params are diff - if ( - response1["choices"][0]["message"]["content"] - != response2["choices"][0]["message"]["content"] - ): # 1 and 2 should be the same - # 1&2 have the exact same input params. This MUST Be a CACHE HIT - print(f"response1: {response1}") - print(f"response2: {response2}") - pytest.fail(f"Error occurred:") - if ( - response1["choices"][0]["message"]["content"] - == response3["choices"][0]["message"]["content"] - ): - # if input params like max_tokens, temperature are diff it should NOT be a cache hit - print(f"response1: {response1}") - print(f"response3: {response3}") - pytest.fail( - f"Response 1 == response 3. Same model, diff params shoudl not cache Error" - f" occurred:" - ) - - assert response1.id == response2.id - assert response1.created == response2.created - assert response1.choices[0].message.content == response2.choices[0].message.content - - # @pytest.mark.skip(reason="AWS Suspended Account") @pytest.mark.parametrize("sync_mode", [True, False]) @pytest.mark.asyncio diff --git a/tests/local_testing/test_disk_cache_unit_tests.py b/tests/local_testing/test_disk_cache_unit_tests.py new file mode 100644 index 000000000000..c777d04ecce3 --- /dev/null +++ b/tests/local_testing/test_disk_cache_unit_tests.py @@ -0,0 +1,11 @@ +from cache_unit_tests import LLMCachingUnitTests +from litellm.caching import LiteLLMCacheType + + +class TestDiskCacheUnitTests(LLMCachingUnitTests): + def get_cache_type(self) -> LiteLLMCacheType: + return LiteLLMCacheType.DISK + + +# if __name__ == "__main__": +# pytest.main([__file__, "-v", "-s"]) diff --git a/tests/local_testing/test_dual_cache.py b/tests/local_testing/test_dual_cache.py index c3f3216d5d79..e81424a9ffed 100644 --- a/tests/local_testing/test_dual_cache.py +++ b/tests/local_testing/test_dual_cache.py @@ -146,7 +146,7 @@ async def test_dual_cache_batch_operations(is_async): # Set values if is_async: - await dual_cache.async_batch_set_cache(cache_list) + await dual_cache.async_set_cache_pipeline(cache_list) else: for key, value in cache_list: dual_cache.set_cache(key, value) diff --git a/tests/local_testing/test_get_llm_provider.py b/tests/local_testing/test_get_llm_provider.py index f7126cec07b2..6654c10c2e23 100644 --- a/tests/local_testing/test_get_llm_provider.py +++ b/tests/local_testing/test_get_llm_provider.py @@ -169,3 +169,11 @@ def test_get_llm_provider_hosted_vllm(): assert custom_llm_provider == "hosted_vllm" assert model == "llama-3.1-70b-instruct" assert dynamic_api_key == "" + + +def test_get_llm_provider_watson_text(): + model, custom_llm_provider, dynamic_api_key, api_base = litellm.get_llm_provider( + model="watsonx_text/watson-text-to-speech", + ) + assert custom_llm_provider == "watsonx_text" + assert model == "watson-text-to-speech" diff --git a/tests/local_testing/test_get_model_list.py b/tests/local_testing/test_get_model_list.py deleted file mode 100644 index 7663eebf5213..000000000000 --- a/tests/local_testing/test_get_model_list.py +++ /dev/null @@ -1,11 +0,0 @@ -import os, sys, traceback - -sys.path.insert( - 0, os.path.abspath("../..") -) # Adds the parent directory to the system path -import litellm -from litellm import get_model_list - -print(get_model_list()) -print(get_model_list()) -# print(litellm.model_list) diff --git a/tests/local_testing/test_langsmith.py b/tests/local_testing/test_langsmith.py index 6a98f244dcbf..ab387e444eff 100644 --- a/tests/local_testing/test_langsmith.py +++ b/tests/local_testing/test_langsmith.py @@ -22,61 +22,6 @@ import time -@pytest.mark.asyncio -async def test_langsmith_queue_logging(): - try: - # Initialize LangsmithLogger - test_langsmith_logger = LangsmithLogger() - - litellm.callbacks = [test_langsmith_logger] - test_langsmith_logger.batch_size = 6 - litellm.set_verbose = True - - # Make multiple calls to ensure we don't hit the batch size - for _ in range(5): - response = await litellm.acompletion( - model="gpt-3.5-turbo", - messages=[{"role": "user", "content": "Test message"}], - max_tokens=10, - temperature=0.2, - mock_response="This is a mock response", - ) - - await asyncio.sleep(3) - - # Check that logs are in the queue - assert len(test_langsmith_logger.log_queue) == 5 - - # Now make calls to exceed the batch size - for _ in range(3): - response = await litellm.acompletion( - model="gpt-3.5-turbo", - messages=[{"role": "user", "content": "Test message"}], - max_tokens=10, - temperature=0.2, - mock_response="This is a mock response", - ) - - # Wait a short time for any asynchronous operations to complete - await asyncio.sleep(1) - - print( - "Length of langsmith log queue: {}".format( - len(test_langsmith_logger.log_queue) - ) - ) - # Check that the queue was flushed after exceeding batch size - assert len(test_langsmith_logger.log_queue) < 5 - - # Clean up - for cb in litellm.callbacks: - if isinstance(cb, LangsmithLogger): - await cb.async_httpx_client.client.aclose() - - except Exception as e: - pytest.fail(f"Error occurred: {e}") - - # test_langsmith_logging() diff --git a/tests/local_testing/test_opentelemetry_unit_tests.py b/tests/local_testing/test_opentelemetry_unit_tests.py deleted file mode 100644 index 530adc6aba59..000000000000 --- a/tests/local_testing/test_opentelemetry_unit_tests.py +++ /dev/null @@ -1,41 +0,0 @@ -# What is this? -## Unit tests for opentelemetry integration - -# What is this? -## Unit test for presidio pii masking -import sys, os, asyncio, time, random -from datetime import datetime -import traceback -from dotenv import load_dotenv - -load_dotenv() -import os -import asyncio - -sys.path.insert( - 0, os.path.abspath("../..") -) # Adds the parent directory to the system path -import pytest -import litellm -from unittest.mock import patch, MagicMock, AsyncMock - - -@pytest.mark.asyncio -async def test_opentelemetry_integration(): - """ - Unit test to confirm the parent otel span is ended - """ - - parent_otel_span = MagicMock() - litellm.callbacks = ["otel"] - - await litellm.acompletion( - model="gpt-3.5-turbo", - messages=[{"role": "user", "content": "Hello, world!"}], - mock_response="Hey!", - metadata={"litellm_parent_otel_span": parent_otel_span}, - ) - - await asyncio.sleep(1) - - parent_otel_span.end.assert_called_once() diff --git a/tests/local_testing/test_utils.py b/tests/local_testing/test_utils.py index 5aa3b610c16b..b3f8208bf82d 100644 --- a/tests/local_testing/test_utils.py +++ b/tests/local_testing/test_utils.py @@ -943,3 +943,24 @@ def test_validate_chat_completion_user_messages(messages, expected_bool): ## Invalid message with pytest.raises(Exception): validate_chat_completion_user_messages(messages=messages) + + +def test_models_by_provider(): + """ + Make sure all providers from model map are in the valid providers list + """ + from litellm import models_by_provider + + providers = set() + for k, v in litellm.model_cost.items(): + if "_" in v["litellm_provider"] and "-" in v["litellm_provider"]: + continue + elif k == "sample_spec": + continue + elif v["litellm_provider"] == "sagemaker": + continue + else: + providers.add(v["litellm_provider"]) + + for provider in providers: + assert provider in models_by_provider.keys() diff --git a/tests/logging_callback_tests/base_test.py b/tests/logging_callback_tests/base_test.py new file mode 100644 index 000000000000..0d1e7dfcf775 --- /dev/null +++ b/tests/logging_callback_tests/base_test.py @@ -0,0 +1,100 @@ +import asyncio +import httpx +import json +import pytest +import sys +from typing import Any, Dict, List +from unittest.mock import MagicMock, Mock, patch +import os + +sys.path.insert( + 0, os.path.abspath("../..") +) # Adds the parent directory to the system path +import litellm +from litellm.exceptions import BadRequestError +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler +from litellm.utils import CustomStreamWrapper +from litellm.types.utils import ModelResponse + +# test_example.py +from abc import ABC, abstractmethod + + +class BaseLoggingCallbackTest(ABC): + """ + Abstract base test class that enforces a common test across all test classes. + """ + + @pytest.fixture + def mock_response_obj(self): + from litellm.types.utils import ( + ModelResponse, + Choices, + Message, + ChatCompletionMessageToolCall, + Function, + Usage, + CompletionTokensDetailsWrapper, + PromptTokensDetailsWrapper, + ) + + # Create a mock response object with the structure you need + return ModelResponse( + id="chatcmpl-ASId3YJWagBpBskWfoNEMPFSkmrEw", + created=1731308157, + model="gpt-4o-mini-2024-07-18", + object="chat.completion", + system_fingerprint="fp_0ba0d124f1", + choices=[ + Choices( + finish_reason="tool_calls", + index=0, + message=Message( + content=None, + role="assistant", + tool_calls=[ + ChatCompletionMessageToolCall( + function=Function( + arguments='{"city": "New York"}', name="get_weather" + ), + id="call_PngsQS5YGmIZKnswhnUOnOVb", + type="function", + ), + ChatCompletionMessageToolCall( + function=Function( + arguments='{"city": "New York"}', name="get_news" + ), + id="call_1zsDThBu0VSK7KuY7eCcJBnq", + type="function", + ), + ], + function_call=None, + ), + ) + ], + usage=Usage( + completion_tokens=46, + prompt_tokens=86, + total_tokens=132, + completion_tokens_details=CompletionTokensDetailsWrapper( + accepted_prediction_tokens=0, + audio_tokens=0, + reasoning_tokens=0, + rejected_prediction_tokens=0, + text_tokens=None, + ), + prompt_tokens_details=PromptTokensDetailsWrapper( + audio_tokens=0, cached_tokens=0, text_tokens=None, image_tokens=None + ), + ), + service_tier=None, + ) + + @abstractmethod + def test_parallel_tool_calls(self, mock_response_obj: ModelResponse): + """ + Check if parallel tool calls are correctly logged by Logging callback + + Relevant issue - https://github.com/BerriAI/litellm/issues/6677 + """ + pass diff --git a/tests/logging_callback_tests/test_langfuse_unit_tests.py b/tests/logging_callback_tests/test_langfuse_unit_tests.py index 20b33f81b52b..c10b6110ccfa 100644 --- a/tests/logging_callback_tests/test_langfuse_unit_tests.py +++ b/tests/logging_callback_tests/test_langfuse_unit_tests.py @@ -212,26 +212,48 @@ def test_get_langfuse_logger_for_request_with_cached_logger(): assert result == cached_logger mock_cache.get_cache.assert_called_once() -@pytest.mark.parametrize("metadata", [ - {'a': 1, 'b': 2, 'c': 3}, - {'a': {'nested_a': 1}, 'b': {'nested_b': 2}}, - {'a': [1, 2, 3], 'b': {4, 5, 6}}, - {'a': (1, 2), 'b': frozenset([3, 4]), 'c': {'d': [5, 6]}}, - {'lock': threading.Lock()}, - {'func': lambda x: x + 1}, - { - 'int': 42, - 'str': 'hello', - 'list': [1, 2, 3], - 'set': {4, 5}, - 'dict': {'nested': 'value'}, - 'non_copyable': threading.Lock(), - 'function': print - }, - ['list', 'not', 'a', 'dict'], - {'timestamp': datetime.now()}, - {}, - None, -]) -def test_langfuse_logger_prepare_metadata(metadata): - global_langfuse_logger._prepare_metadata(metadata) + +@pytest.mark.parametrize( + "metadata, expected_metadata", + [ + ({"a": 1, "b": 2, "c": 3}, {"a": 1, "b": 2, "c": 3}), + ( + {"a": {"nested_a": 1}, "b": {"nested_b": 2}}, + {"a": {"nested_a": 1}, "b": {"nested_b": 2}}, + ), + ({"a": [1, 2, 3], "b": {4, 5, 6}}, {"a": [1, 2, 3], "b": {4, 5, 6}}), + ( + {"a": (1, 2), "b": frozenset([3, 4]), "c": {"d": [5, 6]}}, + {"a": (1, 2), "b": frozenset([3, 4]), "c": {"d": [5, 6]}}, + ), + ({"lock": threading.Lock()}, {}), + ({"func": lambda x: x + 1}, {}), + ( + { + "int": 42, + "str": "hello", + "list": [1, 2, 3], + "set": {4, 5}, + "dict": {"nested": "value"}, + "non_copyable": threading.Lock(), + "function": print, + }, + { + "int": 42, + "str": "hello", + "list": [1, 2, 3], + "set": {4, 5}, + "dict": {"nested": "value"}, + }, + ), + ( + {"list": ["list", "not", "a", "dict"]}, + {"list": ["list", "not", "a", "dict"]}, + ), + ({}, {}), + (None, None), + ], +) +def test_langfuse_logger_prepare_metadata(metadata, expected_metadata): + result = global_langfuse_logger._prepare_metadata(metadata) + assert result == expected_metadata diff --git a/tests/logging_callback_tests/test_langsmith_unit_test.py b/tests/logging_callback_tests/test_langsmith_unit_test.py new file mode 100644 index 000000000000..3e106666f7c8 --- /dev/null +++ b/tests/logging_callback_tests/test_langsmith_unit_test.py @@ -0,0 +1,394 @@ +import io +import os +import sys + + +sys.path.insert(0, os.path.abspath("../..")) + +import asyncio +import gzip +import json +import logging +import time +from unittest.mock import AsyncMock, patch, MagicMock +import pytest +from datetime import datetime, timezone +from litellm.integrations.langsmith import ( + LangsmithLogger, + LangsmithQueueObject, + CredentialsKey, + BatchGroup, +) + +import litellm + + +# Test get_credentials_from_env +@pytest.mark.asyncio +async def test_get_credentials_from_env(): + # Test with direct parameters + logger = LangsmithLogger( + langsmith_api_key="test-key", + langsmith_project="test-project", + langsmith_base_url="http://test-url", + ) + + credentials = logger.get_credentials_from_env( + langsmith_api_key="custom-key", + langsmith_project="custom-project", + langsmith_base_url="http://custom-url", + ) + + assert credentials["LANGSMITH_API_KEY"] == "custom-key" + assert credentials["LANGSMITH_PROJECT"] == "custom-project" + assert credentials["LANGSMITH_BASE_URL"] == "http://custom-url" + + # assert that the default api base is used if not provided + credentials = logger.get_credentials_from_env() + assert credentials["LANGSMITH_BASE_URL"] == "https://api.smith.langchain.com" + + +@pytest.mark.asyncio +async def test_group_batches_by_credentials(): + + logger = LangsmithLogger(langsmith_api_key="test-key") + + # Create test queue objects + queue_obj1 = LangsmithQueueObject( + data={"test": "data1"}, + credentials={ + "LANGSMITH_API_KEY": "key1", + "LANGSMITH_PROJECT": "proj1", + "LANGSMITH_BASE_URL": "url1", + }, + ) + + queue_obj2 = LangsmithQueueObject( + data={"test": "data2"}, + credentials={ + "LANGSMITH_API_KEY": "key1", + "LANGSMITH_PROJECT": "proj1", + "LANGSMITH_BASE_URL": "url1", + }, + ) + + logger.log_queue = [queue_obj1, queue_obj2] + + grouped = logger._group_batches_by_credentials() + + # Check grouping + assert len(grouped) == 1 # Should have one group since credentials are same + key = list(grouped.keys())[0] + assert isinstance(key, CredentialsKey) + assert len(grouped[key].queue_objects) == 2 + + +@pytest.mark.asyncio +async def test_group_batches_by_credentials_multiple_credentials(): + + # Test with multiple different credentials + logger = LangsmithLogger(langsmith_api_key="test-key") + + queue_obj1 = LangsmithQueueObject( + data={"test": "data1"}, + credentials={ + "LANGSMITH_API_KEY": "key1", + "LANGSMITH_PROJECT": "proj1", + "LANGSMITH_BASE_URL": "url1", + }, + ) + + queue_obj2 = LangsmithQueueObject( + data={"test": "data2"}, + credentials={ + "LANGSMITH_API_KEY": "key2", # Different API key + "LANGSMITH_PROJECT": "proj1", + "LANGSMITH_BASE_URL": "url1", + }, + ) + + queue_obj3 = LangsmithQueueObject( + data={"test": "data3"}, + credentials={ + "LANGSMITH_API_KEY": "key1", + "LANGSMITH_PROJECT": "proj2", # Different project + "LANGSMITH_BASE_URL": "url1", + }, + ) + + logger.log_queue = [queue_obj1, queue_obj2, queue_obj3] + + grouped = logger._group_batches_by_credentials() + + # Check grouping + assert len(grouped) == 3 # Should have three groups since credentials differ + for key, batch_group in grouped.items(): + assert isinstance(key, CredentialsKey) + assert len(batch_group.queue_objects) == 1 # Each group should have one object + + +# Test make_dot_order +@pytest.mark.asyncio +async def test_make_dot_order(): + logger = LangsmithLogger(langsmith_api_key="test-key") + run_id = "729cff0e-f30c-4336-8b79-45d6b61c64b4" + dot_order = logger.make_dot_order(run_id) + + print("dot_order=", dot_order) + + # Check format: YYYYMMDDTHHMMSSfffZ + run_id + # Check the timestamp portion (first 23 characters) + timestamp_part = dot_order[:-36] # 36 is length of run_id + assert len(timestamp_part) == 22 + assert timestamp_part[8] == "T" # Check T separator + assert timestamp_part[-1] == "Z" # Check Z suffix + + # Verify timestamp format + try: + # Parse the timestamp portion (removing the Z) + datetime.strptime(timestamp_part[:-1], "%Y%m%dT%H%M%S%f") + except ValueError: + pytest.fail("Timestamp portion is not in correct format") + + # Verify run_id portion + assert dot_order[-36:] == run_id + + +# Test is_serializable +@pytest.mark.asyncio +async def test_is_serializable(): + from litellm.integrations.langsmith import is_serializable + from pydantic import BaseModel + + # Test basic types + assert is_serializable("string") is True + assert is_serializable(123) is True + assert is_serializable({"key": "value"}) is True + + # Test non-serializable types + async def async_func(): + pass + + assert is_serializable(async_func) is False + + class TestModel(BaseModel): + field: str + + assert is_serializable(TestModel(field="test")) is False + + +@pytest.mark.asyncio +async def test_async_send_batch(): + logger = LangsmithLogger(langsmith_api_key="test-key") + + # Mock the httpx client + mock_response = AsyncMock() + mock_response.status_code = 200 + logger.async_httpx_client = AsyncMock() + logger.async_httpx_client.post.return_value = mock_response + + # Add test data to queue + logger.log_queue = [ + LangsmithQueueObject( + data={"test": "data"}, credentials=logger.default_credentials + ) + ] + + await logger.async_send_batch() + + # Verify the API call + logger.async_httpx_client.post.assert_called_once() + call_args = logger.async_httpx_client.post.call_args + assert "runs/batch" in call_args[1]["url"] + assert "x-api-key" in call_args[1]["headers"] + + +@pytest.mark.asyncio +async def test_langsmith_key_based_logging(mocker): + """ + In key based logging langsmith_api_key and langsmith_project are passed directly to litellm.acompletion + """ + try: + # Mock the httpx post request + mock_post = mocker.patch( + "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post" + ) + mock_post.return_value.status_code = 200 + mock_post.return_value.raise_for_status = lambda: None + litellm.set_verbose = True + + litellm.callbacks = [LangsmithLogger()] + response = await litellm.acompletion( + model="gpt-3.5-turbo", + messages=[{"role": "user", "content": "Test message"}], + max_tokens=10, + temperature=0.2, + mock_response="This is a mock response", + langsmith_api_key="fake_key_project2", + langsmith_project="fake_project2", + ) + print("Waiting for logs to be flushed to Langsmith.....") + await asyncio.sleep(15) + + print("done sleeping 15 seconds...") + + # Verify the post request was made with correct parameters + mock_post.assert_called_once() + call_args = mock_post.call_args + + print("call_args", call_args) + + # Check URL contains /runs/batch + assert "/runs/batch" in call_args[1]["url"] + + # Check headers contain the correct API key + assert call_args[1]["headers"]["x-api-key"] == "fake_key_project2" + + # Verify the request body contains the expected data + request_body = call_args[1]["json"] + assert "post" in request_body + assert len(request_body["post"]) == 1 # Should contain one run + + # EXPECTED BODY + expected_body = { + "post": [ + { + "name": "LLMRun", + "run_type": "llm", + "inputs": { + "id": "chatcmpl-82699ee4-7932-4fc0-9585-76abc8caeafa", + "call_type": "acompletion", + "model": "gpt-3.5-turbo", + "messages": [{"role": "user", "content": "Test message"}], + "model_parameters": { + "temperature": 0.2, + "max_tokens": 10, + "extra_body": {}, + }, + }, + "outputs": { + "id": "chatcmpl-82699ee4-7932-4fc0-9585-76abc8caeafa", + "model": "gpt-3.5-turbo", + "choices": [ + { + "finish_reason": "stop", + "index": 0, + "message": { + "content": "This is a mock response", + "role": "assistant", + "tool_calls": None, + "function_call": None, + }, + } + ], + "usage": { + "completion_tokens": 20, + "prompt_tokens": 10, + "total_tokens": 30, + }, + }, + "session_name": "fake_project2", + } + ] + } + + # Print both bodies for debugging + actual_body = call_args[1]["json"] + print("\nExpected body:") + print(json.dumps(expected_body, indent=2)) + print("\nActual body:") + print(json.dumps(actual_body, indent=2)) + + assert len(actual_body["post"]) == 1 + + # Assert only the critical parts we care about + assert actual_body["post"][0]["name"] == expected_body["post"][0]["name"] + assert ( + actual_body["post"][0]["run_type"] == expected_body["post"][0]["run_type"] + ) + assert ( + actual_body["post"][0]["inputs"]["messages"] + == expected_body["post"][0]["inputs"]["messages"] + ) + assert ( + actual_body["post"][0]["inputs"]["model_parameters"] + == expected_body["post"][0]["inputs"]["model_parameters"] + ) + assert ( + actual_body["post"][0]["outputs"]["choices"] + == expected_body["post"][0]["outputs"]["choices"] + ) + assert ( + actual_body["post"][0]["outputs"]["usage"]["completion_tokens"] + == expected_body["post"][0]["outputs"]["usage"]["completion_tokens"] + ) + assert ( + actual_body["post"][0]["outputs"]["usage"]["prompt_tokens"] + == expected_body["post"][0]["outputs"]["usage"]["prompt_tokens"] + ) + assert ( + actual_body["post"][0]["outputs"]["usage"]["total_tokens"] + == expected_body["post"][0]["outputs"]["usage"]["total_tokens"] + ) + assert ( + actual_body["post"][0]["session_name"] + == expected_body["post"][0]["session_name"] + ) + + except Exception as e: + pytest.fail(f"Error occurred: {e}") + + +@pytest.mark.asyncio +async def test_langsmith_queue_logging(): + try: + # Initialize LangsmithLogger + test_langsmith_logger = LangsmithLogger() + + litellm.callbacks = [test_langsmith_logger] + test_langsmith_logger.batch_size = 6 + litellm.set_verbose = True + + # Make multiple calls to ensure we don't hit the batch size + for _ in range(5): + response = await litellm.acompletion( + model="gpt-3.5-turbo", + messages=[{"role": "user", "content": "Test message"}], + max_tokens=10, + temperature=0.2, + mock_response="This is a mock response", + ) + + await asyncio.sleep(3) + + # Check that logs are in the queue + assert len(test_langsmith_logger.log_queue) == 5 + + # Now make calls to exceed the batch size + for _ in range(3): + response = await litellm.acompletion( + model="gpt-3.5-turbo", + messages=[{"role": "user", "content": "Test message"}], + max_tokens=10, + temperature=0.2, + mock_response="This is a mock response", + ) + + # Wait a short time for any asynchronous operations to complete + await asyncio.sleep(1) + + print( + "Length of langsmith log queue: {}".format( + len(test_langsmith_logger.log_queue) + ) + ) + # Check that the queue was flushed after exceeding batch size + assert len(test_langsmith_logger.log_queue) < 5 + + # Clean up + for cb in litellm.callbacks: + if isinstance(cb, LangsmithLogger): + await cb.async_httpx_client.client.aclose() + + except Exception as e: + pytest.fail(f"Error occurred: {e}") diff --git a/tests/logging_callback_tests/test_opentelemetry_unit_tests.py b/tests/logging_callback_tests/test_opentelemetry_unit_tests.py new file mode 100644 index 000000000000..b0d09562c576 --- /dev/null +++ b/tests/logging_callback_tests/test_opentelemetry_unit_tests.py @@ -0,0 +1,58 @@ +# What is this? +## Unit tests for opentelemetry integration + +# What is this? +## Unit test for presidio pii masking +import sys, os, asyncio, time, random +from datetime import datetime +import traceback +from dotenv import load_dotenv + +load_dotenv() +import os +import asyncio + +sys.path.insert( + 0, os.path.abspath("../..") +) # Adds the parent directory to the system path +import pytest +import litellm +from unittest.mock import patch, MagicMock, AsyncMock +from base_test import BaseLoggingCallbackTest +from litellm.types.utils import ModelResponse + + +class TestOpentelemetryUnitTests(BaseLoggingCallbackTest): + def test_parallel_tool_calls(self, mock_response_obj: ModelResponse): + tool_calls = mock_response_obj.choices[0].message.tool_calls + from litellm.integrations.opentelemetry import OpenTelemetry + from litellm.proxy._types import SpanAttributes + + kv_pair_dict = OpenTelemetry._tool_calls_kv_pair(tool_calls) + + assert kv_pair_dict == { + f"{SpanAttributes.LLM_COMPLETIONS}.0.function_call.arguments": '{"city": "New York"}', + f"{SpanAttributes.LLM_COMPLETIONS}.0.function_call.name": "get_weather", + f"{SpanAttributes.LLM_COMPLETIONS}.1.function_call.arguments": '{"city": "New York"}', + f"{SpanAttributes.LLM_COMPLETIONS}.1.function_call.name": "get_news", + } + + @pytest.mark.asyncio + async def test_opentelemetry_integration(self): + """ + Unit test to confirm the parent otel span is ended + """ + + parent_otel_span = MagicMock() + litellm.callbacks = ["otel"] + + await litellm.acompletion( + model="gpt-3.5-turbo", + messages=[{"role": "user", "content": "Hello, world!"}], + mock_response="Hey!", + metadata={"litellm_parent_otel_span": parent_otel_span}, + ) + + await asyncio.sleep(1) + + parent_otel_span.end.assert_called_once() diff --git a/tests/proxy_admin_ui_tests/test_route_check_unit_tests.py b/tests/proxy_admin_ui_tests/test_route_check_unit_tests.py index 9a4ec84671a5..001cc0640ef6 100644 --- a/tests/proxy_admin_ui_tests/test_route_check_unit_tests.py +++ b/tests/proxy_admin_ui_tests/test_route_check_unit_tests.py @@ -147,23 +147,6 @@ def test_key_info_route_allowed(route_checks): ) -def test_key_info_route_forbidden(route_checks): - """ - Internal User is not allowed to access /key/info route for a key they're not using in Authenticated API Key - """ - with pytest.raises(HTTPException) as exc_info: - route_checks.non_proxy_admin_allowed_routes_check( - user_obj=None, - _user_role=LitellmUserRoles.INTERNAL_USER.value, - route="/key/info", - request=MockRequest(query_params={"key": "wrong_key"}), - valid_token=UserAPIKeyAuth(api_key="test_key"), - api_key="test_key", - request_data={}, - ) - assert exc_info.value.status_code == 403 - - def test_user_info_route_allowed(route_checks): """ Internal User is allowed to access /user/info route for their own user_id diff --git a/tests/proxy_unit_tests/test_key_generate_prisma.py b/tests/proxy_unit_tests/test_key_generate_prisma.py index 66b9c7b8f698..78b558cd275d 100644 --- a/tests/proxy_unit_tests/test_key_generate_prisma.py +++ b/tests/proxy_unit_tests/test_key_generate_prisma.py @@ -456,7 +456,10 @@ async def return_body(): print("result from user auth with new key", result) # call /key/info for key - models == "all-proxy-models" - key_info = await info_key_fn(key=generated_key) + key_info = await info_key_fn( + key=generated_key, + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), + ) print("key_info", key_info) models = key_info["info"]["models"] assert models == ["all-team-models"] @@ -1179,7 +1182,12 @@ async def test(): generated_key = key.key # use generated key to auth in - result = await info_key_fn(key=generated_key) + result = await info_key_fn( + key=generated_key, + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, + ), + ) print("result from info_key_fn", result) assert result["key"] == generated_key print("\n info for key=", result["info"]) @@ -1271,7 +1279,12 @@ async def test(): generated_key = key.key # use generated key to auth in - result = await info_key_fn(key=generated_key) + result = await info_key_fn( + key=generated_key, + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, + ), + ) print("result from info_key_fn", result) assert result["key"] == generated_key print("\n info for key=", result["info"]) @@ -1303,7 +1316,12 @@ async def test(): print("response2=", response2) # get info on key after update - result = await info_key_fn(key=generated_key) + result = await info_key_fn( + key=generated_key, + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, + ), + ) print("result from info_key_fn", result) assert result["key"] == generated_key print("\n info for key=", result["info"]) @@ -1989,7 +2007,10 @@ async def test_key_name_null(prisma_client): key = await generate_key_fn(request) print("generated key=", key) generated_key = key.key - result = await info_key_fn(key=generated_key) + result = await info_key_fn( + key=generated_key, + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), + ) print("result from info_key_fn", result) assert result["info"]["key_name"] is None except Exception as e: @@ -2014,7 +2035,10 @@ async def test_key_name_set(prisma_client): request = GenerateKeyRequest() key = await generate_key_fn(request) generated_key = key.key - result = await info_key_fn(key=generated_key) + result = await info_key_fn( + key=generated_key, + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), + ) print("result from info_key_fn", result) assert isinstance(result["info"]["key_name"], str) except Exception as e: @@ -2038,7 +2062,10 @@ async def test_default_key_params(prisma_client): request = GenerateKeyRequest() key = await generate_key_fn(request) generated_key = key.key - result = await info_key_fn(key=generated_key) + result = await info_key_fn( + key=generated_key, + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), + ) print("result from info_key_fn", result) assert result["info"]["max_budget"] == 0.000122 except Exception as e: @@ -2804,7 +2831,10 @@ async def test_generate_key_with_model_tpm_limit(prisma_client): generated_key = key.key # use generated key to auth in - result = await info_key_fn(key=generated_key) + result = await info_key_fn( + key=generated_key, + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), + ) print("result from info_key_fn", result) assert result["key"] == generated_key print("\n info for key=", result["info"]) @@ -2825,7 +2855,10 @@ async def test_generate_key_with_model_tpm_limit(prisma_client): _request._url = URL(url="/update/key") await update_key_fn(data=request, request=_request) - result = await info_key_fn(key=generated_key) + result = await info_key_fn( + key=generated_key, + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), + ) print("result from info_key_fn", result) assert result["key"] == generated_key print("\n info for key=", result["info"]) @@ -2863,7 +2896,10 @@ async def test_generate_key_with_guardrails(prisma_client): generated_key = key.key # use generated key to auth in - result = await info_key_fn(key=generated_key) + result = await info_key_fn( + key=generated_key, + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), + ) print("result from info_key_fn", result) assert result["key"] == generated_key print("\n info for key=", result["info"]) @@ -2882,7 +2918,10 @@ async def test_generate_key_with_guardrails(prisma_client): _request._url = URL(url="/update/key") await update_key_fn(data=request, request=_request) - result = await info_key_fn(key=generated_key) + result = await info_key_fn( + key=generated_key, + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), + ) print("result from info_key_fn", result) assert result["key"] == generated_key print("\n info for key=", result["info"]) diff --git a/tests/proxy_unit_tests/test_proxy_server.py b/tests/proxy_unit_tests/test_proxy_server.py index 76cdf1a541e7..5588d04140f3 100644 --- a/tests/proxy_unit_tests/test_proxy_server.py +++ b/tests/proxy_unit_tests/test_proxy_server.py @@ -1632,6 +1632,139 @@ async def test_add_callback_via_key_litellm_pre_call_utils_gcs_bucket( assert new_data["failure_callback"] == expected_failure_callbacks +@pytest.mark.asyncio +@pytest.mark.parametrize( + "callback_type, expected_success_callbacks, expected_failure_callbacks", + [ + ("success", ["langsmith"], []), + ("failure", [], ["langsmith"]), + ("success_and_failure", ["langsmith"], ["langsmith"]), + ], +) +async def test_add_callback_via_key_litellm_pre_call_utils_langsmith( + prisma_client, callback_type, expected_success_callbacks, expected_failure_callbacks +): + import json + + from fastapi import HTTPException, Request, Response + from starlette.datastructures import URL + + from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request + + setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client) + setattr(litellm.proxy.proxy_server, "master_key", "sk-1234") + await litellm.proxy.proxy_server.prisma_client.connect() + + proxy_config = getattr(litellm.proxy.proxy_server, "proxy_config") + + request = Request(scope={"type": "http", "method": "POST", "headers": {}}) + request._url = URL(url="/chat/completions") + + test_data = { + "model": "azure/chatgpt-v-2", + "messages": [ + {"role": "user", "content": "write 1 sentence poem"}, + ], + "max_tokens": 10, + "mock_response": "Hello world", + "api_key": "my-fake-key", + } + + json_bytes = json.dumps(test_data).encode("utf-8") + + request._body = json_bytes + + data = { + "data": { + "model": "azure/chatgpt-v-2", + "messages": [{"role": "user", "content": "write 1 sentence poem"}], + "max_tokens": 10, + "mock_response": "Hello world", + "api_key": "my-fake-key", + }, + "request": request, + "user_api_key_dict": UserAPIKeyAuth( + token=None, + key_name=None, + key_alias=None, + spend=0.0, + max_budget=None, + expires=None, + models=[], + aliases={}, + config={}, + user_id=None, + team_id=None, + max_parallel_requests=None, + metadata={ + "logging": [ + { + "callback_name": "langsmith", + "callback_type": callback_type, + "callback_vars": { + "langsmith_api_key": "ls-1234", + "langsmith_project": "pr-brief-resemblance-72", + "langsmith_base_url": "https://api.smith.langchain.com", + }, + } + ] + }, + tpm_limit=None, + rpm_limit=None, + budget_duration=None, + budget_reset_at=None, + allowed_cache_controls=[], + permissions={}, + model_spend={}, + model_max_budget={}, + soft_budget_cooldown=False, + litellm_budget_table=None, + org_id=None, + team_spend=None, + team_alias=None, + team_tpm_limit=None, + team_rpm_limit=None, + team_max_budget=None, + team_models=[], + team_blocked=False, + soft_budget=None, + team_model_aliases=None, + team_member_spend=None, + team_metadata=None, + end_user_id=None, + end_user_tpm_limit=None, + end_user_rpm_limit=None, + end_user_max_budget=None, + last_refreshed_at=None, + api_key=None, + user_role=None, + allowed_model_region=None, + parent_otel_span=None, + ), + "proxy_config": proxy_config, + "general_settings": {}, + "version": "0.0.0", + } + + new_data = await add_litellm_data_to_request(**data) + print("NEW DATA: {}".format(new_data)) + + assert "langsmith_api_key" in new_data + assert new_data["langsmith_api_key"] == "ls-1234" + assert "langsmith_project" in new_data + assert new_data["langsmith_project"] == "pr-brief-resemblance-72" + assert "langsmith_base_url" in new_data + assert new_data["langsmith_base_url"] == "https://api.smith.langchain.com" + + if expected_success_callbacks: + assert "success_callback" in new_data + assert new_data["success_callback"] == expected_success_callbacks + + if expected_failure_callbacks: + assert "failure_callback" in new_data + assert new_data["failure_callback"] == expected_failure_callbacks + + @pytest.mark.asyncio async def test_gemini_pass_through_endpoint(): from starlette.datastructures import URL diff --git a/tests/proxy_unit_tests/test_proxy_utils.py b/tests/proxy_unit_tests/test_proxy_utils.py index f3f33bad6e87..2e857808d5c7 100644 --- a/tests/proxy_unit_tests/test_proxy_utils.py +++ b/tests/proxy_unit_tests/test_proxy_utils.py @@ -510,3 +510,23 @@ class ReturnValue(BaseModel): "success_callback": "langfuse", } } + + +def test_prepare_key_update_data(): + from litellm.proxy.management_endpoints.key_management_endpoints import ( + prepare_key_update_data, + ) + from litellm.proxy._types import UpdateKeyRequest + + existing_key_row = MagicMock() + data = UpdateKeyRequest(key="test_key", models=["gpt-4"], duration="120s") + updated_data = prepare_key_update_data(data, existing_key_row) + assert "expires" in updated_data + + data = UpdateKeyRequest(key="test_key", metadata={}) + updated_data = prepare_key_update_data(data, existing_key_row) + assert updated_data["metadata"] == {} + + data = UpdateKeyRequest(key="test_key", metadata=None) + updated_data = prepare_key_update_data(data, existing_key_row) + assert updated_data["metadata"] == None diff --git a/tests/test_keys.py b/tests/test_keys.py index 554a084c908a..a569634bc0c7 100644 --- a/tests/test_keys.py +++ b/tests/test_keys.py @@ -66,6 +66,7 @@ async def generate_key( max_parallel_requests: Optional[int] = None, user_id: Optional[str] = None, team_id: Optional[str] = None, + metadata: Optional[dict] = None, calling_key="sk-1234", ): url = "http://0.0.0.0:4000/key/generate" @@ -82,6 +83,7 @@ async def generate_key( "max_parallel_requests": max_parallel_requests, "user_id": user_id, "team_id": team_id, + "metadata": metadata, } print(f"data: {data}") @@ -136,16 +138,21 @@ async def test_key_gen_bad_key(): pass -async def update_key(session, get_key): +async def update_key(session, get_key, metadata: Optional[dict] = None): """ Make sure only models user has access to are returned """ url = "http://0.0.0.0:4000/key/update" headers = { - "Authorization": f"Bearer sk-1234", + "Authorization": "Bearer sk-1234", "Content-Type": "application/json", } - data = {"key": get_key, "models": ["gpt-4"], "duration": "120s"} + data = {"key": get_key} + + if metadata is not None: + data["metadata"] = metadata + else: + data.update({"models": ["gpt-4"], "duration": "120s"}) async with session.post(url, headers=headers, json=data) as response: status = response.status @@ -276,20 +283,24 @@ async def chat_completion_streaming(session, key, model="gpt-4"): return prompt_tokens, completion_tokens +@pytest.mark.parametrize("metadata", [{"test": "new"}, {}]) @pytest.mark.asyncio -async def test_key_update(): +async def test_key_update(metadata): """ Create key Update key with new model Test key w/ model """ async with aiohttp.ClientSession() as session: - key_gen = await generate_key(session=session, i=0) + key_gen = await generate_key(session=session, i=0, metadata={"test": "test"}) key = key_gen["key"] - await update_key( + assert key_gen["metadata"]["test"] == "test" + updated_key = await update_key( session=session, get_key=key, + metadata=metadata, ) + assert updated_key["metadata"] == metadata await update_proxy_budget(session=session) # resets proxy spend await chat_completion(session=session, key=key) @@ -412,7 +423,7 @@ async def test_key_info(): Get key info - as admin -> 200 - as key itself -> 200 - - as random key -> 403 + - as non existent key -> 404 """ async with aiohttp.ClientSession() as session: key_gen = await generate_key(session=session, i=0) @@ -425,10 +436,9 @@ async def test_key_info(): # as key itself, use the auth param, and no query key needed await get_key_info(session=session, call_key=key) # as random key # - key_gen = await generate_key(session=session, i=0) - random_key = key_gen["key"] - status = await get_key_info(session=session, get_key=key, call_key=random_key) - assert status == 403 + random_key = f"sk-{uuid.uuid4()}" + status = await get_key_info(session=session, get_key=random_key, call_key=key) + assert status == 404 @pytest.mark.asyncio diff --git a/tests/test_team_logging.py b/tests/test_team_logging.py index 97f18b42eedf..cf0fa6354858 100644 --- a/tests/test_team_logging.py +++ b/tests/test_team_logging.py @@ -62,8 +62,8 @@ async def chat_completion(session, key, model="azure-gpt-3.5", request_metadata= @pytest.mark.asyncio -@pytest.mark.flaky(retries=6, delay=1) -async def test_team_logging(): +@pytest.mark.flaky(retries=12, delay=2) +async def test_aaateam_logging(): """ -> Team 1 logs to project 1 -> Create Key diff --git a/ui/litellm-dashboard/src/components/model_dashboard.tsx b/ui/litellm-dashboard/src/components/model_dashboard.tsx index b09df5d7c498..cd915a9be222 100644 --- a/ui/litellm-dashboard/src/components/model_dashboard.tsx +++ b/ui/litellm-dashboard/src/components/model_dashboard.tsx @@ -151,6 +151,7 @@ enum Providers { Cohere = "Cohere", Databricks = "Databricks", Ollama = "Ollama", + xAI = "xAI", } const provider_map: Record = { @@ -166,6 +167,7 @@ const provider_map: Record = { OpenAI_Compatible: "openai", Vertex_AI: "vertex_ai", Databricks: "databricks", + xAI: "xai", Deepseek: "deepseek", Ollama: "ollama",