diff --git a/requirements.txt b/requirements.txt index 1c4e67d5..442ada51 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,33 +1,58 @@ -bcrypt==4.1.2 aiohttp==3.9.3 aiosignal==1.3.1 annotated-types==0.6.0 anthropic==0.19.1 anyio==4.3.0 +appnope==0.1.4 +argon2-cffi==23.1.0 +argon2-cffi-bindings==21.2.0 +arrow==1.3.0 asgiref==3.7.2 +asttokens==2.4.1 +async-lru==2.0.4 async-timeout==4.0.3 attrs==23.2.0 +Babel==2.14.0 backoff==2.2.1 +bcrypt==4.1.2 beautifulsoup4==4.12.3 +bleach==6.1.0 bs4==0.0.2 build==1.0.3 +CacheControl==0.14.0 cachetools==5.3.3 +cattrs==23.2.3 certifi==2024.2.2 +cffi==1.16.0 charset-normalizer==3.3.2 chroma-hnswlib==0.7.3 -chromadb==0.4.23 +chromadb==0.4.24 +cleo==2.1.0 click==8.1.7 +cohere==5.2.5 coloredlogs==15.0.1 +comm==0.2.2 +contextvars==2.4 +crashtest==0.4.1 dataclasses-json==0.6.4 +debugpy==1.8.1 +decorator==5.1.1 +defusedxml==0.7.1 Deprecated==1.2.14 dirtyjson==1.0.8 +distlib==0.3.8 distro==1.9.0 docarray==0.40.0 +dulwich==0.21.7 exceptiongroup==1.2.0 +executing==2.0.1 faiss-cpu==1.7.4 fastapi==0.110.0 +fastavro==1.9.4 +fastjsonschema==2.19.1 filelock==3.13.1 flatbuffers==23.5.26 +fqdn==1.5.1 frozenlist==1.4.1 fsspec==2024.2.0 google-auth==2.28.1 @@ -41,18 +66,46 @@ httpx==0.27.0 huggingface-hub==0.20.3 humanfriendly==10.0 idna==3.6 +immutables==0.20 importlib-metadata==6.11.0 importlib_resources==6.1.2 +iniconfig==2.0.0 +installer==0.7.0 +ipykernel==6.29.3 +ipython==8.22.2 +ipywidgets==8.1.2 +isoduration==20.11.0 +jaraco.classes==3.3.1 +jedi==0.19.1 +Jinja2==3.1.3 joblib==1.3.2 +json5==0.9.24 jsonpatch==1.33 jsonpointer==2.4 +jsonschema==4.21.1 +jsonschema-specifications==2023.12.1 +jupyter==1.0.0 +jupyter-console==6.6.3 +jupyter-events==0.10.0 +jupyter-lsp==2.2.4 +jupyter_client==8.6.1 +jupyter_core==5.7.2 +jupyter_server==2.13.0 +jupyter_server_terminals==0.5.3 +jupyterlab==4.1.5 +jupyterlab_pygments==0.3.0 +jupyterlab_server==2.25.4 +jupyterlab_widgets==3.0.10 +keyring==24.3.1 kubernetes==29.0.0 langchain==0.1.9 langchain-community==0.0.24 -langchain-core==0.1.34 +langchain-core==0.1.27 langchain-openai==0.0.7 langchainhub==0.1.14 langsmith==0.1.9 +-e git+ssh://git@github.com/Scale3-Labs/langtrace-python-sdk.git@a6583c0595c5f1b9e8e68c12c500d16a5c37425d#egg=langtrace_python_sdk +Levenshtein==0.25.0 llama-index==0.10.13.post1 llama-index-agent-openai==0.1.5 llama-index-cli==0.1.5 @@ -70,45 +123,80 @@ llama-index-vector-stores-chroma==0.1.4 llama-parse==0.3.4 llamaindex-py-client==0.1.13 markdown-it-py==3.0.0 +MarkupSafe==2.1.5 marshmallow==3.21.0 +matplotlib-inline==0.1.6 mdurl==0.1.2 +mistune==3.0.2 mmh3==4.1.0 monotonic==1.6 +more-itertools==10.2.0 mpmath==1.3.0 +msgpack==1.0.8 multidict==6.0.5 mypy-extensions==1.0.0 +nbclient==0.10.0 +nbconvert==7.16.3 +nbformat==5.10.3 nest-asyncio==1.6.0 networkx==3.2.1 nltk==3.8.1 +notebook==7.1.2 +notebook_shim==0.2.4 numexpr==2.9.0 numpy==1.26.4 oauthlib==3.2.2 onnxruntime==1.17.1 openai==1.12.0 -opentelemetry-api==1.23.0 -opentelemetry-distro==0.44b0 -opentelemetry-exporter-otlp==1.23.0 -opentelemetry-exporter-otlp-proto-common==1.23.0 -opentelemetry-exporter-otlp-proto-grpc==1.23.0 -opentelemetry-exporter-otlp-proto-http==1.23.0 -opentelemetry-instrumentation==0.44b0 -opentelemetry-instrumentation-asgi==0.44b0 -opentelemetry-instrumentation-fastapi==0.44b0 -opentelemetry-proto==1.23.0 -opentelemetry-sdk==1.23.0 -opentelemetry-semantic-conventions==0.44b0 -opentelemetry-util-http==0.44b0 +opentelemetry-api==1.22.0 +opentelemetry-distro==0.43b0 +opentelemetry-exporter-otlp==1.22.0 +opentelemetry-exporter-otlp-proto-common==1.22.0 +opentelemetry-exporter-otlp-proto-grpc==1.22.0 +opentelemetry-exporter-otlp-proto-http==1.22.0 +opentelemetry-instrumentation==0.43b0 +opentelemetry-instrumentation-aiohttp-client==0.43b0 +opentelemetry-instrumentation-asgi==0.43b0 +opentelemetry-instrumentation-aws-lambda==0.43b0 +opentelemetry-instrumentation-dbapi==0.43b0 +opentelemetry-instrumentation-fastapi==0.43b0 +opentelemetry-instrumentation-logging==0.43b0 +opentelemetry-instrumentation-sqlite3==0.43b0 +opentelemetry-instrumentation-urllib==0.43b0 +opentelemetry-instrumentation-wsgi==0.43b0 +opentelemetry-propagator-aws-xray==1.0.1 +opentelemetry-proto==1.22.0 +opentelemetry-sdk==1.22.0 +opentelemetry-semantic-conventions==0.43b0 +opentelemetry-util-http==0.43b0 orjson==3.9.15 overrides==7.7.0 packaging==23.2 pandas==2.2.1 +pandocfilters==1.5.1 +parea-ai==0.2.107 +parso==0.8.3 +pexpect==4.9.0 pillow==10.2.0 pinecone-client==3.1.0 +pkginfo==1.10.0 +platformdirs==4.2.0 +pluggy==1.4.0 +poetry==1.8.2 +poetry-core==1.9.0 +poetry-plugin-dotenv==0.6.28 +poetry-plugin-export==1.7.1 posthog==3.4.2 +prometheus_client==0.20.0 +prompt-toolkit==3.0.43 protobuf==4.25.3 +psutil==5.9.8 +ptyprocess==0.7.0 pulsar-client==3.4.0 +pure-eval==0.2.2 pyasn1==0.5.1 pyasn1-modules==0.3.0 +pycparser==2.21 pydantic==2.6.2 pydantic_core==2.16.3 Pygments==2.17.2 @@ -117,39 +205,71 @@ PyMuPDFb==1.23.22 pypdf==4.0.2 PyPika==0.48.9 pyproject_hooks==1.0.0 +pysbd==0.3.4 +pytest==8.1.1 python-dateutil==2.8.2 python-dotenv==1.0.1 +python-json-logger==2.0.7 pytz==2024.1 +pyupgrade==3.15.1 PyYAML==6.0.1 +pyzmq==25.1.2 +qtconsole==5.5.1 +QtPy==2.4.1 +rapidfuzz==3.6.2 +redis==5.0.3 +referencing==0.34.0 regex==2023.12.25 requests==2.31.0 requests-oauthlib==1.3.1 +requests-toolbelt==1.0.0 +rfc3339-validator==0.1.4 +rfc3986-validator==0.1.1 rich==13.7.0 +rpds-py==0.18.0 rsa==4.9 +Send2Trash==1.8.2 +shellingham==1.5.4 six==1.16.0 sniffio==1.3.0 soupsieve==2.5 SQLAlchemy==2.0.27 +stack-data==0.6.3 starlette==0.36.3 sympy==1.12 tenacity==8.2.3 +terminado==0.18.1 tiktoken==0.6.0 +tinycss2==1.2.1 +tokenize-rt==5.2.0 tokenizers==0.15.2 tomli==2.0.1 +tomlkit==0.12.4 +tornado==6.4 tqdm==4.66.2 -trace-attributes==1.0.28 +trace-attributes==1.0.32 +traitlets==5.14.2 +trove-classifiers==2024.3.3 typer==0.9.0 +types-python-dateutil==2.9.0.20240316 types-requests==2.31.0.20240218 typing==3.7.4.3 typing-inspect==0.9.0 typing_extensions==4.9.0 tzdata==2024.1 +uri-template==1.3.0 urllib3==2.2.1 uvicorn==0.27.1 uvloop==0.19.0 +virtualenv==20.25.1 watchfiles==0.21.0 +wcwidth==0.2.13 +webcolors==1.13 +webencodings==0.5.1 websocket-client==1.7.0 websockets==12.0 +widgetsnbextension==4.0.10 wrapt==1.16.0 +xattr==1.1.0 yarl==1.9.4 zipp==3.17.0 diff --git a/src/examples/cohere_example/__init__.py b/src/examples/cohere_example/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/src/examples/cohere_example/chat.py b/src/examples/cohere_example/chat.py new file mode 100644 index 00000000..adcf64f5 --- /dev/null +++ b/src/examples/cohere_example/chat.py @@ -0,0 +1,26 @@ +from dotenv import find_dotenv, load_dotenv +import cohere + +from langtrace_python_sdk import langtrace +# from langtrace_python_sdk.utils.with_root_span import with_langtrace_root_span + +_ = load_dotenv(find_dotenv()) + +langtrace.init(batch=False, debug_log_to_console=True, write_to_langtrace_cloud=False) + +co = cohere.Client() + + +# @with_langtrace_root_span("chat_create") +def chat_comp(): + response = co.chat( + chat_history=[ + {"role": "USER", "message": "Who discovered gravity?"}, + {"role": "CHATBOT", "message": "The man who is widely credited with discovering gravity is Sir Isaac Newton"} + ], + message="What is today's news?", + # preamble="answer like yoda", + # perform web search before answering the question. You can also use your own custom connector. + # connectors=[{"id": "web-search"}] + ) + print(response) diff --git a/src/examples/cohere_example/chat_stream.py b/src/examples/cohere_example/chat_stream.py new file mode 100644 index 00000000..5bbc5e7f --- /dev/null +++ b/src/examples/cohere_example/chat_stream.py @@ -0,0 +1,23 @@ +from dotenv import find_dotenv, load_dotenv +import cohere + +from langtrace_python_sdk import langtrace +# from langtrace_python_sdk.utils.with_root_span import with_langtrace_root_span + +_ = load_dotenv(find_dotenv()) + +langtrace.init(batch=False, debug_log_to_console=True, write_to_langtrace_cloud=False) + +co = cohere.Client() + + +# @with_langtrace_root_span("chat_stream") +def chat_stream(): + result = [] + for event in co.chat_stream(message="Tell me a short story in 2 lines"): + if event.event_type == "text-generation": + result.append(event.text) + elif event.event_type == "stream-end": + break + print("".join(result)) + return result diff --git a/src/examples/cohere_example/embed_create.py b/src/examples/cohere_example/embed_create.py new file mode 100644 index 00000000..ad215a9c --- /dev/null +++ b/src/examples/cohere_example/embed_create.py @@ -0,0 +1,21 @@ +from dotenv import find_dotenv, load_dotenv +import cohere + +from langtrace_python_sdk import langtrace +# from langtrace_python_sdk.utils.with_root_span import with_langtrace_root_span + +_ = load_dotenv(find_dotenv()) + +langtrace.init(batch=False, debug_log_to_console=True, write_to_langtrace_cloud=False) + +co = cohere.Client() + + +# @with_langtrace_root_span("embed_create") +def embed_create(): + response = co.embed( + texts=['hello', 'goodbye'], + model='embed-english-v3.0', + input_type='classification' + ) + # print(response) diff --git a/src/langtrace_python_sdk/constants/instrumentation/cohere.py b/src/langtrace_python_sdk/constants/instrumentation/cohere.py new file mode 100644 index 00000000..ad89f04f --- /dev/null +++ b/src/langtrace_python_sdk/constants/instrumentation/cohere.py @@ -0,0 +1,17 @@ +APIS = { + "CHAT_CREATE": { + "URL": "https://api.cohere.ai", + "METHOD": "cohere.client.chat", + "ENDPOINT": "/v1/chat", + }, + "EMBED_CREATE": { + "URL": "https://api.cohere.ai", + "METHOD": "cohere.client.embed", + "ENDPOINT": "/v1/embed", + }, + "CHAT_STREAM": { + "URL": "https://api.cohere.ai", + "METHOD": "cohere.client.chat_stream", + "ENDPOINT": "/v1/messages", + }, +} diff --git a/src/langtrace_python_sdk/constants/instrumentation/common.py b/src/langtrace_python_sdk/constants/instrumentation/common.py index 2e85dfdb..2979542d 100644 --- a/src/langtrace_python_sdk/constants/instrumentation/common.py +++ b/src/langtrace_python_sdk/constants/instrumentation/common.py @@ -16,6 +16,7 @@ "LLAMAINDEX": "LlamaIndex", "OPENAI": "OpenAI", "PINECONE": "Pinecone", + "COHERE": "Cohere", "PPLX": "Perplexity", } diff --git a/src/langtrace_python_sdk/instrumentation/cohere/__init__.py b/src/langtrace_python_sdk/instrumentation/cohere/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/src/langtrace_python_sdk/instrumentation/cohere/instrumentation.py b/src/langtrace_python_sdk/instrumentation/cohere/instrumentation.py new file mode 100644 index 00000000..83eb774e --- /dev/null +++ b/src/langtrace_python_sdk/instrumentation/cohere/instrumentation.py @@ -0,0 +1,53 @@ +""" +Instrumentation for Cohere +""" + +import importlib.metadata +from typing import Collection + +from opentelemetry.instrumentation.instrumentor import BaseInstrumentor +from opentelemetry.trace import get_tracer +from wrapt import wrap_function_wrapper + +from langtrace_python_sdk.instrumentation.cohere.patch import ( + chat_create, + chat_stream, + embed_create +) + +class CohereInstrumentation(BaseInstrumentor): + """ + The CohereInstrumentation class represents the Anthropic instrumentation + """ + + def instrumentation_dependencies(self) -> Collection[str]: + return ["cohere >= 5.0.0"] + + def _instrument(self, **kwargs): + tracer_provider = kwargs.get("tracer_provider") + tracer = get_tracer(__name__, "", tracer_provider) + version = importlib.metadata.version("cohere") + + wrap_function_wrapper( + "cohere.client", + "Client.chat", + chat_create("cohere.client.chat", version, tracer), + ) + + wrap_function_wrapper( + "cohere.client", + "Client.chat_stream", + chat_stream("cohere.client.chat_stream", version, tracer), + ) + + wrap_function_wrapper( + "cohere.client", + "Client.embed", + embed_create("cohere.client.embed", version, tracer), + ) + + def _instrument_module(self, module_name): + pass + + def _uninstrument(self, **kwargs): + pass diff --git a/src/langtrace_python_sdk/instrumentation/cohere/patch.py b/src/langtrace_python_sdk/instrumentation/cohere/patch.py new file mode 100644 index 00000000..ab47aded --- /dev/null +++ b/src/langtrace_python_sdk/instrumentation/cohere/patch.py @@ -0,0 +1,370 @@ +""" +This module contains the patching logic for the Anthropic library.""" + +import json + +from langtrace.trace_attributes import Event, LLMSpanAttributes +from opentelemetry import baggage +from opentelemetry.trace import SpanKind +from opentelemetry.trace.status import Status, StatusCode + +from langtrace_python_sdk.constants.instrumentation.cohere import APIS +from langtrace_python_sdk.constants.instrumentation.common import (LANGTRACE_ADDITIONAL_SPAN_ATTRIBUTES_KEY, SERVICE_PROVIDERS) +from langtrace_python_sdk.utils.llm import estimate_tokens + + +def embed_create(original_method, version, tracer): + """Wrap the `embed_create` method.""" + + def traced_method(wrapped, instance, args, kwargs): + service_provider = SERVICE_PROVIDERS["COHERE"] + extra_attributes = baggage.get_baggage(LANGTRACE_ADDITIONAL_SPAN_ATTRIBUTES_KEY) + + span_attributes = { + "langtrace.sdk.name": "langtrace-python-sdk", + "langtrace.service.name": service_provider, + "langtrace.service.type": "llm", + "langtrace.service.version": version, + "langtrace.version": "1.0.0", + "url.full": APIS["EMBED_CREATE"]["URL"], + "llm.api": APIS["EMBED_CREATE"]["ENDPOINT"], + "llm.model": kwargs.get("model"), + "llm.prompts": "", + "llm.embedding_dataset_id": kwargs.get("dataset_id"), + "llm.embedding_input_type": kwargs.get("input_type"), + "llm.embedding_job_name": kwargs.get("name"), + **(extra_attributes if extra_attributes is not None else {}) + } + + attributes = LLMSpanAttributes(**span_attributes) + + if kwargs.get("user") is not None: + attributes.llm_user = kwargs.get("user") + + span = tracer.start_span( + APIS["EMBED_CREATE"]["METHOD"], kind=SpanKind.CLIENT + ) + for field, value in attributes.model_dump(by_alias=True).items(): + if value is not None: + span.set_attribute(field, value) + try: + # Attempt to call the original method + result = wrapped(*args, **kwargs) + 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_create(original_method, version, tracer): + """Wrap the `chat_create` method.""" + + def traced_method(wrapped, instance, args, kwargs): + service_provider = SERVICE_PROVIDERS["COHERE"] + + message = kwargs.get("message", "") + prompts = json.dumps([ + { + "role": "USER", + "content": message + } + ]) + preamble = kwargs.get("preamble") + if preamble: + prompts = json.dumps( + [{"role": "system", "content": preamble}] + [{"role": "USER", "content": message}] + ) + + chat_history = kwargs.get("chat_history") + if chat_history: + history = [ + { + "message": { + "role": ( + item.get("role") if item.get("role") is not None else "USER" + ), + "content": ( + item.get("message") if item.get("message") is not None else "" + ) + } + } + for item in chat_history + ] + prompts = prompts + json.dumps(history) + + extra_attributes = baggage.get_baggage(LANGTRACE_ADDITIONAL_SPAN_ATTRIBUTES_KEY) + + span_attributes = { + "langtrace.sdk.name": "langtrace-python-sdk", + "langtrace.service.name": service_provider, + "langtrace.service.type": "llm", + "langtrace.service.version": version, + "langtrace.version": "1.0.0", + "url.full": APIS["CHAT_CREATE"]["URL"], + "llm.api": APIS["CHAT_CREATE"]["ENDPOINT"], + "llm.model": kwargs.get("model") if kwargs.get("model") is not None else "command-r", + "llm.stream": False, + "llm.prompts": prompts, + **(extra_attributes if extra_attributes is not None else {}) + } + + attributes = LLMSpanAttributes(**span_attributes) + + if kwargs.get("temperature") is not None: + attributes.llm_temperature = kwargs.get("temperature") + if kwargs.get("max_tokens") is not None: + attributes.max_tokens = kwargs.get("max_tokens") + if kwargs.get("max_input_tokens") is not None: + attributes.max_input_tokens = kwargs.get("max_input_tokens") + if kwargs.get("p") is not None: + attributes.llm_top_p = kwargs.get("p") + if kwargs.get("k") is not None: + attributes.llm_top_p = kwargs.get("k") + if kwargs.get("user") is not None: + attributes.llm_user = kwargs.get("user") + if kwargs.get("conversation_id") is not None: + attributes.conversation_id = kwargs.get("conversation_id") + if kwargs.get("seed") is not None: + attributes.seed = kwargs.get("seed") + if kwargs.get("frequency_penalty") is not None: + attributes.frequency_penalty = kwargs.get("frequency_penalty") + if kwargs.get("presence_penalty") is not None: + attributes.presence_penalty = kwargs.get("presence_penalty") + if kwargs.get("connectors") is not None: + # stringify the list of objects + attributes.llm_connectors = json.dumps(kwargs.get("connectors")) + if kwargs.get("tools") is not None: + # stringify the list of objects + attributes.llm_tools = json.dumps(kwargs.get("tools")) + if kwargs.get("tool_results") is not None: + # stringify the list of objects + attributes.llm_tool_results = json.dumps(kwargs.get("tool_results")) + + span = tracer.start_span( + APIS["CHAT_CREATE"]["METHOD"], kind=SpanKind.CLIENT + ) + + # Set the attributes on the span + for field, value in attributes.model_dump(by_alias=True).items(): + if value is not None: + span.set_attribute(field, value) + try: + # Attempt to call the original method + result = wrapped(*args, **kwargs) + + # Set the response attributes + if (hasattr(result, "generation_id")) and (result.generation_id is not None): + span.set_attribute("llm.generation_id", result.generation_id) + if (hasattr(result, "response_id")) and (result.response_id is not None): + span.set_attribute("llm.response_id", result.response_id) + if (hasattr(result, "is_search_required")) and (result.is_search_required is not None): + span.set_attribute("llm.is_search_required", result.is_search_required) + + if kwargs.get("stream") is False or kwargs.get("stream") is None: + if hasattr(result, "text") and result.text is not None: + if hasattr(result, "chat_history") and result.chat_history is not None: + responses = [ + { + "message": { + "role": ( + item.role if hasattr(item, "role") and item.role is not None else "USER" + ), + "content": ( + item.message if hasattr(item, "message") and item.message is not None else "" + ) + } + } + for item in result.chat_history + ] + span.set_attribute("llm.responses", json.dumps(responses)) + else: + responses = [{ + "message": { + "role": "CHATBOT", + "content": result.text + } + }] + span.set_attribute("llm.responses", json.dumps(responses)) + else: + responses = [] + span.set_attribute("llm.responses", json.dumps(responses)) + + # Get the usage + if hasattr(result, "meta") and result.meta is not None: + if hasattr(result.meta, "billed_units") and result.meta.billed_units is not None: + usage = result.meta.billed_units + if usage is not None: + usage_dict = { + "input_tokens": usage.input_tokens if usage.input_tokens is not None else 0, + "output_tokens": usage.output_tokens if usage.output_tokens is not None else 0, + "total_tokens": usage.input_tokens + usage.output_tokens if usage.input_tokens is not None and usage.output_tokens is not None else 0, + } + span.set_attribute("llm.token.counts", json.dumps(usage_dict)) + span.set_status(StatusCode.OK) + span.end() + return result + else: + # For older version, stream was passed as a parameter + 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.""" + + def traced_method(wrapped, instance, args, kwargs): + service_provider = SERVICE_PROVIDERS["COHERE"] + + message = kwargs.get("message", "") + prompt_tokens = estimate_tokens(message) + prompts = json.dumps([ + { + "role": "USER", + "content": message + } + ]) + preamble = kwargs.get("preamble") + if preamble: + prompts = json.dumps( + [{"role": "system", "content": preamble}] + [{"role": "USER", "content": message}] + ) + + chat_history = kwargs.get("chat_history") + if chat_history: + history = [ + { + "message": { + "role": ( + item.get("role") if item.get("role") is not None else "USER" + ), + "content": ( + item.get("message") if item.get("message") is not None else "" + ) + } + } + for item in chat_history + ] + prompts = prompts + json.dumps(history) + + extra_attributes = baggage.get_baggage(LANGTRACE_ADDITIONAL_SPAN_ATTRIBUTES_KEY) + + span_attributes = { + "langtrace.sdk.name": "langtrace-python-sdk", + "langtrace.service.name": service_provider, + "langtrace.service.type": "llm", + "langtrace.service.version": version, + "langtrace.version": "1.0.0", + "url.full": APIS["CHAT_STREAM"]["URL"], + "llm.api": APIS["CHAT_STREAM"]["ENDPOINT"], + "llm.model": kwargs.get("model") if kwargs.get("model") is not None else "command-r", + "llm.stream": False, + "llm.prompts": prompts, + **(extra_attributes if extra_attributes is not None else {}) + } + + attributes = LLMSpanAttributes(**span_attributes) + + if kwargs.get("temperature") is not None: + attributes.llm_temperature = kwargs.get("temperature") + if kwargs.get("max_tokens") is not None: + attributes.max_tokens = kwargs.get("max_tokens") + if kwargs.get("max_input_tokens") is not None: + attributes.max_input_tokens = kwargs.get("max_input_tokens") + if kwargs.get("p") is not None: + attributes.llm_top_p = kwargs.get("p") + if kwargs.get("k") is not None: + attributes.llm_top_p = kwargs.get("k") + if kwargs.get("user") is not None: + attributes.llm_user = kwargs.get("user") + if kwargs.get("conversation_id") is not None: + attributes.conversation_id = kwargs.get("conversation_id") + if kwargs.get("seed") is not None: + attributes.seed = kwargs.get("seed") + if kwargs.get("frequency_penalty") is not None: + attributes.frequency_penalty = kwargs.get("frequency_penalty") + if kwargs.get("presence_penalty") is not None: + attributes.presence_penalty = kwargs.get("presence_penalty") + if kwargs.get("connectors") is not None: + # stringify the list of objects + attributes.llm_connectors = json.dumps(kwargs.get("connectors")) + if kwargs.get("tools") is not None: + # stringify the list of objects + attributes.llm_tools = json.dumps(kwargs.get("tools")) + if kwargs.get("tool_results") is not None: + # stringify the list of objects + attributes.llm_tool_results = json.dumps(kwargs.get("tool_results")) + + span = tracer.start_span( + APIS["CHAT_CREATE"]["METHOD"], kind=SpanKind.CLIENT + ) + for field, value in attributes.model_dump(by_alias=True).items(): + if value is not None: + span.set_attribute(field, value) + try: + # Attempt to call the original method + result = wrapped(*args, **kwargs) + + result_content = [] + span.add_event(Event.STREAM_START.value) + completion_tokens = 0 + try: + for event in result: + if hasattr(event, "text") and event.text is not None: + completion_tokens += estimate_tokens(event.text) + content = event.text + else: + content = "" + span.add_event( + Event.STREAM_OUTPUT.value, {"response": "".join(content)} + ) + result_content.append(content) + yield event + finally: + + # Finalize span after processing all chunks + span.add_event(Event.STREAM_END.value) + span.set_attribute( + "llm.token.counts", + json.dumps( + { + "input_tokens": prompt_tokens, + "output_tokens": completion_tokens, + "total_tokens": prompt_tokens + completion_tokens, + } + ), + ) + span.set_attribute( + "llm.responses", + json.dumps( + [ + { + "message": { + "role": "CHATBOT", + "content": "".join(result_content), + } + } + ] + ), + ) + span.set_status(StatusCode.OK) + span.end() + + except Exception as error: + span.record_exception(error) + span.set_status(Status(StatusCode.ERROR, str(error))) + span.end() + raise + + return traced_method diff --git a/src/langtrace_python_sdk/instrumentation/openai/patch.py b/src/langtrace_python_sdk/instrumentation/openai/patch.py index 5b5acc94..65c3444d 100644 --- a/src/langtrace_python_sdk/instrumentation/openai/patch.py +++ b/src/langtrace_python_sdk/instrumentation/openai/patch.py @@ -220,12 +220,10 @@ def traced_method(wrapped, instance, args, kwargs): prompt_tokens, function_call=kwargs.get("functions") is not None, ) - except Exception as e: - # Record the exception in the span - span.record_exception(e) - # Set the span status to indicate an error - span.set_status(Status(StatusCode.ERROR, str(e))) - # Reraise the exception to ensure it's not swallowed + + except Exception as error: + span.record_exception(error) + span.set_status(Status(StatusCode.ERROR, str(error))) span.end() raise diff --git a/src/langtrace_python_sdk/langtrace.py b/src/langtrace_python_sdk/langtrace.py index 030094c5..bbe9c3a9 100644 --- a/src/langtrace_python_sdk/langtrace.py +++ b/src/langtrace_python_sdk/langtrace.py @@ -32,6 +32,9 @@ from langtrace_python_sdk.instrumentation.pinecone.instrumentation import ( PineconeInstrumentation, ) +from langtrace_python_sdk.instrumentation.cohere.instrumentation import ( + CohereInstrumentation, +) def init( @@ -77,6 +80,7 @@ def init( langchain_core_instrumentation = LangchainCoreInstrumentation() langchain_community_instrumentation = LangchainCommunityInstrumentation() anthropic_instrumentation = AnthropicInstrumentation() + cohere_instrumentation = CohereInstrumentation() # Call the instrument method with some arguments openai_instrumentation.instrument() @@ -87,3 +91,4 @@ def init( langchain_core_instrumentation.instrument() langchain_community_instrumentation.instrument() anthropic_instrumentation.instrument() + cohere_instrumentation.instrument() diff --git a/src/run_example.py b/src/run_example.py index 4556d3c3..70a174e0 100644 --- a/src/run_example.py +++ b/src/run_example.py @@ -1,7 +1,7 @@ # from examples.langchain_example.basic import basic, load_and_split, rag -from examples.openai.chat_completion import chat_completion +# from examples.openai.chat_completion import chat_completion # from examples.openai import images_generate # from examples.openai.function_calling import function_calling @@ -11,13 +11,20 @@ # from examples.chroma_example.basic import basic # from examples.llamaindex_example.basic import basic # from examples.langchain_example.basic import basic +# from examples.cohere_example.chat import chat_comp +from examples.cohere_example.embed_create import embed_create +# from examples.cohere_example.chat_stream import chat_stream + # from examples.perplexity_example.basic import basic # load_and_split() # rag() # basic() -chat_completion() +# chat_completion() # function_calling() # images_generate() # embeddings_create() # messages_create() +# chat_comp() +embed_create() +# chat_stream()