litellm-mirror/litellm/llms/base.py
Krish Dholakia aa9f1896c6 anthropic prompt caching cost tracking (#5453)
* fix(utils.py): support 'drop_params' for embedding requests

Fixes https://github.com/BerriAI/litellm/issues/5444

* feat(anthropic/cost_calculation.py): Support calculating cost for prompt caching on anthropic

* feat(types/utils.py): allows us to migrate to openai's equivalent, once that comes out

* fix: fix linting errors

* test: mark flaky test
2024-08-31 14:50:12 -07:00

87 lines
2.5 KiB
Python

## This is a template base class to be used for adding new LLM providers via API calls
from typing import Any, Optional, Union
import httpx
import requests
import litellm
class BaseLLM:
_client_session: Optional[httpx.Client] = None
def process_response(
self,
model: str,
response: Union[requests.Response, httpx.Response],
model_response: litellm.utils.ModelResponse,
stream: bool,
logging_obj: Any,
optional_params: dict,
api_key: str,
data: Union[dict, str],
messages: list,
print_verbose,
encoding,
) -> Union[litellm.utils.ModelResponse, litellm.utils.CustomStreamWrapper]:
"""
Helper function to process the response across sync + async completion calls
"""
return model_response
def process_text_completion_response(
self,
model: str,
response: Union[requests.Response, httpx.Response],
model_response: litellm.utils.TextCompletionResponse,
stream: bool,
logging_obj: Any,
optional_params: dict,
api_key: str,
data: Union[dict, str],
messages: list,
print_verbose,
encoding,
) -> Union[litellm.utils.TextCompletionResponse, litellm.utils.CustomStreamWrapper]:
"""
Helper function to process the response across sync + async completion calls
"""
return model_response
def create_client_session(self):
if litellm.client_session:
_client_session = litellm.client_session
else:
_client_session = httpx.Client()
return _client_session
def create_aclient_session(self):
if litellm.aclient_session:
_aclient_session = litellm.aclient_session
else:
_aclient_session = httpx.AsyncClient()
return _aclient_session
def __exit__(self):
if hasattr(self, "_client_session"):
self._client_session.close()
async def __aexit__(self, exc_type, exc_val, exc_tb):
if hasattr(self, "_aclient_session"):
await self._aclient_session.aclose()
def validate_environment(self): # set up the environment required to run the model
pass
def completion(
self, *args, **kwargs
): # logic for parsing in - calling - parsing out model completion calls
pass
def embedding(
self, *args, **kwargs
): # logic for parsing in - calling - parsing out model embedding calls
pass