Skip to content

Commit

Permalink
Revert "[Draft][Tracing] Modify openai span names (#2424)" (#2457)
Browse files Browse the repository at this point in the history
This reverts commit 3ed2542.

# Description

Please add an informative description that covers that changes made by
the pull request and link all relevant issues.

# All Promptflow Contribution checklist:
- [ ] **The pull request does not introduce [breaking changes].**
- [ ] **CHANGELOG is updated for new features, bug fixes or other
significant changes.**
- [ ] **I have read the [contribution guidelines](../CONTRIBUTING.md).**
- [ ] **Create an issue and link to the pull request to get dedicated
review from promptflow team. Learn more: [suggested
workflow](../CONTRIBUTING.md#suggested-workflow).**

## General Guidelines and Best Practices
- [ ] Title of the pull request is clear and informative.
- [ ] There are a small number of commits, each of which have an
informative message. This means that previously merged commits do not
appear in the history of the PR. For more information on cleaning up the
commits in your PR, [see this
page](https://github.com/Azure/azure-powershell/blob/master/documentation/development-docs/cleaning-up-commits.md).

### Testing Guidelines
- [ ] Pull request includes test coverage for the included changes.
  • Loading branch information
crazygao authored Mar 22, 2024
1 parent 3ed2542 commit 60941ef
Show file tree
Hide file tree
Showing 8 changed files with 89 additions and 120 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -21,16 +21,16 @@
IS_LEGACY_OPENAI = version("openai").startswith("0.")


def inject_function_async(args_to_ignore=None, trace_type=TraceType.LLM, name=None):
def inject_function_async(args_to_ignore=None, trace_type=TraceType.LLM):
def decorator(func):
return _traced_async(func, args_to_ignore=args_to_ignore, trace_type=trace_type, name=name)
return _traced_async(func, args_to_ignore=args_to_ignore, trace_type=trace_type)

return decorator


def inject_function_sync(args_to_ignore=None, trace_type=TraceType.LLM, name=None):
def inject_function_sync(args_to_ignore=None, trace_type=TraceType.LLM):
def decorator(func):
return _traced_sync(func, args_to_ignore=args_to_ignore, trace_type=trace_type, name=name)
return _traced_sync(func, args_to_ignore=args_to_ignore, trace_type=trace_type)

return decorator

Expand Down Expand Up @@ -90,67 +90,60 @@ def wrapper(*args, **kwargs):
return wrapper


def inject_async(f, trace_type, name):
def inject_async(f, trace_type):
wrapper_fun = inject_operation_headers(
(inject_function_async(["api_key", "headers", "extra_headers"], trace_type, name)(f))
(inject_function_async(["api_key", "headers", "extra_headers"], trace_type)(f))
)
wrapper_fun._original = f
return wrapper_fun


def inject_sync(f, trace_type, name):
def inject_sync(f, trace_type):
wrapper_fun = inject_operation_headers(
(inject_function_sync(["api_key", "headers", "extra_headers"], trace_type, name)(f))
(inject_function_sync(["api_key", "headers", "extra_headers"], trace_type)(f))
)
wrapper_fun._original = f
return wrapper_fun


def _legacy_openai_apis():
sync_apis = (
("openai", "Completion", "create", TraceType.LLM, "openai_completion_legacy"),
("openai", "ChatCompletion", "create", TraceType.LLM, "openai_chat_legacy"),
("openai", "Embedding", "create", TraceType.EMBEDDING, "openai_embedding_legacy"),
)
async_apis = (
("openai", "Completion", "acreate", TraceType.LLM, "openai_completion_legacy"),
("openai", "ChatCompletion", "acreate", TraceType.LLM, "openai_chat_legacy"),
("openai", "Embedding", "acreate", TraceType.EMBEDDING, "openai_embedding_legacy"),
)
return sync_apis, async_apis


def _openai_apis():
sync_apis = (
("openai.resources.chat", "Completions", "create", TraceType.LLM, "openai_chat"),
("openai.resources", "Completions", "create", TraceType.LLM, "openai_completion"),
("openai.resources", "Embeddings", "create", TraceType.EMBEDDING, "openai_embeddings"),
)
async_apis = (
("openai.resources.chat", "AsyncCompletions", "create", TraceType.LLM, "openai_chat_async"),
("openai.resources", "AsyncCompletions", "create", TraceType.LLM, "openai_completion_async"),
("openai.resources", "AsyncEmbeddings", "create", TraceType.EMBEDDING, "openai_embeddings_async"),
)
return sync_apis, async_apis


def _openai_api_list():
if IS_LEGACY_OPENAI:
sync_apis, async_apis = _legacy_openai_apis()
sync_apis = (
("openai", "Completion", "create", TraceType.LLM),
("openai", "ChatCompletion", "create", TraceType.LLM),
("openai", "Embedding", "create", TraceType.EMBEDDING),
)

async_apis = (
("openai", "Completion", "acreate", TraceType.LLM),
("openai", "ChatCompletion", "acreate", TraceType.LLM),
("openai", "Embedding", "acreate", TraceType.EMBEDDING),
)
else:
sync_apis, async_apis = _openai_apis()
sync_apis = (
("openai.resources.chat", "Completions", "create", TraceType.LLM),
("openai.resources", "Completions", "create", TraceType.LLM),
("openai.resources", "Embeddings", "create", TraceType.EMBEDDING),
)

async_apis = (
("openai.resources.chat", "AsyncCompletions", "create", TraceType.LLM),
("openai.resources", "AsyncCompletions", "create", TraceType.LLM),
("openai.resources", "AsyncEmbeddings", "create", TraceType.EMBEDDING),
)

yield sync_apis, inject_sync
yield async_apis, inject_async


def _generate_api_and_injector(apis):
for apis, injector in apis:
for module_name, class_name, method_name, trace_type, name in apis:
for module_name, class_name, method_name, trace_type in apis:
try:
module = importlib.import_module(module_name)
api = getattr(module, class_name)
if hasattr(api, method_name):
yield api, method_name, trace_type, injector, name
yield api, method_name, trace_type, injector
except AttributeError as e:
# Log the attribute exception with the missing class information
logging.warning(
Expand Down Expand Up @@ -183,10 +176,10 @@ def inject_openai_api():
2. Updates the openai api configs from environment variables.
"""

for api, method, trace_type, injector, name in available_openai_apis_and_injectors():
for api, method, trace_type, injector in available_openai_apis_and_injectors():
# Check if the create method of the openai_api class has already been modified
if not hasattr(getattr(api, method), "_original"):
setattr(api, method, injector(getattr(api, method), trace_type, name))
setattr(api, method, injector(getattr(api, method), trace_type))

if IS_LEGACY_OPENAI:
# For the openai versions lower than 1.0.0, it reads api configs from environment variables only at
Expand All @@ -205,6 +198,6 @@ def recover_openai_api():
"""This function restores the original create methods of the OpenAI API classes
by assigning them back from the _original attributes of the modified methods.
"""
for api, method, _, _, _ in available_openai_apis_and_injectors():
for api, method, _, _ in available_openai_apis_and_injectors():
if hasattr(getattr(api, method), "_original"):
setattr(api, method, getattr(getattr(api, method), "_original"))
37 changes: 7 additions & 30 deletions src/promptflow-tracing/promptflow/tracing/_trace.py
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,7 @@
from ._tracer import Tracer, _create_trace_from_function_call, get_node_name_from_context
from ._utils import get_input_names_for_prompt_template, get_prompt_param_name_from_func, serialize
from .contracts.generator_proxy import GeneratorProxy
from .contracts.trace import TraceType, Trace
from .contracts.trace import TraceType

IS_LEGACY_OPENAI = version("openai").startswith("0.")

Expand Down Expand Up @@ -85,13 +85,13 @@ def enrich_span_with_context(span):
logging.warning(f"Failed to enrich span with context: {e}")


def enrich_span_with_trace(span, trace: Trace):
def enrich_span_with_trace(span, trace):
try:
span.set_attributes(
{
"framework": "promptflow",
"span_type": trace.type.value,
"function": trace.function,
"function": trace.name,
}
)
node_name = get_node_name_from_context()
Expand Down Expand Up @@ -293,11 +293,7 @@ def _traced(


def _traced_async(
func: Callable = None,
*,
args_to_ignore: Optional[List[str]] = None,
trace_type=TraceType.FUNCTION,
name: Optional[str] = None,
func: Callable = None, *, args_to_ignore: Optional[List[str]] = None, trace_type=TraceType.FUNCTION
) -> Callable:
"""
Decorator that adds tracing to an asynchronous function.
Expand All @@ -307,20 +303,14 @@ def _traced_async(
args_to_ignore (Optional[List[str]], optional): A list of argument names to be ignored in the trace.
Defaults to None.
trace_type (TraceType, optional): The type of the trace. Defaults to TraceType.FUNCTION.
name (str, optional): The name of the trace, will set to func name if not provided.
Returns:
Callable: The traced function.
"""

def create_trace(func, args, kwargs):
return _create_trace_from_function_call(
func,
args=args,
kwargs=kwargs,
args_to_ignore=args_to_ignore,
trace_type=trace_type,
name=name,
func, args=args, kwargs=kwargs, args_to_ignore=args_to_ignore, trace_type=trace_type
)

@functools.wraps(func)
Expand Down Expand Up @@ -353,13 +343,7 @@ async def wrapped(*args, **kwargs):
return wrapped


def _traced_sync(
func: Callable = None,
*,
args_to_ignore: Optional[List[str]] = None,
trace_type=TraceType.FUNCTION,
name: Optional[str] = None,
) -> Callable:
def _traced_sync(func: Callable = None, *, args_to_ignore=None, trace_type=TraceType.FUNCTION) -> Callable:
"""
Decorator that adds tracing to a synchronous function.
Expand All @@ -368,21 +352,14 @@ def _traced_sync(
args_to_ignore (Optional[List[str]], optional): A list of argument names to be ignored in the trace.
Defaults to None.
trace_type (TraceType, optional): The type of the trace. Defaults to TraceType.FUNCTION.
name (str, optional): The name of the trace, will set to func name if not provided.
Returns:
Callable: The traced function.
"""

def create_trace(func, args, kwargs):
return _create_trace_from_function_call(
func,
args=args,
kwargs=kwargs,
args_to_ignore=args_to_ignore,
trace_type=trace_type,
name=name,
func, args=args, kwargs=kwargs, args_to_ignore=args_to_ignore, trace_type=trace_type
)

@functools.wraps(func)
Expand Down
10 changes: 4 additions & 6 deletions src/promptflow-tracing/promptflow/tracing/_tracer.py
Original file line number Diff line number Diff line change
Expand Up @@ -137,7 +137,7 @@ def _format_error(error: Exception) -> dict:


def _create_trace_from_function_call(
f, *, args=None, kwargs=None, args_to_ignore: Optional[List[str]] = None, trace_type=TraceType.FUNCTION, name=None,
f, *, args=None, kwargs=None, args_to_ignore: Optional[List[str]] = None, trace_type=TraceType.FUNCTION
):
"""
Creates a trace object from a function call.
Expand All @@ -149,7 +149,6 @@ def _create_trace_from_function_call(
args_to_ignore (Optional[List[str]], optional): A list of argument names to be ignored in the trace.
Defaults to None.
trace_type (TraceType, optional): The type of the trace. Defaults to TraceType.FUNCTION.
name (str, optional): The name of the trace. Defaults to None.
Returns:
Trace: The created trace object.
Expand Down Expand Up @@ -177,17 +176,16 @@ def _create_trace_from_function_call(
for key in args_to_ignore:
all_kwargs.pop(key, None)

function = f.__qualname__
name = f.__qualname__
if trace_type in [TraceType.LLM, TraceType.EMBEDDING] and f.__module__:
function = f"{f.__module__}.{function}"
name = f"{f.__module__}.{name}"

return Trace(
name=name or function, # Use the function name as the trace name if not provided
name=name,
type=trace_type,
start_time=datetime.utcnow().timestamp(),
inputs=all_kwargs,
children=[],
function=function,
)


Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -52,4 +52,3 @@ class Trace:
node_name: Optional[str] = None # The node name of the trace, used for flow level trace
parent_id: str = "" # The parent trace id of the trace
id: str = "" # The trace id
function: str = "" # The function name of the trace
2 changes: 1 addition & 1 deletion src/promptflow-tracing/tests/e2etests/simple_functions.py
Original file line number Diff line number Diff line change
Expand Up @@ -38,7 +38,7 @@ def greetings(user_id):

@trace
async def dummy_llm(prompt: str, model: str):
await asyncio.sleep(0.5)
asyncio.sleep(0.5)
return "dummy_output"


Expand Down
11 changes: 5 additions & 6 deletions src/promptflow-tracing/tests/e2etests/test_tracing.py
Original file line number Diff line number Diff line change
Expand Up @@ -148,18 +148,17 @@ def assert_otel_traces_with_embedding(self, dev_connections, func, inputs, expec
self.validate_openai_tokens(span_list)
for span in span_list:
if span.attributes.get("function", "") in EMBEDDING_FUNCTION_NAMES:
msg = f"Span attributes is not expected: {span.attributes}"
assert "ada" in span.attributes.get("llm.response.model", ""), msg
assert "ada" in span.attributes.get("llm.response.model", "")
embeddings = span.attributes.get("embedding.embeddings", "")
assert "embedding.vector" in embeddings, msg
assert "embedding.text" in embeddings, msg
assert "embedding.vector" in embeddings
assert "embedding.text" in embeddings
if isinstance(inputs["input"], list):
# If the input is a token array, which is list of int, the attribute should contains
# the length of the token array '<len(token_array) dimensional token>'.
assert "dimensional token" in embeddings, msg
assert "dimensional token" in embeddings
else:
# If the input is a string, the attribute should contains the original input string.
assert inputs["input"] in embeddings, msg
assert inputs["input"] in embeddings

def test_otel_trace_with_multiple_functions(self):
execute_function_in_subprocess(self.assert_otel_traces_with_multiple_functions)
Expand Down
Loading

0 comments on commit 60941ef

Please sign in to comment.