From 01dc2af801119508e946e86ab23dd876e25fd44f Mon Sep 17 00:00:00 2001
From: Rohit Kadhe <rohit@scale3labs.com>
Date: Tue, 10 Sep 2024 10:50:18 -0600
Subject: [PATCH 1/9] add sentry integration

---
 .../constants/__init__.py                     |  1 +
 src/langtrace_python_sdk/langtrace.py         | 28 ++++++++++++++++++-
 src/langtrace_python_sdk/utils/__init__.py    |  4 +++
 .../utils/sdk_version_checker.py              |  3 ++
 4 files changed, 35 insertions(+), 1 deletion(-)

diff --git a/src/langtrace_python_sdk/constants/__init__.py b/src/langtrace_python_sdk/constants/__init__.py
index f0dbbdb9..4b77a4b4 100644
--- a/src/langtrace_python_sdk/constants/__init__.py
+++ b/src/langtrace_python_sdk/constants/__init__.py
@@ -1 +1,2 @@
 LANGTRACE_SDK_NAME = "langtrace-python-sdk"
+SENTRY_DSN = "https://7f8eed3a1237fb2efef0f5e96ab407af@o1419498.ingest.us.sentry.io/4507929133056000"
diff --git a/src/langtrace_python_sdk/langtrace.py b/src/langtrace_python_sdk/langtrace.py
index 7be1bf9f..c8169aa3 100644
--- a/src/langtrace_python_sdk/langtrace.py
+++ b/src/langtrace_python_sdk/langtrace.py
@@ -19,6 +19,7 @@
 from typing import Optional
 import importlib.util
 from colorama import Fore
+from langtrace_python_sdk.constants import LANGTRACE_SDK_NAME, SENTRY_DSN
 from opentelemetry import trace
 from opentelemetry.instrumentation.sqlalchemy import SQLAlchemyInstrumentor
 from opentelemetry.sdk.resources import SERVICE_NAME, Resource
@@ -60,8 +61,9 @@
     InstrumentationMethods,
     InstrumentationType,
 )
-from langtrace_python_sdk.utils import check_if_sdk_is_outdated
+from langtrace_python_sdk.utils import check_if_sdk_is_outdated, get_sdk_version
 from langtrace_python_sdk.utils.langtrace_sampler import LangtraceSampler
+import sentry_sdk
 
 
 def init(
@@ -164,6 +166,30 @@ def init(
         provider.add_span_processor(batch_processor_remote)
 
     sys.stdout = sys.__stdout__
+    if (
+        os.environ.get("LANGTRACE_ERROR_REPORTING", "True") == "True"
+        or os.environ.get("LANGTRACE_ERROR_REPORTING", "True") == "true"
+    ):
+        sentry_sdk.init(
+            dsn=SENTRY_DSN,
+            traces_sample_rate=1.0,
+            profiles_sample_rate=1.0,
+        )
+        sdk_options = {
+            "service_name": os.environ.get("OTEL_SERVICE_NAME")
+            or service_name
+            or sys.argv[0],
+            "disable_logging": disable_logging,
+            "disable_instrumentations": disable_instrumentations,
+            "disable_tracing_for_functions": disable_tracing_for_functions,
+            "batch": batch,
+            "write_spans_to_console": write_spans_to_console,
+            "custom_remote_exporter": custom_remote_exporter,
+            "sdk_name": LANGTRACE_SDK_NAME,
+            "sdk_version": get_sdk_version(),
+            "api_host": host,
+        }
+        sentry_sdk.set_context("sdk_init_options", sdk_options)
 
 
 def init_instrumentations(
diff --git a/src/langtrace_python_sdk/utils/__init__.py b/src/langtrace_python_sdk/utils/__init__.py
index df6925f3..43ffa5a7 100644
--- a/src/langtrace_python_sdk/utils/__init__.py
+++ b/src/langtrace_python_sdk/utils/__init__.py
@@ -31,3 +31,7 @@ def set_event_prompt(span: Span, prompt):
 def check_if_sdk_is_outdated():
     SDKVersionChecker().check()
     return
+
+
+def get_sdk_version():
+    return SDKVersionChecker().get_sdk_version()
diff --git a/src/langtrace_python_sdk/utils/sdk_version_checker.py b/src/langtrace_python_sdk/utils/sdk_version_checker.py
index 6b030a27..c708aa82 100644
--- a/src/langtrace_python_sdk/utils/sdk_version_checker.py
+++ b/src/langtrace_python_sdk/utils/sdk_version_checker.py
@@ -44,6 +44,9 @@ def is_outdated(self):
             return self._current_version < latest_version
         return False
 
+    def get_sdk_version(self):
+        return self._current_version
+
     def check(self):
         if self.is_outdated():
             print(

From f45110f969d5d6f369434aaef0aafd1aa7312508 Mon Sep 17 00:00:00 2001
From: Rohit Kadhe <rohit@scale3labs.com>
Date: Tue, 10 Sep 2024 11:01:17 -0600
Subject: [PATCH 2/9] update readme

---
 README.md                             | 5 +++++
 pyproject.toml                        | 1 +
 src/langtrace_python_sdk/langtrace.py | 5 +----
 3 files changed, 7 insertions(+), 4 deletions(-)

diff --git a/README.md b/README.md
index 30d3e6cd..3a65dcab 100644
--- a/README.md
+++ b/README.md
@@ -148,6 +148,10 @@ langtrace.init(custom_remote_exporter=<your_exporter>, batch=<True or False>)
 | `api_host`                 | `Optional[str]`                     | `https://langtrace.ai/`       | The API host for the remote exporter.                                                                                             |
 | `disable_instrumentations` | `Optional[DisableInstrumentations]` | `None`                        | You can pass an object to disable instrumentation for specific vendors ex: `{'only': ['openai']}` or `{'all_except': ['openai']}` |
 
+### Error Reporting to Langtrace
+
+By default all sdk errors are reported to langtrace via Sentry. This can be disabled by setting the following enviroment variable to `False` like so `LANGTRACE_ERROR_REPORTING=False`
+
 ### Additional Customization
 
 - `@with_langtrace_root_span` - this decorator is designed to organize and relate different spans, in a hierarchical manner. When you're performing multiple operations that you want to monitor together as a unit, this function helps by establishing a "parent" (`LangtraceRootSpan` or whatever is passed to `name`) span. Then, any calls to the LLM APIs made within the given function (fn) will be considered "children" of this parent span. This setup is especially useful for tracking the performance or behavior of a group of operations collectively, rather than individually.
@@ -229,6 +233,7 @@ prompt = get_prompt_from_registry(<Registry ID>, options={"prompt_version": 1, "
 ```
 
 ### Opt out of tracing prompt and completion data
+
 By default, prompt and completion data are captured. If you would like to opt out of it, set the following env var,
 
 `TRACE_PROMPT_COMPLETION_DATA=false`
diff --git a/pyproject.toml b/pyproject.toml
index c0b8c990..4875e93b 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -30,6 +30,7 @@ dependencies = [
   'sqlalchemy',
   'fsspec>=2024.6.0',
   "transformers>=4.11.3",
+  "sentry-sdk>=2.14.0",
 ]
 
 requires-python = ">=3.9"
diff --git a/src/langtrace_python_sdk/langtrace.py b/src/langtrace_python_sdk/langtrace.py
index c8169aa3..01c6417d 100644
--- a/src/langtrace_python_sdk/langtrace.py
+++ b/src/langtrace_python_sdk/langtrace.py
@@ -166,10 +166,7 @@ def init(
         provider.add_span_processor(batch_processor_remote)
 
     sys.stdout = sys.__stdout__
-    if (
-        os.environ.get("LANGTRACE_ERROR_REPORTING", "True") == "True"
-        or os.environ.get("LANGTRACE_ERROR_REPORTING", "True") == "true"
-    ):
+    if os.environ.get("LANGTRACE_ERROR_REPORTING", "True") == "True":
         sentry_sdk.init(
             dsn=SENTRY_DSN,
             traces_sample_rate=1.0,

From a9115a08b0f7ab0f9171ae283f9a9f6279f8b711 Mon Sep 17 00:00:00 2001
From: Ali Waleed <ali@scale3labs.com>
Date: Wed, 11 Sep 2024 16:30:26 +0300
Subject: [PATCH 3/9] hotfix for checking package installed

---
 src/langtrace_python_sdk/langtrace.py | 5 ++++-
 src/langtrace_python_sdk/version.py   | 2 +-
 2 files changed, 5 insertions(+), 2 deletions(-)

diff --git a/src/langtrace_python_sdk/langtrace.py b/src/langtrace_python_sdk/langtrace.py
index 01c6417d..0d6be8d1 100644
--- a/src/langtrace_python_sdk/langtrace.py
+++ b/src/langtrace_python_sdk/langtrace.py
@@ -256,4 +256,7 @@ def validate_instrumentations(disable_instrumentations):
 
 
 def is_package_installed(package_name):
-    return importlib.util.find_spec(package_name) is not None
+    import pkg_resources
+
+    installed_packages = {p.key for p in pkg_resources.working_set}
+    return package_name in installed_packages
diff --git a/src/langtrace_python_sdk/version.py b/src/langtrace_python_sdk/version.py
index 534f1f0b..01bdb5c3 100644
--- a/src/langtrace_python_sdk/version.py
+++ b/src/langtrace_python_sdk/version.py
@@ -1 +1 @@
-__version__ = "2.3.10"
+__version__ = "2.3.11"

From 77024fb2fc75a5b65da7670d3ff3317198a82f1a Mon Sep 17 00:00:00 2001
From: Ali Waleed <ali@scale3labs.com>
Date: Wed, 11 Sep 2024 23:15:57 +0300
Subject: [PATCH 4/9] Support Autogen (#242)

* autogen kickstart

* enhance autogen

* run autogen

* finish autogen

* add readme

* disable cache
---
 README.md                                     |  49 +++----
 src/examples/autogen_example/__init__.py      |   8 ++
 src/examples/autogen_example/main.py          |  72 ++++++++++
 .../constants/instrumentation/common.py       |   1 +
 .../instrumentation/__init__.py               |   2 +
 .../instrumentation/autogen/__init__.py       |   3 +
 .../autogen/instrumentation.py                |  42 ++++++
 .../instrumentation/autogen/patch.py          | 132 ++++++++++++++++++
 src/langtrace_python_sdk/langtrace.py         |   2 +
 src/langtrace_python_sdk/utils/__init__.py    |  14 ++
 src/langtrace_python_sdk/utils/llm.py         |  14 +-
 src/run_example.py                            |  10 +-
 12 files changed, 321 insertions(+), 28 deletions(-)
 create mode 100644 src/examples/autogen_example/__init__.py
 create mode 100644 src/examples/autogen_example/main.py
 create mode 100644 src/langtrace_python_sdk/instrumentation/autogen/__init__.py
 create mode 100644 src/langtrace_python_sdk/instrumentation/autogen/instrumentation.py
 create mode 100644 src/langtrace_python_sdk/instrumentation/autogen/patch.py

diff --git a/README.md b/README.md
index 3a65dcab..46a530d3 100644
--- a/README.md
+++ b/README.md
@@ -242,30 +242,31 @@ By default, prompt and completion data are captured. If you would like to opt ou
 
 Langtrace automatically captures traces from the following vendors:
 
-| Vendor       | Type            | Typescript SDK     | Python SDK                      |
-| ------------ | --------------- | ------------------ | ------------------------------- |
-| OpenAI       | LLM             | :white_check_mark: | :white_check_mark:              |
-| Anthropic    | LLM             | :white_check_mark: | :white_check_mark:              |
-| Azure OpenAI | LLM             | :white_check_mark: | :white_check_mark:              |
-| Cohere       | LLM             | :white_check_mark: | :white_check_mark:              |
-| Groq         | LLM             | :x:                | :white_check_mark:              |
-| Perplexity   | LLM             | :white_check_mark: | :white_check_mark:              |
-| Gemini       | LLM             | :x:                | :white_check_mark:              |
-| Mistral      | LLM             | :x:                | :white_check_mark:              |
-| Langchain    | Framework       | :x:                | :white_check_mark:              |
-| LlamaIndex   | Framework       | :white_check_mark: | :white_check_mark:              |
-| Langgraph    | Framework       | :x:                | :white_check_mark:              |
-| DSPy         | Framework       | :x:                | :white_check_mark:              |
-| CrewAI       | Framework       | :x:                | :white_check_mark:              |
-| Ollama       | Framework       | :x:                | :white_check_mark:              |
-| VertexAI     | Framework       | :x:                | :white_check_mark:              |
-| Vercel AI SDK| Framework       | :white_check_mark: | :x:                             |
-| EmbedChain   | Framework       | :x:                | :white_check_mark:              |
-| Pinecone     | Vector Database | :white_check_mark: | :white_check_mark:              |
-| ChromaDB     | Vector Database | :white_check_mark: | :white_check_mark:              |
-| QDrant       | Vector Database | :white_check_mark: | :white_check_mark:              |
-| Weaviate     | Vector Database | :white_check_mark: | :white_check_mark:              |
-| PGVector     | Vector Database | :white_check_mark: | :white_check_mark: (SQLAlchemy) |
+| Vendor        | Type            | Typescript SDK     | Python SDK                      |
+| ------------- | --------------- | ------------------ | ------------------------------- |
+| OpenAI        | LLM             | :white_check_mark: | :white_check_mark:              |
+| Anthropic     | LLM             | :white_check_mark: | :white_check_mark:              |
+| Azure OpenAI  | LLM             | :white_check_mark: | :white_check_mark:              |
+| Cohere        | LLM             | :white_check_mark: | :white_check_mark:              |
+| Groq          | LLM             | :x:                | :white_check_mark:              |
+| Perplexity    | LLM             | :white_check_mark: | :white_check_mark:              |
+| Gemini        | LLM             | :x:                | :white_check_mark:              |
+| Mistral       | LLM             | :x:                | :white_check_mark:              |
+| Langchain     | Framework       | :x:                | :white_check_mark:              |
+| LlamaIndex    | Framework       | :white_check_mark: | :white_check_mark:              |
+| Langgraph     | Framework       | :x:                | :white_check_mark:              |
+| DSPy          | Framework       | :x:                | :white_check_mark:              |
+| CrewAI        | Framework       | :x:                | :white_check_mark:              |
+| Ollama        | Framework       | :x:                | :white_check_mark:              |
+| VertexAI      | Framework       | :x:                | :white_check_mark:              |
+| Vercel AI SDK | Framework       | :white_check_mark: | :x:                             |
+| EmbedChain    | Framework       | :x:                | :white_check_mark:              |
+| Autogen       | Framework       | :x:                | :white_check_mark:              |
+| Pinecone      | Vector Database | :white_check_mark: | :white_check_mark:              |
+| ChromaDB      | Vector Database | :white_check_mark: | :white_check_mark:              |
+| QDrant        | Vector Database | :white_check_mark: | :white_check_mark:              |
+| Weaviate      | Vector Database | :white_check_mark: | :white_check_mark:              |
+| PGVector      | Vector Database | :white_check_mark: | :white_check_mark: (SQLAlchemy) |
 
 ---
 
diff --git a/src/examples/autogen_example/__init__.py b/src/examples/autogen_example/__init__.py
new file mode 100644
index 00000000..e04fb9ef
--- /dev/null
+++ b/src/examples/autogen_example/__init__.py
@@ -0,0 +1,8 @@
+from .main import main as autogen_main
+from .main import comedy_show
+
+
+class AutoGenRunner:
+    def run(self):
+        # autogen_main()
+        comedy_show()
diff --git a/src/examples/autogen_example/main.py b/src/examples/autogen_example/main.py
new file mode 100644
index 00000000..9ed54c17
--- /dev/null
+++ b/src/examples/autogen_example/main.py
@@ -0,0 +1,72 @@
+from langtrace_python_sdk import langtrace
+from autogen import ConversableAgent
+from dotenv import load_dotenv
+from autogen.coding import LocalCommandLineCodeExecutor
+import tempfile
+
+
+load_dotenv()
+langtrace.init(write_spans_to_console=False)
+# agentops.init(api_key=os.getenv("AGENTOPS_API_KEY"))
+# Create a temporary directory to store the code files.
+temp_dir = tempfile.TemporaryDirectory()
+
+
+# Create a local command line code executor.
+executor = LocalCommandLineCodeExecutor(
+    timeout=10,  # Timeout for each code execution in seconds.
+    work_dir=temp_dir.name,  # Use the temporary directory to store the code files.
+)
+
+
+def main():
+
+    agent = ConversableAgent(
+        "chatbot",
+        llm_config={"config_list": [{"model": "gpt-4"}], "cache_seed": None},
+        code_execution_config=False,  # Turn off code execution, by default it is off.
+        function_map=None,  # No registered functions, by default it is None.
+        human_input_mode="NEVER",  # Never ask for human input.
+    )
+
+    reply = agent.generate_reply(
+        messages=[{"content": "Tell me a joke.", "role": "user"}]
+    )
+    return reply
+
+
+def comedy_show():
+    cathy = ConversableAgent(
+        name="cathy",
+        system_message="Your name is Cathy and you are a part of a duo of comedians.",
+        llm_config={
+            "config_list": [{"model": "gpt-4o-mini", "temperature": 0.9}],
+            "cache_seed": None,
+        },
+        description="Cathy is a comedian",
+        max_consecutive_auto_reply=10,
+        code_execution_config={
+            "executor": executor
+        },  # Use the local command line code executor.
+        function_map=None,
+        chat_messages=None,
+        silent=True,
+        default_auto_reply="Sorry, I don't know what to say.",
+        human_input_mode="NEVER",  # Never ask for human input.
+    )
+
+    joe = ConversableAgent(
+        "joe",
+        system_message="Your name is Joe and you are a part of a duo of comedians.",
+        llm_config={
+            "config_list": [{"model": "gpt-4o-mini", "temperature": 0.7}],
+            "cache_seed": None,
+        },
+        human_input_mode="NEVER",  # Never ask for human input.
+    )
+
+    result = joe.initiate_chat(
+        recipient=cathy, message="Cathy, tell me a joke.", max_turns=2
+    )
+
+    return result
diff --git a/src/langtrace_python_sdk/constants/instrumentation/common.py b/src/langtrace_python_sdk/constants/instrumentation/common.py
index 85710ba2..70d92a1b 100644
--- a/src/langtrace_python_sdk/constants/instrumentation/common.py
+++ b/src/langtrace_python_sdk/constants/instrumentation/common.py
@@ -31,6 +31,7 @@
     "GEMINI": "Gemini",
     "MISTRAL": "Mistral",
     "EMBEDCHAIN": "Embedchain",
+    "AUTOGEN": "Autogen",
 }
 
 LANGTRACE_ADDITIONAL_SPAN_ATTRIBUTES_KEY = "langtrace_additional_attributes"
diff --git a/src/langtrace_python_sdk/instrumentation/__init__.py b/src/langtrace_python_sdk/instrumentation/__init__.py
index b0e5fd26..984541dc 100644
--- a/src/langtrace_python_sdk/instrumentation/__init__.py
+++ b/src/langtrace_python_sdk/instrumentation/__init__.py
@@ -14,6 +14,7 @@
 from .weaviate import WeaviateInstrumentation
 from .ollama import OllamaInstrumentor
 from .dspy import DspyInstrumentation
+from .autogen import AutogenInstrumentation
 from .vertexai import VertexAIInstrumentation
 from .gemini import GeminiInstrumentation
 from .mistral import MistralInstrumentation
@@ -37,6 +38,7 @@
     "WeaviateInstrumentation",
     "OllamaInstrumentor",
     "DspyInstrumentation",
+    "AutogenInstrumentation",
     "VertexAIInstrumentation",
     "GeminiInstrumentation",
     "MistralInstrumentation",
diff --git a/src/langtrace_python_sdk/instrumentation/autogen/__init__.py b/src/langtrace_python_sdk/instrumentation/autogen/__init__.py
new file mode 100644
index 00000000..c152b011
--- /dev/null
+++ b/src/langtrace_python_sdk/instrumentation/autogen/__init__.py
@@ -0,0 +1,3 @@
+from .instrumentation import AutogenInstrumentation
+
+__all__ = ["AutogenInstrumentation"]
diff --git a/src/langtrace_python_sdk/instrumentation/autogen/instrumentation.py b/src/langtrace_python_sdk/instrumentation/autogen/instrumentation.py
new file mode 100644
index 00000000..06f89824
--- /dev/null
+++ b/src/langtrace_python_sdk/instrumentation/autogen/instrumentation.py
@@ -0,0 +1,42 @@
+from opentelemetry.instrumentation.instrumentor import BaseInstrumentor
+from opentelemetry.trace import get_tracer
+from wrapt import wrap_function_wrapper as _W
+from importlib_metadata import version as v
+from .patch import patch_generate_reply, patch_initiate_chat
+
+
+class AutogenInstrumentation(BaseInstrumentor):
+    def instrumentation_dependencies(self):
+        return ["autogen >= 0.1.0"]
+
+    def _instrument(self, **kwargs):
+        print("Instrumneting autogen")
+        tracer_provider = kwargs.get("tracer_provider")
+        tracer = get_tracer(__name__, "", tracer_provider)
+        version = v("autogen")
+        # conversable_agent.intiate_chat
+        # conversable_agent.register_function
+        # agent.Agent
+        # AgentCreation
+        # Tools --> Register_for_llm, register_for_execution, register_for_function
+        try:
+            _W(
+                module="autogen.agentchat.conversable_agent",
+                name="ConversableAgent.initiate_chat",
+                wrapper=patch_initiate_chat(
+                    "conversable_agent.initiate_chat", version, tracer
+                ),
+            )
+
+            _W(
+                module="autogen.agentchat.conversable_agent",
+                name="ConversableAgent.generate_reply",
+                wrapper=patch_generate_reply(
+                    "conversable_agent.generate_reply", version, tracer
+                ),
+            )
+        except Exception as e:
+            pass
+
+    def _uninstrument(self, **kwargs):
+        pass
diff --git a/src/langtrace_python_sdk/instrumentation/autogen/patch.py b/src/langtrace_python_sdk/instrumentation/autogen/patch.py
new file mode 100644
index 00000000..ce6c7ff2
--- /dev/null
+++ b/src/langtrace_python_sdk/instrumentation/autogen/patch.py
@@ -0,0 +1,132 @@
+from langtrace_python_sdk.utils.llm import (
+    get_langtrace_attributes,
+    get_extra_attributes,
+    get_span_name,
+    set_span_attributes,
+    get_llm_request_attributes,
+    set_event_completion,
+    set_usage_attributes,
+)
+from langtrace_python_sdk.constants.instrumentation.common import SERVICE_PROVIDERS
+from langtrace.trace_attributes import FrameworkSpanAttributes
+from opentelemetry.trace.status import Status, StatusCode
+from langtrace.trace_attributes import SpanAttributes
+from opentelemetry.trace import Tracer, SpanKind
+
+from langtrace_python_sdk.utils import deduce_args_and_kwargs, set_span_attribute
+import json
+
+
+def patch_initiate_chat(name, version, tracer: Tracer):
+    def traced_method(wrapped, instance, args, kwargs):
+        all_params = deduce_args_and_kwargs(wrapped, *args, **kwargs)
+        all_params["recipient"] = json.dumps(parse_agent(all_params.get("recipient")))
+        span_attributes = {
+            **get_langtrace_attributes(
+                service_provider=SERVICE_PROVIDERS["AUTOGEN"],
+                version=version,
+                vendor_type="framework",
+            ),
+            "sender": json.dumps(parse_agent(instance)),
+            **all_params,
+        }
+        attributes = FrameworkSpanAttributes(**span_attributes)
+
+        with tracer.start_as_current_span(
+            name=get_span_name(name), kind=SpanKind.CLIENT
+        ) as span:
+            try:
+                set_span_attributes(span, attributes)
+                result = wrapped(*args, **kwargs)
+                # set_response_attributes(span, result)
+                return result
+            except Exception as err:
+                # Record the exception in the span
+                span.record_exception(err)
+
+                # Set the span status to indicate an error
+                span.set_status(Status(StatusCode.ERROR, str(err)))
+
+                # Reraise the exception to ensure it's not swallowed
+                raise
+
+    return traced_method
+
+
+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],
+        }
+        service_provider = SERVICE_PROVIDERS["AUTOGEN"]
+
+        span_attributes = {
+            **get_langtrace_attributes(
+                version=version,
+                service_provider=service_provider,
+                vendor_type="framework",
+            ),
+            **get_llm_request_attributes(
+                kwargs,
+                prompts=kwargs.get("messages"),
+            ),
+            **get_extra_attributes(),
+        }
+        attributes = FrameworkSpanAttributes(**span_attributes)
+
+        with tracer.start_as_current_span(
+            name=get_span_name(name), kind=SpanKind.CLIENT
+        ) as span:
+            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:
+                    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)
+                )
+
+                return result
+
+            except Exception as err:
+                # Record the exception in the span
+                span.record_exception(err)
+
+                # Set the span status to indicate an error
+                span.set_status(Status(StatusCode.ERROR, str(err)))
+
+                # Reraise the exception to ensure it's not swallowed
+                raise
+
+    return traced_method
+
+
+def set_response_attributes(span, result):
+    summary = getattr(result, "summary", None)
+    if summary:
+        set_span_attribute(span, "autogen.chat.summary", summary)
+
+
+def parse_agent(agent):
+
+    return {
+        "name": getattr(agent, "name", None),
+        "description": getattr(agent, "description", None),
+        "system_message": str(getattr(agent, "system_message", None)),
+        "silent": getattr(agent, "silent", None),
+        "llm_config": str(getattr(agent, "llm_config", None)),
+        "human_input_mode": getattr(agent, "human_input_mode", None),
+    }
diff --git a/src/langtrace_python_sdk/langtrace.py b/src/langtrace_python_sdk/langtrace.py
index 0d6be8d1..6479cdfb 100644
--- a/src/langtrace_python_sdk/langtrace.py
+++ b/src/langtrace_python_sdk/langtrace.py
@@ -53,6 +53,7 @@
     OpenAIInstrumentation,
     PineconeInstrumentation,
     QdrantInstrumentation,
+    AutogenInstrumentation,
     VertexAIInstrumentation,
     WeaviateInstrumentation,
 )
@@ -141,6 +142,7 @@ def init(
         "google-cloud-aiplatform": VertexAIInstrumentation(),
         "google-generativeai": GeminiInstrumentation(),
         "mistralai": MistralInstrumentation(),
+        "autogen": AutogenInstrumentation(),
     }
 
     init_instrumentations(disable_instrumentations, all_instrumentations)
diff --git a/src/langtrace_python_sdk/utils/__init__.py b/src/langtrace_python_sdk/utils/__init__.py
index 43ffa5a7..d970ba3c 100644
--- a/src/langtrace_python_sdk/utils/__init__.py
+++ b/src/langtrace_python_sdk/utils/__init__.py
@@ -2,6 +2,7 @@
 from .sdk_version_checker import SDKVersionChecker
 from opentelemetry.trace import Span
 from langtrace.trace_attributes import SpanAttributes
+import inspect
 import os
 
 
@@ -28,6 +29,19 @@ def set_event_prompt(span: Span, prompt):
     )
 
 
+def deduce_args_and_kwargs(func, *args, **kwargs):
+    sig = inspect.signature(func)
+    bound_args = sig.bind(*args, **kwargs)
+    bound_args.apply_defaults()
+
+    all_params = {}
+    for param_name, param in sig.parameters.items():
+        if param_name in bound_args.arguments:
+            all_params[param_name] = bound_args.arguments[param_name]
+
+    return all_params
+
+
 def check_if_sdk_is_outdated():
     SDKVersionChecker().check()
     return
diff --git a/src/langtrace_python_sdk/utils/llm.py b/src/langtrace_python_sdk/utils/llm.py
index 21594e58..6d9647e7 100644
--- a/src/langtrace_python_sdk/utils/llm.py
+++ b/src/langtrace_python_sdk/utils/llm.py
@@ -126,7 +126,9 @@ def get_llm_request_attributes(kwargs, prompts=None, model=None, operation_name=
     tools = kwargs.get("tools", None)
     return {
         SpanAttributes.LLM_OPERATION_NAME: operation_name,
-        SpanAttributes.LLM_REQUEST_MODEL: model or kwargs.get("model") or "gpt-3.5-turbo",
+        SpanAttributes.LLM_REQUEST_MODEL: model
+        or kwargs.get("model")
+        or "gpt-3.5-turbo",
         SpanAttributes.LLM_IS_STREAMING: kwargs.get("stream"),
         SpanAttributes.LLM_REQUEST_TEMPERATURE: kwargs.get("temperature"),
         SpanAttributes.LLM_TOP_K: top_k,
@@ -230,7 +232,15 @@ def set_event_completion(span: Span, result_content):
 
 
 def set_span_attributes(span: Span, attributes: Any) -> None:
-    for field, value in attributes.model_dump(by_alias=True).items():
+    from pydantic import BaseModel
+
+    attrs = (
+        attributes.model_dump(by_alias=True)
+        if isinstance(attributes, BaseModel)
+        else attributes
+    )
+
+    for field, value in attrs.items():
         set_span_attribute(span, field, value)
 
 
diff --git a/src/run_example.py b/src/run_example.py
index cddd1e42..864bc8c2 100644
--- a/src/run_example.py
+++ b/src/run_example.py
@@ -2,11 +2,11 @@
 
 ENABLED_EXAMPLES = {
     "anthropic": False,
-    "azureopenai": True,
+    "azureopenai": False,
     "chroma": False,
     "cohere": False,
     "fastapi": False,
-    "langchain": True,
+    "langchain": False,
     "llamaindex": False,
     "hiveagent": False,
     "openai": False,
@@ -16,6 +16,7 @@
     "weaviate": False,
     "ollama": False,
     "groq": False,
+    "autogen": True,
     "vertexai": False,
     "gemini": False,
     "mistral": False,
@@ -93,6 +94,11 @@
     print(Fore.BLUE + "Running Groq example" + Fore.RESET)
     GroqRunner().run()
 
+if ENABLED_EXAMPLES["autogen"]:
+    from examples.autogen_example import AutoGenRunner
+
+    print(Fore.BLUE + "Running Autogen example" + Fore.RESET)
+    AutoGenRunner().run()
 
 if ENABLED_EXAMPLES["vertexai"]:
     from examples.vertexai_example import VertexAIRunner

From c514af466c9ce4e975d84704bd69064800d99b97 Mon Sep 17 00:00:00 2001
From: Rohit Kadhe <113367036+rohit-kadhe@users.noreply.github.com>
Date: Wed, 11 Sep 2024 14:16:20 -0600
Subject: [PATCH 5/9] Support genai and also add token reporting and other data
 points (#345)

* support genai and also add token reporting

* rename example file

* fix

---------

Co-authored-by: Karthik Kalyanaraman <karthik@scale3labs.com>
---
 src/examples/langchain_example/__init__.py    |  2 +
 .../langchain_google_example.py               | 29 ++++++++++
 .../langchain/instrumentation.py              |  9 +--
 .../langchain_core/instrumentation.py         | 14 +++++
 .../instrumentation/langchain_core/patch.py   | 55 +++++++++++++------
 src/langtrace_python_sdk/langtrace.py         |  7 ++-
 6 files changed, 89 insertions(+), 27 deletions(-)
 create mode 100644 src/examples/langchain_example/langchain_google_example.py

diff --git a/src/examples/langchain_example/__init__.py b/src/examples/langchain_example/__init__.py
index f421b6b1..ceaee24c 100644
--- a/src/examples/langchain_example/__init__.py
+++ b/src/examples/langchain_example/__init__.py
@@ -1,3 +1,4 @@
+from examples.langchain_example.langchain_google_genai import basic_google_genai
 from .basic import basic_app, rag, load_and_split
 from langtrace_python_sdk import with_langtrace_root_span
 
@@ -12,6 +13,7 @@ def run(self):
         rag()
         load_and_split()
         basic_graph_tools()
+        basic_google_genai()
 
 
 class GroqRunner:
diff --git a/src/examples/langchain_example/langchain_google_example.py b/src/examples/langchain_example/langchain_google_example.py
new file mode 100644
index 00000000..fc92c8b9
--- /dev/null
+++ b/src/examples/langchain_example/langchain_google_example.py
@@ -0,0 +1,29 @@
+from langchain_core.messages import HumanMessage
+from langchain_google_genai import ChatGoogleGenerativeAI
+from langtrace_python_sdk.utils.with_root_span import with_langtrace_root_span
+from dotenv import find_dotenv, load_dotenv
+from langtrace_python_sdk import langtrace
+
+_ = load_dotenv(find_dotenv())
+
+langtrace.init()
+
+@with_langtrace_root_span("basic_google_genai")
+def basic_google_genai():
+    llm = ChatGoogleGenerativeAI(model="gemini-1.5-flash")
+    # example
+    message = HumanMessage(
+        content=[
+            {
+                "type": "text",
+                "text": "What's in this image?",
+            },
+        ]
+    )
+    message_image = HumanMessage(content="https://picsum.photos/seed/picsum/200/300")
+
+    res = llm.invoke([message, message_image])
+    # print(res)
+
+
+basic_google_genai()
diff --git a/src/langtrace_python_sdk/instrumentation/langchain/instrumentation.py b/src/langtrace_python_sdk/instrumentation/langchain/instrumentation.py
index 4cc4fe8f..7d01906e 100644
--- a/src/langtrace_python_sdk/instrumentation/langchain/instrumentation.py
+++ b/src/langtrace_python_sdk/instrumentation/langchain/instrumentation.py
@@ -81,21 +81,16 @@ def _instrument(self, **kwargs):
         tracer_provider = kwargs.get("tracer_provider")
         tracer = get_tracer(__name__, "", tracer_provider)
         version = importlib.metadata.version("langchain")
-
         wrap_function_wrapper(
             "langchain.agents.agent",
             "RunnableAgent.plan",
-            generic_patch(
-                "RunnableAgent.plan", "plan", tracer, version, True, True
-            ),
+            generic_patch("RunnableAgent.plan", "plan", tracer, version, True, True),
         )
 
         wrap_function_wrapper(
             "langchain.agents.agent",
             "RunnableAgent.aplan",
-            generic_patch(
-                "RunnableAgent.aplan", "plan", tracer, version, True, True
-            ),
+            generic_patch("RunnableAgent.aplan", "plan", tracer, version, True, True),
         )
 
         # modules_to_patch = []
diff --git a/src/langtrace_python_sdk/instrumentation/langchain_core/instrumentation.py b/src/langtrace_python_sdk/instrumentation/langchain_core/instrumentation.py
index e10a6510..a7567ff0 100644
--- a/src/langtrace_python_sdk/instrumentation/langchain_core/instrumentation.py
+++ b/src/langtrace_python_sdk/instrumentation/langchain_core/instrumentation.py
@@ -134,6 +134,20 @@ def _instrument(self, **kwargs):
         ]
 
         modules_to_patch = [
+            (
+                "langchain_core.language_models.chat_models",
+                "chatmodel",
+                generic_patch,
+                True,
+                True,
+            ),
+            (
+                "langchain_core.language_models.base",
+                "language_model",
+                generic_patch,
+                True,
+                True,
+            ),
             ("langchain_core.retrievers", "retriever", generic_patch, True, True),
             ("langchain_core.prompts.chat", "prompt", generic_patch, True, True),
             (
diff --git a/src/langtrace_python_sdk/instrumentation/langchain_core/patch.py b/src/langtrace_python_sdk/instrumentation/langchain_core/patch.py
index 765c738e..ce93c4c5 100644
--- a/src/langtrace_python_sdk/instrumentation/langchain_core/patch.py
+++ b/src/langtrace_python_sdk/instrumentation/langchain_core/patch.py
@@ -57,7 +57,24 @@ def traced_method(wrapped, instance, args, kwargs):
             "langtrace.service.version": version,
             "langtrace.version": v(LANGTRACE_SDK_NAME),
             "langchain.task.name": task,
-            **(extra_attributes if extra_attributes is not None else {}),
+            "gen_ai.request.model": (
+                instance.model if hasattr(instance, "model") else None
+            ),
+            SpanAttributes.LLM_REQUEST_MAX_TOKENS: (
+                instance.max_output_tokens
+                if hasattr(instance, "max_output_tokens")
+                else None
+            ),
+            SpanAttributes.LLM_TOP_K: (
+                instance.top_k if hasattr(instance, "top_k") else None
+            ),
+            SpanAttributes.LLM_REQUEST_TOP_P: (
+                instance.top_p if hasattr(instance, "top_p") else None
+            ),
+            SpanAttributes.LLM_REQUEST_TEMPERATURE: (
+                instance.temperature if hasattr(instance, "temperature") else None
+            ),
+            **(extra_attributes if extra_attributes is not None else {}),  # type: ignore
         }
 
         if trace_input and len(args) > 0:
@@ -79,21 +96,17 @@ def traced_method(wrapped, instance, args, kwargs):
             try:
                 # Attempt to call the original method
                 result = wrapped(*args, **kwargs)
-
                 if trace_output:
                     span.set_attribute("langchain.outputs", to_json_string(result))
-                    if hasattr(result, 'usage'):
-                        prompt_tokens = result.usage.prompt_tokens
-                        completion_tokens = result.usage.completion_tokens
-                        span.set_attribute(SpanAttributes.LLM_USAGE_PROMPT_TOKENS, prompt_tokens)
-                        span.set_attribute(SpanAttributes.LLM_USAGE_COMPLETION_TOKENS, completion_tokens)
-                    elif hasattr(result, 'generations') and len(result.generations) > 0 and len(result.generations[0]) > 0 and hasattr(result.generations[0][0], 'text') and isinstance(result.generations[0][0].text, str):
-                        span.set_attribute(SpanAttributes.LLM_USAGE_COMPLETION_TOKENS, instance.get_num_tokens(result.generations[0][0].text))
-                    elif len(args) > 0 and len(args[0]) > 0 and not hasattr(args[0][0], 'text') and hasattr(instance, 'get_num_tokens'):
-                        span.set_attribute(SpanAttributes.LLM_USAGE_PROMPT_TOKENS, instance.get_num_tokens(args[0][0]))
-                    elif len(args) > 0 and len(args[0]) > 0 and hasattr(args[0][0], 'text') and isinstance(args[0][0].text, str) and hasattr(instance, 'get_num_tokens'):
-                        span.set_attribute(SpanAttributes.LLM_USAGE_PROMPT_TOKENS, instance.get_num_tokens(args[0][0].text))
-
+                    if hasattr(result, "usage_metadata"):
+                        span.set_attribute(
+                            SpanAttributes.LLM_USAGE_PROMPT_TOKENS,
+                            result.usage_metadata["input_tokens"],
+                        )
+                        span.set_attribute(
+                            SpanAttributes.LLM_USAGE_COMPLETION_TOKENS,
+                            result.usage_metadata["output_tokens"],
+                        )
                 span.set_status(StatusCode.OK)
                 return result
             except Exception as err:
@@ -208,9 +221,17 @@ def clean_empty(d):
     if not isinstance(d, (dict, list, tuple)):
         return d
     if isinstance(d, tuple):
-        return tuple(val for val in (clean_empty(val) for val in d) if val != () and val is not None)
+        return tuple(
+            val
+            for val in (clean_empty(val) for val in d)
+            if val != () and val is not None
+        )
     if isinstance(d, list):
-        return [val for val in (clean_empty(val) for val in d) if val != [] and val is not None]
+        return [
+            val
+            for val in (clean_empty(val) for val in d)
+            if val != [] and val is not None
+        ]
     result = {}
     for k, val in d.items():
         if isinstance(val, dict):
@@ -226,7 +247,7 @@ def clean_empty(d):
                 result[k] = val.strip()
         elif isinstance(val, object):
             # some langchain objects have a text attribute
-            val = getattr(val, 'text', None)
+            val = getattr(val, "text", None)
             if val is not None and val.strip() != "":
                 result[k] = val.strip()
     return result
diff --git a/src/langtrace_python_sdk/langtrace.py b/src/langtrace_python_sdk/langtrace.py
index 6479cdfb..86fe0d12 100644
--- a/src/langtrace_python_sdk/langtrace.py
+++ b/src/langtrace_python_sdk/langtrace.py
@@ -129,8 +129,8 @@ def init(
         "embedchain": EmbedchainInstrumentation(),
         "qdrant-client": QdrantInstrumentation(),
         "langchain": LangchainInstrumentation(),
-        "langchain-core": LangchainCoreInstrumentation(),
-        "langchain-community": LangchainCommunityInstrumentation(),
+        "langchain_core": LangchainCoreInstrumentation(),
+        "langchain_community": LangchainCommunityInstrumentation(),
         "langgraph": LanggraphInstrumentation(),
         "anthropic": AnthropicInstrumentation(),
         "cohere": CohereInstrumentation(),
@@ -192,7 +192,8 @@ def init(
 
 
 def init_instrumentations(
-    disable_instrumentations: DisableInstrumentations, all_instrumentations: dict
+    disable_instrumentations: Optional[DisableInstrumentations],
+    all_instrumentations: dict
 ):
     if disable_instrumentations is None:
         for idx, (name, v) in enumerate(all_instrumentations.items()):

From 79d51d83dd3e1537ad503dcabfda3522c761c6f9 Mon Sep 17 00:00:00 2001
From: darshit-s3 <119623510+darshit-s3@users.noreply.github.com>
Date: Thu, 12 Sep 2024 01:49:33 +0530
Subject: [PATCH 6/9] fix: weaviate datetime handling for request and response
 (#346)

* fix: weaviate datetime handling for request and response

* bump version

* bump version

---------

Co-authored-by: Karthik Kalyanaraman <karthik@scale3labs.com>
---
 .../instrumentation/weaviate/patch.py                 | 11 +++++++++--
 src/langtrace_python_sdk/version.py                   |  2 +-
 2 files changed, 10 insertions(+), 3 deletions(-)

diff --git a/src/langtrace_python_sdk/instrumentation/weaviate/patch.py b/src/langtrace_python_sdk/instrumentation/weaviate/patch.py
index 523b047a..5c295ef1 100644
--- a/src/langtrace_python_sdk/instrumentation/weaviate/patch.py
+++ b/src/langtrace_python_sdk/instrumentation/weaviate/patch.py
@@ -15,6 +15,7 @@
 """
 
 import json
+from datetime import datetime
 
 from importlib_metadata import version as v
 from langtrace.trace_attributes import DatabaseSpanAttributes
@@ -48,9 +49,11 @@
 def extract_inputs(args, kwargs):
     extracted_params = {}
     kwargs_without_properties = {
-        k: v for k, v in kwargs.items() if k not in ["properties", "fusion_type"]
+        k: v for k, v in kwargs.items() if k not in ["properties", "fusion_type", "filters"]
     }
     extracted_params.update(extract_input_params(args, kwargs_without_properties))
+    if kwargs.get("filters", None):
+        extracted_params["filters"] = str(kwargs["filters"])
     if kwargs.get("fusion_type", None):
         extracted_params["fusion_type"] = kwargs["fusion_type"].value
     if kwargs.get("properties", None):
@@ -95,9 +98,13 @@ def aggregate_responses(result):
 
 
 def get_response_object_attributes(response_object):
+    def convert_value(value):
+        if isinstance(value, datetime):
+            return value.isoformat()
+        return value
 
     response_attributes = {
-        **response_object.properties,
+        **{k: convert_value(v) for k, v in response_object.properties.items()},
         "uuid": str(response_object.uuid) if hasattr(response_object, "uuid") else None,
         "collection": getattr(response_object, "collection", None),
         "vector": getattr(response_object, "vector", None),
diff --git a/src/langtrace_python_sdk/version.py b/src/langtrace_python_sdk/version.py
index 2559bc66..2aa24fe6 100644
--- a/src/langtrace_python_sdk/version.py
+++ b/src/langtrace_python_sdk/version.py
@@ -1 +1 @@
-__version__ = "2.3.12"
+__version__ = "2.3.13"

From 7ba1161dcb9e5aa482be776d3d6933ec31f40f83 Mon Sep 17 00:00:00 2001
From: Karthik Kalyanaraman <105607645+karthikscale3@users.noreply.github.com>
Date: Wed, 11 Sep 2024 16:23:41 -0700
Subject: [PATCH 7/9] Minor bugfix to langchain instrumentation (#348)

* Bugfix

* bump version
---
 src/langtrace_python_sdk/langtrace.py | 4 ++--
 src/langtrace_python_sdk/version.py   | 2 +-
 2 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/src/langtrace_python_sdk/langtrace.py b/src/langtrace_python_sdk/langtrace.py
index 86fe0d12..b1345d89 100644
--- a/src/langtrace_python_sdk/langtrace.py
+++ b/src/langtrace_python_sdk/langtrace.py
@@ -129,8 +129,8 @@ def init(
         "embedchain": EmbedchainInstrumentation(),
         "qdrant-client": QdrantInstrumentation(),
         "langchain": LangchainInstrumentation(),
-        "langchain_core": LangchainCoreInstrumentation(),
-        "langchain_community": LangchainCommunityInstrumentation(),
+        "langchain-core": LangchainCoreInstrumentation(),
+        "langchain-community": LangchainCommunityInstrumentation(),
         "langgraph": LanggraphInstrumentation(),
         "anthropic": AnthropicInstrumentation(),
         "cohere": CohereInstrumentation(),
diff --git a/src/langtrace_python_sdk/version.py b/src/langtrace_python_sdk/version.py
index 2aa24fe6..bff004f2 100644
--- a/src/langtrace_python_sdk/version.py
+++ b/src/langtrace_python_sdk/version.py
@@ -1 +1 @@
-__version__ = "2.3.13"
+__version__ = "2.3.14"

From ac726b5a1da6a511c4cc81274ba0c186f8074aa2 Mon Sep 17 00:00:00 2001
From: darshit-s3 <119623510+darshit-s3@users.noreply.github.com>
Date: Thu, 12 Sep 2024 19:31:36 +0530
Subject: [PATCH 8/9] fix: handling for openai NOT_GIVEN default arg value
 (#350)

* fix: handling for openai NOT_GIVEN default arg value

* style: fix formating
---
 .../instrumentation/openai/patch.py           | 24 +++++++++++++------
 src/langtrace_python_sdk/version.py           |  2 +-
 2 files changed, 18 insertions(+), 8 deletions(-)

diff --git a/src/langtrace_python_sdk/instrumentation/openai/patch.py b/src/langtrace_python_sdk/instrumentation/openai/patch.py
index b56418fc..bd4f6102 100644
--- a/src/langtrace_python_sdk/instrumentation/openai/patch.py
+++ b/src/langtrace_python_sdk/instrumentation/openai/patch.py
@@ -1,4 +1,5 @@
 import json
+import openai
 from typing import Any, Dict, List, Optional, Callable, Awaitable, Union
 from langtrace.trace_attributes import (
     LLMSpanAttributes,
@@ -40,6 +41,15 @@
 )
 
 
+def filter_valid_attributes(attributes):
+    """Filter attributes where value is not None, not an empty string, and not openai.NOT_GIVEN."""
+    return {
+        key: value
+        for key, value in attributes.items()
+        if value is not None and value != openai.NOT_GIVEN and value != ""
+    }
+
+
 def images_generate(version: str, tracer: Tracer) -> Callable:
     """
     Wrap the `generate` method of the `Images` class to trace it.
@@ -57,7 +67,7 @@ def traced_method(
             **get_extra_attributes(),  # type: ignore
         }
 
-        attributes = LLMSpanAttributes(**span_attributes)
+        attributes = LLMSpanAttributes(**filter_valid_attributes(span_attributes))
 
         with tracer.start_as_current_span(
             name=get_span_name(APIS["IMAGES_GENERATION"]["METHOD"]),
@@ -118,7 +128,7 @@ async def traced_method(
             **get_extra_attributes(),  # type: ignore
         }
 
-        attributes = LLMSpanAttributes(**span_attributes)
+        attributes = LLMSpanAttributes(**filter_valid_attributes(span_attributes))
 
         with tracer.start_as_current_span(
             name=get_span_name(APIS["IMAGES_GENERATION"]["METHOD"]),
@@ -181,7 +191,7 @@ def traced_method(
             **get_extra_attributes(),  # type: ignore
         }
 
-        attributes = LLMSpanAttributes(**span_attributes)
+        attributes = LLMSpanAttributes(**filter_valid_attributes(span_attributes))
 
         with tracer.start_as_current_span(
             name=APIS["IMAGES_EDIT"]["METHOD"],
@@ -268,7 +278,7 @@ def traced_method(
             **get_extra_attributes(),  # type: ignore
         }
 
-        attributes = LLMSpanAttributes(**span_attributes)
+        attributes = LLMSpanAttributes(**filter_valid_attributes(span_attributes))
 
         span = tracer.start_span(
             name=get_span_name(APIS["CHAT_COMPLETION"]["METHOD"]),
@@ -356,7 +366,7 @@ async def traced_method(
             **get_extra_attributes(),  # type: ignore
         }
 
-        attributes = LLMSpanAttributes(**span_attributes)
+        attributes = LLMSpanAttributes(**filter_valid_attributes(span_attributes))
 
         span = tracer.start_span(
             name=get_span_name(APIS["CHAT_COMPLETION"]["METHOD"]),
@@ -438,7 +448,7 @@ def traced_method(
                 [kwargs.get("input", "")]
             )
 
-        attributes = LLMSpanAttributes(**span_attributes)
+        attributes = LLMSpanAttributes(**filter_valid_attributes(span_attributes))
 
         with tracer.start_as_current_span(
             name=get_span_name(APIS["EMBEDDINGS_CREATE"]["METHOD"]),
@@ -487,7 +497,7 @@ async def traced_method(
             **get_extra_attributes(),  # type: ignore
         }
 
-        attributes = LLMSpanAttributes(**span_attributes)
+        attributes = LLMSpanAttributes(**filter_valid_attributes(span_attributes))
 
         encoding_format = kwargs.get("encoding_format")
         if encoding_format is not None:
diff --git a/src/langtrace_python_sdk/version.py b/src/langtrace_python_sdk/version.py
index bff004f2..ab20ff9c 100644
--- a/src/langtrace_python_sdk/version.py
+++ b/src/langtrace_python_sdk/version.py
@@ -1 +1 @@
-__version__ = "2.3.14"
+__version__ = "2.3.15"

From 138c013c157959dc28bbcb30c416c5adebff0225 Mon Sep 17 00:00:00 2001
From: darshit-s3 <119623510+darshit-s3@users.noreply.github.com>
Date: Mon, 16 Sep 2024 23:13:06 +0530
Subject: [PATCH 9/9] fix: tool choice for groq, datetime for cohere (#352)

---
 src/examples/cohere_example/rerank.py         | 21 ++++-
 src/examples/langchain_example/__init__.py    |  4 +-
 .../langchain_example/groq_example.py         | 80 ++++++++++++++++++-
 .../instrumentation/cohere/patch.py           |  5 +-
 src/langtrace_python_sdk/utils/llm.py         |  3 +-
 src/langtrace_python_sdk/utils/misc.py        |  8 ++
 src/langtrace_python_sdk/version.py           |  2 +-
 7 files changed, 113 insertions(+), 10 deletions(-)

diff --git a/src/examples/cohere_example/rerank.py b/src/examples/cohere_example/rerank.py
index de1a21f5..3995feb5 100644
--- a/src/examples/cohere_example/rerank.py
+++ b/src/examples/cohere_example/rerank.py
@@ -1,5 +1,6 @@
 import cohere
 from dotenv import find_dotenv, load_dotenv
+from datetime import datetime
 
 from langtrace_python_sdk import langtrace
 
@@ -16,10 +17,22 @@
 # @with_langtrace_root_span("embed_create")
 def rerank():
     docs = [
-        "Carson City is the capital city of the American state of Nevada.",
-        "The Commonwealth of the Northern Mariana Islands is a group of islands in the Pacific Ocean. Its capital is Saipan.",
-        "Washington, D.C. (also known as simply Washington or D.C., and officially as the District of Columbia) is the capital of the United States. It is a federal district.",
-        "Capital punishment (the death penalty) has existed in the United States since beforethe United States was a country. As of 2017, capital punishment is legal in 30 of the 50 states.",
+        {
+            "text": "Carson City is the capital city of the American state of Nevada.",
+            "date": datetime.now(),
+        },
+        {
+            "text": "The Commonwealth of the Northern Mariana Islands is a group of islands in the Pacific Ocean. Its capital is Saipan.",
+            "date": datetime(2020, 5, 17),
+        },
+        {
+            "text": "Washington, D.C. (also known as simply Washington or D.C., and officially as the District of Columbia) is the capital of the United States. It is a federal district.",
+            "date": datetime(1776, 7, 4),
+        },
+        {
+            "text": "Capital punishment (the death penalty) has existed in the United States since before the United States was a country. As of 2017, capital punishment is legal in 30 of the 50 states.",
+            "date": datetime(2023, 9, 14),
+        },
     ]
 
     response = co.rerank(
diff --git a/src/examples/langchain_example/__init__.py b/src/examples/langchain_example/__init__.py
index ceaee24c..f9621a2c 100644
--- a/src/examples/langchain_example/__init__.py
+++ b/src/examples/langchain_example/__init__.py
@@ -2,7 +2,7 @@
 from .basic import basic_app, rag, load_and_split
 from langtrace_python_sdk import with_langtrace_root_span
 
-from .groq_example import groq_basic, groq_streaming
+from .groq_example import groq_basic, groq_tool_choice, groq_streaming
 from .langgraph_example_tools import basic_graph_tools
 
 
@@ -20,3 +20,5 @@ class GroqRunner:
     @with_langtrace_root_span("Groq")
     def run(self):
         groq_streaming()
+        groq_basic()
+        groq_tool_choice()
diff --git a/src/examples/langchain_example/groq_example.py b/src/examples/langchain_example/groq_example.py
index c16e4851..99a61761 100644
--- a/src/examples/langchain_example/groq_example.py
+++ b/src/examples/langchain_example/groq_example.py
@@ -1,6 +1,6 @@
+import json
+
 from dotenv import find_dotenv, load_dotenv
-from langchain_core.prompts import ChatPromptTemplate
-from langchain_groq import ChatGroq
 from groq import Groq
 
 _ = load_dotenv(find_dotenv())
@@ -30,6 +30,82 @@ def groq_basic():
     return chat_completion
 
 
+def groq_tool_choice():
+
+    user_prompt = "What is 25 * 4 + 10?"
+    MODEL = "llama3-groq-70b-8192-tool-use-preview"
+
+    def calculate(expression):
+        """Evaluate a mathematical expression"""
+        try:
+            result = eval(expression)
+            return json.dumps({"result": result})
+        except:
+            return json.dumps({"error": "Invalid expression"})
+
+    messages = [
+        {
+            "role": "system",
+            "content": "You are a calculator assistant. Use the calculate function to perform mathematical operations and provide the results.",
+        },
+        {
+            "role": "user",
+            "content": user_prompt,
+        },
+    ]
+    tools = [
+        {
+            "type": "function",
+            "function": {
+                "name": "calculate",
+                "description": "Evaluate a mathematical expression",
+                "parameters": {
+                    "type": "object",
+                    "properties": {
+                        "expression": {
+                            "type": "string",
+                            "description": "The mathematical expression to evaluate",
+                        }
+                    },
+                    "required": ["expression"],
+                },
+            },
+        }
+    ]
+    response = client.chat.completions.create(
+        model=MODEL,
+        messages=messages,
+        tools=tools,
+        tool_choice={"type": "function", "function": {"name": "calculate"}},
+        max_tokens=4096,
+    )
+
+    response_message = response.choices[0].message
+    tool_calls = response_message.tool_calls
+    if tool_calls:
+        available_functions = {
+            "calculate": calculate,
+        }
+        messages.append(response_message)
+        for tool_call in tool_calls:
+            function_name = tool_call.function.name
+            function_to_call = available_functions[function_name]
+            function_args = json.loads(tool_call.function.arguments)
+            function_response = function_to_call(
+                expression=function_args.get("expression")
+            )
+            messages.append(
+                {
+                    "tool_call_id": tool_call.id,
+                    "role": "tool",
+                    "name": function_name,
+                    "content": function_response,
+                }
+            )
+        second_response = client.chat.completions.create(model=MODEL, messages=messages)
+        return second_response.choices[0].message.content
+
+
 def groq_streaming():
     chat_completion = client.chat.completions.create(
         messages=[
diff --git a/src/langtrace_python_sdk/instrumentation/cohere/patch.py b/src/langtrace_python_sdk/instrumentation/cohere/patch.py
index c165a10c..38908c3d 100644
--- a/src/langtrace_python_sdk/instrumentation/cohere/patch.py
+++ b/src/langtrace_python_sdk/instrumentation/cohere/patch.py
@@ -27,6 +27,7 @@
 )
 from langtrace.trace_attributes import Event, LLMSpanAttributes
 from langtrace_python_sdk.utils import set_span_attribute
+from langtrace_python_sdk.utils.misc import datetime_encoder
 from opentelemetry.trace import SpanKind
 from opentelemetry.trace.status import Status, StatusCode
 
@@ -50,7 +51,9 @@ def traced_method(wrapped, instance, args, kwargs):
             SpanAttributes.LLM_REQUEST_MODEL: kwargs.get("model") or "command-r-plus",
             SpanAttributes.LLM_URL: APIS["RERANK"]["URL"],
             SpanAttributes.LLM_PATH: APIS["RERANK"]["ENDPOINT"],
-            SpanAttributes.LLM_REQUEST_DOCUMENTS: json.dumps(kwargs.get("documents")),
+            SpanAttributes.LLM_REQUEST_DOCUMENTS: json.dumps(
+                kwargs.get("documents"), cls=datetime_encoder
+            ),
             SpanAttributes.LLM_COHERE_RERANK_QUERY: kwargs.get("query"),
             **get_extra_attributes(),
         }
diff --git a/src/langtrace_python_sdk/utils/llm.py b/src/langtrace_python_sdk/utils/llm.py
index 6d9647e7..e1d6fa9d 100644
--- a/src/langtrace_python_sdk/utils/llm.py
+++ b/src/langtrace_python_sdk/utils/llm.py
@@ -124,6 +124,7 @@ def get_llm_request_attributes(kwargs, prompts=None, model=None, operation_name=
 
     top_p = kwargs.get("p", None) or kwargs.get("top_p", None)
     tools = kwargs.get("tools", None)
+    tool_choice = kwargs.get("tool_choice", None)
     return {
         SpanAttributes.LLM_OPERATION_NAME: operation_name,
         SpanAttributes.LLM_REQUEST_MODEL: model
@@ -141,7 +142,7 @@ def get_llm_request_attributes(kwargs, prompts=None, model=None, operation_name=
         SpanAttributes.LLM_FREQUENCY_PENALTY: kwargs.get("frequency_penalty"),
         SpanAttributes.LLM_REQUEST_SEED: kwargs.get("seed"),
         SpanAttributes.LLM_TOOLS: json.dumps(tools) if tools else None,
-        SpanAttributes.LLM_TOOL_CHOICE: kwargs.get("tool_choice"),
+        SpanAttributes.LLM_TOOL_CHOICE: json.dumps(tool_choice) if tool_choice else None,
         SpanAttributes.LLM_REQUEST_LOGPROPS: kwargs.get("logprobs"),
         SpanAttributes.LLM_REQUEST_LOGITBIAS: kwargs.get("logit_bias"),
         SpanAttributes.LLM_REQUEST_TOP_LOGPROPS: kwargs.get("top_logprobs"),
diff --git a/src/langtrace_python_sdk/utils/misc.py b/src/langtrace_python_sdk/utils/misc.py
index a0d20452..56924cfe 100644
--- a/src/langtrace_python_sdk/utils/misc.py
+++ b/src/langtrace_python_sdk/utils/misc.py
@@ -60,3 +60,11 @@ def is_serializable(value):
 
     # Convert to string representation
     return json.dumps(serializable_args)
+
+
+class datetime_encoder(json.JSONEncoder):
+    def default(self, o):
+        if isinstance(o, datetime):
+            return o.isoformat()
+
+        return json.JSONEncoder.default(self, o)
diff --git a/src/langtrace_python_sdk/version.py b/src/langtrace_python_sdk/version.py
index ab20ff9c..e4f37b40 100644
--- a/src/langtrace_python_sdk/version.py
+++ b/src/langtrace_python_sdk/version.py
@@ -1 +1 @@
-__version__ = "2.3.15"
+__version__ = "2.3.16"