This commit is contained in:
Ishaan Jaff 2025-05-03 11:39:40 -07:00 • committed by GitHub
parent 071c8a2e26
commit a37e4cb0b8
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 26 additions and 12 deletions

View file

@ -59,6 +59,20 @@ Inherits from `StandardLoggingUserAPIKeyMetadata` and adds:
| `spend_logs_metadata` | `Optional[dict]` | Key-value pairs for spend logging |
| `requester_ip_address` | `Optional[str]` | Requester's IP address |
| `requester_metadata` | `Optional[dict]` | Additional requester metadata |
| `vector_store_request_metadata` | `Optional[List[StandardLoggingVectorStoreRequest]]` | Vector store request metadata |
## StandardLoggingVectorStoreRequest
| Field | Type | Description |
|-------|------|-------------|
| vector_store_id | Optional[str] | ID of the vector store |
| custom_llm_provider | Optional[str] | Custom LLM provider the vector store is associated with (e.g., bedrock, openai, anthropic) |
| query | Optional[str] | Query to the vector store |
| vector_store_search_response | Optional[VectorStoreSearchResponse] | OpenAI format vector store search response |
| start_time | Optional[float] | Start time of the vector store request |
| end_time | Optional[float] | End time of the vector store request |
## StandardLoggingAdditionalHeaders

View file

@ -31,8 +31,8 @@ from litellm.types.llms.openai import AllMessageValues, ChatCompletionUserMessag
from litellm.types.utils import StandardLoggingVectorStoreRequest
from litellm.types.vector_stores import (
VectorStoreResultContent,
VectorStoreSearchResponse,
VectorStoreSearchResult,
VectorStorSearchResponse,
)
from litellm.utils import load_credentials_from_list
@ -110,7 +110,7 @@ class BedrockVectorStore(BaseVectorStore, BaseAWSLLM):
#################################################################################################
########## LOGGING for Standard Logging Payload, Langfuse, s3, LiteLLM DB etc. ##################
#################################################################################################
vector_store_search_response: VectorStorSearchResponse = (
vector_store_search_response: VectorStoreSearchResponse = (
self.transform_bedrock_kb_response_to_vector_store_search_response(
bedrock_kb_response=bedrock_kb_response, query=query
)
@ -136,15 +136,15 @@ class BedrockVectorStore(BaseVectorStore, BaseAWSLLM):
self,
bedrock_kb_response: BedrockKBResponse,
query: str,
) -> VectorStorSearchResponse:
) -> VectorStoreSearchResponse:
"""
Transform a BedrockKBResponse to a VectorStorSearchResponse
Transform a BedrockKBResponse to a VectorStoreSearchResponse
"""
retrieval_results: Optional[List[BedrockKBRetrievalResult]] = (
bedrock_kb_response.get("retrievalResults", None)
)
vector_store_search_response: VectorStorSearchResponse = (
VectorStorSearchResponse(search_query=query, data=[])
vector_store_search_response: VectorStoreSearchResponse = (
VectorStoreSearchResponse(search_query=query, data=[])
)
if retrieval_results is None:
return vector_store_search_response

View file

@ -52,9 +52,9 @@ from .llms.openai import (
from .rerank import RerankResponse
if TYPE_CHECKING:
from .vector_stores import VectorStorSearchResponse
from .vector_stores import VectorStoreSearchResponse
else:
VectorStorSearchResponse = Any
VectorStoreSearchResponse = Any
def _generate_id(): # private helper function
@ -1740,7 +1740,7 @@ class StandardLoggingVectorStoreRequest(TypedDict, total=False):
Query to the vector store
"""
vector_store_search_response: Optional[VectorStorSearchResponse]
vector_store_search_response: Optional[VectorStoreSearchResponse]
"""
OpenAI format vector store search response
"""

View file

@ -77,7 +77,7 @@ class VectorStoreSearchResult(TypedDict, total=False):
content: Optional[List[VectorStoreResultContent]]
class VectorStorSearchResponse(TypedDict, total=False):
class VectorStoreSearchResponse(TypedDict, total=False):
"""Response after searching a vector store"""
object: Literal[

View file

@ -23,7 +23,7 @@ from litellm.integrations.vector_stores.bedrock_vector_store import BedrockVecto
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 VectorStorSearchResponse
from litellm.types.vector_stores import VectorStoreSearchResponse
class TestCustomLogger(CustomLogger):
def __init__(self):
@ -262,7 +262,7 @@ async def test_logging_with_knowledge_base_hook(setup_vector_store_registry):
assert vector_store_request_metadata.get("custom_llm_provider") == "bedrock"
vector_store_search_response: VectorStorSearchResponse = vector_store_request_metadata.get("vector_store_search_response")
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