lmnr 0.4.6__py3-none-any.whl → 0.4.8__py3-none-any.whl

lmnr/traceloop_sdk/tests/test_sdk_initialization.py

@@ -0,0 +1,57 @@
+import pytest
+from openai import OpenAI
+from lmnr.traceloop_sdk.decorators import workflow
+
+
+@pytest.fixture
+def openai_client():
+    return OpenAI()
+
+
+@pytest.mark.vcr
+def test_resource_attributes(exporter, openai_client):
+    openai_client.chat.completions.create(
+        model="gpt-3.5-turbo",
+        messages=[{"role": "user", "content": "Tell me a joke about opentelemetry"}],
+    )
+
+    spans = exporter.get_finished_spans()
+    open_ai_span = spans[0]
+    assert open_ai_span.resource.attributes["something"] == "yes"
+    assert open_ai_span.resource.attributes["service.name"] == "test"
+
+
+def test_custom_span_processor(exporter_with_custom_span_processor):
+    @workflow()
+    def run_workflow():
+        pass
+
+    run_workflow()
+
+    spans = exporter_with_custom_span_processor.get_finished_spans()
+    workflow_span = spans[0]
+    assert workflow_span.attributes["custom_span"] == "yes"
+
+
+def test_instruments(exporter_with_custom_instrumentations):
+    @workflow()
+    def run_workflow():
+        pass
+
+    run_workflow()
+
+    spans = exporter_with_custom_instrumentations.get_finished_spans()
+    workflow_span = spans[0]
+    assert workflow_span
+
+
+def test_no_metrics(exporter_with_no_metrics):
+    @workflow()
+    def run_workflow():
+        pass
+
+    run_workflow()
+
+    spans = exporter_with_no_metrics.get_finished_spans()
+    workflow_span = spans[0]
+    assert workflow_span

lmnr/traceloop_sdk/tests/test_tasks.py

@@ -0,0 +1,32 @@
+import json
+import pytest
+
+from langchain_openai import ChatOpenAI
+from lmnr.traceloop_sdk.decorators import task
+from opentelemetry.semconv_ai import SpanAttributes
+
+
+@pytest.mark.vcr
+def test_task_io_serialization_with_langchain(exporter):
+    @task(name="answer_question")
+    def answer_question():
+        chat = ChatOpenAI(temperature=0)
+
+        return chat.invoke("Is Berlin the capital of Germany? Answer with yes or no")
+
+    answer_question()
+
+    spans = exporter.get_finished_spans()
+
+    assert [span.name for span in spans] == [
+        "ChatOpenAI.chat",
+        "answer_question.task",
+    ]
+
+    task_span = next(span for span in spans if span.name == "answer_question.task")
+    assert (
+        json.loads(task_span.attributes.get(SpanAttributes.TRACELOOP_ENTITY_OUTPUT))[
+            "kwargs"
+        ]["content"]
+        == "Yes"
+    )

lmnr/traceloop_sdk/tests/test_workflows.py

@@ -0,0 +1,261 @@
+import json
+
+import pytest
+from openai import OpenAI, AsyncOpenAI
+from opentelemetry.semconv_ai import SpanAttributes
+from lmnr.traceloop_sdk import Traceloop
+from lmnr.traceloop_sdk.decorators import workflow, task, aworkflow, atask
+
+
+@pytest.fixture
+def openai_client():
+    return OpenAI()
+
+
+@pytest.fixture
+def async_openai_client():
+    return AsyncOpenAI()
+
+
+@pytest.mark.vcr
+def test_simple_workflow(exporter, openai_client):
+    @task(name="something_creator", version=2)
+    def create_something(what: str, subject: str):
+        Traceloop.set_prompt(
+            "Tell me a {what} about {subject}", {"what": what, "subject": subject}, 5
+        )
+        completion = openai_client.chat.completions.create(
+            model="gpt-3.5-turbo",
+            messages=[{"role": "user", "content": f"Tell me a {what} about {subject}"}],
+        )
+        return completion.choices[0].message.content
+
+    @workflow(name="pirate_joke_generator", version=1)
+    def joke_workflow():
+        return create_something("joke", subject="OpenTelemetry")
+
+    joke = joke_workflow()
+
+    spans = exporter.get_finished_spans()
+    assert [span.name for span in spans] == [
+        "openai.chat",
+        "something_creator.task",
+        "pirate_joke_generator.workflow",
+    ]
+    open_ai_span = next(span for span in spans if span.name == "openai.chat")
+    assert (
+        open_ai_span.attributes[f"{SpanAttributes.LLM_PROMPTS}.0.content"]
+        == "Tell me a joke about OpenTelemetry"
+    )
+    assert open_ai_span.attributes.get(f"{SpanAttributes.LLM_COMPLETIONS}.0.content")
+    assert (
+        open_ai_span.attributes.get("traceloop.prompt.template")
+        == "Tell me a {what} about {subject}"
+    )
+    assert (
+        open_ai_span.attributes.get("traceloop.prompt.template_variables.what")
+        == "joke"
+    )
+    assert (
+        open_ai_span.attributes.get("traceloop.prompt.template_variables.subject")
+        == "OpenTelemetry"
+    )
+    assert open_ai_span.attributes.get("traceloop.prompt.version") == 5
+
+    workflow_span = next(
+        span for span in spans if span.name == "pirate_joke_generator.workflow"
+    )
+    task_span = next(span for span in spans if span.name == "something_creator.task")
+    assert json.loads(task_span.attributes[SpanAttributes.TRACELOOP_ENTITY_INPUT]) == {
+        "args": ["joke"],
+        "kwargs": {"subject": "OpenTelemetry"},
+    }
+
+    assert (
+        json.loads(task_span.attributes.get(SpanAttributes.TRACELOOP_ENTITY_OUTPUT))
+        == joke
+    )
+    assert task_span.parent.span_id == workflow_span.context.span_id
+    assert (
+        workflow_span.attributes[SpanAttributes.TRACELOOP_ENTITY_NAME]
+        == "pirate_joke_generator"
+    )
+    assert workflow_span.attributes[SpanAttributes.TRACELOOP_ENTITY_VERSION] == 1
+    assert task_span.attributes[SpanAttributes.TRACELOOP_ENTITY_VERSION] == 2
+
+
+@pytest.mark.vcr
+@pytest.mark.asyncio
+async def test_simple_aworkflow(exporter, async_openai_client):
+    @atask(name="something_creator", version=2)
+    async def create_something(what: str, subject: str):
+        Traceloop.set_prompt(
+            "Tell me a {what} about {subject}", {"what": what, "subject": subject}, 5
+        )
+        completion = await async_openai_client.chat.completions.create(
+            model="gpt-3.5-turbo",
+            messages=[{"role": "user", "content": f"Tell me a {what} about {subject}"}],
+        )
+        return completion.choices[0].message.content
+
+    @aworkflow(name="pirate_joke_generator", version=1)
+    async def joke_workflow():
+        return await create_something("joke", subject="OpenTelemetry")
+
+    joke = await joke_workflow()
+
+    spans = exporter.get_finished_spans()
+    assert [span.name for span in spans] == [
+        "openai.chat",
+        "something_creator.task",
+        "pirate_joke_generator.workflow",
+    ]
+    open_ai_span = next(span for span in spans if span.name == "openai.chat")
+    assert (
+        open_ai_span.attributes[f"{SpanAttributes.LLM_PROMPTS}.0.content"]
+        == "Tell me a joke about OpenTelemetry"
+    )
+    assert open_ai_span.attributes.get(f"{SpanAttributes.LLM_COMPLETIONS}.0.content")
+    assert (
+        open_ai_span.attributes.get("traceloop.prompt.template")
+        == "Tell me a {what} about {subject}"
+    )
+    assert (
+        open_ai_span.attributes.get("traceloop.prompt.template_variables.what")
+        == "joke"
+    )
+    assert (
+        open_ai_span.attributes.get("traceloop.prompt.template_variables.subject")
+        == "OpenTelemetry"
+    )
+    assert open_ai_span.attributes.get("traceloop.prompt.version") == 5
+
+    workflow_span = next(
+        span for span in spans if span.name == "pirate_joke_generator.workflow"
+    )
+    task_span = next(span for span in spans if span.name == "something_creator.task")
+    assert json.loads(task_span.attributes[SpanAttributes.TRACELOOP_ENTITY_INPUT]) == {
+        "args": ["joke"],
+        "kwargs": {"subject": "OpenTelemetry"},
+    }
+
+    assert (
+        json.loads(task_span.attributes.get(SpanAttributes.TRACELOOP_ENTITY_OUTPUT))
+        == joke
+    )
+    assert task_span.parent.span_id == workflow_span.context.span_id
+    assert (
+        workflow_span.attributes[SpanAttributes.TRACELOOP_ENTITY_NAME]
+        == "pirate_joke_generator"
+    )
+    assert workflow_span.attributes[SpanAttributes.TRACELOOP_ENTITY_VERSION] == 1
+    assert task_span.attributes[SpanAttributes.TRACELOOP_ENTITY_VERSION] == 2
+
+
+@pytest.mark.vcr
+def test_streaming_workflow(exporter, openai_client):
+
+    @task(name="pirate_joke_generator")
+    def joke_task():
+        response_stream = openai_client.chat.completions.create(
+            model="gpt-3.5-turbo",
+            messages=[
+                {"role": "user", "content": "Tell me a joke about OpenTelemetry"}
+            ],
+            stream=True,
+        )
+        for chunk in response_stream:
+            yield chunk
+
+    @task(name="joke_runner")
+    def joke_runner():
+        res = joke_task()
+        return res
+
+    @workflow(name="joke_manager")
+    def joke_workflow():
+        res = joke_runner()
+        for chunk in res:
+            pass
+
+    joke_workflow()
+
+    spans = exporter.get_finished_spans()
+    assert set([span.name for span in spans]) == set(
+        [
+            "openai.chat",
+            "pirate_joke_generator.task",
+            "joke_runner.task",
+            "joke_manager.workflow",
+        ]
+    )
+    generator_span = next(
+        span for span in spans if span.name == "pirate_joke_generator.task"
+    )
+    runner_span = next(span for span in spans if span.name == "joke_runner.task")
+    manager_span = next(span for span in spans if span.name == "joke_manager.workflow")
+    openai_span = next(span for span in spans if span.name == "openai.chat")
+
+    assert openai_span.parent.span_id == generator_span.context.span_id
+    assert generator_span.parent.span_id == runner_span.context.span_id
+    assert runner_span.parent.span_id == manager_span.context.span_id
+    assert openai_span.end_time <= manager_span.end_time
+
+
+def test_unrelated_entities(exporter):
+    @workflow(name="workflow_1")
+    def workflow_1():
+        return
+
+    @task(name="task_1")
+    def task_1():
+        return
+
+    workflow_1()
+    task_1()
+
+    spans = exporter.get_finished_spans()
+    assert [span.name for span in spans] == ["workflow_1.workflow", "task_1.task"]
+
+    workflow_1_span = spans[0]
+    task_1_span = spans[1]
+
+    assert (
+        workflow_1_span.attributes[SpanAttributes.TRACELOOP_ENTITY_NAME] == "workflow_1"
+    )
+    assert workflow_1_span.attributes[SpanAttributes.TRACELOOP_SPAN_KIND] == "workflow"
+
+    assert task_1_span.attributes[SpanAttributes.TRACELOOP_ENTITY_NAME] == "task_1"
+    assert task_1_span.attributes[SpanAttributes.TRACELOOP_SPAN_KIND] == "task"
+    assert task_1_span.parent is None
+
+
+def test_unserializable_workflow(exporter):
+    @task(name="unserializable_task")
+    def unserializable_task(obj: object):
+        return object()
+
+    @workflow(name="unserializable_workflow")
+    def unserializable_workflow(obj: object):
+        return unserializable_task(obj)
+
+    unserializable_task(object())
+
+    spans = exporter.get_finished_spans()
+    assert [span.name for span in spans] == ["unserializable_task.task"]
+
+
+@pytest.mark.asyncio
+async def test_unserializable_async_workflow(exporter):
+    @atask(name="unserializable_task")
+    async def unserializable_task(obj: object):
+        return object()
+
+    @aworkflow(name="unserializable_workflow")
+    async def unserializable_workflow(obj: object):
+        return await unserializable_task(obj)
+
+    await unserializable_task(object())
+
+    spans = exporter.get_finished_spans()
+    assert [span.name for span in spans] == ["unserializable_task.task"]

lmnr/traceloop_sdk/tracing/__init__.py

@@ -0,0 +1,2 @@
+from lmnr.traceloop_sdk.tracing.context_manager import get_tracer
+from lmnr.traceloop_sdk.tracing.tracing import set_workflow_name

lmnr/traceloop_sdk/tracing/content_allow_list.py

@@ -0,0 +1,24 @@
+# Manages list of associated properties for which content tracing
+# (prompts, vector embeddings, etc.) is allowed.
+class ContentAllowList:
+    def __new__(cls) -> "ContentAllowList":
+        if not hasattr(cls, "instance"):
+            obj = cls.instance = super(ContentAllowList, cls).__new__(cls)
+            obj._allow_list: list[dict] = []
+
+        return cls.instance
+
+    def is_allowed(self, association_properties: dict) -> bool:
+        for allow_list_item in self._allow_list:
+            if all(
+                [
+                    association_properties.get(key) == value
+                    for key, value in allow_list_item.items()
+                ]
+            ):
+                return True
+
+        return False
+
+    def load(self, response_json: dict):
+        self._allow_list = response_json["associationPropertyAllowList"]

lmnr/traceloop_sdk/tracing/context_manager.py

@@ -0,0 +1,13 @@
+from contextlib import contextmanager
+
+from lmnr.traceloop_sdk.tracing.tracing import TracerWrapper
+
+
+@contextmanager
+def get_tracer(flush_on_exit: bool = False):
+    wrapper = TracerWrapper()
+    try:
+        yield wrapper.get_tracer()
+    finally:
+        if flush_on_exit:
+            wrapper.flush()

lmnr/traceloop_sdk/tracing/manual.py

@@ -0,0 +1,57 @@
+from contextlib import contextmanager
+from opentelemetry.semconv_ai import SpanAttributes
+from opentelemetry.trace import Span
+from pydantic import BaseModel
+from lmnr.traceloop_sdk.tracing.context_manager import get_tracer
+
+
+class LLMMessage(BaseModel):
+    role: str
+    content: str
+
+
+class LLMUsage(BaseModel):
+    prompt_tokens: int
+    completion_tokens: int
+    total_tokens: int
+
+
+class LLMSpan:
+    _span: Span = None
+
+    def __init__(self, span: Span):
+        self._span = span
+        pass
+
+    def report_request(self, model: str, messages: list[LLMMessage]):
+        self._span.set_attribute(SpanAttributes.LLM_REQUEST_MODEL, model)
+        for idx, message in enumerate(messages):
+            self._span.set_attribute(
+                f"{SpanAttributes.LLM_PROMPTS}.{idx}.role", message.role
+            )
+            self._span.set_attribute(
+                f"{SpanAttributes.LLM_PROMPTS}.{idx}.content", message.content
+            )
+
+    def report_response(self, model: str, completions: list[str]):
+        self._span.set_attribute(SpanAttributes.LLM_RESPONSE_MODEL, model)
+        for idx, completion in enumerate(completions):
+            self._span.set_attribute(
+                f"{SpanAttributes.LLM_COMPLETIONS}.{idx}.role", "assistant"
+            )
+            self._span.set_attribute(
+                f"{SpanAttributes.LLM_COMPLETIONS}.{idx}", completion
+            )
+
+
+@contextmanager
+def track_llm_call(vendor: str, type: str):
+    with get_tracer() as tracer:
+        with tracer.start_as_current_span(name=f"{vendor}.{type}") as span:
+            span.set_attribute(SpanAttributes.LLM_SYSTEM, vendor)
+            span.set_attribute(SpanAttributes.LLM_REQUEST_TYPE, type)
+            llm_span = LLMSpan(span)
+            try:
+                yield llm_span
+            finally:
+                span.end()