75 lines
3.1 KiB
Python
75 lines
3.1 KiB
Python
from inference import LlamaMessage, replicate_llama_completion
|
|
import re
|
|
|
|
message_format_section = """
|
|
When suggesting code to the user, keep the code as small snippets.
|
|
If the update is meant to be applied to a file, include a line linking the file immediately before the code block like this: "[filename.txt]
|
|
```
|
|
updated segment
|
|
```"
|
|
Omit the file link when the block is not related to a known file
|
|
Code blocks must be under 30 lines. Assistant responses MUST NOT contain horizontal rules.
|
|
When suggesting bash code meant to be directly executed, use the same syntax as above but marking the code block with "[eval]" instead.
|
|
"""
|
|
|
|
|
|
async def gral_message(app, system_prompt, message_history: list, file_map: dict):
|
|
messages = [
|
|
LlamaMessage("system", "\n".join([
|
|
system_prompt,
|
|
message_format_section,
|
|
])),
|
|
]
|
|
|
|
file_parts = []
|
|
for filename, content in file_map.items():
|
|
file_parts.append(f"\n---\n[{filename}]\n\n{content}\n")
|
|
|
|
if file_parts:
|
|
file_lines = [
|
|
"This message contains all files relevant to the prompt: each file section in this message starts with a horizontal rule, followed by the file name in square brackets, followed by the file content.",
|
|
"Only this message contains horizontal rules, assistant responses must not include them.",
|
|
"The prompt and conversation start in the next message after this."
|
|
]
|
|
file_lines.extend(file_parts)
|
|
messages.append(
|
|
LlamaMessage("user", "\n".join(file_lines))
|
|
)
|
|
|
|
messages.extend([LlamaMessage(msg['role'], msg['content']) for msg in message_history])
|
|
|
|
response = await replicate_llama_completion(app, messages)
|
|
return response
|
|
|
|
async def gral_update_file(app, system_prompt, message_history: list, filepath: str, file_content: str):
|
|
messages = [
|
|
LlamaMessage("system", "\n".join([
|
|
system_prompt,
|
|
"Respond with the raw contents of the file updated as concluded in the conversation.",
|
|
])),
|
|
]
|
|
|
|
# History Message
|
|
message_lines = [f"This is the history of the conversation that led to the update of the file {filepath}."]
|
|
message_lines.extend([f"{entry['role']}: {entry['content']}" for entry in message_history if entry.get('content')])
|
|
messages.append(LlamaMessage("user", "\n".join(message_lines)))
|
|
|
|
# Main Instruction Message
|
|
message_lines = [
|
|
f"Write the updated contents of the file: {filepath}, as concluded in the conversation above.",
|
|
f"The raw file contents of the file {filepath} are included following the horizontal rule below."
|
|
"\n---\n",
|
|
file_content,
|
|
]
|
|
messages.append(LlamaMessage("user", "\n".join(message_lines)))
|
|
|
|
new_content = await replicate_llama_completion(app, messages)
|
|
|
|
# Remove backtick fences
|
|
md_code_pattern = re.compile(r"^\s*```[\w\s]*\n([\s\S]*)\n```\s*$", re.MULTILINE)
|
|
md_match = re.search(md_code_pattern, new_content)
|
|
if md_match:
|
|
print("agent used backtick fences")
|
|
new_content = md_match.group(1)
|
|
|
|
return new_content |