Skip to content

Commit 5401641

Browse files
arre-ankitarre-ankit
committed
📦 NEW: Agent Run Tests
1 parent 8dc6990 commit 5401641

File tree

2 files changed

+291
-0
lines changed

2 files changed

+291
-0
lines changed

‎tests/conftest.py‎

Lines changed: 91 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -295,6 +295,97 @@ def mock_responses():
295295
},
296296
"system_fingerprint": "fp_agent1234567890",
297297
},
298+
# Agent run response with structured output
299+
"agent.run.structured": {
300+
"completion": '{"steps": [{"explanation": "Subtract 22 from both sides", "output": "8x = -45"}], "final_answer": "x = -5.625"}',
301+
"output": '{"steps": [{"explanation": "Subtract 22 from both sides", "output": "8x = -45"}], "final_answer": "x = -5.625"}',
302+
"thread_id": "thread_struct123",
303+
"id": "chatcmpl-struct123",
304+
"object": "chat.completion",
305+
"created": timestamp,
306+
"model": "gpt-4",
307+
"choices": [
308+
{
309+
"index": 0,
310+
"message": {
311+
"role": "assistant",
312+
"content": '{"steps": [{"explanation": "Subtract 22 from both sides", "output": "8x = -45"}], "final_answer": "x = -5.625"}',
313+
},
314+
"logprobs": None,
315+
"finish_reason": "stop",
316+
}
317+
],
318+
"usage": {
319+
"prompt_tokens": 60,
320+
"completion_tokens": 80,
321+
"total_tokens": 140,
322+
},
323+
"system_fingerprint": "fp_struct1234567890",
324+
},
325+
# Agent run response with tool calls
326+
"agent.run.tool": {
327+
"completion": None,
328+
"output": None,
329+
"thread_id": "thread_tool123",
330+
"id": "chatcmpl-tool123",
331+
"object": "chat.completion",
332+
"created": timestamp,
333+
"model": "gpt-4-mini",
334+
"choices": [
335+
{
336+
"index": 0,
337+
"message": {
338+
"role": "assistant",
339+
"content": None,
340+
"tool_calls": [
341+
{
342+
"id": "call_123456789",
343+
"type": "function",
344+
"function": {
345+
"name": "send_email",
346+
"arguments": '{"from": "[email protected]", "to": "[email protected]", "subject": "Welcome to Langbase!", "html": "Hello Sam! Welcome to Langbase.", "text": "Hello Sam! Welcome to Langbase."}',
347+
},
348+
}
349+
],
350+
},
351+
"logprobs": None,
352+
"finish_reason": "tool_calls",
353+
}
354+
],
355+
"usage": {
356+
"prompt_tokens": 70,
357+
"completion_tokens": 50,
358+
"total_tokens": 120,
359+
},
360+
"system_fingerprint": "fp_tool1234567890",
361+
},
362+
# Agent run final response after tool execution
363+
"agent.run.tool.final": {
364+
"completion": "✅ Email sent successfully to [email protected]!",
365+
"output": "✅ Email sent successfully to [email protected]!",
366+
"thread_id": "thread_tool123",
367+
"id": "chatcmpl-toolfinal123",
368+
"object": "chat.completion",
369+
"created": timestamp,
370+
"model": "gpt-4-mini",
371+
"choices": [
372+
{
373+
"index": 0,
374+
"message": {
375+
"role": "assistant",
376+
"content": "✅ Email sent successfully to [email protected]!",
377+
},
378+
"logprobs": None,
379+
"finish_reason": "stop",
380+
}
381+
],
382+
"usage": {
383+
"prompt_tokens": 90,
384+
"completion_tokens": 20,
385+
"total_tokens": 110,
386+
},
387+
"system_fingerprint": "fp_toolfinal1234567890",
388+
},
298389
# Error responses
299390
"error_400": {"error": "Bad request", "message": "Invalid parameters"},
300391
"error_401": {"error": "Unauthorized", "message": "Invalid API key"},

‎tests/test_utilities.py‎

Lines changed: 200 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -251,3 +251,203 @@ def test_agent_run_streaming(self, langbase_client, stream_chunks):
251251
request = responses.calls[0].request
252252
validate_response_headers(request.headers, AUTH_AND_JSON_CONTENT_HEADER)
253253
assert json.loads(request.body) == request_body
254+
255+
@responses.activate
256+
def test_agent_run_structured_output(self, langbase_client, mock_responses):
257+
"""Test agent.run method with structured output."""
258+
math_reasoning_schema = {
259+
"type": "object",
260+
"properties": {
261+
"steps": {
262+
"type": "array",
263+
"items": {
264+
"type": "object",
265+
"properties": {
266+
"explanation": {"type": "string"},
267+
"output": {"type": "string"},
268+
},
269+
"required": ["explanation", "output"],
270+
},
271+
},
272+
"final_answer": {"type": "string"},
273+
},
274+
"required": ["steps", "final_answer"],
275+
}
276+
277+
request_body = {
278+
"input": [{"role": "user", "content": "How can I solve 8x + 22 = -23?"}],
279+
"model": "openai:gpt-4.1",
280+
"apiKey": "test-openai-key",
281+
"instructions": "You are a helpful math tutor. Guide the user through the solution step by step.",
282+
"response_format": {
283+
"type": "json_schema",
284+
"json_schema": {
285+
"name": "math_reasoning",
286+
"schema": math_reasoning_schema,
287+
},
288+
},
289+
}
290+
291+
responses.add(
292+
responses.POST,
293+
f"{BASE_URL}{AGENT_RUN_ENDPOINT}",
294+
json=mock_responses["agent.run.structured"],
295+
status=200,
296+
)
297+
298+
result = langbase_client.agent.run(
299+
input=request_body["input"],
300+
model=request_body["model"],
301+
api_key=request_body["apiKey"],
302+
instructions=request_body["instructions"],
303+
response_format=request_body["response_format"],
304+
)
305+
306+
assert result == mock_responses["agent.run.structured"]
307+
assert len(responses.calls) == 1
308+
request = responses.calls[0].request
309+
validate_response_headers(request.headers, AUTH_AND_JSON_CONTENT_HEADER)
310+
assert json.loads(request.body) == request_body
311+
312+
# Verify structured output format
313+
assert "output" in result
314+
output_data = json.loads(result["output"])
315+
assert "steps" in output_data
316+
assert "final_answer" in output_data
317+
assert isinstance(output_data["steps"], list)
318+
assert len(output_data["steps"]) > 0
319+
320+
@responses.activate
321+
def test_agent_run_tool_call(self, langbase_client, mock_responses):
322+
"""Test agent.run method with tool calls."""
323+
send_email_tool_schema = {
324+
"type": "function",
325+
"function": {
326+
"name": "send_email",
327+
"description": "Send an email using Resend API",
328+
"parameters": {
329+
"type": "object",
330+
"required": ["from", "to", "subject", "html", "text"],
331+
"properties": {
332+
"from": {"type": "string"},
333+
"to": {"type": "string"},
334+
"subject": {"type": "string"},
335+
"html": {"type": "string"},
336+
"text": {"type": "string"},
337+
},
338+
"additionalProperties": False,
339+
},
340+
},
341+
}
342+
343+
request_body = {
344+
"input": [{"role": "user", "content": "Send a welcome email to Sam."}],
345+
"model": "openai:gpt-4.1-mini",
346+
"apiKey": "test-openai-key",
347+
"instructions": "You are an email agent. You are given a task to send an email to a recipient. You have the ability to send an email using the send_email tool.",
348+
"tools": [send_email_tool_schema],
349+
}
350+
351+
responses.add(
352+
responses.POST,
353+
f"{BASE_URL}{AGENT_RUN_ENDPOINT}",
354+
json=mock_responses["agent.run.tool"],
355+
status=200,
356+
)
357+
358+
result = langbase_client.agent.run(
359+
input=request_body["input"],
360+
model=request_body["model"],
361+
api_key=request_body["apiKey"],
362+
instructions=request_body["instructions"],
363+
tools=request_body["tools"],
364+
stream=False,
365+
)
366+
367+
assert result == mock_responses["agent.run.tool"]
368+
assert len(responses.calls) == 1
369+
request = responses.calls[0].request
370+
validate_response_headers(request.headers, AUTH_AND_JSON_CONTENT_HEADER)
371+
assert json.loads(request.body) == request_body
372+
373+
# Verify tool call structure
374+
assert "choices" in result
375+
choices = result["choices"]
376+
assert len(choices) > 0
377+
message = choices[0]["message"]
378+
assert "tool_calls" in message
379+
tool_calls = message["tool_calls"]
380+
assert len(tool_calls) > 0
381+
382+
tool_call = tool_calls[0]
383+
assert tool_call["type"] == "function"
384+
assert "function" in tool_call
385+
function = tool_call["function"]
386+
assert function["name"] == "send_email"
387+
assert "arguments" in function
388+
389+
@responses.activate
390+
def test_agent_run_tool_call_final_response(self, langbase_client, mock_responses):
391+
"""Test agent.run method with tool call final response."""
392+
# Simulate messages after tool execution
393+
input_messages = [
394+
{"role": "user", "content": "Send a welcome email to Sam."},
395+
{
396+
"role": "assistant",
397+
"content": None,
398+
"tool_calls": [
399+
{
400+
"id": "call_123456789",
401+
"type": "function",
402+
"function": {
403+
"name": "send_email",
404+
"arguments": '{"from": "[email protected]", "to": "[email protected]", "subject": "Welcome to Langbase!", "html": "Hello Sam! Welcome to Langbase.", "text": "Hello Sam! Welcome to Langbase."}',
405+
},
406+
}
407+
],
408+
},
409+
{
410+
"role": "tool",
411+
"tool_call_id": "call_123456789",
412+
"name": "send_email",
413+
"content": "✅ Email sent successfully to [email protected]!",
414+
},
415+
]
416+
417+
request_body = {
418+
"input": input_messages,
419+
"model": "openai:gpt-4.1-mini",
420+
"apiKey": "test-openai-key",
421+
"instructions": "You are an email sending assistant. Confirm the email has been sent successfully.",
422+
}
423+
424+
responses.add(
425+
responses.POST,
426+
f"{BASE_URL}{AGENT_RUN_ENDPOINT}",
427+
json=mock_responses["agent.run.tool.final"],
428+
status=200,
429+
)
430+
431+
result = langbase_client.agent.run(
432+
input=request_body["input"],
433+
model=request_body["model"],
434+
api_key=request_body["apiKey"],
435+
instructions=request_body["instructions"],
436+
stream=False,
437+
)
438+
439+
assert result == mock_responses["agent.run.tool.final"]
440+
assert len(responses.calls) == 1
441+
request = responses.calls[0].request
442+
validate_response_headers(request.headers, AUTH_AND_JSON_CONTENT_HEADER)
443+
assert json.loads(request.body) == request_body
444+
445+
# Verify final response structure
446+
assert "output" in result
447+
assert result["output"] == "✅ Email sent successfully to [email protected]!"
448+
assert "choices" in result
449+
choices = result["choices"]
450+
assert len(choices) > 0
451+
message = choices[0]["message"]
452+
assert message["role"] == "assistant"
453+
assert message["content"] == "✅ Email sent successfully to [email protected]!"

0 commit comments

Comments
 (0)