fix(traceloop-sdk): add type-checking support with mypy

.github/workflows/ci.yml

@@ -74,6 +74,7 @@ jobs:
         run: npm cache clean --force || true
       - run: npx nx affected -t install --with dev
       - run: npx nx affected -t lint --parallel=3
+      - run: npx nx affected -t type-check --parallel=3
 
   build-packages:
     name: Build Packages

packages/traceloop-sdk/project.json

@@ -50,6 +50,14 @@
         "outputFile": "reports/packages/traceloop-sdk/pylint.txt"
       }
     },
+    "type-check": {
+      "executor": "@nxlv/python:run-commands",
+      "outputs": [],
+      "options": {
+        "command": "poetry run mypy traceloop/sdk",
+        "cwd": "packages/traceloop-sdk"
+      }
+    },
     "test": {
       "executor": "@nxlv/python:run-commands",
       "outputs": [

packages/traceloop-sdk/pyproject.toml

@@ -80,6 +80,10 @@ autopep8 = "^2.2.0"
 flake8 = "7.0.0"
 pytest = "^8.2.2"
 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"
@@ -95,6 +99,44 @@ pandas = ">=1.0.0"
 [tool.poetry.extras]
 datasets = ["pandas"]
 
+[tool.mypy]
+python_version = "3.10"
+warn_return_any = true
+warn_unused_configs = true
+disallow_untyped_defs = true
+disallow_any_unimported = false
+no_implicit_optional = true
+warn_redundant_casts = true
+warn_unused_ignores = true
+warn_no_return = true
+check_untyped_defs = true
+strict_equality = true
+namespace_packages = true
+explicit_package_bases = true
+plugins = ["pydantic.mypy"]
+
+# Blacklist approach - all folders checked except those excluded below
+exclude = [
+    "traceloop/sdk/decorators",
+    "traceloop/sdk/prompts",
+    "traceloop/sdk/tracing",
+    "traceloop/sdk/utils",
+    "traceloop/sdk/__init__.py",
+    "tests/",
+]
+
+[[tool.mypy.overrides]]
+module = [
+    "cuid.*",
+    "posthog.*",
+]
+ignore_missing_imports = true
+
+[pydantic-mypy]
+init_forbid_extra = true
+init_typed = true
+warn_required_dynamic_aliases = true
+
 [build-system]
 requires = ["poetry-core"]
 build-backend = "poetry.core.masonry.api"

packages/traceloop-sdk/tests/test_user_feedback.py

@@ -35,10 +35,10 @@ def test_user_feedback_initialization(mock_http):
     assert feedback._app_name == "test-app"
 
 
-def test_create_basic_feedback(user_feedback, mock_http):
+def test_create_basic_feedback(user_feedback: UserFeedback, mock_http: Mock):
     """Test creating basic user feedback"""
     user_feedback.create(
-        annotation_task="task_123", entity_id="instance_456", tags={"sentiment": "positive"}
+        annotation_task="task_123", entity_instance_id="instance_456", tags={"sentiment": "positive"}
     )
 
     mock_http.post.assert_called_once_with(
@@ -56,11 +56,11 @@ def test_create_basic_feedback(user_feedback, mock_http):
     )
 
 
-def test_create_feedback_complex_tags(user_feedback, mock_http):
+def test_create_feedback_complex_tags(user_feedback: UserFeedback, mock_http: Mock):
     """Test creating user feedback with complex tags"""
     tags = {"sentiment": "positive", "relevance": 0.95, "tones": ["happy", "nice"]}
 
-    user_feedback.create(annotation_task="task_123", entity_id="instance_456", tags=tags)
+    user_feedback.create(annotation_task="task_123", entity_instance_id="instance_456", tags=tags)
 
     mock_http.post.assert_called_once_with(
         "annotation-tasks/task_123/annotations",
@@ -77,13 +77,13 @@ def test_create_feedback_complex_tags(user_feedback, mock_http):
     )
 
 
-def test_create_feedback_parameter_validation(user_feedback):
+def test_create_feedback_parameter_validation(user_feedback: UserFeedback):
     """Test parameter validation for feedback creation"""
     with pytest.raises(ValueError, match="annotation_task is required"):
-        user_feedback.create(annotation_task="", entity_id="instance_456", tags={"sentiment": "positive"})
+        user_feedback.create(annotation_task="", entity_instance_id="instance_456", tags={"sentiment": "positive"})
 
-    with pytest.raises(ValueError, match="entity_id is required"):
-        user_feedback.create(annotation_task="task_123", entity_id="", tags={"sentiment": "positive"})
+    with pytest.raises(ValueError, match="entity_instance_id is required"):
+        user_feedback.create(annotation_task="task_123", entity_instance_id="", tags={"sentiment": "positive"})
 
     with pytest.raises(ValueError, match="tags cannot be empty"):
-        user_feedback.create(annotation_task="task_123", entity_id="instance_456", tags={})
+        user_feedback.create(annotation_task="task_123", entity_instance_id="instance_456", tags={})

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

@@ -21,7 +21,7 @@ def __init__(self, http: HTTPClient, app_name: str, flow: str):
     def create(
         self,
         annotation_task: str,
-        entity_id: str,
+        entity_instance_id: str,
         tags: Dict[str, Any],
     ) -> None:
         """Create an user feedback annotation for a specific task.
@@ -51,15 +51,15 @@ def create(
 
         if not annotation_task:
             raise ValueError("annotation_task is required")
-        if not entity_id:
-            raise ValueError("entity_id is required")
+        if not entity_instance_id:
+            raise ValueError("entity_instance_id is required")
         if not tags:
             raise ValueError("tags cannot be empty")
 
         self._http.post(
             f"annotation-tasks/{annotation_task}/annotations",
             {
-                "entity_instance_id": entity_id,
+                "entity_instance_id": entity_instance_id,
                 "tags": tags,
                 "source": "sdk",
                 "flow": self._flow,

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

@@ -8,36 +8,35 @@ class UserFeedback(BaseAnnotation):
     def __init__(self, http: HTTPClient, app_name: str):
         super().__init__(http, app_name, "user_feedback")
 
-
-def create(
-    self,
-    annotation_task: str,
-    entity_instance_id: str,
-    tags: Dict[str, Any],
-) -> None:
-    """Create an annotation for a specific task.
-
-    Args:
-        annotation_task (str): The ID/slug of the annotation task to report to.
-            Can be found at app.traceloop.com/annotation_tasks/:annotation_task_id
-        entity_instance_id (str): The ID of the specific entity instance being annotated, should be reported
-            in the association properties
-        tags (Dict[str, Any]): Dictionary containing the tags to be reported.
-            Should match the tags defined in the annotation task
-
-    Example:
-        ```python
-        client = Client(api_key="your-key")
-        client.annotation.create(
-            annotation_task="task_123",
-            entity_instance_id="instance_456",
-            tags={
-                "sentiment": "positive",
-                "relevance": 0.95,
-                "tones": ["happy", "nice"]
-            },
-        )
-        ```
-    """
-
-    return BaseAnnotation.create(self, annotation_task, entity_instance_id, tags)
+    def create(
+        self,
+        annotation_task: str,
+        entity_instance_id: str,
+        tags: Dict[str, Any],
+    ) -> None:
+        """Create an annotation for a specific task.
+
+        Args:
+            annotation_task (str): The ID/slug of the annotation task to report to.
+                Can be found at app.traceloop.com/annotation_tasks/:annotation_task_id
+            entity_instance_id (str): The ID of the specific entity instance being annotated, should be reported
+                in the association properties
+            tags (Dict[str, Any]): Dictionary containing the tags to be reported.
+                Should match the tags defined in the annotation task
+
+        Example:
+            ```python
+            client = Client(api_key="your-key")
+            client.annotation.create(
+                annotation_task="task_123",
+                entity_instance_id="instance_456",
+                tags={
+                    "sentiment": "positive",
+                    "relevance": 0.95,
+                    "tones": ["happy", "nice"]
+                },
+            )
+            ```
+        """
+
+        return BaseAnnotation.create(self, annotation_task, entity_instance_id, tags)

packages/traceloop-sdk/traceloop/sdk/client/client.py

@@ -63,4 +63,5 @@ def __init__(
         self.user_feedback = UserFeedback(self._http, self.app_name)
         self.datasets = Datasets(self._http)
         experiment_slug = os.getenv("TRACELOOP_EXP_SLUG")
-        self.experiment = Experiment(self._http, self._async_http, experiment_slug)
+        # TODO: Fix type - Experiment constructor should accept Optional[str]
+        self.experiment = Experiment(self._http, self._async_http, experiment_slug)  # type: ignore[arg-type]

packages/traceloop-sdk/traceloop/sdk/client/http.py

@@ -14,7 +14,7 @@ def __init__(self, base_url: str, api_key: str, version: str):
         self.api_key = api_key
         self.version = version
 
-    def _headers(self):
+    def _headers(self) -> Dict[str, str]:
         return {
             "Authorization": f"Bearer {self.api_key}",
             "X-Traceloop-SDK-Version": self.version,

packages/traceloop-sdk/traceloop/sdk/dataset/dataset.py

@@ -86,7 +86,7 @@ def add_column(self, slug: str, name: str, col_type: ColumnType) -> Column:
         self.columns.append(column)
         return column
 
-    def _create_columns(self, raw_columns: Dict[str, ColumnDefinition]):
+    def _create_columns(self, raw_columns: Dict[str, ColumnDefinition]) -> None:
         """Create Column objects from API response which includes column IDs"""
         for column_slug, column_def in raw_columns.items():
             column = Column(
@@ -98,7 +98,7 @@ def _create_columns(self, raw_columns: Dict[str, ColumnDefinition]):
             )
             self.columns.append(column)
 
-    def _create_rows(self, raw_rows: List[RowObject]):
+    def _create_rows(self, raw_rows: List[RowObject]) -> None:
         for _, row_obj in enumerate(raw_rows):
             row = Row(
                 http=self._http,

packages/traceloop-sdk/traceloop/sdk/datasets/datasets.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/decorators/__init__.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/base.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/evaluator/stream_client.py

@@ -52,13 +52,13 @@ async def wait_for_result(
         except Exception as e:
             raise Exception(f"Unexpected error in SSE stream: {e}")
 
-    async def _handle_sse_response(self, response) -> ExecutionResponse:
+    async def _handle_sse_response(self, response: httpx.Response) -> ExecutionResponse:
         """Handle SSE response: check status and parse result"""
         if response.status_code != 200:
             error_text = await response.aread()
-            raise Exception(
-                f"Failed to stream results: {response.status_code}, body: {error_text}"
-            )
+            # TODO: Fix bytes formatting - should decode error_text or use !r
+            error_msg = f"Failed to stream results: {response.status_code}, body: {error_text}"  # type: ignore[str-bytes-safe]  # noqa: E501
+            raise Exception(error_msg)
 
         response_text = await response.aread()
         return self._parse_sse_result(response_text.decode())