phoenix/litellm
1
0
Fork 0
forked from phoenix/litellm-mirror
Code Pull requests Activity

test(test_function_call_parsing.py): fix test

Browse source
This commit is contained in:
Krrish Dholakia 2024-08-14 09:59:13 -07:00
parent 2b7a64ee28
commit edbe9e0741
1 changed files with 57 additions and 51 deletions
Download patch file Download diff file Expand all files Collapse all files

108
litellm/tests/test_function_call_parsing.py
View file

@ -1,23 +1,27 @@
# What is this? # What is this?
## Test to make sure function call response always works with json.loads() -> no extra parsing required. Relevant issue - https://github.com/BerriAI/litellm/issues/2654 ## Test to make sure function call response always works with json.loads() -> no extra parsing required. Relevant issue - https://github.com/BerriAI/litellm/issues/2654
import sys, os import os
import sys
import traceback import traceback
from dotenv import load_dotenv from dotenv import load_dotenv
load_dotenv() load_dotenv()
import os, io import io
import os
sys.path.insert( sys.path.insert(
0, os.path.abspath("../..") 0, os.path.abspath("../..")
) # Adds the parent directory to the system path ) # Adds the parent directory to the system path
import pytest
import litellm
import json import json
import warnings import warnings
from litellm import completion
from typing import List from typing import List
import pytest
import litellm
from litellm import completion
# Just a stub to keep the sample code simple # Just a stub to keep the sample code simple
class Trade: class Trade:
@ -78,58 +82,60 @@ def trade(model_name: str) -> List[Trade]:
}, },
} }
response = completion( try:
model_name, response = completion(
[ model_name,
{ [
"role": "system", {
"content": """You are an expert asset manager, managing a portfolio. "role": "system",
"content": """You are an expert asset manager, managing a portfolio.
Always use the `trade` function. Make sure that you call it correctly. For example, the following is a valid call: Always use the `trade` function. Make sure that you call it correctly. For example, the following is a valid call:
```
trade({
"orders": [
{"action": "buy", "asset": "BTC", "amount": 0.1},
{"action": "sell", "asset": "ETH", "amount": 0.2}
]
})
```
If there are no trades to make, call `trade` with an empty array:
```
trade({ "orders": [] })
```
""",
},
{
"role": "user",
"content": """Manage the portfolio.
Don't jabber.
This is the current market data:
``` ```
trade({ {market_data}
"orders": [
{"action": "buy", "asset": "BTC", "amount": 0.1},
{"action": "sell", "asset": "ETH", "amount": 0.2}
]
})
``` ```
If there are no trades to make, call `trade` with an empty array: Your portfolio is as follows:
``` ```
trade({ "orders": [] }) {portfolio}
``` ```
""", """.replace(
"{market_data}", "BTC: 64,000 USD\nETH: 3,500 USD"
).replace(
"{portfolio}", "USD: 1000, BTC: 0.1, ETH: 0.2"
),
},
],
tools=[tool_spec],
tool_choice={
"type": "function",
"function": {"name": tool_spec["function"]["name"]}, # type: ignore
}, },
{ )
"role": "user", except litellm.InternalServerError:
"content": """Manage the portfolio. pass
Don't jabber.
This is the current market data:
```
{market_data}
```
Your portfolio is as follows:
```
{portfolio}
```
""".replace(
"{market_data}", "BTC: 64,000 USD\nETH: 3,500 USD"
).replace(
"{portfolio}", "USD: 1000, BTC: 0.1, ETH: 0.2"
),
},
],
tools=[tool_spec],
tool_choice={
"type": "function",
"function": {"name": tool_spec["function"]["name"]}, # type: ignore
},
)
calls = response.choices[0].message.tool_calls calls = response.choices[0].message.tool_calls
trades = [trade for call in calls for trade in parse_call(call)] trades = [trade for call in calls for trade in parse_call(call)]
return trades return trades