Skip to content

Commit

Permalink
Fix returns of OpenAI interface (#108)
Browse files Browse the repository at this point in the history
fix `BaseAPIModel` chat returns

Co-authored-by: wangzy <[email protected]>
  • Loading branch information
braisedpork1964 and wangzy authored Jan 29, 2024
1 parent 39e00f2 commit b5533b0
Showing 1 changed file with 4 additions and 3 deletions.
7 changes: 4 additions & 3 deletions lagent/llms/openai.py
Original file line number Diff line number Diff line change
Expand Up @@ -99,7 +99,7 @@ def chat(
self,
inputs: Union[List[dict], List[List[dict]]],
**gen_params,
) -> List[str]:
) -> Union[str, List[str]]:
"""Generate responses given the contexts.
Args:
Expand All @@ -108,7 +108,7 @@ def chat(
gen_params: additional generation configuration
Returns:
List[str]: A list of generated strings.
Union[str, List[str]]: generated string(s)
"""
assert isinstance(inputs, list)
if isinstance(inputs[0], dict):
Expand All @@ -120,7 +120,8 @@ def chat(
for messages in inputs
]
wait(tasks)
return [task.result() for task in tasks]
ret = [task.result() for task in tasks]
return ret[0] if isinstance(inputs[0], dict) else ret

def _chat(self, messages: List[dict], **gen_params) -> str:
"""Generate completion from a list of templates.
Expand Down

0 comments on commit b5533b0

Please sign in to comment.