Update agent.py

This commit is contained in:
ochafik 2024-03-29 16:19:05 +00:00
parent e874565a13
commit b63f91ade4

View file

@ -24,7 +24,7 @@ def _get_params_schema(fn: Callable, verbose):
def completion_with_tool_usage( def completion_with_tool_usage(
*, *,
response_model: Optional[Union[Json, Type]]=None, response_model: Optional[Union[Json, Type]]=None,
max_tool_iterations: Optional[int]=None, max_iterations: Optional[int]=None,
tools: List[Callable], tools: List[Callable],
endpoint: str, endpoint: str,
messages: List[Message], messages: List[Message],
@ -61,7 +61,7 @@ def completion_with_tool_usage(
] ]
i = 0 i = 0
while (max_tool_iterations is None or i < max_tool_iterations): while (max_iterations is None or i < max_iterations):
request = ChatCompletionRequest( request = ChatCompletionRequest(
messages=messages, messages=messages,
response_format=response_format, response_format=response_format,
@ -117,8 +117,8 @@ def completion_with_tool_usage(
i += 1 i += 1
if max_tool_iterations is not None: if max_iterations is not None:
raise Exception(f"Failed to get a valid response after {max_tool_iterations} tool calls") raise Exception(f"Failed to get a valid response after {max_iterations} tool calls")
def main( def main(
@ -205,7 +205,7 @@ def main(
model="...", model="...",
endpoint=endpoint, endpoint=endpoint,
response_model=response_model, response_model=response_model,
max_tool_iterations=max_tool_iterations, max_iterations=max_iterations,
tools=tool_functions, tools=tool_functions,
auth=auth, auth=auth,
verbose=verbose, verbose=verbose,