diff --git a/src/examples/cohere_example/__init__.py b/src/examples/cohere_example/__init__.py index 5610cf4b..f36a7e7d 100644 --- a/src/examples/cohere_example/__init__.py +++ b/src/examples/cohere_example/__init__.py @@ -1,8 +1,11 @@ from examples.cohere_example.chat import chat_comp +from examples.cohere_example.chatv2 import chat_v2 +from examples.cohere_example.chat_streamv2 import chat_stream_v2 from examples.cohere_example.chat_stream import chat_stream from examples.cohere_example.tools import tool_calling from examples.cohere_example.embed import embed from examples.cohere_example.rerank import rerank +from examples.cohere_example.rerankv2 import rerank_v2 from langtrace_python_sdk import with_langtrace_root_span @@ -10,8 +13,11 @@ class CohereRunner: @with_langtrace_root_span("Cohere") def run(self): + chat_v2() + chat_stream_v2() chat_comp() chat_stream() tool_calling() embed() rerank() + rerank_v2() diff --git a/src/examples/cohere_example/chat_streamv2.py b/src/examples/cohere_example/chat_streamv2.py new file mode 100644 index 00000000..2bce996b --- /dev/null +++ b/src/examples/cohere_example/chat_streamv2.py @@ -0,0 +1,17 @@ +import os +from langtrace_python_sdk import langtrace +import cohere + +langtrace.init(api_key=os.getenv("LANGTRACE_API_KEY")) +co = cohere.ClientV2(api_key=os.getenv("COHERE_API_KEY")) + +def chat_stream_v2(): + res = co.chat_stream( + model="command-r-plus-08-2024", + messages=[{"role": "user", "content": "Write a title for a blog post about API design. Only output the title text"}], + ) + + for event in res: + if event: + if event.type == "content-delta": + print(event.delta.message.content.text) \ No newline at end of file diff --git a/src/examples/cohere_example/chatv2.py b/src/examples/cohere_example/chatv2.py new file mode 100644 index 00000000..26f59745 --- /dev/null +++ b/src/examples/cohere_example/chatv2.py @@ -0,0 +1,21 @@ +import os +from langtrace_python_sdk import langtrace +import cohere + +langtrace.init(api_key=os.getenv("LANGTRACE_API_KEY")) + + +def chat_v2(): + co = cohere.ClientV2(api_key=os.getenv("COHERE_API_KEY")) + + res = co.chat( + model="command-r-plus-08-2024", + messages=[ + { + "role": "user", + "content": "Write a title for a blog post about API design. Only output the title text.", + } + ], + ) + + print(res.message.content[0].text) diff --git a/src/examples/cohere_example/rerankv2.py b/src/examples/cohere_example/rerankv2.py new file mode 100644 index 00000000..55e262cc --- /dev/null +++ b/src/examples/cohere_example/rerankv2.py @@ -0,0 +1,23 @@ +import os +from langtrace_python_sdk import langtrace +import cohere + +langtrace.init(api_key=os.getenv("LANGTRACE_API_KEY")) +co = cohere.ClientV2(api_key=os.getenv("COHERE_API_KEY")) + +docs = [ + "Carson City is the capital city of the American state of Nevada.", + "The Commonwealth of the Northern Mariana Islands is a group of islands in the Pacific Ocean. Its capital is Saipan.", + "Capitalization or capitalisation in English grammar is the use of a capital letter at the start of a word. English usage varies from capitalization in other languages.", + "Washington, D.C. (also known as simply Washington or D.C., and officially as the District of Columbia) is the capital of the United States. It is a federal district.", + "Capital punishment has existed in the United States since beforethe United States was a country. As of 2017, capital punishment is legal in 30 of the 50 states.", +] + +def rerank_v2(): + response = co.rerank( + model="rerank-v3.5", + query="What is the capital of the United States?", + documents=docs, + top_n=3, + ) + print(response) diff --git a/src/langtrace_python_sdk/constants/exporter/langtrace_exporter.py b/src/langtrace_python_sdk/constants/exporter/langtrace_exporter.py index 0261ab6f..3bc4c58c 100644 --- a/src/langtrace_python_sdk/constants/exporter/langtrace_exporter.py +++ b/src/langtrace_python_sdk/constants/exporter/langtrace_exporter.py @@ -1 +1,2 @@ -LANGTRACE_REMOTE_URL = "https://app.langtrace.ai" +LANGTRACE_REMOTE_URL = "https://app.langtrace.ai" +LANGTRACE_SESSION_ID_HEADER = "x-langtrace-session-id" diff --git a/src/langtrace_python_sdk/constants/instrumentation/cohere.py b/src/langtrace_python_sdk/constants/instrumentation/cohere.py index cc38254b..650fa75c 100644 --- a/src/langtrace_python_sdk/constants/instrumentation/cohere.py +++ b/src/langtrace_python_sdk/constants/instrumentation/cohere.py @@ -4,19 +4,39 @@ "METHOD": "cohere.client.chat", "ENDPOINT": "/v1/chat", }, + "CHAT_CREATE_V2": { + "URL": "https://api.cohere.ai", + "METHOD": "cohere.client_v2.chat", + "ENDPOINT": "/v2/chat", + }, "EMBED": { "URL": "https://api.cohere.ai", "METHOD": "cohere.client.embed", "ENDPOINT": "/v1/embed", }, + "EMBED_V2": { + "URL": "https://api.cohere.ai", + "METHOD": "cohere.client_v2.embed", + "ENDPOINT": "/v2/embed", + }, "CHAT_STREAM": { "URL": "https://api.cohere.ai", "METHOD": "cohere.client.chat_stream", - "ENDPOINT": "/v1/messages", + "ENDPOINT": "/v1/chat", + }, + "CHAT_STREAM_V2": { + "URL": "https://api.cohere.ai", + "METHOD": "cohere.client_v2.chat_stream", + "ENDPOINT": "/v2/chat", }, "RERANK": { "URL": "https://api.cohere.ai", "METHOD": "cohere.client.rerank", "ENDPOINT": "/v1/rerank", }, + "RERANK_V2": { + "URL": "https://api.cohere.ai", + "METHOD": "cohere.client_v2.rerank", + "ENDPOINT": "/v2/rerank", + }, } diff --git a/src/langtrace_python_sdk/extensions/langtrace_exporter.py b/src/langtrace_python_sdk/extensions/langtrace_exporter.py index bc26802a..664013db 100644 --- a/src/langtrace_python_sdk/extensions/langtrace_exporter.py +++ b/src/langtrace_python_sdk/extensions/langtrace_exporter.py @@ -9,6 +9,7 @@ from langtrace_python_sdk.constants.exporter.langtrace_exporter import ( LANGTRACE_REMOTE_URL, + LANGTRACE_SESSION_ID_HEADER, ) from colorama import Fore from requests.exceptions import RequestException @@ -51,12 +52,14 @@ class LangTraceExporter(SpanExporter): api_key: str api_host: str disable_logging: bool + session_id: str def __init__( self, api_host, api_key: str = None, disable_logging: bool = False, + session_id: str = None, ) -> None: self.api_key = api_key or os.environ.get("LANGTRACE_API_KEY") self.api_host = ( @@ -65,6 +68,7 @@ def __init__( else api_host ) self.disable_logging = disable_logging + self.session_id = session_id or os.environ.get("LANGTRACE_SESSION_ID") def export(self, spans: typing.Sequence[ReadableSpan]) -> SpanExportResult: """ @@ -82,6 +86,10 @@ def export(self, spans: typing.Sequence[ReadableSpan]) -> SpanExportResult: "User-Agent": "LangtraceExporter", } + # Add session ID if available + if self.session_id: + headers[LANGTRACE_SESSION_ID_HEADER] = self.session_id + # Check if the OTEL_EXPORTER_OTLP_HEADERS environment variable is set otel_headers = os.getenv("OTEL_EXPORTER_OTLP_HEADERS", None) if otel_headers: diff --git a/src/langtrace_python_sdk/instrumentation/cohere/instrumentation.py b/src/langtrace_python_sdk/instrumentation/cohere/instrumentation.py index df433c13..414672ce 100644 --- a/src/langtrace_python_sdk/instrumentation/cohere/instrumentation.py +++ b/src/langtrace_python_sdk/instrumentation/cohere/instrumentation.py @@ -23,6 +23,7 @@ from langtrace_python_sdk.instrumentation.cohere.patch import ( chat_create, + chat_create_v2, chat_stream, embed, rerank, @@ -48,6 +49,18 @@ def _instrument(self, **kwargs): chat_create("cohere.client.chat", version, tracer), ) + wrap_function_wrapper( + "cohere.client_v2", + "ClientV2.chat", + chat_create_v2("cohere.client_v2.chat", version, tracer), + ) + + wrap_function_wrapper( + "cohere.client_v2", + "ClientV2.chat_stream", + chat_create_v2("cohere.client_v2.chat", version, tracer, stream=True), + ) + wrap_function_wrapper( "cohere.client", "Client.chat_stream", @@ -60,12 +73,24 @@ def _instrument(self, **kwargs): embed("cohere.client.embed", version, tracer), ) + wrap_function_wrapper( + "cohere.client_v2", + "ClientV2.embed", + embed("cohere.client.embed", version, tracer, v2=True), + ) + wrap_function_wrapper( "cohere.client", "Client.rerank", rerank("cohere.client.rerank", version, tracer), ) + wrap_function_wrapper( + "cohere.client_v2", + "ClientV2.rerank", + rerank("cohere.client.rerank", version, tracer, v2=True), + ) + def _instrument_module(self, module_name): pass diff --git a/src/langtrace_python_sdk/instrumentation/cohere/patch.py b/src/langtrace_python_sdk/instrumentation/cohere/patch.py index 38908c3d..8b9a8e53 100644 --- a/src/langtrace_python_sdk/instrumentation/cohere/patch.py +++ b/src/langtrace_python_sdk/instrumentation/cohere/patch.py @@ -24,6 +24,7 @@ get_span_name, set_event_completion, set_usage_attributes, + StreamWrapper ) from langtrace.trace_attributes import Event, LLMSpanAttributes from langtrace_python_sdk.utils import set_span_attribute @@ -38,7 +39,7 @@ from langtrace.trace_attributes import SpanAttributes -def rerank(original_method, version, tracer): +def rerank(original_method, version, tracer, v2=False): """Wrap the `rerank` method.""" def traced_method(wrapped, instance, args, kwargs): @@ -49,8 +50,8 @@ def traced_method(wrapped, instance, args, kwargs): **get_llm_request_attributes(kwargs, operation_name="rerank"), **get_llm_url(instance), SpanAttributes.LLM_REQUEST_MODEL: kwargs.get("model") or "command-r-plus", - SpanAttributes.LLM_URL: APIS["RERANK"]["URL"], - SpanAttributes.LLM_PATH: APIS["RERANK"]["ENDPOINT"], + SpanAttributes.LLM_URL: APIS["RERANK" if not v2 else "RERANK_V2"]["URL"], + SpanAttributes.LLM_PATH: APIS["RERANK" if not v2 else "RERANK_V2"]["ENDPOINT"], SpanAttributes.LLM_REQUEST_DOCUMENTS: json.dumps( kwargs.get("documents"), cls=datetime_encoder ), @@ -61,7 +62,7 @@ def traced_method(wrapped, instance, args, kwargs): attributes = LLMSpanAttributes(**span_attributes) span = tracer.start_span( - name=get_span_name(APIS["RERANK"]["METHOD"]), kind=SpanKind.CLIENT + name=get_span_name(APIS["RERANK" if not v2 else "RERANK_V2"]["METHOD"]), kind=SpanKind.CLIENT ) for field, value in attributes.model_dump(by_alias=True).items(): set_span_attribute(span, field, value) @@ -119,7 +120,7 @@ def traced_method(wrapped, instance, args, kwargs): return traced_method -def embed(original_method, version, tracer): +def embed(original_method, version, tracer, v2=False): """Wrap the `embed` method.""" def traced_method(wrapped, instance, args, kwargs): @@ -129,8 +130,8 @@ def traced_method(wrapped, instance, args, kwargs): **get_langtrace_attributes(version, service_provider), **get_llm_request_attributes(kwargs, operation_name="embed"), **get_llm_url(instance), - SpanAttributes.LLM_URL: APIS["EMBED"]["URL"], - SpanAttributes.LLM_PATH: APIS["EMBED"]["ENDPOINT"], + SpanAttributes.LLM_URL: APIS["EMBED" if not v2 else "EMBED_V2"]["URL"], + SpanAttributes.LLM_PATH: APIS["EMBED" if not v2 else "EMBED_V2"]["ENDPOINT"], SpanAttributes.LLM_REQUEST_EMBEDDING_INPUTS: json.dumps( kwargs.get("texts") ), @@ -143,7 +144,7 @@ def traced_method(wrapped, instance, args, kwargs): attributes = LLMSpanAttributes(**span_attributes) span = tracer.start_span( - name=get_span_name(APIS["EMBED"]["METHOD"]), + name=get_span_name(APIS["EMBED" if not v2 else "EMBED_V2"]["METHOD"]), kind=SpanKind.CLIENT, ) for field, value in attributes.model_dump(by_alias=True).items(): @@ -343,6 +344,103 @@ def traced_method(wrapped, instance, args, kwargs): return traced_method +def chat_create_v2(original_method, version, tracer, stream=False): + """Wrap the `chat_create` method for Cohere API v2.""" + + def traced_method(wrapped, instance, args, kwargs): + service_provider = SERVICE_PROVIDERS["COHERE"] + + messages = kwargs.get("messages", []) + if kwargs.get("preamble"): + messages = [{"role": "system", "content": kwargs["preamble"]}] + messages + + span_attributes = { + **get_langtrace_attributes(version, service_provider), + **get_llm_request_attributes(kwargs, prompts=messages), + **get_llm_url(instance), + SpanAttributes.LLM_REQUEST_MODEL: kwargs.get("model") or "command-r-plus", + SpanAttributes.LLM_URL: APIS["CHAT_CREATE_V2"]["URL"], + SpanAttributes.LLM_PATH: APIS["CHAT_CREATE_V2"]["ENDPOINT"], + **get_extra_attributes(), + } + + attributes = LLMSpanAttributes(**span_attributes) + + for attr_name in ["max_input_tokens", "conversation_id", "connectors", "tools", "tool_results"]: + value = kwargs.get(attr_name) + if value is not None: + if attr_name == "max_input_tokens": + attributes.llm_max_input_tokens = str(value) + elif attr_name == "conversation_id": + attributes.conversation_id = value + else: + setattr(attributes, f"llm_{attr_name}", json.dumps(value)) + + span = tracer.start_span( + name=get_span_name(APIS["CHAT_CREATE_V2"]["METHOD"]), + kind=SpanKind.CLIENT + ) + + for field, value in attributes.model_dump(by_alias=True).items(): + set_span_attribute(span, field, value) + + try: + result = wrapped(*args, **kwargs) + + if stream: + return StreamWrapper( + result, + span, + tool_calls=kwargs.get("tools") is not None, + ) + else: + if hasattr(result, "id") and result.id is not None: + span.set_attribute(SpanAttributes.LLM_GENERATION_ID, result.id) + span.set_attribute(SpanAttributes.LLM_RESPONSE_ID, result.id) + + if (hasattr(result, "message") and + hasattr(result.message, "content") and + len(result.message.content) > 0 and + hasattr(result.message.content[0], "text") and + result.message.content[0].text is not None and + result.message.content[0].text != ""): + responses = [{ + "role": result.message.role, + "content": result.message.content[0].text + }] + set_event_completion(span, responses) + if hasattr(result, "tool_calls") and result.tool_calls is not None: + tool_calls = [tool_call.json() for tool_call in result.tool_calls] + span.set_attribute( + SpanAttributes.LLM_TOOL_RESULTS, + json.dumps(tool_calls) + ) + if hasattr(result, "usage") and result.usage is not None: + if (hasattr(result.usage, "billed_units") and + result.usage.billed_units is not None): + usage = result.usage.billed_units + for metric, value in { + "input": usage.input_tokens or 0, + "output": usage.output_tokens or 0, + "total": (usage.input_tokens or 0) + (usage.output_tokens or 0), + }.items(): + span.set_attribute( + f"gen_ai.usage.{metric}_tokens", + int(value) + ) + span.set_status(StatusCode.OK) + span.end() + return result + + except Exception as error: + span.record_exception(error) + span.set_status(Status(StatusCode.ERROR, str(error))) + span.end() + raise + + return traced_method + + def chat_stream(original_method, version, tracer): """Wrap the `messages_stream` method.""" diff --git a/src/langtrace_python_sdk/instrumentation/dspy/patch.py b/src/langtrace_python_sdk/instrumentation/dspy/patch.py index 8a02afcb..65e6ab92 100644 --- a/src/langtrace_python_sdk/instrumentation/dspy/patch.py +++ b/src/langtrace_python_sdk/instrumentation/dspy/patch.py @@ -50,7 +50,7 @@ def traced_method(wrapped, instance, args, kwargs): ), } span_attributes["dspy.optimizer.module.prog"] = json.dumps(prog) - if hasattr(instance, "metric"): + if hasattr(instance, "metric") and getattr(instance, "metric") is not None: span_attributes["dspy.optimizer.metric"] = getattr( instance, "metric" ).__name__ @@ -120,8 +120,6 @@ def traced_method(wrapped, instance, args, kwargs): **get_extra_attributes(), } - trace_checkpoint = os.environ.get("TRACE_DSPY_CHECKPOINT", "true").lower() - if instance.__class__.__name__: span_attributes["dspy.signature.name"] = instance.__class__.__name__ span_attributes["dspy.signature"] = str(instance.signature) @@ -143,9 +141,6 @@ def traced_method(wrapped, instance, args, kwargs): "dspy.signature.result", json.dumps(result.toDict()), ) - if trace_checkpoint == "true": - print(Fore.RED + "Note: DSPy checkpoint tracing is enabled in Langtrace. To disable it, set the env var, TRACE_DSPY_CHECKPOINT to false" + Fore.RESET) - set_span_attribute(span, "dspy.checkpoint", ujson.dumps(instance.dump_state(False), indent=2)) span.set_status(Status(StatusCode.OK)) span.end() diff --git a/src/langtrace_python_sdk/instrumentation/openai/patch.py b/src/langtrace_python_sdk/instrumentation/openai/patch.py index daefeef6..255a7eb7 100644 --- a/src/langtrace_python_sdk/instrumentation/openai/patch.py +++ b/src/langtrace_python_sdk/instrumentation/openai/patch.py @@ -342,8 +342,12 @@ async def traced_method( service_provider = SERVICE_PROVIDERS["PPLX"] elif "azure" in get_base_url(instance): service_provider = SERVICE_PROVIDERS["AZURE"] + elif "groq" in get_base_url(instance): + service_provider = SERVICE_PROVIDERS["GROQ"] elif "x.ai" in get_base_url(instance): service_provider = SERVICE_PROVIDERS["XAI"] + elif "deepseek" in get_base_url(instance): + service_provider = SERVICE_PROVIDERS["DEEPSEEK"] llm_prompts = [] for item in kwargs.get("messages", []): tools = get_tool_calls(item) @@ -431,6 +435,18 @@ def traced_method( kwargs: EmbeddingsCreateKwargs, ) -> Any: service_provider = SERVICE_PROVIDERS["OPENAI"] + base_url = get_base_url(instance) + + if "perplexity" in base_url: + service_provider = SERVICE_PROVIDERS["PPLX"] + elif "azure" in base_url: + service_provider = SERVICE_PROVIDERS["AZURE"] + elif "groq" in base_url: + service_provider = SERVICE_PROVIDERS["GROQ"] + elif "x.ai" in base_url: + service_provider = SERVICE_PROVIDERS["XAI"] + elif "deepseek" in base_url: + service_provider = SERVICE_PROVIDERS["DEEPSEEK"] span_attributes = { **get_langtrace_attributes(version, service_provider, vendor_type="llm"), @@ -469,7 +485,6 @@ def traced_method( kind=SpanKind.CLIENT, context=set_span_in_context(trace.get_current_span()), ) as span: - set_span_attributes(span, attributes) try: # Attempt to call the original method @@ -507,17 +522,27 @@ async def traced_method( ) -> Awaitable[Any]: service_provider = SERVICE_PROVIDERS["OPENAI"] + base_url = get_base_url(instance) + if "perplexity" in base_url: + service_provider = SERVICE_PROVIDERS["PPLX"] + elif "azure" in base_url: + service_provider = SERVICE_PROVIDERS["AZURE"] + elif "groq" in base_url: + service_provider = SERVICE_PROVIDERS["GROQ"] + elif "x.ai" in base_url: + service_provider = SERVICE_PROVIDERS["XAI"] + elif "deepseek" in base_url: + service_provider = SERVICE_PROVIDERS["DEEPSEEK"] span_attributes = { **get_langtrace_attributes(version, service_provider, vendor_type="llm"), **get_llm_request_attributes(kwargs, operation_name="embed"), + **get_llm_url(instance), SpanAttributes.LLM_PATH: APIS["EMBEDDINGS_CREATE"]["ENDPOINT"], SpanAttributes.LLM_REQUEST_DIMENSIONS: kwargs.get("dimensions"), **get_extra_attributes(), # type: ignore } - attributes = LLMSpanAttributes(**filter_valid_attributes(span_attributes)) - encoding_format = kwargs.get("encoding_format") if encoding_format is not None: if not isinstance(encoding_format, list): @@ -530,17 +555,31 @@ async def traced_method( span_attributes[SpanAttributes.LLM_REQUEST_EMBEDDING_INPUTS] = json.dumps( [kwargs.get("input", "")] ) + span_attributes[SpanAttributes.LLM_PROMPTS] = json.dumps( + [ + { + "role": "user", + "content": kwargs.get("input"), + } + ] + ) + + attributes = LLMSpanAttributes(**filter_valid_attributes(span_attributes)) with tracer.start_as_current_span( name=get_span_name(APIS["EMBEDDINGS_CREATE"]["METHOD"]), kind=SpanKind.CLIENT, context=set_span_in_context(trace.get_current_span()), ) as span: - set_span_attributes(span, attributes) try: # Attempt to call the original method result = await wrapped(*args, **kwargs) + usage = getattr(result, "usage", None) + if usage: + set_usage_attributes( + span, {"prompt_tokens": getattr(usage, "prompt_tokens", 0)} + ) span.set_status(StatusCode.OK) return result except Exception as err: diff --git a/src/langtrace_python_sdk/langtrace.py b/src/langtrace_python_sdk/langtrace.py index 1b2c59c8..33e1e998 100644 --- a/src/langtrace_python_sdk/langtrace.py +++ b/src/langtrace_python_sdk/langtrace.py @@ -39,6 +39,7 @@ ) from langtrace_python_sdk.constants.exporter.langtrace_exporter import ( LANGTRACE_REMOTE_URL, + LANGTRACE_SESSION_ID_HEADER, ) from langtrace_python_sdk.instrumentation import ( AnthropicInstrumentation, @@ -98,6 +99,9 @@ def __init__(self, **kwargs): or os.environ.get("LANGTRACE_HEADERS") or os.environ.get("OTEL_EXPORTER_OTLP_HEADERS") ) + self.session_id = kwargs.get("session_id") or os.environ.get( + "LANGTRACE_SESSION_ID" + ) def get_host(config: LangtraceConfig) -> str: @@ -134,15 +138,19 @@ def setup_tracer_provider(config: LangtraceConfig, host: str) -> TracerProvider: def get_headers(config: LangtraceConfig): - if not config.headers: - return { - "x-api-key": config.api_key, - } + headers = { + "x-api-key": config.api_key, + } + + if config.session_id: + headers[LANGTRACE_SESSION_ID_HEADER] = config.session_id if isinstance(config.headers, str): - return parse_env_headers(config.headers, liberal=True) + headers.update(parse_env_headers(config.headers, liberal=True)) + elif config.headers: + headers.update(config.headers) - return config.headers + return headers def get_exporter(config: LangtraceConfig, host: str): @@ -150,9 +158,16 @@ def get_exporter(config: LangtraceConfig, host: str): return config.custom_remote_exporter headers = get_headers(config) - host = f"{host}/api/trace" if host == LANGTRACE_REMOTE_URL else host - if "http" in host.lower() or "https" in host.lower(): - return HTTPExporter(endpoint=host, headers=headers) + exporter_protocol = os.environ.get("OTEL_EXPORTER_OTLP_PROTOCOL", "http") + if "http" in exporter_protocol.lower(): + return HTTPExporter( + endpoint=( + f"{host}/api/trace" + if host == LANGTRACE_REMOTE_URL + else f"{host}/v1/traces" + ), + headers=headers, + ) else: return GRPCExporter(endpoint=host, headers=headers) @@ -215,6 +230,7 @@ def init( service_name: Optional[str] = None, disable_logging: bool = False, headers: Dict[str, str] = {}, + session_id: Optional[str] = None, ): check_if_sdk_is_outdated() @@ -229,6 +245,7 @@ def init( service_name=service_name, disable_logging=disable_logging, headers=headers, + session_id=session_id, ) if config.disable_logging: diff --git a/src/langtrace_python_sdk/utils/llm.py b/src/langtrace_python_sdk/utils/llm.py index b9355124..86cdac3f 100644 --- a/src/langtrace_python_sdk/utils/llm.py +++ b/src/langtrace_python_sdk/utils/llm.py @@ -162,11 +162,12 @@ def get_llm_url(instance): def get_base_url(instance): - return ( + base_url = ( str(instance._client._base_url) if hasattr(instance, "_client") and hasattr(instance._client, "_base_url") else "" ) + return base_url def is_streaming(kwargs): @@ -393,8 +394,19 @@ def build_streaming_response(self, chunk): if hasattr(chunk, "text") and chunk.text is not None: content = [chunk.text] + # CohereV2 + if (hasattr(chunk, "delta") and + chunk.delta is not None and + hasattr(chunk.delta, "message") and + chunk.delta.message is not None and + hasattr(chunk.delta.message, "content") and + chunk.delta.message.content is not None and + hasattr(chunk.delta.message.content, "text") and + chunk.delta.message.content.text is not None): + content = [chunk.delta.message.content.text] + # Anthropic - if hasattr(chunk, "delta") and chunk.delta is not None: + if hasattr(chunk, "delta") and chunk.delta is not None and not hasattr(chunk.delta, "message"): content = [chunk.delta.text] if hasattr(chunk.delta, "text") else [] if isinstance(chunk, dict): @@ -408,7 +420,17 @@ def set_usage_attributes(self, chunk): # Anthropic & OpenAI if hasattr(chunk, "type") and chunk.type == "message_start": - self.prompt_tokens = chunk.message.usage.input_tokens + if hasattr(chunk.message, "usage") and chunk.message.usage is not None: + self.prompt_tokens = chunk.message.usage.input_tokens + + # CohereV2 + if hasattr(chunk, "type") and chunk.type == "message-end": + if (hasattr(chunk, "delta") and chunk.delta is not None and + hasattr(chunk.delta, "usage") and chunk.delta.usage is not None and + hasattr(chunk.delta.usage, "billed_units") and chunk.delta.usage.billed_units is not None): + usage = chunk.delta.usage.billed_units + self.completion_tokens = int(usage.output_tokens) + self.prompt_tokens = int(usage.input_tokens) if hasattr(chunk, "usage") and chunk.usage is not None: if hasattr(chunk.usage, "output_tokens"): diff --git a/src/langtrace_python_sdk/utils/with_root_span.py b/src/langtrace_python_sdk/utils/with_root_span.py index 67e6f66d..f1a5296f 100644 --- a/src/langtrace_python_sdk/utils/with_root_span.py +++ b/src/langtrace_python_sdk/utils/with_root_span.py @@ -61,6 +61,11 @@ def sync_wrapper(*args, **kwargs): span_id = str(span.get_span_context().span_id) trace_id = str(span.get_span_context().trace_id) + # Attach session ID if available + session_id = os.environ.get("LANGTRACE_SESSION_ID") + if session_id: + span.set_attribute("session.id", session_id) + if ( "span_id" in func.__code__.co_varnames and "trace_id" in func.__code__.co_varnames @@ -82,6 +87,12 @@ async def async_wrapper(*args, **kwargs): ) as span: span_id = span.get_span_context().span_id trace_id = span.get_span_context().trace_id + + # Attach session ID if available + session_id = os.environ.get("LANGTRACE_SESSION_ID") + if session_id: + span.set_attribute("session.id", session_id) + if ( "span_id" in func.__code__.co_varnames and "trace_id" in func.__code__.co_varnames diff --git a/src/langtrace_python_sdk/version.py b/src/langtrace_python_sdk/version.py index ef7fb174..4ad21fb4 100644 --- a/src/langtrace_python_sdk/version.py +++ b/src/langtrace_python_sdk/version.py @@ -1 +1 @@ -__version__ = "3.3.14" +__version__ = "3.3.16" diff --git a/src/run_example.py b/src/run_example.py index 9f4cecf2..ab925665 100644 --- a/src/run_example.py +++ b/src/run_example.py @@ -4,7 +4,7 @@ "anthropic": False, "azureopenai": False, "chroma": False, - "cohere": False, + "cohere": True, "fastapi": False, "langchain": False, "llamaindex": False, @@ -21,7 +21,7 @@ "gemini": False, "mistral": False, "awsbedrock": False, - "cerebras": True, + "cerebras": False, } if ENABLED_EXAMPLES["anthropic"]: diff --git a/src/tests/openai/cassettes/test_embeddings_base_url.yaml b/src/tests/openai/cassettes/test_embeddings_base_url.yaml new file mode 100644 index 00000000..56e08927 --- /dev/null +++ b/src/tests/openai/cassettes/test_embeddings_base_url.yaml @@ -0,0 +1,208 @@ +interactions: +- request: + body: '{"input": "Test input", "model": "text-embedding-ada-002", "encoding_format": + "base64"}' + headers: + accept: + - application/json + accept-encoding: + - gzip, deflate + connection: + - keep-alive + content-length: + - '87' + content-type: + - application/json + host: + - api.openai.com + user-agent: + - OpenAI/Python 1.54.4 + x-stainless-arch: + - x64 + x-stainless-async: + - 'false' + x-stainless-lang: + - python + x-stainless-os: + - Linux + x-stainless-package-version: + - 1.54.4 + x-stainless-retry-count: + - '0' + x-stainless-runtime: + - CPython + x-stainless-runtime-version: + - 3.12.5 + method: POST + uri: https://api.openai.com/v1/embeddings + response: + body: + string: !!binary | + H4sIAAAAAAAAA1R6Ww+6Orvn/fspVtYt8wYEpGXdcZaTLYIiTiYTQEUQRA4th5393Sf439kzc6MJ + VrHl6e/09D/+9ddff7dZ9cjHv//56++6HMa//8d27Z6O6d///PU///XXX3/99R+/1/9v5KPJHvd7 + +Sl+w38flp/7Y/77n7+4/77yfwf989fflTcNiI3kjzfmWOQk+XjXsco4dT9Zp1crD3LZULsPODBd + X5kEj2onUQtkz2qO/HqC56TOqcaUSbYwQXkGRThfsHt6FtVKrwwBIbreMdb6TzYx82zB6MLxFEU7 + 4BGeYX3gPv0zdQ7euV/aUipgGp09eimNGNA+PknwGWgU207Y9+t5HErYe09EvUsWZdNL93hplR4l + tUWWAxO3YxSY+ZZL/WjOo9FK+hzeMn5Bsx9TsL6Rn8I44Tk0OUTSaXzlShicFJEeD59sneVamODt + +r1h/Bo+3jLuWw3eUmWh+ulurNv8YsBdlJIe0Zit811NeJDTxKHJAnt9TB8BL5sZw2BdehF9iMye + AcWnPGDjPRveclp2LmCUY4dm2cLRTHhWBG/hFBJQqfm6dmf/DLPGmKnLDIU3c5+pgUQ4+mjnDFd9 + Zgv5DAvtCOmhEzRvqq/9BL1xfmFFArdqag0nBlXFRtjro282nDk+ho9ilLF5YqqKwsa3QRaNCfaa + A1zbafE6wIyORqZEDyJenK4ItjftSw9WtwPLV72msFt5lR74K+uR19tt4NgEDT3yiu2RVzpqYBjY + I6pfn3jtRLdL4WM+xWivkldE4NON4Uc6BVidhZs+67sql569cKGu/HXWtaxgKrJLpGFdeBvRYk4p + B99Pq6f2bs+ty+n7fAOnyS80jF9mttsVbQOPPC6oYyBP51+caEN1FD7YDMAY0YdtPOCzaGx8KJCW + 8Wod5LKraheq9mjvkf0yTGD7Plbv2j7rcVm6sPW7LxIQ1CLhe2Im0GjQoOoqm9Xa6Y8Fvk3DxGiK + T9k0p5YBgXozsUVBWfWsGBfwHp8drOvkDQZgZm8A6/BO1vc30Ze8CHzgz4cXqYKU9KvrcS0gee9j + vxzadVqNhwWukPhI0F7PfmaB/QB9gFTs+ns345ogmmCM+CNiquOnWmonsWFmvhJShoyZzYEUxaC/ + Cohae2OI5usMAqjqUYvVG7xlQ87eUsj35gGj0+MAvo5/U+DjE1LqFUen4vPj3ZLc5Fti5yo8vFm6 + aww7xdeFWteqXAlE+QLl207Hp7u2j1ZXexowPipHahke30+ex+dAZ9oEH8Zrv86dZROoJlpKpN67 + Z6Towwby/pUgfscM2VKBawpkzOX47N/e0bcrrxDsnI9NuN19rObDKXuAu8M4RCw+O4+S/YmXS8UE + aDexejTuMNAge6UaEm52mU0fdZ3AbrnYOJoeRj+TJLLgOJ2OGK37kzcQARlQoUFK1fpcg4H/xAq8 + PXwPG7JTgd98gGF3LLbwKc6G+zSncIbvijTc5QmWNHyWcA8jk3BLoYPFbHYBZET2SyYkNdXsViMP + XhUF2J2KIprDdxXDvEcTdrD08kbHdDTY4+5DfZ8L18ks1A6e+wLh0PBxNk2vhwXPX3uHj3vL1pfJ + eUswpfZMFfFtRbueEVtYVo2FfYeEYI6aSZTgyzSwWXajV1yaRyw1QtlRYz4L2Vg/xjc0Yp6lqhVZ + 3iQUSAHax1Eokp/pOgE027IpxgzG8B30Q2a9cjgX/ht5zSFfZyV3DHDK3QDrFmesvMEs/K/+qcnU + /rqqQxrCLrjw+LjcbX2O93oAO+uKkdjuX9Xs8S2E/eW+o8fxlPWTInEWJCCtSc+aC1hqajBQ7pKc + ZjN0omXN2HA/vOwPka63FUw5e0rkVxeM2KZ8r89FuiL4WM0Ym0w9gG9xunSQp2FBj+Et1L80Sjop + pe5MWA19InozQkPeZcyL6muG9IUGnARn6ehibM6fbKpfzhkykZ1TNbNrsCiVJkHsXZ5Ui191Nkrv + EMLCul0o+vKffmboPQZqgVWsXR5Yn2Tl5MrLeXUJg5ID4BTNcEFXtyrVuoiASdKRAiztWlNjlsZq + LEHdwBd/OBMj3I/ruhh5ILXCpFGHAAWsP7yJGD1Hq2aK0cZPCWCCWafXDb/Wo5i6MNWWBb2iPa7W + azCUQLOzDwF+pni7IT4WoHnGOoGtA7PFr2sEhXdS4ENu+t6MJ3mCi3GcscnYwzox1xH+8AUR3Zn7 + jf8K+B6DHq33StQ/GSk12DcqIvLGb72YOROkKl6xeUyO0XTmmDMwMiXBDvdyKz44J2dQecuA5Ji1 + s/GqtRLgqG7/8AOM32NC4JcZVKwP4F1NTBsV8LHLKT2Zh32/PpJFgXbkxkgKV7maczzx8LN0OpIL + e/FIt7BnwGgfnRDueajWouU5mDGnK/X3We6R2ZNyKXyIPHamhIsmZt4bMOVklkjnuujJfXfNIfHN + FSviu8kGr3He8HTIHlj/Ju9qCdxlgUdDifC1mkudLsY5FLQ+vuHjy7MjYb/EDYRWZWBbKNVoyS3T + gBt/bes7eEvKKz7YObX9469sKs1bCmytkdEOtJU+HXkrBSHNMRH88VkNc8gbcM2og40TL1SrEXcB + PLVFQLb7e+MyAAUKpych4jFT1/ExVQju8WUgfF7x2dr7XAydJYD0hxfzT99QxvMwWiNpHV+a3ILv + 58OjeUjNdXUiowOxdc8oLr4TWPYv1ofo5obYkruymrNjFcJnmdrUGvFzHZVcNWBSP3zCFnraE//M + dtKeHEzCGQHxxlat3/A2pB4+etOiD/wEGtDzzYf0BnXAnMFdC69SvG7zZfuxnVgG2qMLsb2yr2xy + 0vgBLVu6Um9cV50a8GkBZX9SsSk8/XUSEl4D401q8EFsSn2Z8laB+zJ80uMFhvoyTloOVf3UYkN5 + C2CS96olf4qwo8Y4C1H3rW0DHrokxsFRmLKmgIoLpVo1qWWiAXwRe19AdhRd+mwvH33B9ugD2fvq + 2H99w2rDkxaAwH5Q5+6/9Cn0AQNkHuskNrnKW5K5NsCUS28kfM6lvlxkoYBQu/bYMVCvd2udc7DZ + eSci92heO/Yj2HDHly/s5I4LZmkUS6kf3gq+dRe3f3dZbEMntS5IOO19fcdBlpMUGqYb3i8racul + hM1Y6oiVjAaMu6JoQO8cduhVn02w3CrTgNt6EWaKT9FSGbImnfcDou4TKRkx71cCDFavqfJsJp3y + +NGB773G2HpRDXBCwmiAecU5RdlrzmbzCR5A+3gKEfgrq/cP23/86h1FD0X0lrxIEBxL80j16Jp4 + i8+KMTRijqX6Dghg7X0Yw1o/7NF+FN1qNW0t2PNKZWOt6hAgh6VNxanZe1ibbwishgkJbLGmYWUS + A4/eC1cD5dW38EPxW31mWjOGj3tgY3Rc5oqEO5eA/dl6opOgmz1hNS2W27gP0E54+mC5Ot2w3/QL + Wpvw3S/mFHIQiFeD1G44Zgtowhi04VNAAju8wPJo2BxsfEMPUnHJJkXzXchJO4pkIrX6fIKFJE8l + XVDYx7W39PAswg0fiFivr6r+7VfMP0R6jMs2W7+DI4kDwhxV7eO+WmLlaEsXfdhR2xgu0frqjRQW + GoZIMCKtnw77QQGHBslEptdWH8EzKWE3jDtqnhi931Xy6kJxkUeqBympxpAoIrx+55SIftL3K7zE + b0jfdk0VvQ+ytdA1Ddapb9Bzar+9UVvvImyeZ52eteM7Wq6REkBTZRKqeSH1FqVyRek3PvyOYj8r + jsnAqEkxxvZp1qek+uYwIfsbTge5yKZ5VH0wD6qLmOx1yqYIqyVUcqAi8Zi9wBK40gJLcIwwni9S + NCYvc4CpHM8YO+dD9n25xw5ownDDP3+zrLyWQnY5afjQEA90G96Au8DKZPK4pl9N4dxC+Xrv6OGU + +14T7tzhtx/osbAXfQxWPQYkFXwkGsMuW+ire4D2vQuo1semvoh7bpF+8/XvMsyWIrk+YBdSTCrd + Rv1ykxP/p1/pQT/s1pERJVHyJnLGVqU/wWwIsybxNCgQ2GvfjOafUNl3t0Cl7m0vVwThHQLg9ulR + ITB0nWMliqFe5DdsumLVTz4ytD/8l3H56m164C2nRrPgTS+sP/6Qc/95RSx8T/1Ce1sB1+MtpQfu + 3oMl8aZA3h2OFLtXtvMW0S1T2N59Dce+VOkLq7nnPdU/PsWJnG78KXLQqDUWzTFL+4W0Ug4GnuWo + mbtVNdlg3/z8IlrXjHhLd3RLoJysjOpP/10tS3cNYRhZAbagyusE3cIFlnVqYXWV654CRh2k4KCd + 8BHcs6gb94UiP2FKyAj10zrTtG1BsD8+qGWffTB9wtcZKCcj2/QEWOdKBi4M+4D/41c6x3c5WN3g + gPWXlegTe1tLmLY1S/a378tbLu0UwiDkROx8nYe3+ZsQmOR8pcqUBd606XnwMUlHf/6MvtKYg9XD + NumRu2b9KipMC4RSPWB8r6Z1sZUrAc6BQHqQ76w+RRCdIWtev3TT63p3kdkCVk08bPyjeRPz0gsI + 6+BOnWpwq8G/ft8//qQIJZ+1+/nhysQSmcl80wfjc+qg6d3qH394kyY7HVjszEWLsE+i0SBZDvgB + HFDyGF/95FRQA1t9UCv9auskP2QI8kVgyNyjmz463Y6Dy1pW1FzjuJ8AoxdSzS0GmsSDlU3tl3GB + uLAjVRh1349tJhJwXwYHuyetqNaofzDw58/PvVCtAy/qDaSHqcYa2mFvei0PBDrWRPSmEjWaLAVa + sP2+NarzDNG/mx+A36LUMOZfDaDf7BxD521BIjmrm3F1e4vhPbEdAjNlqQaTITl43vsLml8t7VfW + 8W24uFOFphFY0XQRPAken/yeqm9C+9qnAQNf4mMk7M+PcFbfwfp5Fzc/c43WQncV8NwJEIFKhSs9 + 7DskLZCdqOdnhTcvfPCGafth0Y4rq+j7uisLPHVxQN1DpPST2lkFLPOnRFEXffoBfwEDr9J5xcfO + y6rO0cYAZCNl6KHuLxVpPhwD43OfUceqH9FsR18FbH4DFd2uieZAymKpFRaNfNlPl41YNCf4OF9k + uq13Pz8+ZgfOY4yozT97b6bKpYTivnhib9PrvGILk0Tfbk0PnOKB6XzYF9C6hFdsZZa2Lq+01oBa + qiq1iOhny/Nrar+8DYmM98qmk9sksHbantqhePGGnPQcHOm40t/+GJJUl8D9OcfYuNyu/Ur2Jw7W + r12I3/umAms6nB9//Oet+XLRmFtZALDUnbD1lNmV2Mpz+OkffBjlUl9sRw7A5pepNSe8Pmz+R9rW + E/svPc/m+cZK8OzecsKpLvImQnkDumNypqoPG4+oQ/uQNr7DasjfsuWnZ8orsqiLR7la3getAanF + eFRthRrQyuO5/fdT81gJOMmb7wLwwRUOPjb95er9eZ6b/iNCbvr6eo+HB9wdMEXiHDH9UJqnFATK + /Ynmeq7Xr6i3DXRzHmz7c69P4qlRQBeLLemubKcvgXr3QfVwTWytD62fWqkfoA1iSE0FVyv548f3 + XoKskIujWYRJCblXUuL43vj9EqdJAKMmwTRzv7JHIDpPsLGPIjXgoIPJqThNGkLmQzg4VCu9F5om + jxd6/LPfOUIZAxa1TXA4xWK1ytfBhZ/7VcCKO08Z/Q6qJF8OEqIH88P2E/4CCJLKWBFXeCOgTrWg + n/6kjpTdwaizjgZxG4hUiZWmanlRf8OHwNvULefiv/KhxxvEpHtoEZie5SmRlzqu0Bjelu353lz4 + NVyN1MmnqkZBcv/kM9g4dEK/RKNE4FsjA7Uj0Ge0IfwCkjr3kbjLuZW0Q2JBlckUammnSn8L0smV + urpTsbM8FLAE0tGFL4m3MHqomU7W4c1D9kW/1H86LRjsls/hYz3E1Le8cB19mjC/fBgfjtEFrEko + Gftw/+DJB8W+Pty8XgJhZATUWR4FmK7dPoebfqbO/lCsxDUkfm/Z4hUJQnrJfv735zfp4XPWvIX9 + CC58iscL3eofDPv31/3lN2QfKaousI7hwpdXnelxFmYw29FLk77ZisliLGbPO77Lgw88BNg5yR2Y + VS8roavxF6pfZZQBNvwOMOLeGIeNK/ZEk54EGJXPEiFa2eiXf8LafIZoAbkYTcbn1sKM79w/31/P + Oc/8ydeUT2hFdKq+HNBhyKJZHsxKWC4qJ3Ns88LOWeL6UVF3DISEdkiUWjWbNn8mcTK2t/VQ1t3L + NVsQVMmJmlcxBos4PdG+9+6IOqPYVcMzUxPoHZgzWfI31acqb9B+y5fItOHfNCbfAMrcTccx/AzR + uOEj3PiNTDu1if7sB1t7y2hlunc069bJgmt7cLDXBau35ncuB4UkA2pqjaCvLzBD+FtPFs1pv3b3 + rwik9WgiyIMwI1f6LOEv77JfRy9aHG0MwWBDhUgxjMBwETwR7k/Uom59cvXJzd8aFJamoDa9efp8 + vGapdMdHGeutn3jrzQgt+Uvc9k8+Rza+B5fYQNTdyVG14z7TG/786EEqdlHb+4wNVPOrIi56Gvqf + /Grzt9QzqLNywurYYMsDqX48o34uUuBDryzP+HAYfW/VXMHe6ypVsEUvfjYnryP58S9VEn3K1qmb + ecAFgoP1q0yi9UKXP/5708eKLrx6PwGKJmV/8H1dLioPH1+LYqtfhohYX5eDvl9KW94G+ulNXzEw + wr5BPBb8it/4BKqlrhKp5Tp903cBkJKviXHafDyh8hj+T//APwqOzpMlUYCTBBec9ON53fLls2x+ + 1hoxKcj01Wc7F/aaolK7DNR+x4dKA3tNU6kzJefoh/cwmhKe7D+4q0ZtvYiwjb8Btbb8d63yIQHb + 86D2Aj1vbRwHgsQVfZoeoqJayrG14dYPoTY5Mv1wZCsGbv0UxG157fRIVg7y/oUgCd7lbN3wB54V + iHHGliFYyBIof/IaVzucven+yfhfPeKDkO4y8tvvp9wOcGaHDlhMoePhl9gtdbWxi963i/OArd9+ + iXTb36slYopANofyiu1X10cj5K6LdG7WnBq5JvRL0afvfedZKtWOi7gSPRZzuOUlWM90lPGZzpbg + udtBNGsuBuv2BmuwE9H+dOX6jnv43M9fb/mstNL39cXI+YOf6GHLixbXWHiIVhISRquy7E+/wLcv + L2rP0aMi7Id1wZOZDvTnr5Ye5hI4Em1PzS7Wo01v2+DQKvKW7zPeuuUtcFfLNva4PNKnazc/gBjX + I3aPsRuNl5Mt/cnP+M55Ah5Z1xZWixeR/TtUPS4b4lDCgRXi7XlGf/xOAPQHWmKjjtYX2DNw4zc0 + gjvI5sm3O9Ck9WnLV82MBvZTgUxef7Ez7bhobLqEgJo5nv7ww9QNHoHF7XzAiDFnb0WfoQEbXpPl + BKi+9DixYHz+ZkQa5CJa2y8XQzv/VFRnOiPin9+jBlBdhPTn3+ic3zgIr9YdH+3V7n96BZy4WKCH + wzh45c0oQnnT0xRd1M86vhV7+dNfgA4J19Xn/bf0SEVKxAQ8fn4NQQDLO7WlVo0ms3A6uP0eYjV0 + yNaiZTi4GHgmYrtXKy6/8yE88FNMDyge9Gk9rt2fvKfcm6Y3Q9RBiK/0SaTmrGSCHQgWdFEHCFOP + p2i+hfUAN32P/e9ir4PaWSW897X8q8ese+qBLw+XWEEs5qZ1TTSjlE3cMOTGqpw+VWM0AQsc3lQx + TtY638mllbb/j4TNX/7yGeCC9xm7y+kezQlIePDzm7pNkuxrCLPy0z/Y5fQimn/9GtaTMsSh6VrN + gis2EM3nPcX8ywKjA/oUhoYY0OR17DPylgcIh8d6xPbbvegDnnVm7+6XMzYwfG/4yNoQP4ITTdXX + Hfzpl2z+Ffvu2fGW4BzEsLOHFt8G/PCmdkb2b73QfvcuVhpkqQvzk+9uevYc0aw2HlIxFwTb7JPJ + Rut70uTIPN2pw7zPHqGvLgf3d+lga+vv/OpDntzKw26Qu9mEHOcNCUhqrB0o9YbDZOXwh4cx/y77 + lb7OIbyQkMNb/1X/+jSA0GPOEnWXk5wtxWunQXV5p5u/ANXWr5Fge0catsJrHC1hKrx/ef6v3lfB + iLsQDjRukLDp7T/5d8oOLj2c9oO39YcMactvsP2e+myuW8+Fp7YM8DHkF4+iIn9Lm76n2t6s9Ymc + biG89x8ZCeaHrZbPWljypn/+9JPXc3MuYX5CLjW/0dHrDn2SgL9/pwL+819//fW/ficMmvb+qLeD + AeNjHv/930cF/p3e039zHP/nGAIZ0uLx9z//dQLh72/fNt/xf4/t+/EZ/v7nL/7PUYO/x3ZM6//n + 8r+2G/3nv/4PAAAA//8DAEF/1EXeIAAA + headers: + CF-Cache-Status: + - DYNAMIC + CF-RAY: + - 8ef92722e9cb8389-SEA + Connection: + - keep-alive + Content-Encoding: + - gzip + Content-Type: + - application/json + Date: + - Tue, 10 Dec 2024 00:35:05 GMT + Server: + - cloudflare + Set-Cookie: + - __cf_bm=Bc46FmH66u.HvNfo3T4q3vwT9_TGhtm3wFFMbxQcZ0U-1733790905-1.0.1.1-mXshj7zBxHMWpnIh3c8EmaqnQsDr8FCF6kOqoQvhFoytI6b7MXLbqkZNRfXRP0pd_OraJnyZhUzOl8KaZdq9qg; + path=/; expires=Tue, 10-Dec-24 01:05:05 GMT; domain=.api.openai.com; HttpOnly; + Secure; SameSite=None + - _cfuvid=pdqpFgJ05BGtT4Cs8llmvgLkp1kOJXZNJgkv3mk5xnA-1733790905081-0.0.1.1-604800000; + path=/; domain=.api.openai.com; HttpOnly; Secure; SameSite=None + Transfer-Encoding: + - chunked + X-Content-Type-Options: + - nosniff + access-control-allow-origin: + - '*' + access-control-expose-headers: + - X-Request-ID + alt-svc: + - h3=":443"; ma=86400 + openai-model: + - text-embedding-ada-002 + openai-organization: + - scale3-1 + openai-processing-ms: + - '49' + openai-version: + - '2020-10-01' + strict-transport-security: + - max-age=31536000; includeSubDomains; preload + x-ratelimit-limit-requests: + - '10000' + x-ratelimit-limit-tokens: + - '10000000' + x-ratelimit-remaining-requests: + - '9999' + x-ratelimit-remaining-tokens: + - '9999998' + x-ratelimit-reset-requests: + - 6ms + x-ratelimit-reset-tokens: + - 0s + x-request-id: + - req_0ce98b0d1723b6734062d2cd711aa85e + status: + code: 200 + message: OK +version: 1 diff --git a/src/tests/openai/test_embeddings.py b/src/tests/openai/test_embeddings.py index e69de29b..46a3642f 100644 --- a/src/tests/openai/test_embeddings.py +++ b/src/tests/openai/test_embeddings.py @@ -0,0 +1,88 @@ +import pytest +import httpx +from openai import OpenAI +from opentelemetry.trace import SpanKind, StatusCode +from langtrace.trace_attributes import SpanAttributes +from langtrace_python_sdk.constants.instrumentation.openai import APIS +from langtrace_python_sdk.instrumentation.openai import OpenAIInstrumentation +from tests.utils import assert_token_count +from importlib_metadata import version as v + +# Initialize OpenAI instrumentation +instrumentor = OpenAIInstrumentation() +instrumentor.instrument() + +@pytest.mark.vcr() +def test_embeddings_base_url(exporter, openai_client): + input_value = "Test input" + kwargs = { + "input": input_value, + "model": "text-embedding-ada-002", + } + + openai_client.embeddings.create(**kwargs) + spans = exporter.get_finished_spans() + embedding_span = spans[-1] + + attributes = embedding_span.attributes + assert attributes.get(SpanAttributes.LLM_URL) == "https://api.openai.com/v1/" + assert attributes.get(SpanAttributes.LANGTRACE_SERVICE_NAME) == "OpenAI" + + +def test_embeddings_azure_provider(exporter, monkeypatch): + # Mock response data + mock_response = { + "data": [{"embedding": [0.1] * 1536, "index": 0, "object": "embedding"}], + "model": "text-embedding-ada-002", + "object": "list", + "usage": {"prompt_tokens": 5, "total_tokens": 5} + } + + # Create a mock send method for the HTTP client + def mock_send(self, request, **kwargs): + # Create a proper request with headers + headers = { + "authorization": "Bearer test_api_key", + "content-type": "application/json", + } + request = httpx.Request( + method="POST", + url="https://your-resource.azure.openai.com/v1/embeddings", + headers=headers, + ) + + # Create response with proper context + return httpx.Response( + status_code=200, + content=b'{"data": [{"embedding": [0.1, 0.1], "index": 0, "object": "embedding"}], "model": "text-embedding-ada-002", "object": "list", "usage": {"prompt_tokens": 5, "total_tokens": 5}}', + request=request, + headers={"content-type": "application/json"} + ) + + # Create Azure client + azure_client = OpenAI( + api_key="test_api_key", + base_url="https://your-resource.azure.openai.com/v1", + ) + + # Debug prints + print(f"Debug - Azure client type: {type(azure_client)}") + print(f"Debug - Azure client base_url: {azure_client.base_url}") + print(f"Debug - Azure client _client._base_url: {azure_client._client._base_url if hasattr(azure_client, '_client') else 'No _client'}") + + # Patch the HTTP client's send method + monkeypatch.setattr(httpx.Client, "send", mock_send) + + input_value = "Test input" + kwargs = { + "input": input_value, + "model": "text-embedding-ada-002", + } + + azure_client.embeddings.create(**kwargs) + spans = exporter.get_finished_spans() + embedding_span = spans[-1] + + attributes = embedding_span.attributes + assert attributes.get(SpanAttributes.LLM_URL) == "https://your-resource.azure.openai.com/v1/" + assert attributes.get(SpanAttributes.LANGTRACE_SERVICE_NAME) == "Azure" diff --git a/src/tests/test_session_id.py b/src/tests/test_session_id.py new file mode 100644 index 00000000..32adc667 --- /dev/null +++ b/src/tests/test_session_id.py @@ -0,0 +1,59 @@ +import os +import pytest +from opentelemetry.trace import SpanKind +from langtrace_python_sdk.langtrace import LangtraceConfig +from langtrace_python_sdk.extensions.langtrace_exporter import LangTraceExporter +from langtrace_python_sdk.utils.with_root_span import with_langtrace_root_span +from langtrace_python_sdk.constants.exporter.langtrace_exporter import LANGTRACE_SESSION_ID_HEADER + +def test_session_id_from_env(exporter): + # Test session ID from environment variable + test_session_id = "test-session-123" + os.environ["LANGTRACE_SESSION_ID"] = test_session_id + + @with_langtrace_root_span() + def test_function(): + pass + + test_function() + + spans = exporter.get_finished_spans() + assert len(spans) == 1 + span = spans[0] + assert span.attributes.get("session.id") == test_session_id + + # Cleanup + del os.environ["LANGTRACE_SESSION_ID"] + +def test_session_id_in_config(): + # Test session ID through LangtraceConfig + test_session_id = "config-session-456" + config = LangtraceConfig(session_id=test_session_id) + exporter = LangTraceExporter( + api_host="http://test", + api_key="test-key", + session_id=config.session_id + ) + + assert exporter.session_id == test_session_id + +def test_session_id_in_headers(): + # Test session ID in HTTP headers + test_session_id = "header-session-789" + exporter = LangTraceExporter( + api_host="http://test", + api_key="test-key", + session_id=test_session_id + ) + + # Export method adds headers, so we'll check the headers directly + headers = { + "Content-Type": "application/json", + "x-api-key": "test-key", + "User-Agent": "LangtraceExporter", + } + + if test_session_id: + headers[LANGTRACE_SESSION_ID_HEADER] = test_session_id + + assert headers[LANGTRACE_SESSION_ID_HEADER] == test_session_id