chore: add mypy support (#1257)

Author: hassiebp

.github/workflows/ci.yml

@@ -23,6 +23,37 @@ jobs:
         with:
           args: check --config ci.ruff.toml
 
+  type-checking:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v3
+      - name: Set up Python
+        uses: actions/setup-python@v4
+        with:
+          python-version: "3.12"
+      - name: Install poetry
+        uses: abatilo/actions-poetry@v2
+      - name: Setup a local virtual environment
+        run: |
+          poetry config virtualenvs.create true --local
+          poetry config virtualenvs.in-project true --local
+      - uses: actions/cache@v3
+        name: Define a cache for the virtual environment based on the dependencies lock file
+        with:
+          path: ./.venv
+          key: venv-type-check-${{ hashFiles('poetry.lock') }}
+      - uses: actions/cache@v3
+        name: Cache mypy cache
+        with:
+          path: ./.mypy_cache
+          key: mypy-${{ hashFiles('**/*.py', 'pyproject.toml') }}
+          restore-keys: |
+            mypy-
+      - name: Install dependencies
+        run: poetry install --only=main,dev
+      - name: Run mypy type checking
+        run: poetry run mypy langfuse --no-error-summary
+
   ci:
     runs-on: ubuntu-latest
     timeout-minutes: 30
@@ -160,7 +191,7 @@ jobs:
   all-tests-passed:
     # This allows us to have a branch protection rule for tests and deploys with matrix
     runs-on: ubuntu-latest
-    needs: [ci, linting]
+    needs: [ci, linting, type-checking]
     if: always()
     steps:
       - name: Successful deploy

.pre-commit-config.yaml

@@ -2,12 +2,34 @@ repos:
   - repo: https://github.com/astral-sh/ruff-pre-commit
     rev: v0.3.2
     hooks:
-      # Run the linter and fix 
+      # Run the linter and fix
       - id: ruff
-        types_or: [ python, pyi, jupyter ]
-        args: [ --fix,  --config=ci.ruff.toml ]
+        types_or: [python, pyi, jupyter]
+        args: [--fix, --config=ci.ruff.toml]
 
       # Run the formatter.
       - id: ruff-format
-        types_or: [ python, pyi, jupyter ]
+        types_or: [python, pyi, jupyter]
 
+  - repo: https://github.com/pre-commit/mirrors-mypy
+    rev: v1.8.0
+    hooks:
+      - id: mypy
+        additional_dependencies:
+          - types-requests
+          - types-setuptools
+          - httpx
+          - pydantic>=1.10.7
+          - backoff>=1.10.0
+          - openai>=0.27.8
+          - wrapt
+          - packaging>=23.2
+          - opentelemetry-api
+          - opentelemetry-sdk
+          - opentelemetry-exporter-otlp
+          - numpy
+          - langchain>=0.0.309
+          - langchain-core
+          - langgraph
+        args: [--no-error-summary]
+        files: ^langfuse/

langfuse/_client/attributes.py

@@ -12,7 +12,7 @@
 
 import json
 from datetime import datetime
-from typing import Any, Dict, List, Literal, Optional
+from typing import Any, Dict, List, Literal, Optional, Union
 
 from langfuse._utils.serializer import EventSerializer
 from langfuse.model import PromptClient
@@ -68,7 +68,7 @@ def create_trace_attributes(
     metadata: Optional[Any] = None,
     tags: Optional[List[str]] = None,
     public: Optional[bool] = None,
-):
+) -> dict:
     attributes = {
         LangfuseOtelSpanAttributes.TRACE_NAME: name,
         LangfuseOtelSpanAttributes.TRACE_USER_ID: user_id,
@@ -93,7 +93,7 @@ def create_span_attributes(
     level: Optional[SpanLevel] = None,
     status_message: Optional[str] = None,
     version: Optional[str] = None,
-):
+) -> dict:
     attributes = {
         LangfuseOtelSpanAttributes.OBSERVATION_TYPE: "span",
         LangfuseOtelSpanAttributes.OBSERVATION_LEVEL: level,
@@ -122,7 +122,7 @@ def create_generation_attributes(
     usage_details: Optional[Dict[str, int]] = None,
     cost_details: Optional[Dict[str, float]] = None,
     prompt: Optional[PromptClient] = None,
-):
+) -> dict:
     attributes = {
         LangfuseOtelSpanAttributes.OBSERVATION_TYPE: "generation",
         LangfuseOtelSpanAttributes.OBSERVATION_LEVEL: level,
@@ -151,20 +151,20 @@ def create_generation_attributes(
     return {k: v for k, v in attributes.items() if v is not None}
 
 
-def _serialize(obj):
+def _serialize(obj: Any) -> Optional[str]:
     return json.dumps(obj, cls=EventSerializer) if obj is not None else None
 
 
 def _flatten_and_serialize_metadata(
     metadata: Any, type: Literal["observation", "trace"]
-):
+) -> dict:
     prefix = (
         LangfuseOtelSpanAttributes.OBSERVATION_METADATA
         if type == "observation"
         else LangfuseOtelSpanAttributes.TRACE_METADATA
     )
 
-    metadata_attributes = {}
+    metadata_attributes: Dict[str, Union[str, int, None]] = {}
 
     if not isinstance(metadata, dict):
         metadata_attributes[prefix] = _serialize(metadata)