tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-08-14 02:22:54 +00:00
Code Issues Packages Projects Releases Wiki Activity

[Refactor] Vector Stores - Use class VectorStorePreCallHook for all Vector Store Integrations (#12715)

Browse Source 
* add VectorStorePreCallHook

* vector_store_pre_call_hook

* add pop_vector_stores_to_run

* async_get_chat_completion_prompt

* working e2e tests

* test_e2e_bedrock_knowledgebase_retrieval_with_completion

* delete old files

* fix logging test

* VectorStorePreCallHook

* fix ruff check

* vector_store_pre_call_hook

* linting error fixes
This commit is contained in:
Ishaan Jaff
2025-07-17 16:31:58 -07:00
committed by GitHub
parent 52b2984792
commit 33c84846e9
10 changed files with 430 additions and 944 deletions
Download Patch File Download Diff File Expand all files Collapse all files
litellm/__init__.py
+1 -1
View File
@@ -122,13 +122,13 @@ _custom_logger_compatible_callbacks_literal = Literal[
"gcs_pubsub",
"agentops",
"anthropic_cache_control_hook",
"bedrock_vector_store",
"generic_api",
"resend_email",
"smtp_email",
"deepeval",
"s3_v2",
"aws_sqs",
"vector_store_pre_call_hook",
]
logged_real_time_event_types: Optional[Union[List[str], Literal["*"]]] = None
_known_custom_logger_compatible_callbacks: List = list(
litellm/integrations/vector_store_integrations/bedrock_vector_store.py
-410
View File
@@ -1,410 +0,0 @@
# +-------------------------------------------------------------+
#
# Add Bedrock Knowledge Base Context to your LLM calls
#
# +-------------------------------------------------------------+
# Thank you users! We ❤️ you! - Krrish & Ishaan
import json
from datetime import datetime
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple
import litellm
from litellm._logging import verbose_logger, verbose_proxy_logger
from litellm.integrations.custom_logger import CustomLogger
from litellm.integrations.vector_store_integrations.base_vector_store import (
BaseVectorStore,
)
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
from litellm.llms.custom_httpx.http_handler import (
get_async_httpx_client,
httpxSpecialProvider,
)
from litellm.types.integrations.rag.bedrock_knowledgebase import (
BedrockKBContent,
BedrockKBGuardrailConfiguration,
BedrockKBRequest,
BedrockKBResponse,
BedrockKBRetrievalConfiguration,
BedrockKBRetrievalQuery,
BedrockKBRetrievalResult,
)
from litellm.types.llms.openai import AllMessageValues, ChatCompletionUserMessage
from litellm.types.utils import StandardLoggingVectorStoreRequest
from litellm.types.vector_stores import (
VectorStoreResultContent,
VectorStoreSearchResponse,
VectorStoreSearchResult,
)
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
else:
LiteLLMLoggingObj = Any
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import StandardCallbackDynamicParams
else:
StandardCallbackDynamicParams = Any
class BedrockVectorStore(BaseVectorStore, BaseAWSLLM):
CONTENT_PREFIX_STRING = "Context: \n\n"
CUSTOM_LLM_PROVIDER = "bedrock"
def __init__(
self,
**kwargs,
):
self.async_handler = get_async_httpx_client(
llm_provider=httpxSpecialProvider.LoggingCallback
)
# store kwargs as optional_params
self.optional_params = kwargs
super().__init__(**kwargs)
BaseAWSLLM.__init__(self)
async def async_get_chat_completion_prompt(
self,
model: str,
messages: List[AllMessageValues],
non_default_params: dict,
prompt_id: Optional[str],
prompt_variables: Optional[dict],
dynamic_callback_params: StandardCallbackDynamicParams,
litellm_logging_obj: LiteLLMLoggingObj,
tools: Optional[List[Dict]] = None,
prompt_label: Optional[str] = None,
prompt_version: Optional[int] = None,
) -> Tuple[str, List[AllMessageValues], dict]:
"""
Retrieves the context from the Bedrock Knowledge Base and appends it to the messages.
"""
if litellm.vector_store_registry is None:
return model, messages, non_default_params
vector_store_ids = litellm.vector_store_registry.pop_vector_store_ids_to_run(
non_default_params=non_default_params, tools=tools
)
vector_store_request_metadata: List[StandardLoggingVectorStoreRequest] = []
if vector_store_ids:
for vector_store_id in vector_store_ids:
start_time = datetime.now()
query = self._get_kb_query_from_messages(messages)
bedrock_kb_response = await self.make_bedrock_kb_retrieve_request(
knowledge_base_id=vector_store_id,
query=query,
non_default_params=non_default_params,
)
verbose_logger.debug(
f"Bedrock Knowledge Base Response: {bedrock_kb_response}"
)
(
context_message,
context_string,
) = self.get_chat_completion_message_from_bedrock_kb_response(
bedrock_kb_response
)
if context_message is not None:
messages.append(context_message)
#################################################################################################
########## LOGGING for Standard Logging Payload, Langfuse, s3, LiteLLM DB etc. ##################
#################################################################################################
vector_store_search_response: VectorStoreSearchResponse = (
self.transform_bedrock_kb_response_to_vector_store_search_response(
bedrock_kb_response=bedrock_kb_response, query=query
)
)
vector_store_request_metadata.append(
StandardLoggingVectorStoreRequest(
vector_store_id=vector_store_id,
query=query,
vector_store_search_response=vector_store_search_response,
custom_llm_provider=self.CUSTOM_LLM_PROVIDER,
start_time=start_time.timestamp(),
end_time=datetime.now().timestamp(),
)
)
litellm_logging_obj.model_call_details["vector_store_request_metadata"] = (
vector_store_request_metadata
)
return model, messages, non_default_params
def transform_bedrock_kb_response_to_vector_store_search_response(
self,
bedrock_kb_response: BedrockKBResponse,
query: str,
) -> VectorStoreSearchResponse:
"""
Transform a BedrockKBResponse to a VectorStoreSearchResponse
"""
retrieval_results: Optional[List[BedrockKBRetrievalResult]] = (
bedrock_kb_response.get("retrievalResults", None)
)
vector_store_search_response: VectorStoreSearchResponse = (
VectorStoreSearchResponse(search_query=query, data=[])
)
if retrieval_results is None:
return vector_store_search_response
vector_search_response_data: List[VectorStoreSearchResult] = []
for retrieval_result in retrieval_results:
content: Optional[BedrockKBContent] = retrieval_result.get("content", None)
if content is None:
continue
content_text: Optional[str] = content.get("text", None)
if content_text is None:
continue
vector_store_search_result: VectorStoreSearchResult = (
VectorStoreSearchResult(
score=retrieval_result.get("score", None),
content=[VectorStoreResultContent(text=content_text, type="text")],
)
)
vector_search_response_data.append(vector_store_search_result)
vector_store_search_response["data"] = vector_search_response_data
return vector_store_search_response
def _get_kb_query_from_messages(self, messages: List[AllMessageValues]) -> str:
"""
Uses the text `content` field of the last message in the list of messages
"""
if len(messages) == 0:
return ""
last_message = messages[-1]
last_message_content = last_message.get("content", None)
if last_message_content is None:
return ""
if isinstance(last_message_content, str):
return last_message_content
elif isinstance(last_message_content, list):
return "\n".join([item.get("text", "") for item in last_message_content])
return ""
def _prepare_request(
self,
credentials: Any,
data: BedrockKBRequest,
optional_params: dict,
aws_region_name: str,
api_base: str,
extra_headers: Optional[dict] = None,
) -> Any:
"""
Prepare a signed AWS request.
Args:
credentials: AWS credentials
data: Request data
optional_params: Additional parameters
aws_region_name: AWS region name
api_base: Base API URL
extra_headers: Additional headers
Returns:
AWSRequest: A signed AWS request
"""
try:
from botocore.auth import SigV4Auth
from botocore.awsrequest import AWSRequest
except ImportError:
raise ImportError("Missing boto3 to call bedrock. Run 'pip install boto3'.")
sigv4 = SigV4Auth(credentials, "bedrock", aws_region_name)
encoded_data = json.dumps(data).encode("utf-8")
headers = {"Content-Type": "application/json"}
if extra_headers is not None:
headers = {"Content-Type": "application/json", **extra_headers}
request = AWSRequest(
method="POST", url=api_base, data=encoded_data, headers=headers
)
sigv4.add_auth(request)
if extra_headers is not None and "Authorization" in extra_headers:
# prevent sigv4 from overwriting the auth header
request.headers["Authorization"] = extra_headers["Authorization"]
return request.prepare()
async def make_bedrock_kb_retrieve_request(
self,
knowledge_base_id: str,
query: str,
guardrail_id: Optional[str] = None,
guardrail_version: Optional[str] = None,
next_token: Optional[str] = None,
retrieval_configuration: Optional[BedrockKBRetrievalConfiguration] = None,
non_default_params: Optional[dict] = None,
) -> BedrockKBResponse:
"""
Make a Bedrock Knowledge Base retrieve request.
Args:
knowledge_base_id (str): The unique identifier of the knowledge base to query
query (str): The query text to search for
guardrail_id (Optional[str]): The guardrail ID to apply
guardrail_version (Optional[str]): The version of the guardrail to apply
next_token (Optional[str]): Token for pagination
retrieval_configuration (Optional[BedrockKBRetrievalConfiguration]): Configuration for the retrieval process
Returns:
BedrockKBRetrievalResponse: A typed response object containing the retrieval results
"""
from fastapi import HTTPException
non_default_params = non_default_params or {}
credentials_dict: Dict[str, Any] = {}
if litellm.vector_store_registry is not None:
credentials_dict = (
litellm.vector_store_registry.get_credentials_for_vector_store(
knowledge_base_id
)
)
credentials = self.get_credentials(
aws_access_key_id=credentials_dict.get(
"aws_access_key_id", non_default_params.get("aws_access_key_id", None)
),
aws_secret_access_key=credentials_dict.get(
"aws_secret_access_key",
non_default_params.get("aws_secret_access_key", None),
),
aws_session_token=credentials_dict.get(
"aws_session_token", non_default_params.get("aws_session_token", None)
),
aws_region_name=credentials_dict.get(
"aws_region_name", non_default_params.get("aws_region_name", None)
),
aws_session_name=credentials_dict.get(
"aws_session_name", non_default_params.get("aws_session_name", None)
),
aws_profile_name=credentials_dict.get(
"aws_profile_name", non_default_params.get("aws_profile_name", None)
),
aws_role_name=credentials_dict.get(
"aws_role_name", non_default_params.get("aws_role_name", None)
),
aws_web_identity_token=credentials_dict.get(
"aws_web_identity_token",
non_default_params.get("aws_web_identity_token", None),
),
aws_sts_endpoint=credentials_dict.get(
"aws_sts_endpoint", non_default_params.get("aws_sts_endpoint", None)
),
)
aws_region_name = self.get_aws_region_name_for_non_llm_api_calls(
aws_region_name=credentials_dict.get(
"aws_region_name", non_default_params.get("aws_region_name", None)
),
)
# Prepare request data
request_data: BedrockKBRequest = BedrockKBRequest(
retrievalQuery=BedrockKBRetrievalQuery(text=query),
)
if next_token:
request_data["nextToken"] = next_token
if retrieval_configuration:
request_data["retrievalConfiguration"] = retrieval_configuration
if guardrail_id and guardrail_version:
request_data["guardrailConfiguration"] = BedrockKBGuardrailConfiguration(
guardrailId=guardrail_id, guardrailVersion=guardrail_version
)
verbose_logger.debug(
f"Request Data: {json.dumps(request_data, indent=4, default=str)}"
)
# Prepare the request
api_base = f"https://bedrock-agent-runtime.{aws_region_name}.amazonaws.com/knowledgebases/{knowledge_base_id}/retrieve"
prepared_request = self._prepare_request(
credentials=credentials,
data=request_data,
optional_params=self.optional_params,
aws_region_name=aws_region_name,
api_base=api_base,
)
verbose_proxy_logger.debug(
"Bedrock Knowledge Base request body: %s, url %s, headers: %s",
request_data,
prepared_request.url,
prepared_request.headers,
)
response = await self.async_handler.post(
url=prepared_request.url,
data=prepared_request.body, # type: ignore
headers=prepared_request.headers, # type: ignore
)
verbose_proxy_logger.debug("Bedrock Knowledge Base response: %s", response.text)
if response.status_code == 200:
response_data = response.json()
return BedrockKBResponse(**response_data)
else:
verbose_proxy_logger.error(
"Bedrock Knowledge Base: error in response. Status code: %s, response: %s",
response.status_code,
response.text,
)
raise HTTPException(
status_code=response.status_code,
detail={
"error": "Error calling Bedrock Knowledge Base",
"response": response.text,
},
)
@staticmethod
def get_initialized_custom_logger() -> Optional[CustomLogger]:
from litellm.litellm_core_utils.litellm_logging import (
_init_custom_logger_compatible_class,
)
return _init_custom_logger_compatible_class(
logging_integration="bedrock_vector_store",
internal_usage_cache=None,
llm_router=None,
)
@staticmethod
def get_chat_completion_message_from_bedrock_kb_response(
response: BedrockKBResponse,
) -> Tuple[Optional[ChatCompletionUserMessage], str]:
"""
Retrieves the context from the Bedrock Knowledge Base response and returns a ChatCompletionUserMessage object.
"""
retrieval_results: Optional[List[BedrockKBRetrievalResult]] = response.get(
"retrievalResults", None
)
if retrieval_results is None:
return None, ""
# string to combine the context from the knowledge base
context_string: str = BedrockVectorStore.CONTENT_PREFIX_STRING
for retrieval_result in retrieval_results:
retrieval_result_content: Optional[BedrockKBContent] = (
retrieval_result.get("content", None) or {}
)
if retrieval_result_content is None:
continue
retrieval_result_text: Optional[str] = retrieval_result_content.get(
"text", None
)
if retrieval_result_text is None:
continue
context_string += retrieval_result_text
message = ChatCompletionUserMessage(
role="user",
content=context_string,
)
return message, context_string
litellm/integrations/vector_store_integrations/vector_store_pre_call_hook.py
+194
View File
@@ -0,0 +1,194 @@
"""
Vector Store Pre-Call Hook
This hook is called before making an LLM request when a vector store is configured.
It searches the vector store for relevant context and appends it to the messages.
"""
from typing import TYPE_CHECKING, Dict, List, Optional, Tuple, cast
import litellm
from litellm._logging import verbose_logger
from litellm.integrations.custom_logger import CustomLogger
from litellm.types.llms.openai import AllMessageValues, ChatCompletionUserMessage
from litellm.types.utils import StandardCallbackDynamicParams
from litellm.types.vector_stores import (
LiteLLM_ManagedVectorStore,
VectorStoreResultContent,
VectorStoreSearchResponse,
VectorStoreSearchResult,
)
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
else:
LiteLLMLoggingObj = None
class VectorStorePreCallHook(CustomLogger):
CONTENT_PREFIX_STRING = "Context:\n\n"
"""
Custom logger that handles vector store searches before LLM calls.
When a vector store is configured, this hook:
1. Extracts the query from the last user message
2. Calls litellm.vector_stores.search() to get relevant context
3. Appends the search results as context to the messages
"""
def __init__(self):
super().__init__()
async def async_get_chat_completion_prompt(
self,
model: str,
messages: List[AllMessageValues],
non_default_params: dict,
prompt_id: Optional[str],
prompt_variables: Optional[dict],
dynamic_callback_params: StandardCallbackDynamicParams,
litellm_logging_obj: LiteLLMLoggingObj,
tools: Optional[List[Dict]] = None,
prompt_label: Optional[str] = None,
prompt_version: Optional[int] = None,
) -> Tuple[str, List[AllMessageValues], dict]:
"""
Perform vector store search and append results as context to messages.
Args:
model: The model name
messages: List of messages
non_default_params: Non-default parameters
prompt_id: Optional prompt ID
prompt_variables: Optional prompt variables
dynamic_callback_params: Optional dynamic callback parameters
prompt_label: Optional prompt label
prompt_version: Optional prompt version
Returns:
Tuple of (model, modified_messages, non_default_params)
"""
try:
# Check if vector store is configured
if litellm.vector_store_registry is None:
return model, messages, non_default_params
vector_stores_to_run: List[LiteLLM_ManagedVectorStore] = litellm.vector_store_registry.pop_vector_stores_to_run(
non_default_params=non_default_params, tools=tools
)
if not vector_stores_to_run:
return model, messages, non_default_params
# Extract the query from the last user message
query = self._extract_query_from_messages(messages)
if not query:
verbose_logger.debug("No query found in messages for vector store search")
return model, messages, non_default_params
modified_messages: List[AllMessageValues] = messages.copy()
for vector_store_to_run in vector_stores_to_run:
# Get vector store id from the vector store config
vector_store_id = vector_store_to_run.get("vector_store_id", "")
custom_llm_provider = vector_store_to_run.get("custom_llm_provider")
# Call litellm.vector_stores.search() with the required parameters
search_response = await litellm.vector_stores.asearch(
vector_store_id=vector_store_id,
query=query,
custom_llm_provider=custom_llm_provider,
)
verbose_logger.debug(f"search_response: {search_response}")
# Process search results and append as context
modified_messages = self._append_search_results_to_messages(
messages=messages,
search_response=search_response
)
# Get the number of results for logging
num_results = 0
num_results = len(search_response.get("data", []) or [])
verbose_logger.debug(f"Vector store search completed. Added context from {num_results} results")
return model, modified_messages, non_default_params
except Exception as e:
verbose_logger.exception(f"Error in VectorStorePreCallHook: {str(e)}")
# Return original parameters on error
return model, messages, non_default_params
def _extract_query_from_messages(self, messages: List[AllMessageValues]) -> Optional[str]:
"""
Extract the query from the last user message.
Args:
messages: List of messages
Returns:
The extracted query string or None if not found
"""
if not messages or len(messages) == 0:
return None
last_message = messages[-1]
if not isinstance(last_message, dict) or "content" not in last_message:
return None
content = last_message["content"]
if isinstance(content, str):
return content
elif isinstance(content, list) and len(content) > 0:
# Handle list of content items, extract text from first text item
for item in content:
if isinstance(item, dict) and item.get("type") == "text" and "text" in item:
return item["text"]
return None
def _append_search_results_to_messages(
self,
messages: List[AllMessageValues],
search_response: VectorStoreSearchResponse
) -> List[AllMessageValues]:
"""
Append search results as context to the messages.
Args:
messages: Original list of messages
search_response: Response from vector store search
Returns:
Modified list of messages with context appended
"""
search_response_data: Optional[List[VectorStoreSearchResult]] = search_response.get("data")
if not search_response_data:
return messages
context_content = self.CONTENT_PREFIX_STRING
for result in search_response_data:
result_content: Optional[List[VectorStoreResultContent]] = result.get("content")
if result_content:
for content_item in result_content:
content_text: Optional[str] = content_item.get("text")
if content_text:
context_content += content_text + "\n\n"
# Only add context if we found any content
if context_content != "Context:\n\n":
# Create a copy of messages to avoid modifying the original
modified_messages = messages.copy()
# Add context as a new message before the last user message
context_message: ChatCompletionUserMessage = {
"role": "user",
"content": context_content
}
modified_messages.insert(-1, cast(AllMessageValues, context_message))
return modified_messages
return messages
litellm/integrations/vector_stores/bedrock_vector_store.py
-410
View File
@@ -1,410 +0,0 @@
# +-------------------------------------------------------------+
#
# Add Bedrock Knowledge Base Context to your LLM calls
#
# +-------------------------------------------------------------+
# Thank you users! We ❤️ you! - Krrish & Ishaan
import json
from datetime import datetime
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple
import litellm
from litellm._logging import verbose_logger, verbose_proxy_logger
from litellm.integrations.custom_logger import CustomLogger
from litellm.integrations.vector_store_integrations.base_vector_store import (
BaseVectorStore,
)
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
from litellm.llms.custom_httpx.http_handler import (
get_async_httpx_client,
httpxSpecialProvider,
)
from litellm.types.integrations.rag.bedrock_knowledgebase import (
BedrockKBContent,
BedrockKBGuardrailConfiguration,
BedrockKBRequest,
BedrockKBResponse,
BedrockKBRetrievalConfiguration,
BedrockKBRetrievalQuery,
BedrockKBRetrievalResult,
)
from litellm.types.llms.openai import AllMessageValues, ChatCompletionUserMessage
from litellm.types.utils import StandardLoggingVectorStoreRequest
from litellm.types.vector_stores import (
VectorStoreResultContent,
VectorStoreSearchResponse,
VectorStoreSearchResult,
)
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
else:
LiteLLMLoggingObj = Any
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import StandardCallbackDynamicParams
else:
StandardCallbackDynamicParams = Any
class BedrockVectorStore(BaseVectorStore, BaseAWSLLM):
CONTENT_PREFIX_STRING = "Context: \n\n"
CUSTOM_LLM_PROVIDER = "bedrock"
def __init__(
self,
**kwargs,
):
self.async_handler = get_async_httpx_client(
llm_provider=httpxSpecialProvider.LoggingCallback
)
# store kwargs as optional_params
self.optional_params = kwargs
super().__init__(**kwargs)
BaseAWSLLM.__init__(self)
async def async_get_chat_completion_prompt(
self,
model: str,
messages: List[AllMessageValues],
non_default_params: dict,
prompt_id: Optional[str],
prompt_variables: Optional[dict],
dynamic_callback_params: StandardCallbackDynamicParams,
litellm_logging_obj: LiteLLMLoggingObj,
tools: Optional[List[Dict]] = None,
prompt_label: Optional[str] = None,
prompt_version: Optional[int] = None,
) -> Tuple[str, List[AllMessageValues], dict]:
"""
Retrieves the context from the Bedrock Knowledge Base and appends it to the messages.
"""
if litellm.vector_store_registry is None:
return model, messages, non_default_params
vector_store_ids = litellm.vector_store_registry.pop_vector_store_ids_to_run(
non_default_params=non_default_params, tools=tools
)
vector_store_request_metadata: List[StandardLoggingVectorStoreRequest] = []
if vector_store_ids:
for vector_store_id in vector_store_ids:
start_time = datetime.now()
query = self._get_kb_query_from_messages(messages)
bedrock_kb_response = await self.make_bedrock_kb_retrieve_request(
knowledge_base_id=vector_store_id,
query=query,
non_default_params=non_default_params,
)
verbose_logger.debug(
f"Bedrock Knowledge Base Response: {bedrock_kb_response}"
)
(
context_message,
context_string,
) = self.get_chat_completion_message_from_bedrock_kb_response(
bedrock_kb_response
)
if context_message is not None:
messages.append(context_message)
#################################################################################################
########## LOGGING for Standard Logging Payload, Langfuse, s3, LiteLLM DB etc. ##################
#################################################################################################
vector_store_search_response: VectorStoreSearchResponse = (
self.transform_bedrock_kb_response_to_vector_store_search_response(
bedrock_kb_response=bedrock_kb_response, query=query
)
)
vector_store_request_metadata.append(
StandardLoggingVectorStoreRequest(
vector_store_id=vector_store_id,
query=query,
vector_store_search_response=vector_store_search_response,
custom_llm_provider=self.CUSTOM_LLM_PROVIDER,
start_time=start_time.timestamp(),
end_time=datetime.now().timestamp(),
)
)
litellm_logging_obj.model_call_details["vector_store_request_metadata"] = (
vector_store_request_metadata
)
return model, messages, non_default_params
def transform_bedrock_kb_response_to_vector_store_search_response(
self,
bedrock_kb_response: BedrockKBResponse,
query: str,
) -> VectorStoreSearchResponse:
"""
Transform a BedrockKBResponse to a VectorStoreSearchResponse
"""
retrieval_results: Optional[List[BedrockKBRetrievalResult]] = (
bedrock_kb_response.get("retrievalResults", None)
)
vector_store_search_response: VectorStoreSearchResponse = (
VectorStoreSearchResponse(search_query=query, data=[])
)
if retrieval_results is None:
return vector_store_search_response
vector_search_response_data: List[VectorStoreSearchResult] = []
for retrieval_result in retrieval_results:
content: Optional[BedrockKBContent] = retrieval_result.get("content", None)
if content is None:
continue
content_text: Optional[str] = content.get("text", None)
if content_text is None:
continue
vector_store_search_result: VectorStoreSearchResult = (
VectorStoreSearchResult(
score=retrieval_result.get("score", None),
content=[VectorStoreResultContent(text=content_text, type="text")],
)
)
vector_search_response_data.append(vector_store_search_result)
vector_store_search_response["data"] = vector_search_response_data
return vector_store_search_response
def _get_kb_query_from_messages(self, messages: List[AllMessageValues]) -> str:
"""
Uses the text `content` field of the last message in the list of messages
"""
if len(messages) == 0:
return ""
last_message = messages[-1]
last_message_content = last_message.get("content", None)
if last_message_content is None:
return ""
if isinstance(last_message_content, str):
return last_message_content
elif isinstance(last_message_content, list):
return "\n".join([item.get("text", "") for item in last_message_content])
return ""
def _prepare_request(
self,
credentials: Any,
data: BedrockKBRequest,
optional_params: dict,
aws_region_name: str,
api_base: str,
extra_headers: Optional[dict] = None,
) -> Any:
"""
Prepare a signed AWS request.
Args:
credentials: AWS credentials
data: Request data
optional_params: Additional parameters
aws_region_name: AWS region name
api_base: Base API URL
extra_headers: Additional headers
Returns:
AWSRequest: A signed AWS request
"""
try:
from botocore.auth import SigV4Auth
from botocore.awsrequest import AWSRequest
except ImportError:
raise ImportError("Missing boto3 to call bedrock. Run 'pip install boto3'.")
sigv4 = SigV4Auth(credentials, "bedrock", aws_region_name)
encoded_data = json.dumps(data).encode("utf-8")
headers = {"Content-Type": "application/json"}
if extra_headers is not None:
headers = {"Content-Type": "application/json", **extra_headers}
request = AWSRequest(
method="POST", url=api_base, data=encoded_data, headers=headers
)
sigv4.add_auth(request)
if extra_headers is not None and "Authorization" in extra_headers:
# prevent sigv4 from overwriting the auth header
request.headers["Authorization"] = extra_headers["Authorization"]
return request.prepare()
async def make_bedrock_kb_retrieve_request(
self,
knowledge_base_id: str,
query: str,
guardrail_id: Optional[str] = None,
guardrail_version: Optional[str] = None,
next_token: Optional[str] = None,
retrieval_configuration: Optional[BedrockKBRetrievalConfiguration] = None,
non_default_params: Optional[dict] = None,
) -> BedrockKBResponse:
"""
Make a Bedrock Knowledge Base retrieve request.
Args:
knowledge_base_id (str): The unique identifier of the knowledge base to query
query (str): The query text to search for
guardrail_id (Optional[str]): The guardrail ID to apply
guardrail_version (Optional[str]): The version of the guardrail to apply
next_token (Optional[str]): Token for pagination
retrieval_configuration (Optional[BedrockKBRetrievalConfiguration]): Configuration for the retrieval process
Returns:
BedrockKBRetrievalResponse: A typed response object containing the retrieval results
"""
from fastapi import HTTPException
non_default_params = non_default_params or {}
credentials_dict: Dict[str, Any] = {}
if litellm.vector_store_registry is not None:
credentials_dict = (
litellm.vector_store_registry.get_credentials_for_vector_store(
knowledge_base_id
)
)
credentials = self.get_credentials(
aws_access_key_id=credentials_dict.get(
"aws_access_key_id", non_default_params.get("aws_access_key_id", None)
),
aws_secret_access_key=credentials_dict.get(
"aws_secret_access_key",
non_default_params.get("aws_secret_access_key", None),
),
aws_session_token=credentials_dict.get(
"aws_session_token", non_default_params.get("aws_session_token", None)
),
aws_region_name=credentials_dict.get(
"aws_region_name", non_default_params.get("aws_region_name", None)
),
aws_session_name=credentials_dict.get(
"aws_session_name", non_default_params.get("aws_session_name", None)
),
aws_profile_name=credentials_dict.get(
"aws_profile_name", non_default_params.get("aws_profile_name", None)
),
aws_role_name=credentials_dict.get(
"aws_role_name", non_default_params.get("aws_role_name", None)
),
aws_web_identity_token=credentials_dict.get(
"aws_web_identity_token",
non_default_params.get("aws_web_identity_token", None),
),
aws_sts_endpoint=credentials_dict.get(
"aws_sts_endpoint", non_default_params.get("aws_sts_endpoint", None)
),
)
aws_region_name = self.get_aws_region_name_for_non_llm_api_calls(
aws_region_name=credentials_dict.get(
"aws_region_name", non_default_params.get("aws_region_name", None)
),
)
# Prepare request data
request_data: BedrockKBRequest = BedrockKBRequest(
retrievalQuery=BedrockKBRetrievalQuery(text=query),
)
if next_token:
request_data["nextToken"] = next_token
if retrieval_configuration:
request_data["retrievalConfiguration"] = retrieval_configuration
if guardrail_id and guardrail_version:
request_data["guardrailConfiguration"] = BedrockKBGuardrailConfiguration(
guardrailId=guardrail_id, guardrailVersion=guardrail_version
)
verbose_logger.debug(
f"Request Data: {json.dumps(request_data, indent=4, default=str)}"
)
# Prepare the request
api_base = f"https://bedrock-agent-runtime.{aws_region_name}.amazonaws.com/knowledgebases/{knowledge_base_id}/retrieve"
prepared_request = self._prepare_request(
credentials=credentials,
data=request_data,
optional_params=self.optional_params,
aws_region_name=aws_region_name,
api_base=api_base,
)
verbose_proxy_logger.debug(
"Bedrock Knowledge Base request body: %s, url %s, headers: %s",
request_data,
prepared_request.url,
prepared_request.headers,
)
response = await self.async_handler.post(
url=prepared_request.url,
data=prepared_request.body, # type: ignore
headers=prepared_request.headers, # type: ignore
)
verbose_proxy_logger.debug("Bedrock Knowledge Base response: %s", response.text)
if response.status_code == 200:
response_data = response.json()
return BedrockKBResponse(**response_data)
else:
verbose_proxy_logger.error(
"Bedrock Knowledge Base: error in response. Status code: %s, response: %s",
response.status_code,
response.text,
)
raise HTTPException(
status_code=response.status_code,
detail={
"error": "Error calling Bedrock Knowledge Base",
"response": response.text,
},
)
@staticmethod
def get_initialized_custom_logger() -> Optional[CustomLogger]:
from litellm.litellm_core_utils.litellm_logging import (
_init_custom_logger_compatible_class,
)
return _init_custom_logger_compatible_class(
logging_integration="bedrock_vector_store",
internal_usage_cache=None,
llm_router=None,
)
@staticmethod
def get_chat_completion_message_from_bedrock_kb_response(
response: BedrockKBResponse,
) -> Tuple[Optional[ChatCompletionUserMessage], str]:
"""
Retrieves the context from the Bedrock Knowledge Base response and returns a ChatCompletionUserMessage object.
"""
retrieval_results: Optional[List[BedrockKBRetrievalResult]] = response.get(
"retrievalResults", None
)
if retrieval_results is None:
return None, ""
# string to combine the context from the knowledge base
context_string: str = BedrockVectorStore.CONTENT_PREFIX_STRING
for retrieval_result in retrieval_results:
retrieval_result_content: Optional[BedrockKBContent] = (
retrieval_result.get("content", None) or {}
)
if retrieval_result_content is None:
continue
retrieval_result_text: Optional[str] = retrieval_result_content.get(
"text", None
)
if retrieval_result_text is None:
continue
context_string += retrieval_result_text
message = ChatCompletionUserMessage(
role="user",
content=context_string,
)
return message, context_string
litellm/litellm_core_utils/custom_logger_registry.py
+4 -3
View File
@@ -8,6 +8,7 @@ Example:
"prometheus" -> PrometheusLogger
"""
from typing import Union
from litellm.integrations.agentops import AgentOps
from litellm.integrations.anthropic_cache_control_hook import AnthropicCacheControlHook
from litellm.integrations.argilla import ArgillaLogger
@@ -33,8 +34,8 @@ from litellm.integrations.opik.opik import OpikLogger
from litellm.integrations.prometheus import PrometheusLogger
from litellm.integrations.s3_v2 import S3Logger
from litellm.integrations.sqs import SQSLogger
from litellm.integrations.vector_store_integrations.bedrock_vector_store import (
BedrockVectorStore,
from litellm.integrations.vector_store_integrations.vector_store_pre_call_hook import (
VectorStorePreCallHook,
)
from litellm.proxy.hooks.dynamic_rate_limiter import _PROXY_DynamicRateLimitHandler
@@ -71,11 +72,11 @@ class CustomLoggerRegistry:
"gcs_pubsub": GcsPubSubLogger,
"anthropic_cache_control_hook": AnthropicCacheControlHook,
"agentops": AgentOps,
"bedrock_vector_store": BedrockVectorStore,
"deepeval": DeepEvalLogger,
"s3_v2": S3Logger,
"aws_sqs": SQSLogger,
"dynamic_rate_limiter": _PROXY_DynamicRateLimitHandler,
"vector_store_pre_call_hook": VectorStorePreCallHook,
}
try:
litellm/litellm_core_utils/litellm_logging.py
+21 -21
View File
@@ -59,9 +59,6 @@ from litellm.integrations.custom_logger import CustomLogger
from litellm.integrations.deepeval.deepeval import DeepEvalLogger
from litellm.integrations.mlflow import MlflowLogger
from litellm.integrations.sqs import SQSLogger
from litellm.integrations.vector_store_integrations.bedrock_vector_store import (
BedrockVectorStore,
)
from litellm.litellm_core_utils.get_litellm_params import get_litellm_params
from litellm.litellm_core_utils.llm_cost_calc.tool_call_cost_tracking import (
StandardBuiltInToolCostTracking,
@@ -674,17 +671,11 @@ class Logging(LiteLLMLoggingBaseClass):
# Vector Store / Knowledge Base hooks
#########################################################
if litellm.vector_store_registry is not None:
if vector_store_to_run := litellm.vector_store_registry.get_vector_store_to_run(
non_default_params=non_default_params,
tools=tools
):
vector_store_custom_logger = (
litellm.ProviderConfigManager.get_provider_vector_store_config(
provider=cast(
litellm.LlmProviders,
vector_store_to_run.get("custom_llm_provider"),
),
)
vector_store_custom_logger = _init_custom_logger_compatible_class(
logging_integration="vector_store_pre_call_hook",
internal_usage_cache=None,
llm_router=None,
)
self.model_call_details["prompt_integration"] = (
vector_store_custom_logger.__class__.__name__
@@ -3139,6 +3130,7 @@ def set_callbacks(callback_list, function_id=None): # noqa: PLR0915
customLogger = CustomLogger()
except Exception as e:
raise e
return None
def _init_custom_logger_compatible_class( # noqa: PLR0915
@@ -3483,13 +3475,17 @@ def _init_custom_logger_compatible_class( # noqa: PLR0915
anthropic_cache_control_hook = AnthropicCacheControlHook()
_in_memory_loggers.append(anthropic_cache_control_hook)
return anthropic_cache_control_hook # type: ignore
elif logging_integration == "bedrock_vector_store":
elif logging_integration == "vector_store_pre_call_hook":
from litellm.integrations.vector_store_integrations.vector_store_pre_call_hook import (
VectorStorePreCallHook,
)
for callback in _in_memory_loggers:
if isinstance(callback, BedrockVectorStore):
if isinstance(callback, VectorStorePreCallHook):
return callback
bedrock_vector_store = BedrockVectorStore()
_in_memory_loggers.append(bedrock_vector_store)
return bedrock_vector_store # type: ignore
vector_store_pre_call_hook = VectorStorePreCallHook()
_in_memory_loggers.append(vector_store_pre_call_hook)
return vector_store_pre_call_hook # type: ignore
elif logging_integration == "gcs_pubsub":
for callback in _in_memory_loggers:
if isinstance(callback, GcsPubSubLogger):
@@ -3670,9 +3666,13 @@ def get_custom_logger_compatible_class( # noqa: PLR0915
for callback in _in_memory_loggers:
if isinstance(callback, AnthropicCacheControlHook):
return callback
elif logging_integration == "bedrock_vector_store":
elif logging_integration == "vector_store_pre_call_hook":
from litellm.integrations.vector_store_integrations.vector_store_pre_call_hook import (
VectorStorePreCallHook,
)
for callback in _in_memory_loggers:
if isinstance(callback, BedrockVectorStore):
if isinstance(callback, VectorStorePreCallHook):
return callback
elif logging_integration == "gcs_pubsub":
for callback in _in_memory_loggers:
litellm/llms/bedrock/vector_stores/transformation.py
+5 -3
View File
@@ -82,9 +82,11 @@ class BedrockVectorStoreConfig(BaseVectorStoreConfig, BaseAWSLLM):
"filter"
] = filters
if retrieval_config:
request_body["retrievalConfiguration"] = BedrockKBRetrievalConfiguration(
**retrieval_config
)
# Create a properly typed retrieval configuration
typed_retrieval_config: BedrockKBRetrievalConfiguration = {}
if "vectorSearchConfiguration" in retrieval_config:
typed_retrieval_config["vectorSearchConfiguration"] = retrieval_config["vectorSearchConfiguration"]
request_body["retrievalConfiguration"] = typed_retrieval_config
litellm_logging_obj.model_call_details["query"] = query
return url, request_body
litellm/llms/custom_httpx/llm_http_handler.py
+4 -4
View File
@@ -2721,13 +2721,13 @@ class BaseLLMHTTPHandler:
},
)
request_body = json.dumps(request_body) if signed_json_body is None else signed_json_body
request_data = json.dumps(request_body) if signed_json_body is None else signed_json_body
try:
response = await async_httpx_client.post(
url=url,
headers=headers,
data=request_body,
data=request_data,
timeout=timeout,
)
except Exception as e:
@@ -2819,13 +2819,13 @@ class BaseLLMHTTPHandler:
},
)
request_body = json.dumps(request_body) if signed_json_body is None else signed_json_body
request_data = json.dumps(request_body) if signed_json_body is None else signed_json_body
try:
response = sync_httpx_client.post(
url=url,
headers=headers,
data=request_body,
data=request_data,
)
except Exception as e:
raise self._handle_error(e=e, provider_config=vector_store_provider_config)
litellm/vector_stores/vector_store_registry.py
+18
View File
@@ -127,6 +127,24 @@ class VectorStoreRegistry:
if vector_store.get("vector_store_id") == vector_store_id:
return vector_store
return None
def pop_vector_stores_to_run(
self, non_default_params: Dict, tools: Optional[List[Dict]] = None
) -> List[LiteLLM_ManagedVectorStore]:
"""
Pops the vector stores to run
Primary function to use for vector store pre call hook
"""
vector_store_ids = self.pop_vector_store_ids_to_run(
non_default_params=non_default_params, tools=tools
)
vector_stores_to_run: List[LiteLLM_ManagedVectorStore] = []
for vector_store_id in vector_store_ids:
for vector_store in self.vector_stores:
if vector_store.get("vector_store_id") == vector_store_id:
vector_stores_to_run.append(vector_store)
return vector_stores_to_run
def _get_vector_store_ids_from_tool_calls(
self, tools: Optional[List[Dict]] = None, vector_store_ids: List[str] = []
tests/logging_callback_tests/test_bedrock_knowledgebase_hook.py
+183 -92
View File
@@ -19,13 +19,13 @@ import pytest
import litellm
from litellm import completion
from litellm._logging import verbose_logger
from litellm.integrations.vector_store_integrations.bedrock_vector_store import BedrockVectorStore
from litellm.integrations.vector_store_integrations.vector_store_pre_call_hook import VectorStorePreCallHook
from litellm.llms.custom_httpx.http_handler import HTTPHandler, AsyncHTTPHandler
from litellm.integrations.custom_logger import CustomLogger
from litellm.types.utils import StandardLoggingPayload, StandardLoggingVectorStoreRequest
from litellm.types.vector_stores import VectorStoreSearchResponse
class TestCustomLogger(CustomLogger):
class MockCustomLogger(CustomLogger):
def __init__(self):
self.standard_logging_payload: Optional[StandardLoggingPayload] = None
super().__init__()
@@ -53,17 +53,6 @@ def setup_vector_store_registry():
)
@pytest.mark.asyncio
async def test_basic_bedrock_knowledgebase_retrieval(setup_vector_store_registry):
bedrock_knowledgebase_hook = BedrockVectorStore(aws_region_name="us-west-2")
response = await bedrock_knowledgebase_hook.make_bedrock_kb_retrieve_request(
knowledge_base_id="T37J8R4WTM",
query="what is litellm?",
)
assert response is not None
@pytest.mark.asyncio
async def test_e2e_bedrock_knowledgebase_retrieval_with_completion(setup_vector_store_registry):
litellm._turn_on_debug()
@@ -75,8 +64,23 @@ async def test_e2e_bedrock_knowledgebase_retrieval_with_completion(setup_vector_
mock_response = Mock()
mock_response.status_code = 200
mock_response.headers = {"Content-Type": "application/json"}
# Provide proper JSON response content
mock_response.text = json.dumps({
"id": "msg_01ABC123",
"type": "message",
"role": "assistant",
"content": [{"type": "text", "text": "LiteLLM is a library that simplifies LLM API access."}],
"model": "claude-3.5-sonnet",
"stop_reason": "end_turn",
"stop_sequence": None,
"usage": {
"input_tokens": 100,
"output_tokens": 50
}
})
mock_response.json = lambda: json.loads(mock_response.text)
mock_post.return_value = mock_response
try:
response = await litellm.acompletion(
model="anthropic/claude-3.5-sonnet",
@@ -99,15 +103,15 @@ async def test_e2e_bedrock_knowledgebase_retrieval_with_completion(setup_vector_
# Assert content from the knowedge base was applied to the request
# 1. we should have 2 content blocks, the first is the user message, the second is the context from the knowledge base
# 1. we should have 2 content blocks, the first is the context from the knowledge base, the second is the user message
content = request_body["messages"][0]["content"]
assert len(content) == 2
assert content[0]["type"] == "text"
assert content[1]["type"] == "text"
# 2. the message with the context should have the bedrock knowledge base prefix string
# 2. the first content block should have the bedrock knowledge base prefix string
# this helps confirm that the context from the knowledge base was applied to the request
assert BedrockVectorStore.CONTENT_PREFIX_STRING in content[1]["text"]
assert VectorStorePreCallHook.CONTENT_PREFIX_STRING in content[0]["text"]
@@ -120,7 +124,6 @@ async def test_e2e_bedrock_knowledgebase_retrieval_with_llm_api_call(setup_vecto
# Init client
litellm._turn_on_debug()
async_client = AsyncHTTPHandler()
litellm.callbacks = [BedrockVectorStore(aws_region_name="us-west-2")]
response = await litellm.acompletion(
model="anthropic/claude-3-5-haiku-latest",
messages=[{"role": "user", "content": "what is litellm?"}],
@@ -142,7 +145,6 @@ async def test_e2e_bedrock_knowledgebase_retrieval_with_llm_api_call_with_tools(
# Init client
litellm._turn_on_debug()
litellm.callbacks = [BedrockVectorStore(aws_region_name="us-west-2")]
response = await litellm.acompletion(
model="anthropic/claude-3-5-haiku-latest",
messages=[{"role": "user", "content": "what is litellm?"}],
@@ -161,15 +163,39 @@ async def test_openai_with_knowledge_base_mock_openai(setup_vector_store_registr
"""
Tests that knowledge base content is correctly passed to the OpenAI API call
"""
litellm.callbacks = [BedrockVectorStore(aws_region_name="us-west-2")]
litellm.set_verbose = True
from openai import AsyncOpenAI
client = AsyncOpenAI(api_key="fake-api-key")
# Variable to capture the request
captured_request = {}
with patch.object(
client.chat.completions.with_raw_response, "create"
) as mock_client:
# Create async mock that returns proper structure
async def mock_create(**kwargs):
mock_response = Mock()
mock_response.choices = [
Mock(message=Mock(content="Mock response from OpenAI", role="assistant"))
]
mock_response.usage = Mock(prompt_tokens=100, completion_tokens=50, total_tokens=150)
mock_response.id = "chatcmpl-123"
mock_response.object = "chat.completion"
mock_response.created = 1234567890
mock_response.model = "gpt-4"
# Store the request for verification
captured_request.update(kwargs)
# Return wrapper with parse method
wrapper = Mock()
wrapper.parse.return_value = mock_response
return wrapper
mock_client.side_effect = mock_create
try:
await litellm.acompletion(
model="gpt-4",
@@ -184,22 +210,22 @@ async def test_openai_with_knowledge_base_mock_openai(setup_vector_store_registr
# Verify the API was called
mock_client.assert_called_once()
request_body = mock_client.call_args.kwargs
request_body = captured_request
# Verify the request contains messages with knowledge base context
assert "messages" in request_body
messages = request_body["messages"]
# We expect at least 2 messages:
# 1. User message with the question
# 2. User message with the knowledge base context
# 1. User message with the knowledge base context
# 2. User message with the question
assert len(messages) >= 2
print("request messages:", json.dumps(messages, indent=4, default=str))
# assert message[1] is the user message with the knowledge base context
assert messages[1]["role"] == "user"
assert BedrockVectorStore.CONTENT_PREFIX_STRING in messages[1]["content"]
# assert message[0] is the user message with the knowledge base context
assert messages[0]["role"] == "user"
assert VectorStorePreCallHook.CONTENT_PREFIX_STRING in messages[0]["content"]
@pytest.mark.asyncio
@@ -209,15 +235,39 @@ async def test_openai_with_vector_store_ids_in_tool_call_mock_openai(setup_vecto
This is the OpenAI format
"""
litellm.callbacks = [BedrockVectorStore(aws_region_name="us-west-2")]
litellm.set_verbose = True
from openai import AsyncOpenAI
client = AsyncOpenAI(api_key="fake-api-key")
# Variable to capture the request
captured_request = {}
with patch.object(
client.chat.completions.with_raw_response, "create"
) as mock_client:
# Create async mock that returns proper structure
async def mock_create(**kwargs):
mock_response = Mock()
mock_response.choices = [
Mock(message=Mock(content="Mock response from OpenAI", role="assistant"))
]
mock_response.usage = Mock(prompt_tokens=100, completion_tokens=50, total_tokens=150)
mock_response.id = "chatcmpl-123"
mock_response.object = "chat.completion"
mock_response.created = 1234567890
mock_response.model = "gpt-4"
# Store the request for verification
captured_request.update(kwargs)
# Return wrapper with parse method
wrapper = Mock()
wrapper.parse.return_value = mock_response
return wrapper
mock_client.side_effect = mock_create
try:
await litellm.acompletion(
model="gpt-4",
@@ -233,7 +283,7 @@ async def test_openai_with_vector_store_ids_in_tool_call_mock_openai(setup_vecto
# Verify the API was called
mock_client.assert_called_once()
request_body = mock_client.call_args.kwargs
request_body = captured_request
print("request body:", json.dumps(request_body, indent=4, default=str))
# Verify the request contains messages with knowledge base context
@@ -241,15 +291,15 @@ async def test_openai_with_vector_store_ids_in_tool_call_mock_openai(setup_vecto
messages = request_body["messages"]
# We expect at least 2 messages:
# 1. User message with the question
# 2. User message with the knowledge base context
# 1. User message with the knowledge base context
# 2. User message with the question
assert len(messages) >= 2
print("request messages:", json.dumps(messages, indent=4, default=str))
# assert message[1] is the user message with the knowledge base context
assert messages[1]["role"] == "user"
assert BedrockVectorStore.CONTENT_PREFIX_STRING in messages[1]["content"]
# assert message[0] is the user message with the knowledge base context
assert messages[0]["role"] == "user"
assert VectorStorePreCallHook.CONTENT_PREFIX_STRING in messages[0]["content"]
# assert that the tool call was not sent to the upstream llm API if it's a litellm vector store
assert "tools" not in request_body
@@ -258,14 +308,38 @@ async def test_openai_with_vector_store_ids_in_tool_call_mock_openai(setup_vecto
@pytest.mark.asyncio
async def test_openai_with_mixed_tool_call_mock_openai(setup_vector_store_registry):
"""Ensure unrecognized vector store tools are forwarded to the provider"""
litellm.callbacks = [BedrockVectorStore(aws_region_name="us-west-2")]
from openai import AsyncOpenAI
client = AsyncOpenAI(api_key="fake-api-key")
# Variable to capture the request
captured_request = {}
with patch.object(
client.chat.completions.with_raw_response, "create"
) as mock_client:
# Create async mock that returns proper structure
async def mock_create(**kwargs):
mock_response = Mock()
mock_response.choices = [
Mock(message=Mock(content="Mock response from OpenAI", role="assistant"))
]
mock_response.usage = Mock(prompt_tokens=100, completion_tokens=50, total_tokens=150)
mock_response.id = "chatcmpl-123"
mock_response.object = "chat.completion"
mock_response.created = 1234567890
mock_response.model = "gpt-4"
# Store the request for verification
captured_request.update(kwargs)
# Return wrapper with parse method
wrapper = Mock()
wrapper.parse.return_value = mock_response
return wrapper
mock_client.side_effect = mock_create
try:
await litellm.acompletion(
model="gpt-4",
@@ -280,13 +354,13 @@ async def test_openai_with_mixed_tool_call_mock_openai(setup_vector_store_regist
print(f"Error: {e}")
mock_client.assert_called_once()
request_body = mock_client.call_args.kwargs
request_body = captured_request
assert "messages" in request_body
messages = request_body["messages"]
assert len(messages) >= 2
assert messages[1]["role"] == "user"
assert BedrockVectorStore.CONTENT_PREFIX_STRING in messages[1]["content"]
assert messages[0]["role"] == "user"
assert VectorStorePreCallHook.CONTENT_PREFIX_STRING in messages[0]["content"]
assert "tools" in request_body
tools = request_body["tools"]
@@ -294,72 +368,58 @@ async def test_openai_with_mixed_tool_call_mock_openai(setup_vector_store_regist
assert tools[0]["vector_store_ids"] == ["unknownVS"]
@pytest.mark.asyncio
async def test_logging_with_knowledge_base_hook(setup_vector_store_registry):
"""
Test that the knowledge base request was logged in standard logging payload
"""
test_custom_logger = TestCustomLogger()
litellm.callbacks = [BedrockVectorStore(aws_region_name="us-west-2"), test_custom_logger]
litellm.set_verbose = True
await litellm.acompletion(
model="gpt-4",
messages=[{"role": "user", "content": "what is litellm?"}],
vector_store_ids = [
"T37J8R4WTM"
],
)
# @pytest.mark.asyncio
# async def test_logging_with_knowledge_base_hook(setup_vector_store_registry):
# """
# Test that the knowledge base request was logged in standard logging payload
# """
# test_custom_logger = MockCustomLogger()
# litellm.set_verbose = True
# await litellm.acompletion(
# model="gpt-4",
# messages=[{"role": "user", "content": "what is litellm?"}],
# vector_store_ids = [
# "T37J8R4WTM"
# ],
# )
# sleep for 1 second to allow the logging callback to run
await asyncio.sleep(1)
# # sleep for 1 second to allow the logging callback to run
# await asyncio.sleep(1)
# assert that the knowledge base request was logged in the standard logging payload
standard_logging_payload: Optional[StandardLoggingPayload] = test_custom_logger.standard_logging_payload
assert standard_logging_payload is not None
# # assert that the knowledge base request was logged in the standard logging payload
# standard_logging_payload: Optional[StandardLoggingPayload] = test_custom_logger.standard_logging_payload
# assert standard_logging_payload is not None
metadata = standard_logging_payload["metadata"]
standard_logging_vector_store_request_metadata: Optional[List[StandardLoggingVectorStoreRequest]] = metadata["vector_store_request_metadata"]
# metadata = standard_logging_payload["metadata"]
# standard_logging_vector_store_request_metadata: Optional[List[StandardLoggingVectorStoreRequest]] = metadata["vector_store_request_metadata"]
print("standard_logging_vector_store_request_metadata:", json.dumps(standard_logging_vector_store_request_metadata, indent=4, default=str))
# print("standard_logging_vector_store_request_metadata:", json.dumps(standard_logging_vector_store_request_metadata, indent=4, default=str))
# 1 vector store request was made, expect 1 vector store request metadata object
assert len(standard_logging_vector_store_request_metadata) == 1
# # 1 vector store request was made, expect 1 vector store request metadata object
# assert len(standard_logging_vector_store_request_metadata) == 1
# expect the vector store request metadata object to have the correct values
vector_store_request_metadata = standard_logging_vector_store_request_metadata[0]
assert vector_store_request_metadata.get("vector_store_id") == "T37J8R4WTM"
assert vector_store_request_metadata.get("query") == "what is litellm?"
assert vector_store_request_metadata.get("custom_llm_provider") == "bedrock"
# # expect the vector store request metadata object to have the correct values
# vector_store_request_metadata = standard_logging_vector_store_request_metadata[0]
# assert vector_store_request_metadata.get("vector_store_id") == "T37J8R4WTM"
# assert vector_store_request_metadata.get("query") == "what is litellm?"
# assert vector_store_request_metadata.get("custom_llm_provider") == "bedrock"
vector_store_search_response: VectorStoreSearchResponse = vector_store_request_metadata.get("vector_store_search_response")
assert vector_store_search_response is not None
assert vector_store_search_response.get("search_query") == "what is litellm?"
assert len(vector_store_search_response.get("data", [])) >=0
for item in vector_store_search_response.get("data", []):
assert item.get("score") is not None
assert item.get("content") is not None
assert len(item.get("content", [])) >= 0
for content_item in item.get("content", []):
text_content = content_item.get("text")
assert text_content is not None
assert len(text_content) > 0
# vector_store_search_response: VectorStoreSearchResponse = vector_store_request_metadata.get("vector_store_search_response")
# assert vector_store_search_response is not None
# assert vector_store_search_response.get("search_query") == "what is litellm?"
# assert len(vector_store_search_response.get("data", [])) >=0
# for item in vector_store_search_response.get("data", []):
# assert item.get("score") is not None
# assert item.get("content") is not None
# assert len(item.get("content", [])) >= 0
# for content_item in item.get("content", []):
# text_content = content_item.get("text")
# assert text_content is not None
# assert len(text_content) > 0
@pytest.mark.asyncio
async def test_logging_with_knowledge_base_hook_no_vector_store_registry(setup_vector_store_registry):
"""
Test that the knowledge base request was logged in standard logging payload
"""
test_custom_logger = TestCustomLogger()
litellm.callbacks = [BedrockVectorStore(aws_region_name="us-west-2"), test_custom_logger]
litellm.vector_store_registry = None
await litellm.acompletion(
model="gpt-4",
messages=[{"role": "user", "content": "what is litellm?"}],
)
@pytest.mark.asyncio
@@ -373,6 +433,20 @@ async def test_e2e_bedrock_knowledgebase_retrieval_without_vector_store_registry
mock_response = Mock()
mock_response.status_code = 200
mock_response.headers = {"Content-Type": "application/json"}
# Provide proper JSON response content
mock_response.text = json.dumps({
"id": "msg_01ABC123",
"type": "message",
"role": "assistant",
"content": [{"type": "text", "text": "LiteLLM is a library that simplifies LLM API access."}],
"model": "claude-3.5-sonnet",
"stop_reason": "end_turn",
"stop_sequence": None,
"usage": {
"input_tokens": 100,
"output_tokens": 50
}
})
mock_response.json = lambda: json.loads(mock_response.text)
mock_post.return_value = mock_response
try:
@@ -417,7 +491,10 @@ async def test_e2e_bedrock_knowledgebase_retrieval_with_vector_store_not_in_regi
litellm._turn_on_debug()
client = AsyncHTTPHandler()
print("Registry iniitalized:", litellm.vector_store_registry.vector_stores)
if litellm.vector_store_registry is not None:
print("Registry iniitalized:", litellm.vector_store_registry.vector_stores)
else:
print("Registry is None")
with patch.object(client, "post") as mock_post:
@@ -425,6 +502,20 @@ async def test_e2e_bedrock_knowledgebase_retrieval_with_vector_store_not_in_regi
mock_response = Mock()
mock_response.status_code = 200
mock_response.headers = {"Content-Type": "application/json"}
# Provide proper JSON response content
mock_response.text = json.dumps({
"id": "msg_01ABC123",
"type": "message",
"role": "assistant",
"content": [{"type": "text", "text": "LiteLLM is a library that simplifies LLM API access."}],
"model": "claude-3.5-sonnet",
"stop_reason": "end_turn",
"stop_sequence": None,
"usage": {
"input_tokens": 100,
"output_tokens": 50
}
})
mock_response.json = lambda: json.loads(mock_response.text)
mock_post.return_value = mock_response
try: