b700cfac64
- Removed Portkey from the equation, and defined explicit monitoring using Langchain native code - Optimization of Business Event logging
52 lines
1.6 KiB
Python
52 lines
1.6 KiB
Python
from langchain_openai import OpenAIEmbeddings
|
|
from typing import List, Any
|
|
import time
|
|
from common.utils.business_event_context import current_event
|
|
|
|
|
|
class TrackedOpenAIEmbeddings(OpenAIEmbeddings):
|
|
def __init__(self, *args, **kwargs):
|
|
super().__init__(*args, **kwargs)
|
|
|
|
def embed_documents(self, texts: list[str]) -> list[list[float]]:
|
|
start_time = time.time()
|
|
result = super().embed_documents(texts)
|
|
end_time = time.time()
|
|
|
|
# Estimate token usage (OpenAI uses tiktoken for this)
|
|
import tiktoken
|
|
enc = tiktoken.encoding_for_model(self.model)
|
|
total_tokens = sum(len(enc.encode(text)) for text in texts)
|
|
|
|
metrics = {
|
|
'total_tokens': total_tokens,
|
|
'prompt_tokens': total_tokens, # For embeddings, all tokens are prompt tokens
|
|
'completion_tokens': 0,
|
|
'time_elapsed': end_time - start_time,
|
|
'interaction_type': 'Embedding',
|
|
}
|
|
current_event.log_llm_metrics(metrics)
|
|
|
|
return result
|
|
|
|
def embed_query(self, text: str) -> List[float]:
|
|
start_time = time.time()
|
|
result = super().embed_query(text)
|
|
end_time = time.time()
|
|
|
|
# Estimate token usage
|
|
import tiktoken
|
|
enc = tiktoken.encoding_for_model(self.model)
|
|
total_tokens = len(enc.encode(text))
|
|
|
|
metrics = {
|
|
'total_tokens': total_tokens,
|
|
'prompt_tokens': total_tokens,
|
|
'completion_tokens': 0,
|
|
'time_elapsed': end_time - start_time,
|
|
'interaction_type': 'Embedding',
|
|
}
|
|
current_event.log_llm_metrics(metrics)
|
|
|
|
return result
|