[FEATURE] Send provider in the usage callback (#75)

* Send provider in the usage callback Signed-off-by: Deepak <[email protected]> * Updated lock file Signed-off-by: Deepak <[email protected]> --------- Signed-off-by: Deepak <[email protected]>
Zipstack · Jul 19, 2024 · c8d93f8 · c8d93f8
1 parent 04aa44d
commit c8d93f8
Show file tree

Hide file tree

Showing 6 changed files with 10 additions and 6 deletions.
diff --git a/pdm.lock b/pdm.lock
diff --git a/pyproject.toml b/pyproject.toml
@@ -12,7 +12,7 @@ dependencies = [
     "python-magic~=0.4.27",
     "python-dotenv==1.0.0",
     # LLM Triad
-    "unstract-adapters~=0.21.0",
+    "unstract-adapters~=0.22.0",
     "llama-index==0.10.38",
     "tiktoken~=0.4.0",
     "transformers==4.37.0",

diff --git a/src/unstract/sdk/__init__.py b/src/unstract/sdk/__init__.py
@@ -1,4 +1,4 @@
-__version__ = "0.37.0"
+__version__ = "0.38.0"
 
 
 def get_sdk_version():

diff --git a/src/unstract/sdk/audit.py b/src/unstract/sdk/audit.py
@@ -69,6 +69,7 @@ def push_usage_data(
         execution_id = kwargs.get("execution_id", "")
         adapter_instance_id = kwargs.get("adapter_instance_id", "")
         run_id = kwargs.get("run_id", "")
+        provider = kwargs.get("provider", "")
         llm_usage_reason = ""
         if event_type == "llm":
             llm_usage_reason = kwargs.get("llm_usage_reason", "")
@@ -80,6 +81,7 @@ def push_usage_data(
             "usage_type": event_type,
             "llm_usage_reason": llm_usage_reason,
             "model_name": model_name,
+            "provider": provider,
             "embedding_tokens": token_counter.total_embedding_token_count,
             "prompt_tokens": token_counter.prompt_llm_token_count,
             "completion_tokens": token_counter.completion_llm_token_count,

diff --git a/src/unstract/sdk/embedding.py b/src/unstract/sdk/embedding.py
@@ -71,6 +71,7 @@ def _get_embedding(self) -> BaseEmbedding:
             ][Common.ADAPTER]
             embedding_metadata = embedding_config_data.get(Common.ADAPTER_METADATA)
             embedding_adapter_class = embedding_adapter(embedding_metadata)
+            self._usage_kwargs["provider"] = embedding_adapter_class.get_provider()
             return embedding_adapter_class.get_embedding_instance()
         except Exception as e:
             self._tool.stream_log(

diff --git a/src/unstract/sdk/llm.py b/src/unstract/sdk/llm.py
@@ -106,6 +106,7 @@ def _get_llm(self, adapter_instance_id: str) -> LlamaIndexLLM:
             ]
             llm_metadata = llm_config_data.get(Common.ADAPTER_METADATA)
             llm_adapter_class: LLMAdapter = llm_adapter(llm_metadata)
+            self._usage_kwargs["provider"] = llm_adapter_class.get_provider()
             llm_instance: LLM = llm_adapter_class.get_llm_instance()
             return llm_instance
         except Exception as e: