mirror of https://github.com/aliasrobotics/cai.git
109 lines
3.4 KiB
Python
109 lines
3.4 KiB
Python
import pytest
|
|
import tempfile
|
|
import shutil
|
|
import statistics
|
|
from pathlib import Path
|
|
from typing import List, Union
|
|
from cai import CAI, Agent
|
|
|
|
|
|
class TestFunctionCallBenchmarksBasic:
|
|
@pytest.fixture(autouse=True)
|
|
def setup_teardown(self):
|
|
self.test_dir = tempfile.mkdtemp()
|
|
yield
|
|
shutil.rmtree(self.test_dir)
|
|
|
|
def _test_function_call(self, model_name: str, is_multi: bool) -> bool:
|
|
client = CAI()
|
|
|
|
test_agent = Agent(
|
|
model=model_name,
|
|
name="Function Call Demo Agent",
|
|
instructions="You are an agent that demonstrates function calling capabilities.",
|
|
parallel_tool_calls=is_multi
|
|
)
|
|
|
|
def add_numbers(a: Union[int, str], b: Union[int, str]) -> str:
|
|
try:
|
|
return str(int(a) + int(b))
|
|
except (ValueError, TypeError):
|
|
return "error"
|
|
|
|
def analyze_text(text: str) -> dict:
|
|
try:
|
|
words = str(text).split()
|
|
return {
|
|
"word_count": str(len(words)),
|
|
"char_count": str(len(str(text)))
|
|
}
|
|
except (ValueError, TypeError):
|
|
return {
|
|
"word_count": "error",
|
|
"char_count": "error"
|
|
}
|
|
|
|
test_agent.functions.extend([add_numbers, analyze_text])
|
|
|
|
if is_multi:
|
|
prompt = """Please:
|
|
1. Add 5 and 7
|
|
2. Analyze text: 'Hello world'"""
|
|
else:
|
|
prompt = "Add 5 and 7"
|
|
|
|
try:
|
|
response = client.run(agent=test_agent, messages=[
|
|
{"role": "user", "content": prompt}], debug=False)
|
|
success = "error" not in response.messages[-1]["content"]
|
|
test_type = "Multi" if is_multi else "Single"
|
|
if success:
|
|
print(f"\n✅ {test_type}-function test passed for {model_name}")
|
|
else:
|
|
print(f"\n❌ {test_type}-function test failed for {model_name}")
|
|
return success
|
|
except Exception as e:
|
|
print(f"\n❌ Test failed for {model_name}: {str(e)}")
|
|
return False
|
|
|
|
@pytest.mark.parametrize("model", [
|
|
"dwightfoster03/functionary-small-v3.1",
|
|
"llama3.1:8b",
|
|
"qwen2.5:14b",
|
|
"qwen2.5:32b",
|
|
"marco-o1:7b-fp16"
|
|
])
|
|
def test_function_call_benchmark(self, model):
|
|
results = []
|
|
print(f"\nTesting {model}")
|
|
single_results = []
|
|
multi_results = []
|
|
|
|
for i in range(3):
|
|
print(f"Iteration {i + 1}:")
|
|
single_results.append(self._test_function_call(model, False))
|
|
multi_results.append(self._test_function_call(model, True))
|
|
|
|
single_avg = statistics.mean(single_results)
|
|
multi_avg = statistics.mean(multi_results)
|
|
overall = (single_avg + multi_avg) / 2
|
|
|
|
print(f"Results for {model}:")
|
|
print(f"Single: {single_avg * 100:.1f}%")
|
|
print(f"Multi: {multi_avg * 100:.1f}%")
|
|
print(f"Overall: {overall * 100:.1f}%")
|
|
|
|
results.append({
|
|
"model": model,
|
|
"single": single_avg,
|
|
"multi": multi_avg
|
|
})
|
|
|
|
total = statistics.mean([
|
|
(r["single"] + r["multi"]) / 2
|
|
for r in results
|
|
])
|
|
|
|
print(f"\nOverall success rate: {total * 100:.1f}%")
|
|
assert total >= 0.5, "Success rate below 50%"
|