feat(traceloop-sdk): enhance type annotations and add new dependencies

- Updated poetry.lock to include pandas-stubs and types-pytz for improved type checking.
- Modified project.json to specify the mypy command for type checking the sdk directory.
- Enhanced type annotations across various modules for better type safety and clarity.
- Refactored function signatures to include return types for improved type hinting.

Author: doronkopit5

packages/traceloop-sdk/poetry.lock

@@ -54,7 +54,7 @@
       "executor": "@nxlv/python:run-commands",
       "outputs": [],
       "options": {
-        "command": "poetry run mypy",
+        "command": "poetry run mypy traceloop/sdk",
         "cwd": "packages/traceloop-sdk"
       }
     },

packages/traceloop-sdk/project.json

@@ -83,6 +83,7 @@ pytest-sugar = "1.0.0"
 mypy = "^1.18.2"
 types-requests = "^2.31.0"
 types-colorama = "^0.4.15"
+pandas-stubs = "*"
 
 [tool.poetry.group.test.dependencies]
 openai = "^1.31.1"
@@ -114,16 +115,7 @@ namespace_packages = true
 explicit_package_bases = true
 plugins = ["pydantic.mypy"]
 
-# Target specific folders for strict type checking
-files = [
-    "traceloop/sdk/annotation",
-    "traceloop/sdk/client",
-    "traceloop/sdk/dataset",
-    "traceloop/sdk/evaluator",
-    "traceloop/sdk/experiment",
-]
-
-# Exclude other folders (blacklist approach - new folders will be checked by default)
+# Blacklist approach - all folders checked except those excluded below
 exclude = [
     "traceloop/sdk/decorators",
     "traceloop/sdk/prompts",

packages/traceloop-sdk/pyproject.toml

@@ -8,7 +8,6 @@ class UserFeedback(BaseAnnotation):
     def __init__(self, http: HTTPClient, app_name: str):
         super().__init__(http, app_name, "user_feedback")
 
-
     def create(
         self,
         annotation_task: str,

packages/traceloop-sdk/traceloop/sdk/annotation/user_feedback.py

@@ -1,5 +1,5 @@
 import csv
-from typing import List, Optional
+from typing import List, Optional, cast
 from pathlib import Path
 
 try:
@@ -81,6 +81,10 @@ def from_csv(
 
             reader = csv.DictReader(csvfile, delimiter=delimiter)
 
+            # TODO: Handle None case for fieldnames more gracefully
+            if reader.fieldnames is None:
+                raise ValueError("CSV file has no headers")
+
             for field_name in reader.fieldnames:
                 columns_definition.append(
                     ColumnDefinition(
@@ -138,8 +142,9 @@ def from_dataframe(
                 )
             )
 
+        # TODO: Pandas returns Hashable keys, should ensure they're strings
         rows = [
-            {self._slugify(k): v for k, v in row.items()}
+            {self._slugify(str(k)): v for k, v in row.items()}
             for row in df.to_dict(orient="records")
         ]
 
@@ -160,14 +165,14 @@ def get_version_csv(self, slug: str, version: str) -> str:
         result = self._http.get(f"datasets/{slug}/versions/{version}")
         if result is None:
             raise Exception(f"Failed to get dataset {slug} by version {version}")
-        return result
+        return cast(str, result)
 
     def get_version_jsonl(self, slug: str, version: str) -> str:
         """Get a specific version of a dataset as a JSONL string"""
         result = self._http.get(f"datasets/{slug}/versions/{version}/jsonl")
         if result is None:
             raise Exception(f"Failed to get dataset {slug} by version {version}")
-        return result
+        return cast(str, result)
 
     def _create_dataset(self, input: CreateDatasetRequest) -> CreateDatasetResponse:
         """Create new dataset"""

packages/traceloop-sdk/traceloop/sdk/datasets/datasets.py

@@ -1,4 +1,4 @@
-from typing import Optional, TypeVar, Callable, Any, ParamSpec, Awaitable
+from typing import Optional, TypeVar, Callable
 import warnings
 
 from opentelemetry.semconv_ai import TraceloopSpanKindValues
@@ -8,9 +8,7 @@
     entity_method,
 )
 
-P = ParamSpec("P")
-R = TypeVar("R")
-F = TypeVar("F", bound=Callable[P, R | Awaitable[R]])
+F = TypeVar("F")
 
 
 def task(

packages/traceloop-sdk/traceloop/sdk/decorators/__init__.py

@@ -7,8 +7,6 @@
     Callable,
     Any,
     cast,
-    ParamSpec,
-    Awaitable,
 )
 import inspect
 import warnings
@@ -29,10 +27,7 @@
 from traceloop.sdk.utils import camel_to_snake
 from traceloop.sdk.utils.json_encoder import JSONEncoder
 
-P = ParamSpec("P")
-
-R = TypeVar("R")
-F = TypeVar("F", bound=Callable[P, R | Awaitable[R]])
+F = TypeVar("F")
 
 
 def _truncate_json_if_needed(json_str: str) -> str:

packages/traceloop-sdk/traceloop/sdk/decorators/base.py

@@ -6,7 +6,7 @@
 import requests
 
 from threading import Thread, Event
-from typing import Dict
+from typing import Dict, Any
 from tenacity import (
     RetryError,
     retry,
@@ -51,7 +51,7 @@ def __init__(self, base_url: str, api_key: str):
             ),
         )
 
-    def run(self):
+    def run(self) -> None:
         refresh_data(
             self._base_url,
             self._api_key,
@@ -61,10 +61,10 @@ def run(self):
         self._exit_monitor.start()
         self._poller_thread.start()
 
-    def post(self, api: str, body: Dict[str, str]):
+    def post(self, api: str, body: Dict[str, str]) -> None:
         post_url(f"{self._base_url}/v1/traceloop/{api}", self._api_key, body)
 
-    def api_post(self, api: str, body: Dict[str, typing.Any]):
+    def api_post(self, api: str, body: Dict[str, typing.Any]) -> None:
         try:
             post_url(f"{self._base_url}/v2/{api}", self._api_key, body)
         except Exception as e:
@@ -83,7 +83,7 @@ def __init__(
         super().__init__(lambda e: check_http_error(e))
 
 
-def check_http_error(e):
+def check_http_error(e: BaseException) -> bool:
     return isinstance(e, requests.exceptions.HTTPError) and (
         500 <= e.response.status_code < 600
     )
@@ -94,7 +94,7 @@ def check_http_error(e):
     stop=stop_after_attempt(MAX_RETRIES),
     retry=RetryIfServerError(),
 )
-def fetch_url(url: str, api_key: str):
+def fetch_url(url: str, api_key: str) -> Any:
     response = requests.get(
         url,
         headers={
@@ -114,7 +114,7 @@ def fetch_url(url: str, api_key: str):
         return response.json()
 
 
-def post_url(url: str, api_key: str, body: Dict[str, typing.Any]):
+def post_url(url: str, api_key: str, body: Dict[str, typing.Any]) -> None:
     response = requests.post(
         url,
         headers={
@@ -135,7 +135,7 @@ def thread_func(
     api_key: str,
     stop_polling_event: Event,
     seconds_interval: float = 5.0,
-):
+) -> None:
     while not stop_polling_event.is_set():
         try:
             refresh_data(base_url, api_key, prompt_registry, content_allow_list)
@@ -151,15 +151,15 @@ def refresh_data(
     api_key: str,
     prompt_registry: PromptRegistry,
     content_allow_list: ContentAllowList,
-):
+) -> None:
     response = fetch_url(f"{base_url}/v1/traceloop/prompts", api_key)
     prompt_registry.load(response)
 
     response = fetch_url(f"{base_url}/v1/traceloop/pii/tracing-allow-list", api_key)
     content_allow_list.load(response)
 
 
-def monitor_exit(exit_event: Event):
+def monitor_exit(exit_event: Event) -> None:
     main_thread = threading.main_thread()
     main_thread.join()
     exit_event.set()

packages/traceloop-sdk/traceloop/sdk/fetcher.py

@@ -6,22 +6,26 @@
 
 
 class ImageUploader:
-    def __init__(self, base_url, api_key):
+    def __init__(self, base_url: str, api_key: str) -> None:
         self.base_url = base_url
         self.api_key = api_key
         self.logger = logging.getLogger(__name__)
 
-    def upload_base64_image(self, trace_id, span_id, image_name, image_file):
-        asyncio.run(self.aupload_image_file(trace_id, span_id, image_name, image_file))
+    def upload_base64_image(
+        self, trace_id: str, span_id: str, image_name: str, image_file: str
+    ) -> None:
+        asyncio.run(self.aupload_base64_image(trace_id, span_id, image_name, image_file))
 
-    async def aupload_base64_image(self, trace_id, span_id, image_name, image_file):
+    async def aupload_base64_image(
+        self, trace_id: str, span_id: str, image_name: str, image_file: str
+    ) -> str:
         url = self._get_image_url(trace_id, span_id, image_name)
 
         await self._async_upload(url, image_file)
 
         return url
 
-    def _get_image_url(self, trace_id, span_id, image_name):
+    def _get_image_url(self, trace_id: str, span_id: str, image_name: str) -> str:
         response = requests.post(
             f"{self.base_url}/v2/traces/{trace_id}/spans/{span_id}/images",
             json={
@@ -33,9 +37,9 @@ def _get_image_url(self, trace_id, span_id, image_name):
             },
         )
 
-        return response.json()["url"]
+        return response.json()["url"]  # type: ignore[no-any-return]
 
-    async def _async_upload(self, url, base64_image):
+    async def _async_upload(self, url: str, base64_image: str) -> None:
         headers = {
             "Authorization": f"Bearer {self.api_key}",
             "Content-Type": "application/json",

packages/traceloop-sdk/traceloop/sdk/images/image_uploader.py

@@ -1,5 +1,5 @@
 import logging
-from typing import Dict
+from typing import Dict, Optional, Any
 
 from opentelemetry.exporter.otlp.proto.grpc._log_exporter import (
     OTLPLogExporter as GRPCExporter,
@@ -15,13 +15,13 @@
 
 
 class LoggerWrapper(object):
-    resource_attributes: dict = {}
-    endpoint: str = None
+    resource_attributes: Dict[Any, Any] = {}
+    endpoint: Optional[str] = None
     headers: Dict[str, str] = {}
-    __logging_exporter: LogExporter = None
-    __logging_provider: LoggerProvider = None
+    __logging_exporter: Optional[LogExporter] = None
+    __logging_provider: Optional[LoggerProvider] = None
 
-    def __new__(cls, exporter: LogExporter = None) -> "LoggerWrapper":
+    def __new__(cls, exporter: Optional[LogExporter] = None) -> "LoggerWrapper":
         if not hasattr(cls, "instance"):
             obj = cls.instance = super(LoggerWrapper, cls).__new__(cls)
             if not LoggerWrapper.endpoint:
@@ -58,7 +58,7 @@ def init_logging_exporter(endpoint: str, headers: Dict[str, str]) -> LogExporter
 
 
 def init_logging_provider(
-    exporter: LogExporter, resource_attributes: dict = None
+    exporter: LogExporter, resource_attributes: Optional[Dict[Any, Any]] = None
 ) -> LoggerProvider:
     resource = (
         Resource.create(resource_attributes)