From 92f4faaba49d10787c44561d5d340c6a4f170713 Mon Sep 17 00:00:00 2001 From: Munir Abdinur Date: Mon, 10 Feb 2025 14:54:33 -0500 Subject: [PATCH] chore(llmobs): migrate to core api --- .../internal/openai/_endpoint_hooks.py | 3 +- ddtrace/contrib/internal/openai/utils.py | 5 +- .../_evaluators/ragas/answer_relevancy.py | 3 +- .../_evaluators/ragas/context_precision.py | 3 +- .../llmobs/_evaluators/ragas/faithfulness.py | 3 +- ddtrace/llmobs/_integrations/anthropic.py | 3 +- ddtrace/llmobs/_integrations/base.py | 3 +- ddtrace/llmobs/_integrations/bedrock.py | 5 +- ddtrace/llmobs/_integrations/gemini.py | 3 +- ddtrace/llmobs/_integrations/langchain.py | 49 +++-- ddtrace/llmobs/_integrations/langgraph.py | 13 +- ddtrace/llmobs/_integrations/openai.py | 21 +- ddtrace/llmobs/_integrations/vertexai.py | 3 +- ddtrace/llmobs/_llmobs.py | 84 +++---- ddtrace/llmobs/_utils.py | 19 +- ddtrace/llmobs/decorators.py | 5 +- tests/llmobs/test_llmobs.py | 131 +++++------ tests/llmobs/test_llmobs_service.py | 207 +++++++++--------- 18 files changed, 289 insertions(+), 274 deletions(-) diff --git a/ddtrace/contrib/internal/openai/_endpoint_hooks.py b/ddtrace/contrib/internal/openai/_endpoint_hooks.py index 786bb67f919..eb8822d4fdb 100644 --- a/ddtrace/contrib/internal/openai/_endpoint_hooks.py +++ b/ddtrace/contrib/internal/openai/_endpoint_hooks.py @@ -8,6 +8,7 @@ from ddtrace.contrib.internal.openai.utils import _loop_handler from ddtrace.contrib.internal.openai.utils import _process_finished_stream from ddtrace.contrib.internal.openai.utils import _tag_tool_calls +from ddtrace.internal import core from ddtrace.internal.utils.version import parse_version @@ -248,7 +249,7 @@ def _record_request(self, pin, integration, instance, span, args, kwargs): if kwargs.get("stream_options", {}).get("include_usage", None) is not None: # Only perform token chunk auto-extraction if this option is not explicitly set return - span._set_ctx_item("_dd.auto_extract_token_chunk", True) + core.set_item("_dd.auto_extract_token_chunk", True) stream_options = kwargs.get("stream_options", {}) stream_options["include_usage"] = True kwargs["stream_options"] = stream_options diff --git a/ddtrace/contrib/internal/openai/utils.py b/ddtrace/contrib/internal/openai/utils.py index 0217b1e61d2..a203e50ecb7 100644 --- a/ddtrace/contrib/internal/openai/utils.py +++ b/ddtrace/contrib/internal/openai/utils.py @@ -8,6 +8,7 @@ import wrapt +from ddtrace.internal import core from ddtrace.internal.logger import get_logger from ddtrace.llmobs._utils import _get_attr @@ -87,7 +88,7 @@ def __next__(self): def _extract_token_chunk(self, chunk): """Attempt to extract the token chunk (last chunk in the stream) from the streamed response.""" - if not self._dd_span._get_ctx_item("_dd.auto_extract_token_chunk"): + if not core.get_item("_dd.auto_extract_token_chunk"): return choices = getattr(chunk, "choices") if not choices: @@ -153,7 +154,7 @@ async def __anext__(self): async def _extract_token_chunk(self, chunk): """Attempt to extract the token chunk (last chunk in the stream) from the streamed response.""" - if not self._dd_span._get_ctx_item("_dd.auto_extract_token_chunk"): + if not core.get_item("_dd.auto_extract_token_chunk"): return choices = getattr(chunk, "choices") if not choices: diff --git a/ddtrace/llmobs/_evaluators/ragas/answer_relevancy.py b/ddtrace/llmobs/_evaluators/ragas/answer_relevancy.py index 5fd6e6b7c0c..a5e79af001b 100644 --- a/ddtrace/llmobs/_evaluators/ragas/answer_relevancy.py +++ b/ddtrace/llmobs/_evaluators/ragas/answer_relevancy.py @@ -3,6 +3,7 @@ from typing import Tuple from typing import Union +from ddtrace.internal import core from ddtrace.internal.logger import get_logger from ddtrace.llmobs._constants import EVALUATION_SPAN_METADATA from ddtrace.llmobs._constants import IS_EVALUATION_SPAN @@ -85,7 +86,7 @@ def evaluate(self, span_event: dict) -> Tuple[Union[float, str], Optional[dict]] with self.llmobs_service.workflow( "dd-ragas.answer_relevancy", ml_app=_get_ml_app_for_ragas_trace(span_event) ) as ragas_ar_workflow: - ragas_ar_workflow._set_ctx_item(IS_EVALUATION_SPAN, True) + core.set_item(IS_EVALUATION_SPAN, True) try: evaluation_metadata[EVALUATION_SPAN_METADATA] = self.llmobs_service.export_span(span=ragas_ar_workflow) diff --git a/ddtrace/llmobs/_evaluators/ragas/context_precision.py b/ddtrace/llmobs/_evaluators/ragas/context_precision.py index 13ccb1d593a..8d805ed7d73 100644 --- a/ddtrace/llmobs/_evaluators/ragas/context_precision.py +++ b/ddtrace/llmobs/_evaluators/ragas/context_precision.py @@ -3,6 +3,7 @@ from typing import Tuple from typing import Union +from ddtrace.internal import core from ddtrace.internal.logger import get_logger from ddtrace.llmobs._constants import EVALUATION_KIND_METADATA from ddtrace.llmobs._constants import EVALUATION_SPAN_METADATA @@ -83,7 +84,7 @@ def evaluate(self, span_event: dict) -> Tuple[Union[float, str], Optional[dict]] with self.llmobs_service.workflow( "dd-ragas.context_precision", ml_app=_get_ml_app_for_ragas_trace(span_event) ) as ragas_cp_workflow: - ragas_cp_workflow._set_ctx_item(IS_EVALUATION_SPAN, True) + core.set_item(IS_EVALUATION_SPAN, True) try: evaluation_metadata[EVALUATION_SPAN_METADATA] = self.llmobs_service.export_span(span=ragas_cp_workflow) diff --git a/ddtrace/llmobs/_evaluators/ragas/faithfulness.py b/ddtrace/llmobs/_evaluators/ragas/faithfulness.py index 2c413f2cec7..07eef34999e 100644 --- a/ddtrace/llmobs/_evaluators/ragas/faithfulness.py +++ b/ddtrace/llmobs/_evaluators/ragas/faithfulness.py @@ -5,6 +5,7 @@ from typing import Tuple from typing import Union +from ddtrace.internal import core from ddtrace.internal.logger import get_logger from ddtrace.llmobs._constants import EVALUATION_KIND_METADATA from ddtrace.llmobs._constants import EVALUATION_SPAN_METADATA @@ -97,7 +98,7 @@ def evaluate(self, span_event: dict) -> Tuple[Union[float, str], Optional[dict]] with self.llmobs_service.workflow( "dd-ragas.faithfulness", ml_app=_get_ml_app_for_ragas_trace(span_event) ) as ragas_faithfulness_workflow: - ragas_faithfulness_workflow._set_ctx_item(IS_EVALUATION_SPAN, True) + core.set_item(IS_EVALUATION_SPAN, True) try: evaluation_metadata[EVALUATION_SPAN_METADATA] = self.llmobs_service.export_span( span=ragas_faithfulness_workflow diff --git a/ddtrace/llmobs/_integrations/anthropic.py b/ddtrace/llmobs/_integrations/anthropic.py index bb4f96e7814..22d7e860eab 100644 --- a/ddtrace/llmobs/_integrations/anthropic.py +++ b/ddtrace/llmobs/_integrations/anthropic.py @@ -5,6 +5,7 @@ from typing import List from typing import Optional +from ddtrace.internal import core from ddtrace.internal.logger import get_logger from ddtrace.llmobs._constants import INPUT_MESSAGES from ddtrace.llmobs._constants import METADATA @@ -67,7 +68,7 @@ def _llmobs_set_tags( if not span.error and response is not None: output_messages = self._extract_output_message(response) - span._set_ctx_items( + core.set_items( { SPAN_KIND: "llm", MODEL_NAME: span.get_tag("anthropic.request.model") or "", diff --git a/ddtrace/llmobs/_integrations/base.py b/ddtrace/llmobs/_integrations/base.py index a098c899014..a0fee5f08c7 100644 --- a/ddtrace/llmobs/_integrations/base.py +++ b/ddtrace/llmobs/_integrations/base.py @@ -11,6 +11,7 @@ from ddtrace.constants import _SPAN_MEASURED_KEY from ddtrace.contrib.internal.trace_utils import int_service from ddtrace.ext import SpanTypes +from ddtrace.internal import core from ddtrace.internal.agent import get_stats_url from ddtrace.internal.dogstatsd import get_dogstatsd_client from ddtrace.internal.hostname import get_hostname @@ -137,7 +138,7 @@ def trace(self, pin: Pin, operation_id: str, submit_to_llmobs: bool = False, **k # The LLMObs parent ID tag is not set at span start time. We need to manually set the parent ID tag now # in these cases to avoid conflicting with the later propagated tags. parent_id = _get_llmobs_parent_id(span) or "undefined" - span._set_ctx_item(PARENT_ID_KEY, str(parent_id)) + core.set_item(PARENT_ID_KEY, str(parent_id)) telemetry_writer.add_count_metric( namespace=TELEMETRY_NAMESPACE.MLOBS, name="span.start", diff --git a/ddtrace/llmobs/_integrations/bedrock.py b/ddtrace/llmobs/_integrations/bedrock.py index cbc1456fc24..0dfbd84b88a 100644 --- a/ddtrace/llmobs/_integrations/bedrock.py +++ b/ddtrace/llmobs/_integrations/bedrock.py @@ -3,6 +3,7 @@ from typing import List from typing import Optional +from ddtrace.internal import core from ddtrace.internal.logger import get_logger from ddtrace.llmobs._constants import INPUT_MESSAGES from ddtrace.llmobs._constants import METADATA @@ -36,7 +37,7 @@ def _llmobs_set_tags( """Extract prompt/response tags from a completion and set them as temporary "_ml_obs.*" tags.""" if span.get_tag(PROPAGATED_PARENT_ID_KEY) is None: parent_id = _get_llmobs_parent_id(span) or "undefined" - span._set_ctx_item(PARENT_ID_KEY, parent_id) + core.set_item(PARENT_ID_KEY, parent_id) parameters = {} if span.get_tag("bedrock.request.temperature"): parameters["temperature"] = float(span.get_tag("bedrock.request.temperature") or 0.0) @@ -48,7 +49,7 @@ def _llmobs_set_tags( output_messages = [{"content": ""}] if not span.error and response is not None: output_messages = self._extract_output_message(response) - span._set_ctx_items( + core.set_items( { SPAN_KIND: "llm", MODEL_NAME: span.get_tag("bedrock.request.model") or "", diff --git a/ddtrace/llmobs/_integrations/gemini.py b/ddtrace/llmobs/_integrations/gemini.py index 0407ec7188b..c314a567bd6 100644 --- a/ddtrace/llmobs/_integrations/gemini.py +++ b/ddtrace/llmobs/_integrations/gemini.py @@ -4,6 +4,7 @@ from typing import List from typing import Optional +from ddtrace.internal import core from ddtrace.internal.utils import get_argument_value from ddtrace.llmobs._constants import INPUT_MESSAGES from ddtrace.llmobs._constants import METADATA @@ -51,7 +52,7 @@ def _llmobs_set_tags( if not span.error and response is not None: output_messages = self._extract_output_message(response) - span._set_ctx_items( + core.set_items( { SPAN_KIND: "llm", MODEL_NAME: span.get_tag("google_generativeai.request.model") or "", diff --git a/ddtrace/llmobs/_integrations/langchain.py b/ddtrace/llmobs/_integrations/langchain.py index d380c6ab7a8..89af50df28c 100644 --- a/ddtrace/llmobs/_integrations/langchain.py +++ b/ddtrace/llmobs/_integrations/langchain.py @@ -6,6 +6,7 @@ from typing import Optional from typing import Union +from ddtrace.internal import core from ddtrace.internal.logger import get_logger from ddtrace.internal.utils import ArgumentError from ddtrace.internal.utils import get_argument_value @@ -128,7 +129,7 @@ def _llmobs_set_metadata(self, span: Span, model_provider: Optional[str] = None) if max_tokens is not None and max_tokens != "None": metadata["max_tokens"] = int(max_tokens) if metadata: - span._set_ctx_item(METADATA, metadata) + core.set_item(METADATA, metadata) def _llmobs_set_tags_from_llm( self, span: Span, args: List[Any], kwargs: Dict[str, Any], completions: Any, is_workflow: bool = False @@ -146,7 +147,7 @@ def _llmobs_set_tags_from_llm( else: input_messages = [{"content": str(prompt)} for prompt in prompts] - span._set_ctx_items( + core.set_items( { SPAN_KIND: "workflow" if is_workflow else "llm", MODEL_NAME: span.get_tag(MODEL) or "", @@ -156,7 +157,7 @@ def _llmobs_set_tags_from_llm( ) if span.error: - span._set_ctx_item(output_tag_key, [{"content": ""}]) + core.set_item(output_tag_key, [{"content": ""}]) return if stream: message_content = [{"content": completions}] # single completion for streams @@ -170,8 +171,8 @@ def _llmobs_set_tags_from_llm( OUTPUT_TOKENS_METRIC_KEY: output_tokens, TOTAL_TOKENS_METRIC_KEY: total_tokens, } - span._set_ctx_item(METRICS, metrics) - span._set_ctx_item(output_tag_key, message_content) + core.set_item(METRICS, metrics) + core.set_item(output_tag_key, message_content) def _llmobs_set_tags_from_chat_model( self, @@ -181,7 +182,7 @@ def _llmobs_set_tags_from_chat_model( chat_completions: Any, is_workflow: bool = False, ) -> None: - span._set_ctx_items( + core.set_items( { SPAN_KIND: "workflow" if is_workflow else "llm", MODEL_NAME: span.get_tag(MODEL) or "", @@ -207,17 +208,17 @@ def _llmobs_set_tags_from_chat_model( ) role = getattr(message, "role", ROLE_MAPPING.get(message.type, "")) input_messages.append({"content": str(content), "role": str(role)}) - span._set_ctx_item(input_tag_key, input_messages) + core.set_item(input_tag_key, input_messages) if span.error: - span._set_ctx_item(output_tag_key, [{"content": ""}]) + core.set_item(output_tag_key, [{"content": ""}]) return output_messages = [] if stream: content = chat_completions.content role = chat_completions.__class__.__name__.replace("MessageChunk", "").lower() # AIMessageChunk --> ai - span._set_ctx_item(output_tag_key, [{"content": content, "role": ROLE_MAPPING.get(role, "")}]) + core.set_item(output_tag_key, [{"content": content, "role": ROLE_MAPPING.get(role, "")}]) return input_tokens, output_tokens, total_tokens = 0, 0, 0 @@ -253,7 +254,7 @@ def _llmobs_set_tags_from_chat_model( output_tokens = sum(v["output_tokens"] for v in tokens_per_choice_run_id.values()) total_tokens = sum(v["total_tokens"] for v in tokens_per_choice_run_id.values()) - span._set_ctx_item(output_tag_key, output_messages) + core.set_item(output_tag_key, output_messages) if not is_workflow and total_tokens > 0: metrics = { @@ -261,7 +262,7 @@ def _llmobs_set_tags_from_chat_model( OUTPUT_TOKENS_METRIC_KEY: output_tokens, TOTAL_TOKENS_METRIC_KEY: total_tokens, } - span._set_ctx_item(METRICS, metrics) + core.set_item(METRICS, metrics) def _extract_tool_calls(self, chat_completion_msg: Any) -> List[Dict[str, Any]]: """Extracts tool calls from a langchain chat completion.""" @@ -315,7 +316,7 @@ def _llmobs_set_meta_tags_from_chain(self, span: Span, args, kwargs, outputs: An formatted_outputs = "" if not span.error and outputs is not None: formatted_outputs = format_langchain_io(outputs) - span._set_ctx_items({SPAN_KIND: "workflow", INPUT_VALUE: formatted_inputs, OUTPUT_VALUE: formatted_outputs}) + core.set_items({SPAN_KIND: "workflow", INPUT_VALUE: formatted_inputs, OUTPUT_VALUE: formatted_outputs}) def _llmobs_set_meta_tags_from_embedding( self, @@ -325,7 +326,7 @@ def _llmobs_set_meta_tags_from_embedding( output_embedding: Union[List[float], List[List[float]], None], is_workflow: bool = False, ) -> None: - span._set_ctx_items( + core.set_items( { SPAN_KIND: "workflow" if is_workflow else "embedding", MODEL_NAME: span.get_tag(MODEL) or "", @@ -346,16 +347,16 @@ def _llmobs_set_meta_tags_from_embedding( ): if is_workflow: formatted_inputs = format_langchain_io(input_texts) - span._set_ctx_item(input_tag_key, formatted_inputs) + core.set_item(input_tag_key, formatted_inputs) else: if isinstance(input_texts, str): input_texts = [input_texts] input_documents = [Document(text=str(doc)) for doc in input_texts] - span._set_ctx_item(input_tag_key, input_documents) + core.set_item(input_tag_key, input_documents) except TypeError: log.warning("Failed to serialize embedding input data to JSON") if span.error or output_embedding is None: - span._set_ctx_item(output_tag_key, "") + core.set_item(output_tag_key, "") return try: if isinstance(output_embedding[0], float): @@ -367,7 +368,7 @@ def _llmobs_set_meta_tags_from_embedding( output_values = output_embedding embeddings_count = len(output_embedding) embedding_dim = len(output_values[0]) - span._set_ctx_item( + core.set_item( output_tag_key, "[{} embedding(s) returned with size {}]".format(embeddings_count, embedding_dim), ) @@ -382,7 +383,7 @@ def _llmobs_set_meta_tags_from_similarity_search( output_documents: Union[List[Any], None], is_workflow: bool = False, ) -> None: - span._set_ctx_items( + core.set_items( { SPAN_KIND: "workflow" if is_workflow else "retrieval", MODEL_NAME: span.get_tag(MODEL) or "", @@ -392,12 +393,12 @@ def _llmobs_set_meta_tags_from_similarity_search( input_query = get_argument_value(args, kwargs, 0, "query") if input_query is not None: formatted_inputs = format_langchain_io(input_query) - span._set_ctx_item(INPUT_VALUE, formatted_inputs) + core.set_item(INPUT_VALUE, formatted_inputs) if span.error or not output_documents or not isinstance(output_documents, list): - span._set_ctx_item(OUTPUT_VALUE, "") + core.set_item(OUTPUT_VALUE, "") return if is_workflow: - span._set_ctx_item(OUTPUT_VALUE, "[{} document(s) retrieved]".format(len(output_documents))) + core.set_item(OUTPUT_VALUE, "[{} document(s) retrieved]".format(len(output_documents))) return documents = [] for d in output_documents: @@ -406,9 +407,9 @@ def _llmobs_set_meta_tags_from_similarity_search( metadata = getattr(d, "metadata", {}) doc["name"] = metadata.get("name", doc["id"]) documents.append(doc) - span._set_ctx_item(OUTPUT_DOCUMENTS, format_langchain_io(documents)) + core.set_item(OUTPUT_DOCUMENTS, format_langchain_io(documents)) # we set the value as well to ensure that the UI would display it in case the span was the root - span._set_ctx_item(OUTPUT_VALUE, "[{} document(s) retrieved]".format(len(documents))) + core.set_item(OUTPUT_VALUE, "[{} document(s) retrieved]".format(len(documents))) def _llmobs_set_meta_tags_from_tool(self, span: Span, tool_inputs: Dict[str, Any], tool_output: object) -> None: metadata = json.loads(str(span.get_tag(METADATA))) if span.get_tag(METADATA) else {} @@ -423,7 +424,7 @@ def _llmobs_set_meta_tags_from_tool(self, span: Span, tool_inputs: Dict[str, Any formatted_outputs = "" if not span.error and tool_output is not None: formatted_outputs = format_langchain_io(tool_output) - span._set_ctx_items( + core.set_items( { SPAN_KIND: "tool", METADATA: metadata, diff --git a/ddtrace/llmobs/_integrations/langgraph.py b/ddtrace/llmobs/_integrations/langgraph.py index 08f3943e57e..818ec878e8a 100644 --- a/ddtrace/llmobs/_integrations/langgraph.py +++ b/ddtrace/llmobs/_integrations/langgraph.py @@ -4,6 +4,7 @@ from typing import Optional from ddtrace.ext import SpanTypes +from ddtrace.internal import core from ddtrace.internal.utils import get_argument_value from ddtrace.llmobs._constants import INPUT_VALUE from ddtrace.llmobs._constants import NAME @@ -45,9 +46,9 @@ def _llmobs_set_tags( invoked_node_span_links = invoked_node.get("span_links") if invoked_node_span_links is not None: span_links = invoked_node_span_links - current_span_links = span._get_ctx_item(SPAN_LINKS) or [] + current_span_links = core.get_item(SPAN_LINKS) or [] - span._set_ctx_items( + core.set_items( { SPAN_KIND: "agent" if operation == "graph" else "task", INPUT_VALUE: format_langchain_io(inputs), @@ -90,10 +91,10 @@ def _handle_finished_graph(self, graph_span, finished_tasks, is_subgraph_node): {**self._graph_nodes_by_task_id[task_id]["span"], "attributes": {"from": "output", "to": "output"}} for task_id in finished_tasks.keys() ] - graph_span_span_links = graph_span._get_ctx_item(SPAN_LINKS) or [] - graph_span._set_ctx_item(SPAN_LINKS, graph_span_span_links + output_span_links) + graph_span_span_links = core.get_item(SPAN_LINKS) or [] + core.set_item(SPAN_LINKS, graph_span_span_links + output_span_links) if graph_caller_span is not None and not is_subgraph_node: - graph_caller_span_links = graph_caller_span._get_ctx_item(SPAN_LINKS) or [] + graph_caller_span_links = core.get_item(SPAN_LINKS) or [] span_links = [ { "span_id": str(graph_span.span_id) or "undefined", @@ -101,7 +102,7 @@ def _handle_finished_graph(self, graph_span, finished_tasks, is_subgraph_node): "attributes": {"from": "output", "to": "output"}, } ] - graph_caller_span._set_ctx_item(SPAN_LINKS, graph_caller_span_links + span_links) + core.set_item(SPAN_LINKS, graph_caller_span_links + span_links) return def _link_task_to_parent(self, task_id, task, finished_task_names_to_ids): diff --git a/ddtrace/llmobs/_integrations/openai.py b/ddtrace/llmobs/_integrations/openai.py index eb01a679191..ad892ba2285 100644 --- a/ddtrace/llmobs/_integrations/openai.py +++ b/ddtrace/llmobs/_integrations/openai.py @@ -5,6 +5,7 @@ from typing import Optional from typing import Tuple +from ddtrace.internal import core from ddtrace.internal.constants import COMPONENT from ddtrace.internal.utils.version import parse_version from ddtrace.llmobs._constants import INPUT_DOCUMENTS @@ -116,7 +117,7 @@ def _llmobs_set_tags( elif operation == "embedding": self._llmobs_set_meta_tags_from_embedding(span, kwargs, response) metrics = self._extract_llmobs_metrics_tags(span, response) - span._set_ctx_items( + core.set_items( {SPAN_KIND: span_kind, MODEL_NAME: model_name or "", MODEL_PROVIDER: model_provider, METRICS: metrics} ) @@ -131,7 +132,7 @@ def _llmobs_set_meta_tags_from_completion(span: Span, kwargs: Dict[str, Any], co if not span.error and completions: choices = getattr(completions, "choices", completions) output_messages = [{"content": _get_attr(choice, "text", "")} for choice in choices] - span._set_ctx_items( + core.set_items( { INPUT_MESSAGES: [{"content": str(p)} for p in prompt], METADATA: parameters, @@ -146,10 +147,10 @@ def _llmobs_set_meta_tags_from_chat(span: Span, kwargs: Dict[str, Any], messages for m in kwargs.get("messages", []): input_messages.append({"content": str(_get_attr(m, "content", "")), "role": str(_get_attr(m, "role", ""))}) parameters = {k: v for k, v in kwargs.items() if k not in ("model", "messages", "tools", "functions")} - span._set_ctx_items({INPUT_MESSAGES: input_messages, METADATA: parameters}) + core.set_items({INPUT_MESSAGES: input_messages, METADATA: parameters}) if span.error or not messages: - span._set_ctx_item(OUTPUT_MESSAGES, [{"content": ""}]) + core.set_item(OUTPUT_MESSAGES, [{"content": ""}]) return if isinstance(messages, list): # streamed response output_messages = [] @@ -167,7 +168,7 @@ def _llmobs_set_meta_tags_from_chat(span: Span, kwargs: Dict[str, Any], messages for tool_call in tool_calls ] output_messages.append(message) - span._set_ctx_item(OUTPUT_MESSAGES, output_messages) + core.set_item(OUTPUT_MESSAGES, output_messages) return choices = _get_attr(messages, "choices", []) output_messages = [] @@ -196,7 +197,7 @@ def _llmobs_set_meta_tags_from_chat(span: Span, kwargs: Dict[str, Any], messages output_messages.append({"content": content, "role": role, "tool_calls": tool_calls_info}) continue output_messages.append({"content": content, "role": role}) - span._set_ctx_item(OUTPUT_MESSAGES, output_messages) + core.set_item(OUTPUT_MESSAGES, output_messages) @staticmethod def _llmobs_set_meta_tags_from_embedding(span: Span, kwargs: Dict[str, Any], resp: Any) -> None: @@ -212,16 +213,14 @@ def _llmobs_set_meta_tags_from_embedding(span: Span, kwargs: Dict[str, Any], res input_documents = [] for doc in embedding_inputs: input_documents.append(Document(text=str(doc))) - span._set_ctx_items({METADATA: metadata, INPUT_DOCUMENTS: input_documents}) + core.set_items({METADATA: metadata, INPUT_DOCUMENTS: input_documents}) if span.error: return if encoding_format == "float": embedding_dim = len(resp.data[0].embedding) - span._set_ctx_item( - OUTPUT_VALUE, "[{} embedding(s) returned with size {}]".format(len(resp.data), embedding_dim) - ) + core.set_item(OUTPUT_VALUE, "[{} embedding(s) returned with size {}]".format(len(resp.data), embedding_dim)) return - span._set_ctx_item(OUTPUT_VALUE, "[{} embedding(s) returned]".format(len(resp.data))) + core.set_item(OUTPUT_VALUE, "[{} embedding(s) returned]".format(len(resp.data))) @staticmethod def _extract_llmobs_metrics_tags(span: Span, resp: Any) -> Dict[str, Any]: diff --git a/ddtrace/llmobs/_integrations/vertexai.py b/ddtrace/llmobs/_integrations/vertexai.py index db40ac15b19..e6f8d1c11a7 100644 --- a/ddtrace/llmobs/_integrations/vertexai.py +++ b/ddtrace/llmobs/_integrations/vertexai.py @@ -4,6 +4,7 @@ from typing import List from typing import Optional +from ddtrace.internal import core from ddtrace.internal.utils import ArgumentError from ddtrace.internal.utils import get_argument_value from ddtrace.llmobs._constants import INPUT_MESSAGES @@ -57,7 +58,7 @@ def _llmobs_set_tags( if not span.error and response is not None: output_messages = self._extract_output_message(response) - span._set_ctx_items( + core.set_items( { SPAN_KIND: "llm", MODEL_NAME: span.get_tag("vertexai.request.model") or "", diff --git a/ddtrace/llmobs/_llmobs.py b/ddtrace/llmobs/_llmobs.py index 14e12d3151e..ae91b17329d 100644 --- a/ddtrace/llmobs/_llmobs.py +++ b/ddtrace/llmobs/_llmobs.py @@ -131,7 +131,7 @@ def _submit_llmobs_span(self, span: Span) -> None: "Error generating LLMObs span event for span %s, likely due to malformed span", span, exc_info=True ) finally: - if not span_event or not span._get_ctx_item(SPAN_KIND) == "llm" or _is_evaluation_span(span): + if not span_event or not core.get_item(SPAN_KIND) == "llm" or _is_evaluation_span(span): return if self._evaluator_runner: self._evaluator_runner.enqueue(span_event, span) @@ -139,28 +139,28 @@ def _submit_llmobs_span(self, span: Span) -> None: @classmethod def _llmobs_span_event(cls, span: Span) -> Dict[str, Any]: """Span event object structure.""" - span_kind = span._get_ctx_item(SPAN_KIND) + span_kind = core.get_item(SPAN_KIND) if not span_kind: raise KeyError("Span kind not found in span context") meta: Dict[str, Any] = {"span.kind": span_kind, "input": {}, "output": {}} - if span_kind in ("llm", "embedding") and span._get_ctx_item(MODEL_NAME) is not None: - meta["model_name"] = span._get_ctx_item(MODEL_NAME) - meta["model_provider"] = (span._get_ctx_item(MODEL_PROVIDER) or "custom").lower() - meta["metadata"] = span._get_ctx_item(METADATA) or {} - if span_kind == "llm" and span._get_ctx_item(INPUT_MESSAGES) is not None: - meta["input"]["messages"] = span._get_ctx_item(INPUT_MESSAGES) - if span._get_ctx_item(INPUT_VALUE) is not None: - meta["input"]["value"] = safe_json(span._get_ctx_item(INPUT_VALUE), ensure_ascii=False) - if span_kind == "llm" and span._get_ctx_item(OUTPUT_MESSAGES) is not None: - meta["output"]["messages"] = span._get_ctx_item(OUTPUT_MESSAGES) - if span_kind == "embedding" and span._get_ctx_item(INPUT_DOCUMENTS) is not None: - meta["input"]["documents"] = span._get_ctx_item(INPUT_DOCUMENTS) - if span._get_ctx_item(OUTPUT_VALUE) is not None: - meta["output"]["value"] = safe_json(span._get_ctx_item(OUTPUT_VALUE), ensure_ascii=False) - if span_kind == "retrieval" and span._get_ctx_item(OUTPUT_DOCUMENTS) is not None: - meta["output"]["documents"] = span._get_ctx_item(OUTPUT_DOCUMENTS) - if span._get_ctx_item(INPUT_PROMPT) is not None: - prompt_json_str = span._get_ctx_item(INPUT_PROMPT) + if span_kind in ("llm", "embedding") and core.get_item(MODEL_NAME) is not None: + meta["model_name"] = core.get_item(MODEL_NAME) + meta["model_provider"] = (core.get_item(MODEL_PROVIDER) or "custom").lower() + meta["metadata"] = core.get_item(METADATA) or {} + if span_kind == "llm" and core.get_item(INPUT_MESSAGES) is not None: + meta["input"]["messages"] = core.get_item(INPUT_MESSAGES) + if core.get_item(INPUT_VALUE) is not None: + meta["input"]["value"] = safe_json(core.get_item(INPUT_VALUE), ensure_ascii=False) + if span_kind == "llm" and core.get_item(OUTPUT_MESSAGES) is not None: + meta["output"]["messages"] = core.get_item(OUTPUT_MESSAGES) + if span_kind == "embedding" and core.get_item(INPUT_DOCUMENTS) is not None: + meta["input"]["documents"] = core.get_item(INPUT_DOCUMENTS) + if core.get_item(OUTPUT_VALUE) is not None: + meta["output"]["value"] = safe_json(core.get_item(OUTPUT_VALUE), ensure_ascii=False) + if span_kind == "retrieval" and core.get_item(OUTPUT_DOCUMENTS) is not None: + meta["output"]["documents"] = core.get_item(OUTPUT_DOCUMENTS) + if core.get_item(INPUT_PROMPT) is not None: + prompt_json_str = core.get_item(INPUT_PROMPT) if span_kind != "llm": log.warning( "Dropping prompt on non-LLM span kind, annotating prompts is only supported for LLM span kinds." @@ -179,10 +179,10 @@ def _llmobs_span_event(cls, span: Span) -> Dict[str, Any]: meta.pop("input") if not meta["output"]: meta.pop("output") - metrics = span._get_ctx_item(METRICS) or {} + metrics = core.get_item(METRICS) or {} ml_app = _get_ml_app(span) - span._set_ctx_item(ML_APP, ml_app) + core.set_item(ML_APP, ml_app) parent_id = str(_get_llmobs_parent_id(span) or "undefined") llmobs_span_event = { @@ -198,12 +198,12 @@ def _llmobs_span_event(cls, span: Span) -> Dict[str, Any]: } session_id = _get_session_id(span) if session_id is not None: - span._set_ctx_item(SESSION_ID, session_id) + core.set_item(SESSION_ID, session_id) llmobs_span_event["session_id"] = session_id llmobs_span_event["tags"] = cls._llmobs_tags(span, ml_app, session_id) - span_links = span._get_ctx_item(SPAN_LINKS) + span_links = core.get_item(SPAN_LINKS) if isinstance(span_links, list): llmobs_span_event["span_links"] = span_links @@ -228,7 +228,7 @@ def _llmobs_tags(span: Span, ml_app: str, session_id: Optional[str] = None) -> L tags["session_id"] = session_id if _is_evaluation_span(span): tags[constants.RUNNER_IS_INTEGRATION_SPAN_TAG] = "ragas" - existing_tags = span._get_ctx_item(TAGS) + existing_tags = core.get_item(TAGS) if existing_tags is not None: tags.update(existing_tags) return ["{}:{}".format(k, v) for k, v in tags.items()] @@ -529,23 +529,23 @@ def _start_span( if name is None: name = operation_kind span = self.tracer.trace(name, resource=operation_kind, span_type=SpanTypes.LLM) - span._set_ctx_item(SPAN_KIND, operation_kind) + core.set_item(SPAN_KIND, operation_kind) if model_name is not None: - span._set_ctx_item(MODEL_NAME, model_name) + core.set_item(MODEL_NAME, model_name) if model_provider is not None: - span._set_ctx_item(MODEL_PROVIDER, model_provider) + core.set_item(MODEL_PROVIDER, model_provider) session_id = session_id if session_id is not None else _get_session_id(span) if session_id is not None: - span._set_ctx_item(SESSION_ID, session_id) + core.set_item(SESSION_ID, session_id) if ml_app is None: ml_app = _get_ml_app(span) - span._set_ctx_item(ML_APP, ml_app) + core.set_item(ML_APP, ml_app) if span.get_tag(PROPAGATED_PARENT_ID_KEY) is None: # For non-distributed traces or spans in the first service of a distributed trace, # The LLMObs parent ID tag is not set at span start time. We need to manually set the parent ID tag now # in these cases to avoid conflicting with the later propagated tags. parent_id = _get_llmobs_parent_id(span) or "undefined" - span._set_ctx_item(PARENT_ID_KEY, str(parent_id)) + core.set_item(PARENT_ID_KEY, str(parent_id)) return span @classmethod @@ -775,7 +775,7 @@ def annotate( log.warning("span tags must be a dictionary of string key - primitive value pairs.") else: cls._set_dict_attribute(span, TAGS, tags) - span_kind = span._get_ctx_item(SPAN_KIND) + span_kind = core.get_item(SPAN_KIND) if _name is not None: span.name = _name if prompt is not None: @@ -807,7 +807,7 @@ def _tag_llm_io(cls, span, input_messages=None, output_messages=None): if not isinstance(input_messages, Messages): input_messages = Messages(input_messages) if input_messages.messages: - span._set_ctx_item(INPUT_MESSAGES, input_messages.messages) + core.set_item(INPUT_MESSAGES, input_messages.messages) except TypeError: log.warning("Failed to parse input messages.", exc_info=True) if output_messages is None: @@ -817,7 +817,7 @@ def _tag_llm_io(cls, span, input_messages=None, output_messages=None): output_messages = Messages(output_messages) if not output_messages.messages: return - span._set_ctx_item(OUTPUT_MESSAGES, output_messages.messages) + core.set_item(OUTPUT_MESSAGES, output_messages.messages) except TypeError: log.warning("Failed to parse output messages.", exc_info=True) @@ -831,12 +831,12 @@ def _tag_embedding_io(cls, span, input_documents=None, output_text=None): if not isinstance(input_documents, Documents): input_documents = Documents(input_documents) if input_documents.documents: - span._set_ctx_item(INPUT_DOCUMENTS, input_documents.documents) + core.set_item(INPUT_DOCUMENTS, input_documents.documents) except TypeError: log.warning("Failed to parse input documents.", exc_info=True) if output_text is None: return - span._set_ctx_item(OUTPUT_VALUE, str(output_text)) + core.set_item(OUTPUT_VALUE, str(output_text)) @classmethod def _tag_retrieval_io(cls, span, input_text=None, output_documents=None): @@ -844,7 +844,7 @@ def _tag_retrieval_io(cls, span, input_text=None, output_documents=None): Will be mapped to span's `meta.{input,output}.text` fields. """ if input_text is not None: - span._set_ctx_item(INPUT_VALUE, str(input_text)) + core.set_item(INPUT_VALUE, str(input_text)) if output_documents is None: return try: @@ -852,7 +852,7 @@ def _tag_retrieval_io(cls, span, input_text=None, output_documents=None): output_documents = Documents(output_documents) if not output_documents.documents: return - span._set_ctx_item(OUTPUT_DOCUMENTS, output_documents.documents) + core.set_item(OUTPUT_DOCUMENTS, output_documents.documents) except TypeError: log.warning("Failed to parse output documents.", exc_info=True) @@ -862,9 +862,9 @@ def _tag_text_io(cls, span, input_value=None, output_value=None): Will be mapped to span's `meta.{input,output}.values` fields. """ if input_value is not None: - span._set_ctx_item(INPUT_VALUE, str(input_value)) + core.set_item(INPUT_VALUE, str(input_value)) if output_value is not None: - span._set_ctx_item(OUTPUT_VALUE, str(output_value)) + core.set_item(OUTPUT_VALUE, str(output_value)) @staticmethod def _set_dict_attribute(span: Span, key, value: Dict[str, Any]) -> None: @@ -872,9 +872,9 @@ def _set_dict_attribute(span: Span, key, value: Dict[str, Any]) -> None: If the attribute is already set on the span, the new dict with be merged with the existing dict. """ - existing_value = span._get_ctx_item(key) or {} + existing_value = core.get_item(key) or {} existing_value.update(value) - span._set_ctx_item(key, existing_value) + core.set_item(key, existing_value) @classmethod def submit_evaluation_for( diff --git a/ddtrace/llmobs/_utils.py b/ddtrace/llmobs/_utils.py index 8861820002c..3bd285ca63f 100644 --- a/ddtrace/llmobs/_utils.py +++ b/ddtrace/llmobs/_utils.py @@ -7,6 +7,7 @@ import ddtrace from ddtrace import config from ddtrace.ext import SpanTypes +from ddtrace.internal import core from ddtrace.internal.logger import get_logger from ddtrace.llmobs._constants import GEMINI_APM_SPAN_NAME from ddtrace.llmobs._constants import INTERNAL_CONTEXT_VARIABLE_KEYS @@ -112,8 +113,8 @@ def _get_llmobs_parent_id(span: Span) -> Optional[str]: """Return the span ID of the nearest LLMObs-type span in the span's ancestor tree. In priority order: manually set parent ID tag, nearest LLMObs ancestor, local root's propagated parent ID tag. """ - if span._get_ctx_item(PARENT_ID_KEY): - return span._get_ctx_item(PARENT_ID_KEY) + if core.get_item(PARENT_ID_KEY): + return core.get_item(PARENT_ID_KEY) nearest_llmobs_ancestor = _get_nearest_llmobs_ancestor(span) if nearest_llmobs_ancestor: return str(nearest_llmobs_ancestor.span_id) @@ -126,7 +127,7 @@ def _get_span_name(span: Span) -> str: elif span.name == OPENAI_APM_SPAN_NAME and span.resource != "": client_name = span.get_tag("openai.request.client") or "OpenAI" return "{}.{}".format(client_name, span.resource) - return span._get_ctx_item(NAME) or span.name + return core.get_item(NAME) or span.name def _is_evaluation_span(span: Span) -> bool: @@ -134,12 +135,12 @@ def _is_evaluation_span(span: Span) -> bool: Return whether or not a span is an evaluation span by checking the span's nearest LLMObs span ancestor. Default to 'False' """ - is_evaluation_span = span._get_ctx_item(IS_EVALUATION_SPAN) + is_evaluation_span = core.get_item(IS_EVALUATION_SPAN) if is_evaluation_span: return is_evaluation_span llmobs_parent = _get_nearest_llmobs_ancestor(span) while llmobs_parent: - is_evaluation_span = llmobs_parent._get_ctx_item(IS_EVALUATION_SPAN) + is_evaluation_span = core.get_item(IS_EVALUATION_SPAN) if is_evaluation_span: return is_evaluation_span llmobs_parent = _get_nearest_llmobs_ancestor(llmobs_parent) @@ -151,12 +152,12 @@ def _get_ml_app(span: Span) -> str: Return the ML app name for a given span, by checking the span's nearest LLMObs span ancestor. Default to the global config LLMObs ML app name otherwise. """ - ml_app = span._get_ctx_item(ML_APP) + ml_app = core.get_item(ML_APP) if ml_app: return ml_app llmobs_parent = _get_nearest_llmobs_ancestor(span) while llmobs_parent: - ml_app = llmobs_parent._get_ctx_item(ML_APP) + ml_app = core.get_item(ML_APP) if ml_app is not None: return ml_app llmobs_parent = _get_nearest_llmobs_ancestor(llmobs_parent) @@ -168,12 +169,12 @@ def _get_session_id(span: Span) -> Optional[str]: Return the session ID for a given span, by checking the span's nearest LLMObs span ancestor. Default to the span's trace ID. """ - session_id = span._get_ctx_item(SESSION_ID) + session_id = core.get_item(SESSION_ID) if session_id: return session_id llmobs_parent = _get_nearest_llmobs_ancestor(span) while llmobs_parent: - session_id = llmobs_parent._get_ctx_item(SESSION_ID) + session_id = core.get_item(SESSION_ID) if session_id is not None: return session_id llmobs_parent = _get_nearest_llmobs_ancestor(llmobs_parent) diff --git a/ddtrace/llmobs/decorators.py b/ddtrace/llmobs/decorators.py index 7e61f9b4e18..65a59704259 100644 --- a/ddtrace/llmobs/decorators.py +++ b/ddtrace/llmobs/decorators.py @@ -5,6 +5,7 @@ from typing import Callable from typing import Optional +from ddtrace.internal import core from ddtrace.internal.compat import iscoroutinefunction from ddtrace.internal.compat import isgeneratorfunction from ddtrace.internal.logger import get_logger @@ -192,7 +193,7 @@ async def wrapper(*args, **kwargs): _automatic_io_annotation and resp and operation_kind != "retrieval" - and span._get_ctx_item(OUTPUT_VALUE) is None + and core.get_item(OUTPUT_VALUE) is None ): LLMObs.annotate(span=span, output_data=resp) return resp @@ -240,7 +241,7 @@ def wrapper(*args, **kwargs): _automatic_io_annotation and resp and operation_kind != "retrieval" - and span._get_ctx_item(OUTPUT_VALUE) is None + and core.get_item(OUTPUT_VALUE) is None ): LLMObs.annotate(span=span, output_data=resp) return resp diff --git a/tests/llmobs/test_llmobs.py b/tests/llmobs/test_llmobs.py index 086e867c787..996714d026c 100644 --- a/tests/llmobs/test_llmobs.py +++ b/tests/llmobs/test_llmobs.py @@ -1,6 +1,7 @@ import pytest from ddtrace.ext import SpanTypes +from ddtrace.internal import core from ddtrace.llmobs import _constants as const from ddtrace.llmobs._utils import _get_llmobs_parent_id from ddtrace.llmobs._utils import _get_session_id @@ -11,16 +12,16 @@ class TestMLApp: @pytest.mark.parametrize("llmobs_env", [{"DD_LLMOBS_ML_APP": ""}]) def test_tag_defaults_to_env_var(self, tracer, llmobs_env, llmobs_events): """Test that no ml_app defaults to the environment variable DD_LLMOBS_ML_APP.""" - with tracer.trace("root_llm_span", span_type=SpanTypes.LLM) as llm_span: - llm_span._set_ctx_item(const.SPAN_KIND, "llm") + with tracer.trace("root_llm_span", span_type=SpanTypes.LLM): + core.set_item(const.SPAN_KIND, "llm") assert "ml_app:" in llmobs_events[0]["tags"] @pytest.mark.parametrize("llmobs_env", [{"DD_LLMOBS_ML_APP": ""}]) def test_tag_overrides_env_var(self, tracer, llmobs_env, llmobs_events): """Test that when ml_app is set on the span, it overrides the environment variable DD_LLMOBS_ML_APP.""" - with tracer.trace("root_llm_span", span_type=SpanTypes.LLM) as llm_span: - llm_span._set_ctx_item(const.SPAN_KIND, "llm") - llm_span._set_ctx_item(const.ML_APP, "test-ml-app") + with tracer.trace("root_llm_span", span_type=SpanTypes.LLM): + core.set_item(const.SPAN_KIND, "llm") + core.set_item(const.ML_APP, "test-ml-app") assert "ml_app:test-ml-app" in llmobs_events[0]["tags"] def test_propagates_ignore_non_llmobs_spans(self, tracer, llmobs_events): @@ -28,14 +29,14 @@ def test_propagates_ignore_non_llmobs_spans(self, tracer, llmobs_events): Test that when ml_app is not set, we propagate from nearest LLMObs ancestor even if there are non-LLMObs spans in between. """ - with tracer.trace("root_llm_span", span_type=SpanTypes.LLM) as llm_span: - llm_span._set_ctx_item(const.SPAN_KIND, "llm") - llm_span._set_ctx_item(const.ML_APP, "test-ml-app") + with tracer.trace("root_llm_span", span_type=SpanTypes.LLM): + core.set_item(const.SPAN_KIND, "llm") + core.set_item(const.ML_APP, "test-ml-app") with tracer.trace("child_span"): - with tracer.trace("llm_grandchild_span", span_type=SpanTypes.LLM) as grandchild_span: - grandchild_span._set_ctx_item(const.SPAN_KIND, "llm") - with tracer.trace("great_grandchild_span", span_type=SpanTypes.LLM) as great_grandchild_span: - great_grandchild_span._set_ctx_item(const.SPAN_KIND, "llm") + with tracer.trace("llm_grandchild_span", span_type=SpanTypes.LLM): + core.set_item(const.SPAN_KIND, "llm") + with tracer.trace("great_grandchild_span", span_type=SpanTypes.LLM): + core.set_item(const.SPAN_KIND, "llm") assert len(llmobs_events) == 3 for llmobs_event in llmobs_events: assert "ml_app:test-ml-app" in llmobs_event["tags"] @@ -62,8 +63,8 @@ def test_propagate_from_ancestors(self, tracer): Test that session_id is propagated from the nearest LLMObs span in the span's ancestor tree if no session_id is not set on the span itself. """ - with tracer.trace("root_llm_span", span_type=SpanTypes.LLM) as root_span: - root_span._set_ctx_item(const.SESSION_ID, "test_session_id") + with tracer.trace("root_llm_span", span_type=SpanTypes.LLM): + core.set_item(const.SESSION_ID, "test_session_id") with tracer.trace("child_span"): with tracer.trace("llm_span", span_type=SpanTypes.LLM) as llm_span: pass @@ -71,11 +72,11 @@ def test_propagate_from_ancestors(self, tracer): def test_if_set_manually(self, tracer): """Test that session_id is extracted from the span if it is already set manually.""" - with tracer.trace("root_llm_span", span_type=SpanTypes.LLM) as root_span: - root_span._set_ctx_item(const.SESSION_ID, "test_session_id") + with tracer.trace("root_llm_span", span_type=SpanTypes.LLM): + core.set_item(const.SESSION_ID, "test_session_id") with tracer.trace("child_span"): with tracer.trace("llm_span", span_type=SpanTypes.LLM) as llm_span: - llm_span._set_ctx_item(const.SESSION_ID, "test_different_session_id") + core.set_item(const.SESSION_ID, "test_different_session_id") assert _get_session_id(llm_span) == "test_different_session_id" def test_propagates_ignore_non_llmobs_spans(self, tracer, llmobs_events): @@ -83,14 +84,14 @@ def test_propagates_ignore_non_llmobs_spans(self, tracer, llmobs_events): Test that when session_id is not set, we propagate from nearest LLMObs ancestor even if there are non-LLMObs spans in between. """ - with tracer.trace("root_llm_span", span_type=SpanTypes.LLM) as llm_span: - llm_span._set_ctx_item(const.SPAN_KIND, "llm") - llm_span._set_ctx_item(const.SESSION_ID, "session-123") + with tracer.trace("root_llm_span", span_type=SpanTypes.LLM): + core.set_item(const.SPAN_KIND, "llm") + core.set_item(const.SESSION_ID, "session-123") with tracer.trace("child_span"): - with tracer.trace("llm_grandchild_span", span_type=SpanTypes.LLM) as grandchild_span: - grandchild_span._set_ctx_item(const.SPAN_KIND, "llm") - with tracer.trace("great_grandchild_span", span_type=SpanTypes.LLM) as great_grandchild_span: - great_grandchild_span._set_ctx_item(const.SPAN_KIND, "llm") + with tracer.trace("llm_grandchild_span", span_type=SpanTypes.LLM): + core.set_item(const.SPAN_KIND, "llm") + with tracer.trace("great_grandchild_span", span_type=SpanTypes.LLM): + core.set_item(const.SPAN_KIND, "llm") llm_event, grandchild_event, great_grandchild_event = llmobs_events assert llm_event["session_id"] == "session-123" @@ -100,81 +101,81 @@ def test_propagates_ignore_non_llmobs_spans(self, tracer, llmobs_events): def test_input_value_is_set(tracer, llmobs_events): """Test that input value is set on the span event if they are present on the span.""" - with tracer.trace("root_llm_span", span_type=SpanTypes.LLM) as llm_span: - llm_span._set_ctx_item(const.SPAN_KIND, "llm") - llm_span._set_ctx_item(const.INPUT_VALUE, "value") + with tracer.trace("root_llm_span", span_type=SpanTypes.LLM): + core.set_item(const.SPAN_KIND, "llm") + core.set_item(const.INPUT_VALUE, "value") assert llmobs_events[0]["meta"]["input"]["value"] == "value" def test_input_messages_are_set(tracer, llmobs_events): """Test that input messages are set on the span event if they are present on the span.""" - with tracer.trace("root_llm_span", span_type=SpanTypes.LLM) as llm_span: - llm_span._set_ctx_item(const.SPAN_KIND, "llm") - llm_span._set_ctx_item(const.INPUT_MESSAGES, [{"content": "message", "role": "user"}]) + with tracer.trace("root_llm_span", span_type=SpanTypes.LLM): + core.set_item(const.SPAN_KIND, "llm") + core.set_item(const.INPUT_MESSAGES, [{"content": "message", "role": "user"}]) assert llmobs_events[0]["meta"]["input"]["messages"] == [{"content": "message", "role": "user"}] def test_output_messages_are_set(tracer, llmobs_events): """Test that output messages are set on the span event if they are present on the span.""" - with tracer.trace("root_llm_span", span_type=SpanTypes.LLM) as llm_span: - llm_span._set_ctx_item(const.SPAN_KIND, "llm") - llm_span._set_ctx_item(const.OUTPUT_MESSAGES, [{"content": "message", "role": "user"}]) + with tracer.trace("root_llm_span", span_type=SpanTypes.LLM): + core.set_item(const.SPAN_KIND, "llm") + core.set_item(const.OUTPUT_MESSAGES, [{"content": "message", "role": "user"}]) assert llmobs_events[0]["meta"]["output"]["messages"] == [{"content": "message", "role": "user"}] def test_output_value_is_set(tracer, llmobs_events): """Test that output value is set on the span event if they are present on the span.""" - with tracer.trace("root_llm_span", span_type=SpanTypes.LLM) as llm_span: - llm_span._set_ctx_item(const.SPAN_KIND, "llm") - llm_span._set_ctx_item(const.OUTPUT_VALUE, "value") + with tracer.trace("root_llm_span", span_type=SpanTypes.LLM): + core.set_item(const.SPAN_KIND, "llm") + core.set_item(const.OUTPUT_VALUE, "value") assert llmobs_events[0]["meta"]["output"]["value"] == "value" def test_prompt_is_set(tracer, llmobs_events): """Test that prompt is set on the span event if they are present on the span.""" - with tracer.trace("root_llm_span", span_type=SpanTypes.LLM) as llm_span: - llm_span._set_ctx_item(const.SPAN_KIND, "llm") - llm_span._set_ctx_item(const.INPUT_PROMPT, {"variables": {"var1": "var2"}}) + with tracer.trace("root_llm_span", span_type=SpanTypes.LLM): + core.set_item(const.SPAN_KIND, "llm") + core.set_item(const.INPUT_PROMPT, {"variables": {"var1": "var2"}}) assert llmobs_events[0]["meta"]["input"]["prompt"] == {"variables": {"var1": "var2"}} def test_prompt_is_not_set_for_non_llm_spans(tracer, llmobs_events): """Test that prompt is NOT set on the span event if the span is not an LLM span.""" - with tracer.trace("task_span", span_type=SpanTypes.LLM) as task_span: - task_span._set_ctx_item(const.SPAN_KIND, "task") - task_span._set_ctx_item(const.INPUT_VALUE, "ival") - task_span._set_ctx_item(const.INPUT_PROMPT, {"variables": {"var1": "var2"}}) + with tracer.trace("task_span", span_type=SpanTypes.LLM): + core.set_item(const.SPAN_KIND, "task") + core.set_item(const.INPUT_VALUE, "ival") + core.set_item(const.INPUT_PROMPT, {"variables": {"var1": "var2"}}) assert llmobs_events[0]["meta"]["input"].get("prompt") is None def test_metadata_is_set(tracer, llmobs_events): """Test that metadata is set on the span event if it is present on the span.""" - with tracer.trace("root_llm_span", span_type=SpanTypes.LLM) as llm_span: - llm_span._set_ctx_item(const.SPAN_KIND, "llm") - llm_span._set_ctx_item(const.METADATA, {"key": "value"}) + with tracer.trace("root_llm_span", span_type=SpanTypes.LLM): + core.set_item(const.SPAN_KIND, "llm") + core.set_item(const.METADATA, {"key": "value"}) assert llmobs_events[0]["meta"]["metadata"] == {"key": "value"} def test_metrics_are_set(tracer, llmobs_events): """Test that metadata is set on the span event if it is present on the span.""" - with tracer.trace("root_llm_span", span_type=SpanTypes.LLM) as llm_span: - llm_span._set_ctx_item(const.SPAN_KIND, "llm") - llm_span._set_ctx_item(const.METRICS, {"tokens": 100}) + with tracer.trace("root_llm_span", span_type=SpanTypes.LLM): + core.set_item(const.SPAN_KIND, "llm") + core.set_item(const.METRICS, {"tokens": 100}) assert llmobs_events[0]["metrics"] == {"tokens": 100} def test_langchain_span_name_is_set_to_class_name(tracer, llmobs_events): """Test span names for langchain auto-instrumented spans is set correctly.""" - with tracer.trace(const.LANGCHAIN_APM_SPAN_NAME, resource="expected_name", span_type=SpanTypes.LLM) as llm_span: - llm_span._set_ctx_item(const.SPAN_KIND, "llm") + with tracer.trace(const.LANGCHAIN_APM_SPAN_NAME, resource="expected_name", span_type=SpanTypes.LLM): + core.set_item(const.SPAN_KIND, "llm") assert llmobs_events[0]["name"] == "expected_name" def test_error_is_set(tracer, llmobs_events): """Test that error is set on the span event if it is present on the span.""" with pytest.raises(ValueError): - with tracer.trace("root_llm_span", span_type=SpanTypes.LLM) as llm_span: - llm_span._set_ctx_item(const.SPAN_KIND, "llm") + with tracer.trace("root_llm_span", span_type=SpanTypes.LLM): + core.set_item(const.SPAN_KIND, "llm") raise ValueError("error") span_event = llmobs_events[0] assert span_event["meta"]["error.message"] == "error" @@ -184,9 +185,9 @@ def test_error_is_set(tracer, llmobs_events): def test_model_provider_defaults_to_custom(tracer, llmobs_events): """Test that model provider defaults to "custom" if not provided.""" - with tracer.trace("root_llm_span", span_type=SpanTypes.LLM) as llm_span: - llm_span._set_ctx_item(const.SPAN_KIND, "llm") - llm_span._set_ctx_item(const.MODEL_NAME, "model_name") + with tracer.trace("root_llm_span", span_type=SpanTypes.LLM): + core.set_item(const.SPAN_KIND, "llm") + core.set_item(const.MODEL_NAME, "model_name") span_event = llmobs_events[0] assert span_event["meta"]["model_name"] == "model_name" assert span_event["meta"]["model_provider"] == "custom" @@ -194,9 +195,9 @@ def test_model_provider_defaults_to_custom(tracer, llmobs_events): def test_model_not_set_if_not_llm_kind_span(tracer, llmobs_events): """Test that model name and provider not set if non-LLM span.""" - with tracer.trace("root_workflow_span", span_type=SpanTypes.LLM) as span: - span._set_ctx_item(const.SPAN_KIND, "workflow") - span._set_ctx_item(const.MODEL_NAME, "model_name") + with tracer.trace("root_workflow_span", span_type=SpanTypes.LLM): + core.set_item(const.SPAN_KIND, "workflow") + core.set_item(const.MODEL_NAME, "model_name") span_event = llmobs_events[0] assert "model_name" not in span_event["meta"] assert "model_provider" not in span_event["meta"] @@ -204,10 +205,10 @@ def test_model_not_set_if_not_llm_kind_span(tracer, llmobs_events): def test_model_and_provider_are_set(tracer, llmobs_events): """Test that model and provider are set on the span event if they are present on the LLM-kind span.""" - with tracer.trace("root_llm_span", span_type=SpanTypes.LLM) as llm_span: - llm_span._set_ctx_item(const.SPAN_KIND, "llm") - llm_span._set_ctx_item(const.MODEL_NAME, "model_name") - llm_span._set_ctx_item(const.MODEL_PROVIDER, "model_provider") + with tracer.trace("root_llm_span", span_type=SpanTypes.LLM): + core.set_item(const.SPAN_KIND, "llm") + core.set_item(const.MODEL_NAME, "model_name") + core.set_item(const.MODEL_PROVIDER, "model_provider") span_event = llmobs_events[0] assert span_event["meta"]["model_name"] == "model_name" assert span_event["meta"]["model_provider"] == "model_provider" @@ -227,10 +228,10 @@ def test_malformed_span_logs_error_instead_of_raising(tracer, llmobs_events, moc def test_only_generate_span_events_from_llmobs_spans(tracer, llmobs_events): """Test that we only generate LLMObs span events for LLM span types.""" with tracer.trace("root_llm_span", service="tests.llmobs", span_type=SpanTypes.LLM) as root_span: - root_span._set_ctx_item(const.SPAN_KIND, "llm") + core.set_item(const.SPAN_KIND, "llm") with tracer.trace("child_span"): with tracer.trace("llm_span", span_type=SpanTypes.LLM) as grandchild_span: - grandchild_span._set_ctx_item(const.SPAN_KIND, "llm") + core.set_item(const.SPAN_KIND, "llm") expected_grandchild_llmobs_span = _expected_llmobs_llm_span_event(grandchild_span, "llm") expected_grandchild_llmobs_span["parent_id"] = str(root_span.span_id) diff --git a/tests/llmobs/test_llmobs_service.py b/tests/llmobs/test_llmobs_service.py index 2fe3e1fbfab..be8e78b24df 100644 --- a/tests/llmobs/test_llmobs_service.py +++ b/tests/llmobs/test_llmobs_service.py @@ -8,6 +8,7 @@ import ddtrace from ddtrace.ext import SpanTypes +from ddtrace.internal import core from ddtrace.internal.service import ServiceStatus from ddtrace.llmobs import LLMObs as llmobs_service from ddtrace.llmobs._constants import INPUT_DOCUMENTS @@ -213,16 +214,16 @@ def test_start_span_uses_kind_as_default_name(llmobs): def test_start_span_with_session_id(llmobs): - with llmobs.llm(model_name="test_model", session_id="test_session_id") as span: - assert span._get_ctx_item(SESSION_ID) == "test_session_id" - with llmobs.tool(session_id="test_session_id") as span: - assert span._get_ctx_item(SESSION_ID) == "test_session_id" - with llmobs.task(session_id="test_session_id") as span: - assert span._get_ctx_item(SESSION_ID) == "test_session_id" - with llmobs.workflow(session_id="test_session_id") as span: - assert span._get_ctx_item(SESSION_ID) == "test_session_id" - with llmobs.agent(session_id="test_session_id") as span: - assert span._get_ctx_item(SESSION_ID) == "test_session_id" + with llmobs.llm(model_name="test_model", session_id="test_session_id"): + assert core.get_item(SESSION_ID) == "test_session_id" + with llmobs.tool(session_id="test_session_id"): + assert core.get_item(SESSION_ID) == "test_session_id" + with llmobs.task(session_id="test_session_id"): + assert core.get_item(SESSION_ID) == "test_session_id" + with llmobs.workflow(session_id="test_session_id"): + assert core.get_item(SESSION_ID) == "test_session_id" + with llmobs.agent(session_id="test_session_id"): + assert core.get_item(SESSION_ID) == "test_session_id" def test_session_id_becomes_top_level_field(llmobs, llmobs_events): @@ -238,9 +239,9 @@ def test_llm_span(llmobs, llmobs_events): assert span.name == "test_llm_call" assert span.resource == "llm" assert span.span_type == "llm" - assert span._get_ctx_item(SPAN_KIND) == "llm" - assert span._get_ctx_item(MODEL_NAME) == "test_model" - assert span._get_ctx_item(MODEL_PROVIDER) == "test_provider" + assert core.get_item(SPAN_KIND) == "llm" + assert core.get_item(MODEL_NAME) == "test_model" + assert core.get_item(MODEL_PROVIDER) == "test_provider" assert len(llmobs_events) == 1 assert llmobs_events[0] == _expected_llmobs_llm_span_event( span, "llm", model_name="test_model", model_provider="test_provider" @@ -249,7 +250,7 @@ def test_llm_span(llmobs, llmobs_events): def test_llm_span_no_model_sets_default(llmobs, llmobs_events): with llmobs.llm(name="test_llm_call", model_provider="test_provider") as span: - assert span._get_ctx_item(MODEL_NAME) == "custom" + assert core.get_item(MODEL_NAME) == "custom" assert len(llmobs_events) == 1 assert llmobs_events[0] == _expected_llmobs_llm_span_event( span, "llm", model_name="custom", model_provider="test_provider" @@ -261,9 +262,9 @@ def test_default_model_provider_set_to_custom(llmobs): assert span.name == "test_llm_call" assert span.resource == "llm" assert span.span_type == "llm" - assert span._get_ctx_item(SPAN_KIND) == "llm" - assert span._get_ctx_item(MODEL_NAME) == "test_model" - assert span._get_ctx_item(MODEL_PROVIDER) == "custom" + assert core.get_item(SPAN_KIND) == "llm" + assert core.get_item(MODEL_NAME) == "test_model" + assert core.get_item(MODEL_PROVIDER) == "custom" def test_tool_span(llmobs, llmobs_events): @@ -271,7 +272,7 @@ def test_tool_span(llmobs, llmobs_events): assert span.name == "test_tool" assert span.resource == "tool" assert span.span_type == "llm" - assert span._get_ctx_item(SPAN_KIND) == "tool" + assert core.get_item(SPAN_KIND) == "tool" assert len(llmobs_events) == 1 assert llmobs_events[0] == _expected_llmobs_non_llm_span_event(span, "tool") @@ -281,7 +282,7 @@ def test_task_span(llmobs, llmobs_events): assert span.name == "test_task" assert span.resource == "task" assert span.span_type == "llm" - assert span._get_ctx_item(SPAN_KIND) == "task" + assert core.get_item(SPAN_KIND) == "task" assert len(llmobs_events) == 1 assert llmobs_events[0] == _expected_llmobs_non_llm_span_event(span, "task") @@ -291,7 +292,7 @@ def test_workflow_span(llmobs, llmobs_events): assert span.name == "test_workflow" assert span.resource == "workflow" assert span.span_type == "llm" - assert span._get_ctx_item(SPAN_KIND) == "workflow" + assert core.get_item(SPAN_KIND) == "workflow" assert len(llmobs_events) == 1 assert llmobs_events[0] == _expected_llmobs_non_llm_span_event(span, "workflow") @@ -301,14 +302,14 @@ def test_agent_span(llmobs, llmobs_events): assert span.name == "test_agent" assert span.resource == "agent" assert span.span_type == "llm" - assert span._get_ctx_item(SPAN_KIND) == "agent" + assert core.get_item(SPAN_KIND) == "agent" assert len(llmobs_events) == 1 assert llmobs_events[0] == _expected_llmobs_llm_span_event(span, "agent") def test_embedding_span_no_model_sets_default(llmobs, llmobs_events): with llmobs.embedding(name="test_embedding", model_provider="test_provider") as span: - assert span._get_ctx_item(MODEL_NAME) == "custom" + assert core.get_item(MODEL_NAME) == "custom" assert len(llmobs_events) == 1 assert llmobs_events[0] == _expected_llmobs_llm_span_event( span, "embedding", model_name="custom", model_provider="test_provider" @@ -320,9 +321,9 @@ def test_embedding_default_model_provider_set_to_custom(llmobs): assert span.name == "test_embedding" assert span.resource == "embedding" assert span.span_type == "llm" - assert span._get_ctx_item(SPAN_KIND) == "embedding" - assert span._get_ctx_item(MODEL_NAME) == "test_model" - assert span._get_ctx_item(MODEL_PROVIDER) == "custom" + assert core.get_item(SPAN_KIND) == "embedding" + assert core.get_item(MODEL_NAME) == "test_model" + assert core.get_item(MODEL_PROVIDER) == "custom" def test_embedding_span(llmobs, llmobs_events): @@ -330,9 +331,9 @@ def test_embedding_span(llmobs, llmobs_events): assert span.name == "test_embedding" assert span.resource == "embedding" assert span.span_type == "llm" - assert span._get_ctx_item(SPAN_KIND) == "embedding" - assert span._get_ctx_item(MODEL_NAME) == "test_model" - assert span._get_ctx_item(MODEL_PROVIDER) == "test_provider" + assert core.get_item(SPAN_KIND) == "embedding" + assert core.get_item(MODEL_NAME) == "test_model" + assert core.get_item(MODEL_PROVIDER) == "test_provider" assert len(llmobs_events) == 1 assert llmobs_events[0] == _expected_llmobs_llm_span_event( span, "embedding", model_name="test_model", model_provider="test_provider" @@ -361,14 +362,14 @@ def test_annotate_finished_span_does_nothing(llmobs, mock_llmobs_logs): def test_annotate_metadata(llmobs): with llmobs.llm(model_name="test_model", name="test_llm_call", model_provider="test_provider") as span: llmobs.annotate(span=span, metadata={"temperature": 0.5, "max_tokens": 20, "top_k": 10, "n": 3}) - assert span._get_ctx_item(METADATA) == {"temperature": 0.5, "max_tokens": 20, "top_k": 10, "n": 3} + assert core.get_item(METADATA) == {"temperature": 0.5, "max_tokens": 20, "top_k": 10, "n": 3} def test_annotate_metadata_updates(llmobs): with llmobs.llm(model_name="test_model", name="test_llm_call", model_provider="test_provider") as span: llmobs.annotate(span=span, metadata={"temperature": 0.5, "max_tokens": 20, "top_k": 10, "n": 3}) llmobs.annotate(span=span, metadata={"temperature": 1, "logit_bias": [{"1": 2}]}) - assert span._get_ctx_item(METADATA) == { + assert core.get_item(METADATA) == { "temperature": 1, "max_tokens": 20, "top_k": 10, @@ -380,7 +381,7 @@ def test_annotate_metadata_updates(llmobs): def test_annotate_metadata_wrong_type_raises_warning(llmobs, mock_llmobs_logs): with llmobs.llm(model_name="test_model", name="test_llm_call", model_provider="test_provider") as span: llmobs.annotate(span=span, metadata="wrong_metadata") - assert span._get_ctx_item(METADATA) is None + assert core.get_item(METADATA) is None mock_llmobs_logs.warning.assert_called_once_with("metadata must be a dictionary") mock_llmobs_logs.reset_mock() @@ -388,13 +389,13 @@ def test_annotate_metadata_wrong_type_raises_warning(llmobs, mock_llmobs_logs): def test_annotate_tag(llmobs): with llmobs.llm(model_name="test_model", name="test_llm_call", model_provider="test_provider") as span: llmobs.annotate(span=span, tags={"test_tag_name": "test_tag_value", "test_numeric_tag": 10}) - assert span._get_ctx_item(TAGS) == {"test_tag_name": "test_tag_value", "test_numeric_tag": 10} + assert core.get_item(TAGS) == {"test_tag_name": "test_tag_value", "test_numeric_tag": 10} def test_annotate_tag_wrong_type(llmobs, mock_llmobs_logs): with llmobs.llm(model_name="test_model", name="test_llm_call", model_provider="test_provider") as span: llmobs.annotate(span=span, tags=12345) - assert span._get_ctx_item(TAGS) is None + assert core.get_item(TAGS) is None mock_llmobs_logs.warning.assert_called_once_with( "span tags must be a dictionary of string key - primitive value pairs." ) @@ -403,63 +404,63 @@ def test_annotate_tag_wrong_type(llmobs, mock_llmobs_logs): def test_annotate_input_string(llmobs): with llmobs.llm(model_name="test_model") as llm_span: llmobs.annotate(span=llm_span, input_data="test_input") - assert llm_span._get_ctx_item(INPUT_MESSAGES) == [{"content": "test_input"}] + assert core.get_item(INPUT_MESSAGES) == [{"content": "test_input"}] with llmobs.task() as task_span: llmobs.annotate(span=task_span, input_data="test_input") - assert task_span._get_ctx_item(INPUT_VALUE) == "test_input" + assert core.get_item(INPUT_VALUE) == "test_input" with llmobs.tool() as tool_span: llmobs.annotate(span=tool_span, input_data="test_input") - assert tool_span._get_ctx_item(INPUT_VALUE) == "test_input" + assert core.get_item(INPUT_VALUE) == "test_input" with llmobs.workflow() as workflow_span: llmobs.annotate(span=workflow_span, input_data="test_input") - assert workflow_span._get_ctx_item(INPUT_VALUE) == "test_input" + assert core.get_item(INPUT_VALUE) == "test_input" with llmobs.agent() as agent_span: llmobs.annotate(span=agent_span, input_data="test_input") - assert agent_span._get_ctx_item(INPUT_VALUE) == "test_input" + assert core.get_item(INPUT_VALUE) == "test_input" with llmobs.retrieval() as retrieval_span: llmobs.annotate(span=retrieval_span, input_data="test_input") - assert retrieval_span._get_ctx_item(INPUT_VALUE) == "test_input" + assert core.get_item(INPUT_VALUE) == "test_input" def test_annotate_numeric_io(llmobs): with llmobs.task() as task_span: llmobs.annotate(span=task_span, input_data=0, output_data=0) - assert task_span._get_ctx_item(INPUT_VALUE) == "0" - assert task_span._get_ctx_item(OUTPUT_VALUE) == "0" + assert core.get_item(INPUT_VALUE) == "0" + assert core.get_item(OUTPUT_VALUE) == "0" with llmobs.task() as task_span: llmobs.annotate(span=task_span, input_data=1.23, output_data=1.23) - assert task_span._get_ctx_item(INPUT_VALUE) == "1.23" - assert task_span._get_ctx_item(OUTPUT_VALUE) == "1.23" + assert core.get_item(INPUT_VALUE) == "1.23" + assert core.get_item(OUTPUT_VALUE) == "1.23" def test_annotate_input_serializable_value(llmobs): with llmobs.task() as task_span: llmobs.annotate(span=task_span, input_data=["test_input"]) - assert task_span._get_ctx_item(INPUT_VALUE) == str(["test_input"]) + assert core.get_item(INPUT_VALUE) == str(["test_input"]) with llmobs.tool() as tool_span: llmobs.annotate(span=tool_span, input_data={"test_input": "hello world"}) - assert tool_span._get_ctx_item(INPUT_VALUE) == str({"test_input": "hello world"}) + assert core.get_item(INPUT_VALUE) == str({"test_input": "hello world"}) with llmobs.workflow() as workflow_span: llmobs.annotate(span=workflow_span, input_data=("asd", 123)) - assert workflow_span._get_ctx_item(INPUT_VALUE) == str(("asd", 123)) + assert core.get_item(INPUT_VALUE) == str(("asd", 123)) with llmobs.agent() as agent_span: llmobs.annotate(span=agent_span, input_data="test_input") - assert agent_span._get_ctx_item(INPUT_VALUE) == "test_input" + assert core.get_item(INPUT_VALUE) == "test_input" with llmobs.retrieval() as retrieval_span: llmobs.annotate(span=retrieval_span, input_data=[0, 1, 2, 3, 4]) - assert retrieval_span._get_ctx_item(INPUT_VALUE) == str([0, 1, 2, 3, 4]) + assert core.get_item(INPUT_VALUE) == str([0, 1, 2, 3, 4]) def test_annotate_input_llm_message(llmobs): with llmobs.llm(model_name="test_model") as span: llmobs.annotate(span=span, input_data=[{"content": "test_input", "role": "human"}]) - assert span._get_ctx_item(INPUT_MESSAGES) == [{"content": "test_input", "role": "human"}] + assert core.get_item(INPUT_MESSAGES) == [{"content": "test_input", "role": "human"}] def test_annotate_input_llm_message_wrong_type(llmobs, mock_llmobs_logs): with llmobs.llm(model_name="test_model") as span: llmobs.annotate(span=span, input_data=[{"content": object()}]) - assert span._get_ctx_item(INPUT_MESSAGES) is None + assert core.get_item(INPUT_MESSAGES) is None mock_llmobs_logs.warning.assert_called_once_with("Failed to parse input messages.", exc_info=True) @@ -475,13 +476,13 @@ def test_llmobs_annotate_incorrect_message_content_type_raises_warning(llmobs, m def test_annotate_document_str(llmobs): with llmobs.embedding(model_name="test_model") as span: llmobs.annotate(span=span, input_data="test_document_text") - documents = span._get_ctx_item(INPUT_DOCUMENTS) + documents = core.get_item(INPUT_DOCUMENTS) assert documents assert len(documents) == 1 assert documents[0]["text"] == "test_document_text" with llmobs.retrieval() as span: llmobs.annotate(span=span, output_data="test_document_text") - documents = span._get_ctx_item(OUTPUT_DOCUMENTS) + documents = core.get_item(OUTPUT_DOCUMENTS) assert documents assert len(documents) == 1 assert documents[0]["text"] == "test_document_text" @@ -490,13 +491,13 @@ def test_annotate_document_str(llmobs): def test_annotate_document_dict(llmobs): with llmobs.embedding(model_name="test_model") as span: llmobs.annotate(span=span, input_data={"text": "test_document_text"}) - documents = span._get_ctx_item(INPUT_DOCUMENTS) + documents = core.get_item(INPUT_DOCUMENTS) assert documents assert len(documents) == 1 assert documents[0]["text"] == "test_document_text" with llmobs.retrieval() as span: llmobs.annotate(span=span, output_data={"text": "test_document_text"}) - documents = span._get_ctx_item(OUTPUT_DOCUMENTS) + documents = core.get_item(OUTPUT_DOCUMENTS) assert documents assert len(documents) == 1 assert documents[0]["text"] == "test_document_text" @@ -508,7 +509,7 @@ def test_annotate_document_list(llmobs): span=span, input_data=[{"text": "test_document_text"}, {"text": "text", "name": "name", "score": 0.9, "id": "id"}], ) - documents = span._get_ctx_item(INPUT_DOCUMENTS) + documents = core.get_item(INPUT_DOCUMENTS) assert documents assert len(documents) == 2 assert documents[0]["text"] == "test_document_text" @@ -521,7 +522,7 @@ def test_annotate_document_list(llmobs): span=span, output_data=[{"text": "test_document_text"}, {"text": "text", "name": "name", "score": 0.9, "id": "id"}], ) - documents = span._get_ctx_item(OUTPUT_DOCUMENTS) + documents = core.get_item(OUTPUT_DOCUMENTS) assert documents assert len(documents) == 2 assert documents[0]["text"] == "test_document_text" @@ -588,72 +589,72 @@ def test_annotate_incorrect_document_field_type_raises_warning(llmobs, mock_llmo def test_annotate_output_string(llmobs): with llmobs.llm(model_name="test_model") as llm_span: llmobs.annotate(span=llm_span, output_data="test_output") - assert llm_span._get_ctx_item(OUTPUT_MESSAGES) == [{"content": "test_output"}] + assert core.get_item(OUTPUT_MESSAGES) == [{"content": "test_output"}] with llmobs.embedding(model_name="test_model") as embedding_span: llmobs.annotate(span=embedding_span, output_data="test_output") - assert embedding_span._get_ctx_item(OUTPUT_VALUE) == "test_output" + assert core.get_item(OUTPUT_VALUE) == "test_output" with llmobs.task() as task_span: llmobs.annotate(span=task_span, output_data="test_output") - assert task_span._get_ctx_item(OUTPUT_VALUE) == "test_output" + assert core.get_item(OUTPUT_VALUE) == "test_output" with llmobs.tool() as tool_span: llmobs.annotate(span=tool_span, output_data="test_output") - assert tool_span._get_ctx_item(OUTPUT_VALUE) == "test_output" + assert core.get_item(OUTPUT_VALUE) == "test_output" with llmobs.workflow() as workflow_span: llmobs.annotate(span=workflow_span, output_data="test_output") - assert workflow_span._get_ctx_item(OUTPUT_VALUE) == "test_output" + assert core.get_item(OUTPUT_VALUE) == "test_output" with llmobs.agent() as agent_span: llmobs.annotate(span=agent_span, output_data="test_output") - assert agent_span._get_ctx_item(OUTPUT_VALUE) == "test_output" + assert core.get_item(OUTPUT_VALUE) == "test_output" def test_annotate_output_serializable_value(llmobs): with llmobs.embedding(model_name="test_model") as embedding_span: llmobs.annotate(span=embedding_span, output_data=[[0, 1, 2, 3], [4, 5, 6, 7]]) - assert embedding_span._get_ctx_item(OUTPUT_VALUE) == str([[0, 1, 2, 3], [4, 5, 6, 7]]) + assert core.get_item(OUTPUT_VALUE) == str([[0, 1, 2, 3], [4, 5, 6, 7]]) with llmobs.task() as task_span: llmobs.annotate(span=task_span, output_data=["test_output"]) - assert task_span._get_ctx_item(OUTPUT_VALUE) == str(["test_output"]) + assert core.get_item(OUTPUT_VALUE) == str(["test_output"]) with llmobs.tool() as tool_span: llmobs.annotate(span=tool_span, output_data={"test_output": "hello world"}) - assert tool_span._get_ctx_item(OUTPUT_VALUE) == str({"test_output": "hello world"}) + assert core.get_item(OUTPUT_VALUE) == str({"test_output": "hello world"}) with llmobs.workflow() as workflow_span: llmobs.annotate(span=workflow_span, output_data=("asd", 123)) - assert workflow_span._get_ctx_item(OUTPUT_VALUE) == str(("asd", 123)) + assert core.get_item(OUTPUT_VALUE) == str(("asd", 123)) with llmobs.agent() as agent_span: llmobs.annotate(span=agent_span, output_data="test_output") - assert agent_span._get_ctx_item(OUTPUT_VALUE) == "test_output" + assert core.get_item(OUTPUT_VALUE) == "test_output" def test_annotate_output_llm_message(llmobs): with llmobs.llm(model_name="test_model") as llm_span: llmobs.annotate(span=llm_span, output_data=[{"content": "test_output", "role": "human"}]) - assert llm_span._get_ctx_item(OUTPUT_MESSAGES) == [{"content": "test_output", "role": "human"}] + assert core.get_item(OUTPUT_MESSAGES) == [{"content": "test_output", "role": "human"}] def test_annotate_output_llm_message_wrong_type(llmobs, mock_llmobs_logs): with llmobs.llm(model_name="test_model") as llm_span: llmobs.annotate(span=llm_span, output_data=[{"content": object()}]) - assert llm_span._get_ctx_item(OUTPUT_MESSAGES) is None + assert core.get_item(OUTPUT_MESSAGES) is None mock_llmobs_logs.warning.assert_called_once_with("Failed to parse output messages.", exc_info=True) def test_annotate_metrics(llmobs): with llmobs.llm(model_name="test_model") as span: llmobs.annotate(span=span, metrics={"input_tokens": 10, "output_tokens": 20, "total_tokens": 30}) - assert span._get_ctx_item(METRICS) == {"input_tokens": 10, "output_tokens": 20, "total_tokens": 30} + assert core.get_item(METRICS) == {"input_tokens": 10, "output_tokens": 20, "total_tokens": 30} def test_annotate_metrics_updates(llmobs): with llmobs.llm(model_name="test_model") as span: llmobs.annotate(span=span, metrics={"input_tokens": 10, "output_tokens": 20}) llmobs.annotate(span=span, metrics={"input_tokens": 20, "total_tokens": 40}) - assert span._get_ctx_item(METRICS) == {"input_tokens": 20, "output_tokens": 20, "total_tokens": 40} + assert core.get_item(METRICS) == {"input_tokens": 20, "output_tokens": 20, "total_tokens": 40} def test_annotate_metrics_wrong_type(llmobs, mock_llmobs_logs): with llmobs.llm(model_name="test_model") as llm_span: llmobs.annotate(span=llm_span, metrics=12345) - assert llm_span._get_ctx_item(METRICS) is None + assert core.get_item(METRICS) is None mock_llmobs_logs.warning.assert_called_once_with( "metrics must be a dictionary of string key - numeric value pairs." ) @@ -671,7 +672,7 @@ def test_annotate_prompt_dict(llmobs): "id": "test_prompt", }, ) - assert span._get_ctx_item(INPUT_PROMPT) == { + assert core.get_item(INPUT_PROMPT) == { "template": "{var1} {var3}", "variables": {"var1": "var1", "var2": "var3"}, "version": "1.0.0", @@ -694,7 +695,7 @@ def test_annotate_prompt_dict_with_context_var_keys(llmobs): "rag_query_variables": ["user_input"], }, ) - assert span._get_ctx_item(INPUT_PROMPT) == { + assert core.get_item(INPUT_PROMPT) == { "template": "{var1} {var3}", "variables": {"var1": "var1", "var2": "var3"}, "version": "1.0.0", @@ -717,7 +718,7 @@ def test_annotate_prompt_typed_dict(llmobs): rag_query_variables=["user_input"], ), ) - assert span._get_ctx_item(INPUT_PROMPT) == { + assert core.get_item(INPUT_PROMPT) == { "template": "{var1} {var3}", "variables": {"var1": "var1", "var2": "var3"}, "version": "1.0.0", @@ -730,7 +731,7 @@ def test_annotate_prompt_typed_dict(llmobs): def test_annotate_prompt_wrong_type(llmobs, mock_llmobs_logs): with llmobs.llm(model_name="test_model") as span: llmobs.annotate(span=span, prompt="prompt") - assert span._get_ctx_item(INPUT_PROMPT) is None + assert core.get_item(INPUT_PROMPT) is None mock_llmobs_logs.warning.assert_called_once_with("Failed to validate prompt with error: ", exc_info=True) mock_llmobs_logs.reset_mock() @@ -1536,8 +1537,8 @@ def test_llmobs_with_evaluator_runner(llmobs, mock_llmobs_evaluator_runner): def test_llmobs_with_evaluator_runner_does_not_enqueue_evaluation_spans(mock_llmobs_evaluator_runner, llmobs): - with llmobs.agent(name="test") as agent: - agent._set_ctx_item(IS_EVALUATION_SPAN, True) + with llmobs.agent(name="test"): + core.set_item(IS_EVALUATION_SPAN, True) with llmobs.llm(model_name="test_model"): pass time.sleep(0.1) @@ -1563,14 +1564,14 @@ def test_llmobs_with_evaluation_runner_does_not_enqueue_non_llm_spans(mock_llmob def test_annotation_context_modifies_span_tags(llmobs): with llmobs.annotation_context(tags={"foo": "bar"}): - with llmobs.agent(name="test_agent") as span: - assert span._get_ctx_item(TAGS) == {"foo": "bar"} + with llmobs.agent(name="test_agent"): + assert core.get_item(TAGS) == {"foo": "bar"} def test_annotation_context_modifies_prompt(llmobs): with llmobs.annotation_context(prompt={"template": "test_template"}): - with llmobs.llm(name="test_agent", model_name="test") as span: - assert span._get_ctx_item(INPUT_PROMPT) == { + with llmobs.llm(name="test_agent", model_name="test"): + assert core.get_item(INPUT_PROMPT) == { "template": "test_template", "_dd_context_variable_keys": ["context"], "_dd_query_variable_keys": ["question"], @@ -1586,15 +1587,15 @@ def test_annotation_context_modifies_name(llmobs): def test_annotation_context_finished_context_does_not_modify_tags(llmobs): with llmobs.annotation_context(tags={"foo": "bar"}): pass - with llmobs.agent(name="test_agent") as span: - assert span._get_ctx_item(TAGS) is None + with llmobs.agent(name="test_agent"): + assert core.get_item(TAGS) is None def test_annotation_context_finished_context_does_not_modify_prompt(llmobs): with llmobs.annotation_context(prompt={"template": "test_template"}): pass - with llmobs.llm(name="test_agent", model_name="test") as span: - assert span._get_ctx_item(INPUT_PROMPT) is None + with llmobs.llm(name="test_agent", model_name="test"): + assert core.get_item(INPUT_PROMPT) is None def test_annotation_context_finished_context_does_not_modify_name(llmobs): @@ -1607,8 +1608,8 @@ def test_annotation_context_finished_context_does_not_modify_name(llmobs): def test_annotation_context_nested(llmobs): with llmobs.annotation_context(tags={"foo": "bar", "boo": "bar"}): with llmobs.annotation_context(tags={"foo": "baz"}): - with llmobs.agent(name="test_agent") as span: - assert span._get_ctx_item(TAGS) == {"foo": "baz", "boo": "bar"} + with llmobs.agent(name="test_agent"): + assert core.get_item(TAGS) == {"foo": "baz", "boo": "bar"} def test_annotation_context_nested_overrides_name(llmobs): @@ -1624,8 +1625,8 @@ def test_annotation_context_nested_maintains_trace_structure(llmobs, llmobs_even with llmobs.agent(name="parent_span") as parent_span: with llmobs.annotation_context(tags={"foo": "baz"}): with llmobs.workflow(name="child_span") as child_span: - assert child_span._get_ctx_item(TAGS) == {"foo": "baz", "boo": "bar"} - assert parent_span._get_ctx_item(TAGS) == {"foo": "bar", "boo": "bar"} + assert core.get_item(TAGS) == {"foo": "baz", "boo": "bar"} + assert core.get_item(TAGS) == {"foo": "bar", "boo": "bar"} assert len(llmobs_events) == 2 parent_span, child_span = llmobs_events[1], llmobs_events[0] @@ -1669,7 +1670,7 @@ def context_one(): with llmobs.annotation_context(name="expected_agent", tags={"foo": "bar"}): with llmobs.agent(name="test_agent") as span: event.wait() - agent_has_correct_tags = span._get_ctx_item(TAGS) == {"foo": "bar"} + agent_has_correct_tags = core.get_item(TAGS) == {"foo": "bar"} agent_has_correct_name = span.name == "expected_agent" # thread which registers an annotation context for 0.5 seconds @@ -1680,7 +1681,7 @@ def context_two(): with llmobs.annotation_context(name="expected_tool"): with llmobs.tool(name="test_tool") as tool_span: event.wait() - tool_does_not_have_tags = tool_span._get_ctx_item(TAGS) is None + tool_does_not_have_tags = core.get_item(TAGS) is None tool_has_correct_name = tool_span.name == "expected_tool" thread_one = threading.Thread(target=context_one) @@ -1690,7 +1691,7 @@ def context_two(): with llmobs.agent(name="test_agent") as span: assert span.name == "test_agent" - assert span._get_ctx_item(TAGS) is None + assert core.get_item(TAGS) is None event.set() thread_one.join() @@ -1706,14 +1707,14 @@ def context_two(): async def test_annotation_context_async_modifies_span_tags(llmobs): async with llmobs.annotation_context(tags={"foo": "bar"}): - with llmobs.agent(name="test_agent") as span: - assert span._get_ctx_item(TAGS) == {"foo": "bar"} + with llmobs.agent(name="test_agent"): + assert core.get_item(TAGS) == {"foo": "bar"} async def test_annotation_context_async_modifies_prompt(llmobs): async with llmobs.annotation_context(prompt={"template": "test_template"}): - with llmobs.llm(name="test_agent", model_name="test") as span: - assert span._get_ctx_item(INPUT_PROMPT) == { + with llmobs.llm(name="test_agent", model_name="test"): + assert core.get_item(INPUT_PROMPT) == { "template": "test_template", "_dd_context_variable_keys": ["context"], "_dd_query_variable_keys": ["question"], @@ -1729,15 +1730,15 @@ async def test_annotation_context_async_modifies_name(llmobs): async def test_annotation_context_async_finished_context_does_not_modify_tags(llmobs): async with llmobs.annotation_context(tags={"foo": "bar"}): pass - with llmobs.agent(name="test_agent") as span: - assert span._get_ctx_item(TAGS) is None + with llmobs.agent(name="test_agent"): + assert core.get_item(TAGS) is None async def test_annotation_context_async_finished_context_does_not_modify_prompt(llmobs): async with llmobs.annotation_context(prompt={"template": "test_template"}): pass - with llmobs.llm(name="test_agent", model_name="test") as span: - assert span._get_ctx_item(INPUT_PROMPT) is None + with llmobs.llm(name="test_agent", model_name="test"): + assert core.get_item(INPUT_PROMPT) is None async def test_annotation_context_finished_context_async_does_not_modify_name(llmobs): @@ -1750,8 +1751,8 @@ async def test_annotation_context_finished_context_async_does_not_modify_name(ll async def test_annotation_context_async_nested(llmobs): async with llmobs.annotation_context(tags={"foo": "bar", "boo": "bar"}): async with llmobs.annotation_context(tags={"foo": "baz"}): - with llmobs.agent(name="test_agent") as span: - assert span._get_ctx_item(TAGS) == {"foo": "baz", "boo": "bar"} + with llmobs.agent(name="test_agent"): + assert core.get_item(TAGS) == {"foo": "baz", "boo": "bar"} def test_service_enable_starts_evaluator_runner_when_evaluators_exist():