type checks

Author: Andrei Bratu

poetry.lock

@@ -1,3 +1,6 @@
+[project]
+name = "humanloop"
+
 [tool.poetry]
 name = "humanloop"
 version = "0.8.24"
@@ -27,12 +30,7 @@ packages = [
     { include = "humanloop", from = "src"}
 ]
 
-[tool.poetry.group.dev.dependencies]
-python-dotenv = "^1.0.1"
 
-chromadb = "^0.6.3"
-pandas = "^2.2.3"
-pyarrow = "^19.0.0"
 [project.urls]
 Repository = 'https://github.com/humanloop/humanloop-python'
 
@@ -70,6 +68,9 @@ python-dotenv = "^1.0.1"
 replicate = "^1.0.3"
 ruff = "^0.5.6"
 types-jsonschema = "^4.23.0.20240813"
+chromadb="^0.3.8"
+pandas = "^2.2.3"
+pyarrow = "^19.0.0"
 
 [tool.pytest.ini_options]
 testpaths = [ "tests" ]
@@ -84,4 +85,4 @@ line-length = 120
 
 [build-system]
 requires = ["poetry-core"]
-build-backend = "poetry.core.masonry.api"
+build-backend = "poetry.core.masonry.api"

pyproject.toml

@@ -5,7 +5,7 @@
 from queue import Empty as EmptyQueue
 from queue import Queue
 from threading import Thread
-from typing import Any, Optional
+from typing import Any, Optional, Sequence
 
 from opentelemetry.sdk.trace import ReadableSpan
 from opentelemetry.sdk.trace.export import SpanExporter, SpanExportResult
@@ -94,7 +94,7 @@ def __init__(
         # Flow Log Span ID mapping to children Spans that must be uploaded first
         self._spans_left_in_trace: dict[int, set[int]] = {}
 
-    def export(self, spans: list[ReadableSpan]) -> SpanExportResult:
+    def export(self, spans: Sequence[ReadableSpan]) -> SpanExportResult:
         if self._shutdown:
             logger.warning("[HumanloopSpanExporter] Shutting down, not accepting new spans")
             return SpanExportResult.FAILURE
@@ -205,15 +205,14 @@ def _export_span_dispatch(self, span: ReadableSpan) -> None:
             file_type,
         )
 
-        match file_type:
-            case "prompt":
-                self._export_prompt_span(span=span)
-            case "tool":
-                self._export_tool_span(span=span)
-            case "flow":
-                self._export_flow_span(span=span)
-            case _:
-                raise NotImplementedError(f"Unknown span type: {hl_file}")
+        if file_type == "prompt":
+            self._export_prompt_span(span=span)
+        if file_type == "tool":
+            self._export_tool_span(span=span)
+        if file_type == "flow":
+            self._export_flow_span(span=span)
+        else:
+            raise NotImplementedError(f"Unknown span type: {hl_file}")
 
     def _export_prompt_span(self, span: ReadableSpan) -> None:
         file_object: dict[str, Any] = read_from_opentelemetry_span(
@@ -305,7 +304,7 @@ def _export_flow_span(self, span: ReadableSpan) -> None:
         # passed by the Processor. Each uploaded child in the trace
         # will check if it's the last one and mark the Flow Log as complete
         try:
-            prerequisites: list[int] = read_from_opentelemetry_span(
+            prerequisites: list[int] = read_from_opentelemetry_span(  # type: ignore
                 span=span,
                 key=HUMANLOOP_FLOW_PREREQUISITES_KEY,
             )

src/humanloop/otel/exporter.py

@@ -244,7 +244,7 @@ def pseudo_to_list(sub_dict):
 
 def is_llm_provider_call(span: ReadableSpan) -> bool:
     """Determines if the span was created by an Instrumentor for LLM provider clients."""
-    if not hasattr(span, "instrumentation_scope"):
+    if not hasattr(span, "instrumentation_scope") or span.instrumentation_scope is None:
         return False
     span_instrumentor_name = span.instrumentation_scope.name
     # Match against the prefix of the Instrumentor name since

src/humanloop/otel/helpers.py

@@ -1,6 +1,7 @@
 from dataclasses import dataclass
 import logging
 from collections import defaultdict
+from typing import Optional
 
 from opentelemetry.sdk.trace import ReadableSpan
 from opentelemetry.sdk.trace.export import SimpleSpanProcessor, SpanExporter
@@ -81,7 +82,10 @@ def on_end(self, span: ReadableSpan) -> None:
         span_id = span.context.span_id
         if is_humanloop_span(span=span):
             if not self._must_wait(span):
-                self._send_to_exporter(span, self._dependencies[span.context.span_id])
+                self._send_to_exporter(
+                    span=span,
+                    dependencies=[dependency.span for dependency in self._dependencies[span.context.span_id]],
+                )
             else:
                 # Must wait for dependencies
                 self._waiting[span_id] = span
@@ -93,7 +97,10 @@ def on_end(self, span: ReadableSpan) -> None:
 
             waiting_span = self._get_waiting_parent(span)
             if waiting_span is not None:
-                self._send_to_exporter(span, self._dependencies[span])
+                self._send_to_exporter(
+                    span=span,
+                    dependencies=[dependency.span for dependency in self._dependencies[span.context.span_id]],
+                )
             return
 
         # Be unopinionated and pass all other spans to Exporter
@@ -106,18 +113,20 @@ def _must_wait(self, span: ReadableSpan) -> bool:
             return False
         return True
 
-    def _get_waiting_parent(self, span: ReadableSpan) -> ReadableSpan | None:
-        parent_span_id = span.parent.span_id
+    def _get_waiting_parent(self, span: ReadableSpan) -> Optional[ReadableSpan]:
+        # We know this span has a parent, need to satisfy the type checker
+        parent_span_id = span.parent.span_id  # type: ignore
         if parent_span_id in self._waiting:
             if all([dependency.finished for dependency in self._dependencies[parent_span_id]]):
                 waiting_span = self._waiting[parent_span_id]
                 del self._dependencies[parent_span_id]
-                del waiting_span[parent_span_id]
+                del self._waiting[parent_span_id]
                 return waiting_span
         return None
 
     def _add_dependency_to_await(self, span: ReadableSpan):
-        parent_span_id = span.parent.span_id if span.parent else None
+        # We know this span has a parent, need to satisfy the type checker
+        parent_span_id = span.parent.span_id  # type: ignore
         if self._is_dependency(span):
             self._dependencies[parent_span_id].append(DependantSpan(span=span, finished=False))
 
@@ -138,7 +147,8 @@ def _track_flow_traces(self, span: ReadableSpan):
 
     def _mark_dependency_arrival(self, span: ReadableSpan):
         span_id = span.context.span_id
-        parent_span_id = span.parent.span_id
+        # We know this span has a parent, need to satisfy type checker
+        parent_span_id = span.parent.span_id  # type: ignore
         self._dependencies[parent_span_id] = [
             dependency if dependency.span.context.span_id != span_id else DependantSpan(span=span, finished=True)
             for dependency in self._dependencies[parent_span_id]
@@ -160,28 +170,27 @@ def _send_to_exporter(
             # Processing specific to each Humanloop File type
             file_type = span.attributes[HUMANLOOP_FILE_TYPE_KEY]  # type: ignore
             span_id = span.context.span_id
-            match file_type:
-                case "prompt":
-                    enhance_prompt_span(
-                        prompt_span=span,
-                        dependencies=[dependency.span for dependency in dependencies],
-                    )
-                case "tool":
-                    # No extra processing needed
-                    pass
-                case "flow":
-                    trace = self._spans_to_complete_flow_trace.get(span_id, [])
-                    write_to_opentelemetry_span(
-                        span=span,
-                        key=HUMANLOOP_FLOW_PREREQUISITES_KEY,
-                        value=trace,
-                    )
-                case _:
-                    logger.error(
-                        "[HumanloopSpanProcessor] Unknown Humanloop File span %s %s",
-                        span_id,
-                        span.name,
-                    )
+            if file_type == "prompt":
+                enhance_prompt_span(
+                    prompt_span=span,
+                    dependencies=dependencies,
+                )
+            elif file_type == "tool":
+                # No extra processing needed
+                pass
+            elif file_type == "flow":
+                trace = self._spans_to_complete_flow_trace.get(span_id, [])
+                write_to_opentelemetry_span(
+                    span=span,
+                    key=HUMANLOOP_FLOW_PREREQUISITES_KEY,
+                    value=trace,
+                )
+            else:
+                logger.error(
+                    "[HumanloopSpanProcessor] Unknown Humanloop File span %s %s",
+                    span_id,
+                    span.name,
+                )
 
         self.span_exporter.export([span])
 

src/humanloop/otel/processor/__init__.py

@@ -1,3 +1,4 @@
+import logging
 from typing import Any
 from opentelemetry.sdk.trace import ReadableSpan
 from pydantic import ValidationError as PydanticValidationError
@@ -10,6 +11,8 @@
 )
 from humanloop.types.prompt_kernel_request import PromptKernelRequest
 
+logger = logging.getLogger("humanloop.sdk")
+
 
 def enhance_prompt_span(prompt_span: ReadableSpan, dependencies: list[ReadableSpan]):
     """Add information from the LLM provider span to the Prompt span.
@@ -57,7 +60,7 @@ def _enrich_prompt_kernel(prompt_span: ReadableSpan, llm_provider_call_span: Rea
 
     try:
         # Validate the Prompt Kernel
-        PromptKernelRequest.model_validate(obj=prompt)
+        PromptKernelRequest.model_validate(obj=prompt)  # type: ignore
     except PydanticValidationError as e:
         logger.error(
             "[HumanloopSpanProcessor] Could not validate Prompt Kernel extracted from span: %s %s. Error: %s",

src/humanloop/otel/processor/prompts.py

@@ -1,5 +1,5 @@
 from dataclasses import dataclass
-from typing import Callable
+from typing import Any, Callable
 
 import pytest
 
@@ -44,7 +44,7 @@ def surfer_agent_scenario(
     client = OpenAI(api_key=api_keys.openai)
 
     @humanloop_client.tool(path=get_test_path("Calculator"))
-    def calculator(operation: str, num1: int, num2: int) -> str:
+    def calculator(operation: str, num1: int, num2: int) -> float:
         """Do arithmetic operations on two numbers."""
         if operation == "add":
             return num1 + num2
@@ -70,8 +70,8 @@ def pick_random_number():
             calculator.json_schema,
         ],
     )
-    def call_agent(messages: list[str]) -> str:
-        output = client.chat.completions.create(
+    def call_agent(messages: list[dict[str, Any]]) -> str:  # type: ignore [call-arg]
+        output = client.chat.completions.create(  # type: ignore [call-overload]
             model="gpt-4o-mini",
             messages=messages,
             # Use .json_schema property on decorated functions to easily access

tests/integration/chat_agent/conftest.py

@@ -8,7 +8,7 @@
 import time
 from typing import ContextManager, TextIO
 from unittest.mock import MagicMock, patch
-from src.humanloop import Humanloop
+from humanloop import Humanloop
 from tests.integration.chat_agent.conftest import SurferAgentScenario
 from tests.integration.conftest import DirectoryIdentifiers
 
@@ -27,7 +27,7 @@ def test_scenario_runs(
         "exit",
     ]
     mocked_input.side_effect = scenario_io
-    with capture_stdout() as console_output:
+    with capture_stdout() as console_output:  # type: ignore [operator]
         surfer_agent_scenario.agent_chat_workflow()
 
     time.sleep(5)
@@ -46,7 +46,9 @@ def test_scenario_runs(
     # List will not pass the children to the trace_children attribute
     assert len(flow_log.trace_children) == 0
     response = humanloop_client.logs.get(flow_log.id)
-    assert response["trace_status"] == "complete"
+    if not isinstance(response, dict):
+        response = response.dict()
+    assert response["trace_status"] == "complete"  # type: ignore [attr-defined]
     assert len(response["trace_children"]) == 2
     messages = response["trace_children"][1]["messages"]
     assert len(messages) == 4

tests/integration/chat_agent/test_chat_agent.py

@@ -10,7 +10,7 @@
 from typing import Generator, TextIO
 import pytest
 from dotenv import load_dotenv
-from src.humanloop import Humanloop
+from humanloop import Humanloop
 
 
 @dataclass
@@ -49,9 +49,13 @@ def root_integration_directory(humanloop_client: Humanloop) -> Generator[str, No
 
 @pytest.fixture(scope="session")
 def api_keys() -> APIKeys:
+    openai_key = os.getenv("OPENAI_API_KEY")
+    humanloop_key = os.getenv("HUMANLOOP_API_KEY")
+    if openai_key is None or humanloop_key is None:
+        raise ValueError("API keys are not set in .env file")
     api_keys = APIKeys(
-        openai=os.getenv("OPENAI_API_KEY"),
-        humanloop=os.getenv("HUMANLOOP_API_KEY"),
+        openai=openai_key,
+        humanloop=humanloop_key,
     )
     for key, value in asdict(api_keys).items():
         if value is None:
@@ -68,17 +72,18 @@ def _directory_cleanup(directory_id: str, humanloop_client: Humanloop):
     response = humanloop_client.directories.get(directory_id)
     for file in response.files:
         file_id = file.id
-        match file.type:
-            case "prompt":
-                client = humanloop_client.prompts
-            case "tool":
-                client = humanloop_client.tools
-            case "dataset":
-                client = humanloop_client.datasets
-            case "evaluator":
-                client = humanloop_client.evaluators
-            case "flow":
-                client = humanloop_client.flows
+        if file.type == "prompt":
+            client = humanloop_client.prompts  # type: ignore [assignment]
+        elif file.type == "tool":
+            client = humanloop_client.tools  # type: ignore [assignment]
+        elif file.type == "dataset":
+            client = humanloop_client.datasets  # type: ignore [assignment]
+        elif file.type == "evaluator":
+            client = humanloop_client.evaluators  # type: ignore [assignment]
+        elif file.type == "flow":
+            client = humanloop_client.flows  # type: ignore [assignment]
+        else:
+            raise NotImplementedError(f"Unknown HL file type {file.type}")
         client.delete(file_id)
 
     for subdirectory in response.subdirectories:

tests/integration/conftest.py

@@ -47,17 +47,17 @@ class MedQAScenario:
 
 @pytest.fixture()
 def evaluate_medqa_scenario(
-    humanloop_client: Humanloop,
+    humanloop_client: "Humanloop",
     get_test_path: Callable[[str], str],
     api_keys: APIKeys,
     medqa_knowledge_base_path: str,
     medqa_dataset_path: str,
 ) -> MedQAScenario:
     import inspect
 
-    from chromadb import chromadb
+    from chromadb import chromadb  # type: ignore
     from openai import OpenAI
-    import pandas as pd
+    import pandas as pd  # type: ignore
 
     chroma = chromadb.Client()
     collection = chroma.get_or_create_collection(name="MedQA")