🐛 Bug Report: `aentity_class` Missing Required Positional Argument: `version`

[Carlososuna11]

Which component is this bug for?

Traceloop SDK

📜 Description

When attempting to use the @aagent decorator from the traceloop-sdk==0.26.5, a TypeError is raised indicating that aentity_class is missing a required positional argument: version.

👟 Reproduction steps

1. Implement the @aagent decorator in your code as shown below:

  from traceloop.sdk.decorators import aagent
  from some_module import AbstractAgent, get_logger

  logger = get_logger("CompletionAgent")

  @aagent(
      name="CompletionAgent",
      version=1,
      method_name="generate_sse_response",
  )
  class CompletionAgent(AbstractAgent):
      def start_llm(self, request: GenericAgentRequestSchema):
          self.llm = LLMManager().get_generator_from_request(request)
          logger.debug(f"Starting LLM generator for {request.type}")
          self.request = request

2. Run the application.
3. Observe the following error:

👍 Expected behavior

It should properly initialize the CompletionAgent class using the @aagent decorator from the traceloop-sdk library version 0.26.5, without raising any errors. The @aagent decorator should correctly pass all the required parameters, including name, version, and method_name, to the aentity_class function. The decorated class should be able to instantiate and execute methods as expected, enabling the functionality of the CompletionAgent without any issues related to missing arguments.

👎 Actual Behavior with Screenshots

It actually raises a TypeError, indicating that the aentity_class function is missing a required positional argument: version. This error occurs when the @aagent decorator tries to apply the atask function, which fails to pass the version parameter to the aentity_class function. As a result, the CompletionAgent class cannot be initialized properly, and the expected functionality is interrupted by this error.

🤖 Python Version

Python 3.10

📃 Provide any additional context for the Bug.

Suspected Cause

The issue seems to be in the atask function within the traceloop SDK. The version parameter is not being passed to the aentity_class function when method_name is provided.

Relevant Code Snippet

The problematic code appears to be here:

def atask(
    name: Optional[str] = None,
    version: Optional[int] = None,
    method_name: Optional[str] = None,
    tlp_span_kind: Optional[TraceloopSpanKindValues] = TraceloopSpanKindValues.TASK,
):
    if method_name is None:
        return aentity_method(name=name, version=version, tlp_span_kind=tlp_span_kind)
    else:
        return aentity_class(
            name=name, method_name=method_name, tlp_span_kind=tlp_span_kind
        )

The version parameter is not passed to aentity_class in the else block.

Suggested Fix

Modify the else block to include the version parameter:

def atask(
    name: Optional[str] = None,
    version: Optional[int] = None,
    method_name: Optional[str] = None,
    tlp_span_kind: Optional[TraceloopSpanKindValues] = TraceloopSpanKindValues.TASK,
):
    if method_name is None:
        return aentity_method(name=name, version=version, tlp_span_kind=tlp_span_kind)
    else:
        return aentity_class(
            name=name, version=version, method_name=method_name, tlp_span_kind=tlp_span_kind
        )

👀 Have you spent some time to check if this bug has been raised before?

• [X] I checked and didn't find similar issue

Are you willing to submit PR?

Yes I am willing to submit a PR!

[nirga]

Thanks for reporting @Carlososuna11! I was about to fix it but saw you're willing to submit a PR, so go ahead :)

[Carlososuna11]

ok! It's the #1816

[nirga]

Fixed with #1816