[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]>

Author: Deepak-Kesavan

pdm.lock

@@ -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",

pyproject.toml

@@ -1,4 +1,4 @@
-__version__ = "0.37.0"
+__version__ = "0.38.0"
 
 
 def get_sdk_version():

src/unstract/sdk/__init__.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,

src/unstract/sdk/audit.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(

src/unstract/sdk/embedding.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: