Changed llm stream output logic and removed unnecessary outputs to UI

This commit is contained in:
matija-ilijas
2024-08-27 19:59:13 +02:00
parent eb184469f8
commit 85ff210c1c
19 changed files with 43 additions and 51 deletions

View File

@@ -126,7 +126,7 @@ class Architect(BaseAgent):
"""
await self.send_message("Selecting starter templates ...")
llm = self.get_llm()
llm = self.get_llm(stream_output=True)
convo = (
AgentConvo(self)
.template(
@@ -168,7 +168,7 @@ class Architect(BaseAgent):
await self.send_message("Picking technologies to use ...")
llm = self.get_llm()
llm = self.get_llm(stream_output=True)
convo = (
AgentConvo(self)
.template(
@@ -271,7 +271,7 @@ class Architect(BaseAgent):
# If template has no options, no need to ask LLM for anything
return NoOptions()
llm = self.get_llm()
llm = self.get_llm(stream_output=True)
convo = (
AgentConvo(self)
.template(

View File

@@ -34,7 +34,6 @@ class BaseAgent:
"""
self.ui_source = AgentSource(self.display_name, self.agent_type)
self.ui = ui
self.stream_output = True
self.state_manager = state_manager
self.process_manager = process_manager
self.prev_response = prev_response
@@ -106,13 +105,11 @@ class BaseAgent:
Handle streamed response from the LLM.
Serves as a callback to `AgentBase.llm()` so it can stream the responses to the UI.
This can be turned on/off on a pe-request basis by setting `BaseAgent.stream_output`
to True or False.
:param content: Response content.
"""
if self.stream_output:
await self.ui.send_stream_chunk(content, source=self.ui_source)
await self.ui.send_stream_chunk(content, source=self.ui_source)
if content is None:
await self.ui.send_message("", source=self.ui_source)
@@ -150,7 +147,7 @@ class BaseAgent:
return False
def get_llm(self, name=None) -> Callable:
def get_llm(self, name=None, stream_output=False) -> Callable:
"""
Get a new instance of the agent-specific LLM client.
@@ -170,7 +167,8 @@ class BaseAgent:
llm_config = config.llm_for_agent(name)
client_class = BaseLLMClient.for_provider(llm_config.provider)
llm_client = client_class(llm_config, stream_handler=self.stream_handler, error_handler=self.error_handler)
stream_handler = self.stream_handler if stream_output else None
llm_client = client_class(llm_config, stream_handler=stream_handler, error_handler=self.error_handler)
async def client(convo, **kwargs) -> Any:
"""

View File

@@ -66,7 +66,7 @@ class BugHunter(BaseAgent):
return await self.start_pair_programming()
async def get_bug_reproduction_instructions(self):
llm = self.get_llm()
llm = self.get_llm(stream_output=True)
convo = AgentConvo(self).template(
"get_bug_reproduction_instructions",
current_task=self.current_state.current_task,
@@ -79,7 +79,7 @@ class BugHunter(BaseAgent):
self.next_state.current_iteration["bug_reproduction_description"] = bug_reproduction_instructions
async def check_logs(self, logs_message: str = None):
llm = self.get_llm(CHECK_LOGS_AGENT_NAME)
llm = self.get_llm(CHECK_LOGS_AGENT_NAME, stream_output=True)
convo = self.generate_iteration_convo_so_far()
human_readable_instructions = await llm(convo, temperature=0.5)
@@ -91,17 +91,17 @@ class BugHunter(BaseAgent):
)
.require_schema(HuntConclusionOptions)
)
llm = self.get_llm()
llm = self.get_llm(stream_output=True)
hunt_conclusion = await llm(convo, parser=JSONParser(HuntConclusionOptions), temperature=0)
if hunt_conclusion.conclusion == magic_words.PROBLEM_IDENTIFIED:
# if no need for logs, implement iteration same as before
self.set_data_for_next_hunting_cycle(human_readable_instructions, IterationStatus.AWAITING_BUG_FIX)
await self.send_message("The bug is found - I'm attempting to fix it.")
await self.send_message("Found the bug - I'm attempting to fix it ...")
else:
# if logs are needed, add logging steps
self.set_data_for_next_hunting_cycle(human_readable_instructions, IterationStatus.AWAITING_LOGGING)
await self.send_message("Adding more logs to identify the bug.")
await self.send_message("Adding more logs to identify the bug ...")
self.next_state.flag_iterations_as_modified()
return AgentResponse.done(self)
@@ -147,7 +147,7 @@ class BugHunter(BaseAgent):
buttons["continue"] = "Continue"
buttons["done"] = "Bug is fixed"
backend_logs = await self.ask_question(
"Please do exactly what you did in the last iteration, paste **BACKEND** logs here and click CONTINUE.",
"Please do exactly what you did in the last iteration, paste the BACKEND logs here and click CONTINUE.",
buttons=buttons,
default="continue",
hint="Instructions for testing:\n\n"
@@ -161,7 +161,7 @@ class BugHunter(BaseAgent):
self.next_state.flag_iterations_as_modified()
else:
frontend_logs = await self.ask_question(
"Please paste **frontend** logs here and click CONTINUE.",
"Please paste the FRONTEND logs here and click CONTINUE.",
buttons={"continue": "Continue", "done": "Bug is fixed"},
default="continue",
hint="Instructions for testing:\n\n"
@@ -188,7 +188,7 @@ class BugHunter(BaseAgent):
return AgentResponse.done(self)
async def start_pair_programming(self):
llm = self.get_llm()
llm = self.get_llm(stream_output=True)
convo = self.generate_iteration_convo_so_far(True)
if len(convo.messages) > 1:
convo.remove_last_x_messages(1)
@@ -280,12 +280,12 @@ class BugHunter(BaseAgent):
human_hint = await self.ask_question(human_hint_label)
convo = convo.template("instructions_from_human_hint", human_hint=human_hint.text)
await self.ui.start_important_stream()
llm = self.get_llm(CHECK_LOGS_AGENT_NAME)
llm = self.get_llm(CHECK_LOGS_AGENT_NAME, stream_output=True)
human_readable_instructions = await llm(convo, temperature=0.5)
human_approval = await self.ask_question(
"Can I implement this solution?", buttons={"yes": "Yes", "no": "No"}, buttons_only=True
)
llm = self.get_llm()
llm = self.get_llm(stream_output=True)
if human_approval.button == "yes":
self.set_data_for_next_hunting_cycle(
human_readable_instructions, IterationStatus.AWAITING_BUG_FIX
@@ -344,7 +344,7 @@ class BugHunter(BaseAgent):
self.next_state.current_iteration["status"] = new_status
async def continue_on(self, convo, button_value, user_response):
llm = self.get_llm()
llm = self.get_llm(stream_output=True)
convo = convo.template("continue_on")
continue_on = await llm(convo, temperature=0.5)
return continue_on

View File

@@ -108,8 +108,6 @@ class CodeMonkey(BaseAgent):
continue
log.debug(f"Describing file {file.path}")
await self.send_message(f"Describing file {file.path} ...")
convo = (
AgentConvo(self)
.template(

View File

@@ -179,21 +179,21 @@ class CodeReviewer(BaseAgent):
)
if len(hunks_to_apply) == len(hunks):
await self.send_message("Applying entire change")
# await self.send_message("Applying entire change")
log.info(f"Applying entire change to {file_name}")
return new_content, None
elif len(hunks_to_apply) == 0:
if hunks_to_rework:
await self.send_message(
f"Requesting rework for {len(hunks_to_rework)} changes with reason: {llm_response.review_notes}"
)
# await self.send_message(
# f"Requesting rework for {len(hunks_to_rework)} changes with reason: {llm_response.review_notes}"
# )
log.info(f"Requesting rework for {len(hunks_to_rework)} changes to {file_name} (0 hunks to apply)")
return old_content, review_log
else:
# If everything can be safely ignored, it's probably because the files already implement the changes
# from previous tasks (which can happen often). Insisting on a change here is likely to cause problems.
await self.send_message(f"Rejecting entire change with reason: {llm_response.review_notes}")
# await self.send_message(f"Rejecting entire change with reason: {llm_response.review_notes}")
log.info(f"Rejecting entire change to {file_name} with reason: {llm_response.review_notes}")
return old_content, None

View File

@@ -217,7 +217,7 @@ class Developer(RelevantFilesMixin, BaseAgent):
current_task_index = self.current_state.tasks.index(current_task)
llm = self.get_llm(TASK_BREAKDOWN_AGENT_NAME)
llm = self.get_llm(TASK_BREAKDOWN_AGENT_NAME, stream_output=True)
convo = AgentConvo(self).template(
"breakdown",
task=current_task,
@@ -236,7 +236,6 @@ class Developer(RelevantFilesMixin, BaseAgent):
self.next_state.flag_tasks_as_modified()
llm = self.get_llm()
await self.send_message("Breaking down the task into steps ...")
convo.assistant(response).template("parse_task").require_schema(TaskSteps)
response: TaskSteps = await llm(convo, parser=JSONParser(TaskSteps), temperature=0)
@@ -302,8 +301,7 @@ class Developer(RelevantFilesMixin, BaseAgent):
buttons["skip"] = "Skip Task"
description = self.current_state.current_task["description"]
await self.send_message("Starting new task with description:")
await self.send_message(description)
await self.send_message("Starting new task with description:\n\n" + description)
user_response = await self.ask_question(
"Do you want to execute the above task?",
buttons=buttons,

View File

@@ -85,7 +85,7 @@ class ErrorHandler(BaseAgent):
log.info("Skipping command error debug (requested by user)")
return AgentResponse.done(self)
llm = self.get_llm()
llm = self.get_llm(stream_output=True)
convo = AgentConvo(self).template(
"debug",
task_steps=self.current_state.steps,

View File

@@ -54,7 +54,6 @@ class Executor(BaseAgent):
output_handler=self.output_handler,
exit_handler=self.exit_handler,
)
self.stream_output = True
def for_step(self, step):
# FIXME: not needed, refactor to use self.current_state.current_step
@@ -142,7 +141,7 @@ class Executor(BaseAgent):
async def check_command_output(
self, cmd: str, timeout: Optional[int], stdout: str, stderr: str, status_code: int
) -> CommandResult:
llm = self.get_llm()
llm = self.get_llm(stream_output=True)
convo = (
AgentConvo(self)
.template(

View File

@@ -85,7 +85,7 @@ class ExternalDocumentation(BaseAgent):
if not available_docsets:
return {}
llm = self.get_llm()
llm = self.get_llm(stream_output=True)
convo = (
AgentConvo(self)
.template(
@@ -109,7 +109,7 @@ class ExternalDocumentation(BaseAgent):
queries = {}
await self.send_message("Getting relevant documentation for the following topics:")
for k, short_desc in docsets.items():
llm = self.get_llm()
llm = self.get_llm(stream_output=True)
convo = (
AgentConvo(self)
.template(

View File

@@ -54,7 +54,7 @@ class Importer(BaseAgent):
await self.state_manager.commit()
async def analyze_project(self):
llm = self.get_llm()
llm = self.get_llm(stream_output=True)
self.send_message("Inspecting most important project files ...")

View File

@@ -42,7 +42,7 @@ class IterationPromptMixin:
:param bug_hunting_cycles: Data about logs that need to be added to the code (optional).
:return: The generated solution to the problem.
"""
llm = self.get_llm(TROUBLESHOOTER_BUG_REPORT)
llm = self.get_llm(TROUBLESHOOTER_BUG_REPORT, stream_output=True)
convo = AgentConvo(self).template(
"iteration",
user_feedback=user_feedback,
@@ -63,8 +63,6 @@ class RelevantFilesMixin:
self, user_feedback: Optional[str] = None, solution_description: Optional[str] = None
) -> AgentResponse:
log.debug("Getting relevant files for the current task")
await self.send_message("Figuring out which project files are relevant for the next task ...")
done = False
relevant_files = set()
llm = self.get_llm(GET_RELEVANT_FILES_AGENT_NAME)

View File

@@ -46,7 +46,7 @@ class ProblemSolver(IterationPromptMixin, BaseAgent):
return await self.try_alternative_solutions()
async def generate_alternative_solutions(self):
llm = self.get_llm()
llm = self.get_llm(stream_output=True)
convo = (
AgentConvo(self)
.template(

View File

@@ -21,7 +21,7 @@ class TaskReviewer(BaseAgent):
"""
log.debug(f"Reviewing code changes for task {self.current_state.current_task['description']}")
await self.send_message("Reviewing the task implementation ...")
# await self.send_message("Reviewing the task implementation ...")
all_feedbacks = [
iteration["user_feedback"].replace("```", "").strip()
for iteration in self.current_state.iterations

View File

@@ -156,7 +156,7 @@ class TechLead(BaseAgent):
log.debug(f"Planning tasks for the epic: {epic['name']}")
await self.send_message("Starting to create the action plan for development ...")
llm = self.get_llm(TECH_LEAD_PLANNING)
llm = self.get_llm(TECH_LEAD_PLANNING, stream_output=True)
convo = (
AgentConvo(self)
.template(
@@ -220,7 +220,7 @@ class TechLead(BaseAgent):
log.debug(f"Updating development plan for {epic['name']}")
await self.send_message("Updating development plan ...")
llm = self.get_llm()
llm = self.get_llm(stream_output=True)
convo = (
AgentConvo(self)
.template(

View File

@@ -52,7 +52,7 @@ class TechnicalWriter(BaseAgent):
async def create_readme(self):
await self.send_message("Creating README ...")
llm = self.get_llm()
llm = self.get_llm(stream_output=True)
convo = AgentConvo(self).template("create_readme")
llm_response: str = await llm(convo)
await self.state_manager.save_file("README.md", llm_response)

View File

@@ -174,7 +174,7 @@ class Troubleshooter(IterationPromptMixin, RelevantFilesMixin, BaseAgent):
route_files = await self._get_route_files()
llm = self.get_llm()
llm = self.get_llm(stream_output=True)
convo = self._get_task_convo().template(
"define_user_review_goal", task=self.current_state.current_task, route_files=route_files
)
@@ -304,7 +304,7 @@ class Troubleshooter(IterationPromptMixin, RelevantFilesMixin, BaseAgent):
:return: Additional questions and answers to generate a better bug report.
"""
additional_qa = []
llm = self.get_llm()
llm = self.get_llm(stream_output=True)
convo = (
AgentConvo(self)
.template(

View File

@@ -45,3 +45,4 @@ Expected result: Form is submitted, page is reloaded and "Thank you" message is
---end_of_example---
If nothing needs to be tested for this task, instead of outputting the steps, just output a single word: DONE
Make sure you do not output any duplicate steps.

View File

@@ -63,7 +63,7 @@ async def test_get_llm(mock_BaseLLMClient):
mock_client = AsyncMock(return_value=("response", "log"))
mock_OpenAIClient.return_value = mock_client
llm = agent.get_llm()
llm = agent.get_llm(stream_output=True)
mock_BaseLLMClient.for_provider.assert_called_once_with("openai")

View File

@@ -65,7 +65,7 @@ async def test_ask_for_feature(agentcontext):
assert sm.current_state.epics[1]["completed"] is False
@pytest.mark.asyncio
@pytest.mark.skip(reason="Temporary")
async def test_plan_epic(agentcontext):
"""
If called and there's an incomplete epic, the TechLead agent should plan the epic.
@@ -102,7 +102,7 @@ async def test_plan_epic(agentcontext):
assert sm.current_state.tasks[1]["description"] == "Task 2"
@pytest.mark.asyncio
@pytest.mark.skip(reason="Temporary")
async def test_update_epic(agentcontext):
"""
Updating the current epic's dev plan according to the current task iterations.