From 2b723a53f666b185bcc10658eae67de38f30cb3b Mon Sep 17 00:00:00 2001
From: Ali Waleed <ali@scale3labs.com>
Date: Tue, 7 Jan 2025 14:15:08 +0200
Subject: [PATCH] handle `llm_config` different cases

---
 .../instrumentation/autogen/patch.py          | 46 +++++++++++++------
 src/langtrace_python_sdk/version.py           |  2 +-
 2 files changed, 33 insertions(+), 15 deletions(-)

diff --git a/src/langtrace_python_sdk/instrumentation/autogen/patch.py b/src/langtrace_python_sdk/instrumentation/autogen/patch.py
index ce6c7ff2..2cb44341 100644
--- a/src/langtrace_python_sdk/instrumentation/autogen/patch.py
+++ b/src/langtrace_python_sdk/instrumentation/autogen/patch.py
@@ -56,12 +56,8 @@ def traced_method(wrapped, instance, args, kwargs):
 def patch_generate_reply(name, version, tracer: Tracer):
 
     def traced_method(wrapped, instance, args, kwargs):
-
         llm_config = instance.llm_config
-        kwargs = {
-            **kwargs,
-            **llm_config.get("config_list")[0],
-        }
+        kwargs = parse_kwargs(kwargs, llm_config)
         service_provider = SERVICE_PROVIDERS["AUTOGEN"]
 
         span_attributes = {
@@ -84,20 +80,25 @@ def traced_method(wrapped, instance, args, kwargs):
             try:
 
                 result = wrapped(*args, **kwargs)
-                
+
                 # if caching is disabled, return result as langtrace will instrument the rest.
-                if "cache_seed" in llm_config and llm_config.get("cache_seed") is None:
+                if (
+                    llm_config
+                    and "cache_seed" in llm_config
+                    and llm_config.get("cache_seed") is None
+                ):
                     return result
 
                 set_span_attributes(span, attributes)
                 set_event_completion(span, [{"role": "assistant", "content": result}])
-                total_cost, response_model = list(instance.get_total_usage().keys())
-                set_span_attribute(
-                    span, SpanAttributes.LLM_RESPONSE_MODEL, response_model
-                )
-                set_usage_attributes(
-                    span, instance.get_total_usage().get(response_model)
-                )
+                if llm_config:
+                    total_cost, response_model = list(instance.get_total_usage().keys())
+                    set_span_attribute(
+                        span, SpanAttributes.LLM_RESPONSE_MODEL, response_model
+                    )
+                    set_usage_attributes(
+                        span, instance.get_total_usage().get(response_model)
+                    )
 
                 return result
 
@@ -130,3 +131,20 @@ def parse_agent(agent):
         "llm_config": str(getattr(agent, "llm_config", None)),
         "human_input_mode": getattr(agent, "human_input_mode", None),
     }
+
+
+def parse_kwargs(kwargs, llm_config):
+    # Handle cases where llm_config is False or None
+    if not llm_config:
+        return kwargs
+
+    if isinstance(llm_config, dict) and "config_list" in llm_config:
+        return {
+            **kwargs,
+            **llm_config.get("config_list")[0],
+        }
+    else:
+        return {
+            **kwargs,
+            **llm_config,
+        }
diff --git a/src/langtrace_python_sdk/version.py b/src/langtrace_python_sdk/version.py
index bb2175a9..eb7bf495 100644
--- a/src/langtrace_python_sdk/version.py
+++ b/src/langtrace_python_sdk/version.py
@@ -1 +1 @@
-__version__ = "3.3.18"
+__version__ = "3.3.19"