From 331a717bb4edf5ca223b070a4e3ba2d3a8c44012 Mon Sep 17 00:00:00 2001 From: Matthieu Olenga Date: Wed, 13 Nov 2024 13:01:04 +0100 Subject: [PATCH] feat: add better typing for decorator/context manager fucn --- literalai/client.py | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/literalai/client.py b/literalai/client.py index c3743cd..8401e96 100644 --- a/literalai/client.py +++ b/literalai/client.py @@ -1,5 +1,5 @@ import os -from typing import Any, Dict, List, Optional, Union +from typing import Any, Callable, Dict, List, Optional, overload, TypeVar, Union from literalai.api import AsyncLiteralAPI, LiteralAPI from literalai.callback.langchain_callback import get_langchain_callback @@ -26,6 +26,8 @@ from literalai.requirements import check_all_requirements +T = TypeVar("T", bound=Callable) + class BaseLiteralClient: """ @@ -144,9 +146,15 @@ def langchain_callback( **kwargs, ) + @overload + def thread(self, original_function: Callable[..., T], **kwargs) -> Callable[..., T]: ... + + @overload + def thread(self, **kwargs) -> ThreadContextManager: ... + def thread( self, - original_function=None, + original_function: Optional[Callable] = None, *, thread_id: Optional[str] = None, name: Optional[str] = None,