diff --git a/chat/notbadai_chat/__init__.py b/chat/notbadai_chat/__init__.py index 201cd10..a7bde2f 100644 --- a/chat/notbadai_chat/__init__.py +++ b/chat/notbadai_chat/__init__.py @@ -1,70 +1,37 @@ from string import Template from pathlib import Path - -from notbadai_ide import api, START_METADATA, END_METADATA - +from notbadai_ide import api from .common.llm import call_llm -from .common.utils import parse_prompt -from .common.prompt import build_context - -module_dir = Path(__file__).parent - - -def get_prompt_template(template_path: str, **kwargs) -> str: - path = module_dir / f'{template_path}.md' - with open(str(path)) as f: - template = Template(f.read()) - - return template.substitute(kwargs) +from .common.formatting import markdown_section, markdown_code_block def start(): """Main extension function that handles chat interactions with the AI assistant.""" - command, model, prompt = parse_prompt() - selection = api.get_selection() - chat_history = api.get_chat_history() - prompt = api.get_prompt() - - api.chat(f'{START_METADATA}model: {model}, command: {command}{END_METADATA}') - - if command == '': - api.chat(f'{START_METADATA}Without context{END_METADATA}') - messages = [ - {'role': 'system', 'content': get_prompt_template('chat.system', model=model)}, - *[m.to_dict() for m in chat_history], - {'role': 'user', 'content': prompt}, - ] - elif command == 'here': - context = build_context() - - api.chat(f'{START_METADATA}With context: {len(context) :,} characters,' - f' selection: {bool(selection)}{END_METADATA}') - api.log(context) - messages = [ - {'role': 'system', 'content': get_prompt_template('chat.system', model=model)}, - {'role': 'user', 'content': context}, - *[m.to_dict() for m in chat_history], - {'role': 'user', 'content': prompt}, - ] - elif command == 'context': - context = build_context() - - api.chat(f'{START_METADATA}With context: {len(context) :,} characters,' - f' selection: {bool(selection)}{END_METADATA}') - # api.log(context) - messages = [ - {'role': 'system', 'content': get_prompt_template('chat.system', model=model)}, - {'role': 'user', 'content': context}, - *[m.to_dict() for m in chat_history], - {'role': 'user', 'content': prompt}, - ] - else: - raise ValueError(f'Unknown command: {command}') - - api.log(f'messages {len(messages)}') - api.log(f'prompt {prompt}') - # api.log(context) - - content = call_llm(model, messages) - - api.log(content) + model = 'qwen' + context = [] + + open_files = [f for f in api.get_repo_files() if f.is_open] + if open_files: + context.append(markdown_section("Relevant files", "\n\n".join( + f'Path: `{f.path}`\n\n{markdown_code_block(f.get_content())}' for f in open_files))) + + if current_file := api.get_current_file(): + file_content = markdown_code_block(current_file.get_content()) + context.append(markdown_section("Current File", f"Path: `{current_file.path}`\n\n{file_content}")) + + if terminal := api.get_current_terminal().get_snapshot(): + context.append(markdown_section("Terminal output", markdown_code_block(terminal[-40000:]))) + + if selection := api.get_selection().strip(): + context.append(markdown_section("Selection", f"This is the code snippet " + f"that I'm referring to\n\n{markdown_code_block(selection)}")) + + with open(Path(__file__).parent / 'chat.system.md') as f: + system_prompt = Template(f.read()).substitute(model=model) + + messages = [ + {'role': 'system', 'content': system_prompt}, + {'role': 'user', 'content': "\n\n".join(context)}, + {'role': 'user', 'content': api.get_prompt()}, + ] + call_llm(model, messages)