add log10_tags context manager and with_tags decorator

examples/logging/session_openai.py

@@ -1,8 +1,23 @@
-from log10.load import OpenAI, log10_session
+from log10.load import OpenAI, log10_session, log10_tags, with_log10_tags
 
 
 client = OpenAI()
 
+
+@with_log10_tags(["decorator-tags", "decorator-tags-2"])
+def completion_with_tags():
+    completion = client.chat.completions.create(
+        model="gpt-4o",
+        messages=[
+            {
+                "role": "user",
+                "content": "Hello?",
+            },
+        ],
+    )
+    print(completion.choices[0].message)
+
+
 with log10_session(tags=["log10-io/examples"]):
     completion = client.chat.completions.create(
         model="gpt-3.5-turbo",
@@ -15,13 +30,42 @@
     )
     print(completion.choices[0].message)
 
+    with log10_tags(["extra_tag_1", "extra_tag_2"]):
+        completion = client.chat.completions.create(
+            model="gpt-3.5-turbo",
+            messages=[
+                {
+                    "role": "user",
+                    "content": "Hello again, are you there?",
+                },
+            ],
+        )
+        print(completion.choices[0].message)
+
     completion = client.chat.completions.create(
         model="gpt-3.5-turbo",
         messages=[
             {
                 "role": "user",
-                "content": "Hello again, are you there?",
+                "content": "Hello again and again?",
             },
         ],
     )
     print(completion.choices[0].message)
+
+    completion_with_tags()
+
+# add a test with log10_tags and log10_session, where log10_session is nested inside log10_tags
+with log10_tags(["outer-tag-1", "outer-tag-2"]):
+    with log10_session(tags=["inner-tag-1", "inner-tag-2"]):
+        completion = client.chat.completions.create(
+            model="gpt-3.5-turbo",
+            messages=[
+                {
+                    "role": "user",
+                    "content": "Hello again and again?",
+                },
+            ],
+        )
+        print(completion.choices[0].message)
+    completion_with_tags()

log10/load.py

@@ -120,33 +120,82 @@ def get_log10_session_tags():
     return tags_var.get()
 
 
+class TagsManager:
+    def __init__(self, tags: list[str] = None):
+        self.tags = self._validate_tags(tags) or []
+
+    @staticmethod
+    def _validate_tags(tags: list[str] | None) -> list[str]:
+        if tags is None:
+            return None
+
+        if not isinstance(tags, list):
+            logger.warning(
+                f"Invalid tags format: expected list, got {type(tags).__name__}. Tags will be omitted from the log."
+            )
+            return None
+
+        validated_tags = []
+        for tag in tags:
+            if not isinstance(tag, str):
+                logger.warning(
+                    f"Invalid tag type: expected str, got {type(tag).__name__}. This tag will be omitted: {repr(tag)}"
+                )
+                continue
+            validated_tags.append(tag)
+        return validated_tags
+
+    def _enter(self):
+        current_tags = tags_var.get()
+        new_tags = current_tags + self.tags
+
+        self.tags_token = tags_var.set(new_tags)
+        return self.tags_token
+
+    def _exit(self, exc_type, exc_value, traceback):
+        tags_var.reset(self.tags_token)
+
+    def __enter__(self):
+        return self._enter()
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        self._exit(exc_type, exc_value, traceback)
+
+    async def __aenter__(self):
+        return self._enter()
+
+    async def __aexit__(self, exc_type, exc_value, traceback):
+        self._exit(exc_type, exc_value, traceback)
+
+
 class log10_session:
     def __init__(self, tags=None):
-        self.tags = tags
+        self.tags_manager = TagsManager(tags)
 
-    def __enter__(self):
+    def _enter(self):
         self.session_id_token = session_id_var.set(get_session_id())
         self.last_completion_response_token = last_completion_response_var.set(None)
-        self.tags_token = tags_var.set(self.tags)
+        self.tags_manager._enter()
 
         return self
 
-    def __exit__(self, exc_type, exc_value, traceback):
+    def _exit(self, exc_type, exc_value, traceback):
         session_id_var.reset(self.session_id_token)
         last_completion_response_var.reset(self.last_completion_response_token)
-        tags_var.reset(self.tags_token)
+        self.tags_manager._exit(exc_type, exc_value, traceback)
+
+    def __enter__(self):
+        return self._enter()
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        self._exit(exc_type, exc_value, traceback)
         return
 
     async def __aenter__(self):
-        self.session_id_token = session_id_var.set(get_session_id())
-        self.last_completion_response_token = last_completion_response_var.set(None)
-        self.tags_token = tags_var.set(self.tags)
-        return self
+        return self._enter()
 
     async def __aexit__(self, exc_type, exc_value, traceback):
-        session_id_var.reset(self.session_id_token)
-        last_completion_response_var.reset(self.last_completion_response_token)
-        tags_var.reset(self.tags_token)
+        self._exit(exc_type, exc_value, traceback)
         return
 
     def last_completion_url(self):
@@ -166,6 +215,60 @@ def last_completion_id(self):
         return response["completionID"]
 
 
+@contextmanager
+def log10_tags(tags: list[str]):
+    """
+    A context manager that adds tags to the current session.
+    This could be used with log10_session to add extra tags to the session.
+    Example:
+    >>> from log10.load import log10_tags
+    >>> with log10_tags(["tag1", "tag2"]):
+    >>>     completion = client.chat.completions.create(
+    >>>         model="gpt-4o",
+    >>>         messages=[
+    >>>             {
+    >>>                 "role": "user",
+    >>>                 "content": "Hello?",
+    >>>             },
+    >>>         ],
+    >>>     )
+    >>>     print(completion.choices[0].message)
+    """
+    tags_manager = TagsManager(tags)
+    with tags_manager:
+        yield
+
+
+def with_log10_tags(tags: list[str]):
+    """
+    A decorator that adds tags to a function call.
+    Example:
+    >>> from log10.load import with_log10_tags
+    >>> @with_log10_tags(["decorator-tags", "decorator-tags-2"])
+    >>> def completion_with_tags():
+    >>>     completion = client.chat.completions.create(
+    >>>         model="gpt-4o",
+    >>>         messages=[
+    >>>             {
+    >>>                 "role": "user",
+    >>>                 "content": "Hello?",
+    >>>             },
+    >>>         ],
+    >>>     )
+    >>>     print(completion.choices[0].message)
+    """
+
+    def decorator(func):
+        @functools.wraps(func)
+        def wrapper(*args, **kwargs):
+            with log10_tags(tags):
+                return func(*args, **kwargs)
+
+        return wrapper
+
+    return decorator
+
+
 @contextmanager
 def timed_block(block_name):
     if DEBUG:

tests/test_load.py

@@ -0,0 +1,106 @@
+import pytest
+
+from log10.load import get_log10_session_tags, log10_session, log10_tags, with_log10_tags
+
+
+def test_log10_tags():
+    # Test single tag
+    with log10_tags(["test_tag"]):
+        assert get_log10_session_tags() == ["test_tag"]
+    assert get_log10_session_tags() == []
+
+    # Test multiple tags
+    with log10_tags(["tag1", "tag2"]):
+        assert get_log10_session_tags() == ["tag1", "tag2"]
+    assert get_log10_session_tags() == []
+
+    # Test nested tags
+    with log10_tags(["outer"]):
+        assert get_log10_session_tags() == ["outer"]
+        with log10_tags(["inner"]):
+            assert get_log10_session_tags() == ["outer", "inner"]
+        assert get_log10_session_tags() == ["outer"]
+
+    # Test that tags are cleared after context
+    assert get_log10_session_tags() == []
+
+
+def test_log10_session():
+    # Test session with no tags
+    with log10_session():
+        assert get_log10_session_tags() == []
+
+    # Test session with tags
+    with log10_session(tags=["session_tag"]):
+        assert get_log10_session_tags() == ["session_tag"]
+
+
+def test_log10_tags_session():
+    # Test nested session and tags
+    with log10_session(tags=["outer_session"]):
+        assert get_log10_session_tags() == ["outer_session"]
+        with log10_tags(["inner_tag"]):
+            assert get_log10_session_tags() == ["outer_session", "inner_tag"]
+        assert get_log10_session_tags() == ["outer_session"]
+    assert get_log10_session_tags() == []
+
+    with log10_tags(["outer_tag"]):
+        assert get_log10_session_tags() == ["outer_tag"]
+        with log10_session(tags=["inner_session"]):
+            assert get_log10_session_tags() == ["outer_tag", "inner_session"]
+        assert get_log10_session_tags() == ["outer_tag"]
+    assert get_log10_session_tags() == []
+
+
+@pytest.mark.asyncio
+async def test_log10_session_async():
+    # Test async session with tags
+    async with log10_session(tags=["async_session"]):
+        assert get_log10_session_tags() == ["async_session"]
+
+    # Test that tags are cleared after async session
+    assert get_log10_session_tags() == []
+
+
+def test_log10_tags_invalid_input():
+    # Test with non-list input
+    with log10_tags("not_a_list"):
+        assert get_log10_session_tags() == []
+
+    # Test with non-string tags
+    with log10_tags(["valid", 123, {"invalid": "tag"}]):
+        assert get_log10_session_tags() == ["valid"]
+
+
+def test_log10_session_invalid_input():
+    # Test with non-list tags
+    with log10_session(tags="not_a_list"):
+        assert get_log10_session_tags() == []
+
+    # Test with non-string tags
+    with log10_session(tags=["valid", 123, {"invalid": "tag"}]):
+        assert get_log10_session_tags() == ["valid"]
+
+
+def test_with_log10_tags_decorator():
+    @with_log10_tags(["decorator_tag1", "decorator_tag2"])
+    def decorated_function():
+        return get_log10_session_tags()
+
+    # Test that the decorator adds tags
+    assert decorated_function() == ["decorator_tag1", "decorator_tag2"]
+
+    # Test that tags are cleared after the decorated function call
+    assert get_log10_session_tags() == []
+
+    # Test nested decorators
+    @with_log10_tags(["outer_tag"])
+    def outer_function():
+        @with_log10_tags(["inner_tag"])
+        def inner_function():
+            return get_log10_session_tags()
+
+        return inner_function()
+
+    assert outer_function() == ["outer_tag", "inner_tag"]
+    assert get_log10_session_tags() == []