nits

Author: Andrei Bratu

src/humanloop/overload.py

@@ -264,7 +264,7 @@ def overload_client(
         def log_wrapper(self: T, **kwargs) -> LogResponseType:
             return _overload_log(self, file_syncer, use_local_files, **kwargs)
 
-        client.log = types.MethodType(log_wrapper, client)  # type: ignore [method-assign, union-attr, assignment]
+        client.log = types.MethodType(log_wrapper, client)  # type: ignore [method-assign]
 
     # Overload call method for Prompt and Agent clients
     if _get_file_type_from_client(client) in FileSyncer.SERIALIZABLE_FILE_TYPES:

tests/custom/integration/test_decorators.py

@@ -132,14 +132,17 @@ def test_flow_decorator_populates_output_message(
     try:
         humanloop_client = get_humanloop_client()
 
+        # GIVEN a flow that returns a ChatMessage like dict
         @humanloop_client.flow(path=f"{sdk_test_dir}/test_flow_log_output_message")
         def my_flow(question: str) -> dict[str, Any]:
             return {"role": "user", "content": question}
 
+        # WHEN the flow is called
         assert "france" in my_flow("What is the capital of the France?")["content"].lower()
 
         time.sleep(5)
 
+        # THEN the Flow is created and the Log has output_message populated
         flow_response = humanloop_client.files.retrieve_by_path(path=f"{sdk_test_dir}/test_flow_log_output_message")
         assert flow_response is not None
         flow_logs_response = humanloop_client.logs.list(file_id=flow_response.id, page=1, size=50)
@@ -161,12 +164,13 @@ async def test_async_flow_decorator(
 ):
     humanloop_client = get_humanloop_client()
 
+    # GIVEN an async flow that returns a string
     @humanloop_client.a_flow(path=f"{sdk_test_dir}/test_async_flow")
     async def my_flow(question: str) -> str:
         return "baz!"
 
-    # THEN the output is the one expected
-    assert "baz!" in await my_flow("test")
+    # WHEN the flow is called
+    assert "baz!" == await my_flow("test")
 
     # Wait for the flow and log to propagate to integration backend
     await asyncio.sleep(3)
@@ -191,12 +195,15 @@ async def test_async_tool_decorator(
 ):
     humanloop_client = get_humanloop_client()
 
+    # GIVEN an async tool that returns a string
     @humanloop_client.a_tool(path=f"{sdk_test_dir}/test_async_tool")
     async def my_tool(question: str) -> str:
         return "baz!"
 
+    # THEN the tool has a json_schema
     assert hasattr(my_tool, "json_schema")
 
+    # WHEN the tool is called
     await my_tool("test")
 
     # Wait for the flow and log to propagate to integration backend
@@ -223,6 +230,7 @@ async def test_async_prompt_decorator(
 ):
     humanloop_client = get_humanloop_client()
 
+    # GIVEN an async prompt that calls OpenAI
     @humanloop_client.a_prompt(path=f"{sdk_test_dir}/test_async_prompt")
     async def my_prompt(question: str) -> str:
         openai_client = AsyncOpenAI(api_key=openai_key)
@@ -235,7 +243,8 @@ async def my_prompt(question: str) -> str:
         assert response.choices[0].message.content is not None
         return response.choices[0].message.content
 
-    # THEN the output is the one expected
+    # WHEN the prompt is called
+    # THEN the output is not null
     output = await my_prompt("What is the capital of the France?")
     assert output is not None
 
@@ -249,6 +258,7 @@ async def my_prompt(question: str) -> str:
     # THEN a Log exists on the File
     prompt_logs_response = humanloop_client.logs.list(file_id=prompt_file_response.id, page=1, size=50)
     assert prompt_logs_response.items is not None and len(prompt_logs_response.items) == 1
+    # THEN output_message matches the one intercepted from OpenAI response
     assert prompt_logs_response.items[0].output_message.content == output  # type: ignore [union-attr]
 
 
@@ -259,6 +269,7 @@ async def test_async_flow_decorator_with_trace(
 ):
     humanloop_client = get_humanloop_client()
 
+    # GIVEN async flow and prompt decorators
     @humanloop_client.a_prompt(path=f"{sdk_test_dir}/test_async_prompt_with_trace")
     async def my_prompt(question: str) -> str:
         openai_client = AsyncOpenAI(api_key=openai_key)
@@ -273,20 +284,26 @@ async def my_prompt(question: str) -> str:
 
     @humanloop_client.a_flow(path=f"{sdk_test_dir}/test_async_flow_with_trace")
     async def my_flow(question: str) -> str:
-        return await my_prompt("test")
+        return await my_prompt(question="test")
 
+    # WHEN the flow is called
     await my_flow("test")
 
     # Wait for the flow and log to propagate to integration backend
     await asyncio.sleep(3)
 
+    # THEN both files exist on Humanloop
     flow_file_response = humanloop_client.files.retrieve_by_path(path=f"{sdk_test_dir}/test_async_flow_with_trace")
     assert flow_file_response is not None
+    prompt_file_response = humanloop_client.files.retrieve_by_path(path=f"{sdk_test_dir}/test_async_prompt_with_trace")
+    assert prompt_file_response is not None
 
+    # THEN a Log exists on the File
     flow_logs_response = humanloop_client.logs.list(file_id=flow_file_response.id, page=1, size=50)
     assert flow_logs_response.items is not None and len(flow_logs_response.items) == 1
     assert flow_logs_response.items[0].output is not None and flow_logs_response.items[0].inputs == {"question": "test"}
     flow_log_with_trace_response = humanloop_client.logs.get(id=flow_logs_response.items[0].id)
+    # THEN a Prompt Log is added to the Flow Log trace
     assert (
         flow_log_with_trace_response["trace_children"] is not None  # type: ignore [index]
         and len(flow_log_with_trace_response["trace_children"]) == 1  # type: ignore [index]