rename our completion as a drop-in replacement of litellm completion (#2509)

This commit is contained in:
Engel Nyst
2024-06-19 05:25:25 +02:00
committed by GitHub
parent b2307db010
commit 80fe13f4be
10 changed files with 30 additions and 28 deletions

View File

@@ -38,7 +38,7 @@ class SWEAgent(Agent):
self.cur_line: int = 0
def _think_act(self, messages: list[dict]) -> tuple[Action, str]:
resp = self.llm.do_completion(
resp = self.llm.completion(
messages=messages,
temperature=0.05,
)

View File

@@ -208,7 +208,7 @@ class BrowsingAgent(Agent):
prompt = get_prompt(error_prefix, cur_axtree_txt, prev_action_str)
messages.append({'role': 'user', 'content': prompt})
logger.info(prompt)
response = self.llm.do_completion(
response = self.llm.completion(
messages=messages,
temperature=0.0,
stop=[')```', ')\n```'],

View File

@@ -221,7 +221,7 @@ class CodeActAgent(Agent):
f'\n\nENVIRONMENT REMINDER: You have {state.max_iterations - state.iteration} turns left to complete the task.'
)
response = self.llm.do_completion(
response = self.llm.completion(
messages=messages,
stop=[
'</execute_ipython>',

View File

@@ -173,7 +173,7 @@ class CodeActSWEAgent(Agent):
f'\n\nENVIRONMENT REMINDER: You have {state.max_iterations - state.iteration} turns left to complete the task.'
)
response = self.llm.do_completion(
response = self.llm.completion(
messages=messages,
stop=[
'</execute_ipython>',

View File

@@ -64,7 +64,7 @@ class MicroAgent(Agent):
latest_user_message=state.get_current_user_intent(),
)
messages = [{'content': prompt, 'role': 'user'}]
resp = self.llm.do_completion(messages=messages)
resp = self.llm.completion(messages=messages)
action_resp = resp['choices'][0]['message']['content']
state.num_of_chars += len(prompt) + len(action_resp)
action = parse_response(action_resp)

View File

@@ -181,7 +181,7 @@ class MonologueAgent(Agent):
]
# format all as a single message, a monologue
resp = self.llm.do_completion(messages=messages)
resp = self.llm.completion(messages=messages)
# keep track of max_chars fallback option
state.num_of_chars += len(prompt) + len(

View File

@@ -47,7 +47,7 @@ class PlannerAgent(Agent):
return AgentFinishAction()
prompt = get_prompt(state)
messages = [{'content': prompt, 'role': 'user'}]
resp = self.llm.do_completion(messages=messages)
resp = self.llm.completion(messages=messages)
state.num_of_chars += len(prompt) + len(
resp['choices'][0]['message']['content']
)

View File

@@ -189,17 +189,31 @@ class LLM:
after=attempt_on_error,
)
def wrapper(*args, **kwargs):
"""
Wrapper for the litellm completion function. Logs the input and output of the completion function.
"""
# some callers might just send the messages directly
if 'messages' in kwargs:
messages = kwargs['messages']
else:
messages = args[1]
# log the prompt
debug_message = ''
for message in messages:
debug_message += message_separator + message['content']
llm_prompt_logger.debug(debug_message)
# call the completion function
resp = completion_unwrapped(*args, **kwargs)
# log the response
message_back = resp['choices'][0]['message']['content']
llm_response_logger.debug(message_back)
# post-process to log costs
self._post_completion(resp)
return resp
self._completion = wrapper # type: ignore
@@ -208,20 +222,12 @@ class LLM:
def completion(self):
"""
Decorator for the litellm completion function.
"""
return self._completion
def do_completion(self, *args, **kwargs):
"""
Wrapper for the litellm completion function.
Check the complete documentation at https://litellm.vercel.app/docs/completion
"""
resp = self._completion(*args, **kwargs)
self.post_completion(resp)
return resp
return self._completion
def post_completion(self, response: str) -> None:
def _post_completion(self, response: str) -> None:
"""
Post-process the completion response.
"""

View File

@@ -16,7 +16,7 @@ class MemoryCondenser:
try:
messages = [{'content': summarize_prompt, 'role': 'user'}]
resp = llm.do_completion(messages=messages)
resp = llm.completion(messages=messages)
summary_response = resp['choices'][0]['message']['content']
return summary_response
except Exception as e:

View File

@@ -35,9 +35,7 @@ def test_coder_agent_with_summary():
"""
mock_llm = MagicMock()
content = json.dumps({'action': 'finish', 'args': {}})
mock_llm.do_completion.return_value = {
'choices': [{'message': {'content': content}}]
}
mock_llm.completion.return_value = {'choices': [{'message': {'content': content}}]}
coder_agent = Agent.get_cls('CoderAgent')(llm=mock_llm)
assert coder_agent is not None
@@ -49,8 +47,8 @@ def test_coder_agent_with_summary():
state = State(history=history, inputs={'summary': summary})
coder_agent.step(state)
mock_llm.do_completion.assert_called_once()
_, kwargs = mock_llm.do_completion.call_args
mock_llm.completion.assert_called_once()
_, kwargs = mock_llm.completion.call_args
prompt = kwargs['messages'][0]['content']
assert task in prompt
assert "Here's a summary of the codebase, as it relates to this task" in prompt
@@ -64,9 +62,7 @@ def test_coder_agent_without_summary():
"""
mock_llm = MagicMock()
content = json.dumps({'action': 'finish', 'args': {}})
mock_llm.do_completion.return_value = {
'choices': [{'message': {'content': content}}]
}
mock_llm.completion.return_value = {'choices': [{'message': {'content': content}}]}
coder_agent = Agent.get_cls('CoderAgent')(llm=mock_llm)
assert coder_agent is not None
@@ -77,8 +73,8 @@ def test_coder_agent_without_summary():
state = State(history=history)
coder_agent.step(state)
mock_llm.do_completion.assert_called_once()
_, kwargs = mock_llm.do_completion.call_args
mock_llm.completion.assert_called_once()
_, kwargs = mock_llm.completion.call_args
prompt = kwargs['messages'][0]['content']
assert task in prompt
assert "Here's a summary of the codebase, as it relates to this task" not in prompt