File size: 2,030 Bytes
d26280a |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 |
import os
import tempfile
from langchain.schema import HumanMessage, SystemMessage
from gpt_engineer.core.ai import AI
from gpt_engineer.core.base_memory import BaseMemory
from gpt_engineer.core.chat_to_files import overwrite_code_with_edits
from gpt_engineer.core.default.disk_memory import DiskMemory
from gpt_engineer.core.default.paths import IMPROVE_LOG_FILE
from gpt_engineer.core.default.steps import curr_fn, setup_sys_prompt_existing_code
from gpt_engineer.core.files_dict import FilesDict
from gpt_engineer.core.preprompts_holder import PrepromptsHolder
from gpt_engineer.tools.experimental.code_vector_repository import CodeVectorRepository
def improve_automatic_file_selection(
ai: AI,
prompt: str,
code: FilesDict,
memory: BaseMemory,
preprompts_holder: PrepromptsHolder,
):
code_vector_repository = CodeVectorRepository()
# ToDo: Replace this hacky way to get the right langchain document format
temp_dir = tempfile.mkdtemp()
temp_saver = DiskMemory(temp_dir)
for file, content in code.items():
temp_saver[file] = content
code_vector_repository.load_from_directory(temp_dir)
relevant_documents = code_vector_repository.relevent_code_chunks(prompt)
relevant_code = FilesDict()
for doc in relevant_documents:
file_path = os.path.relpath(doc.metadata["filename"], temp_dir)
relevant_code[file_path] = code[file_path]
print(
"Relevant documents to be modified are: "
+ "\n".join(sorted(relevant_code.keys()))
)
preprompts = preprompts_holder.get_preprompts()
messages = [
SystemMessage(content=setup_sys_prompt_existing_code(preprompts)),
]
# Add files as input
messages.append(HumanMessage(content=f"{relevant_code.to_chat()}"))
messages.append(HumanMessage(content=f"Request: {prompt}"))
messages = ai.next(messages, step_name=curr_fn())
chat = messages[-1].content.strip()
overwrite_code_with_edits(chat, code)
memory[IMPROVE_LOG_FILE] = chat
return code
|