写点什么

5 分钟上手!用 #Workspace 为 screenshot-to-code 项目增加新功能

作者:豆包MarsCode
  • 2024-11-25
    北京
  • 本文字数:3503 字

    阅读完需:约 11 分钟

在之前的文章中,我们教过大家如何 5 分钟实现网站复刻,同时为大家介绍了 screenshot-to-code 这个开源项目,以及如何基于该项目增加调用 Gemini 模型,最终实现上传图片生成代码的效果。


当时就有不少小伙伴提出疑问,怎么快速找到需要修改的地方呢?

在此之前,我们要先花一些时间去了解阅读源码才能找到相关函数,但在今天,给大家介绍一个豆包MarsCode 的利器:#Workspace,有了它,妈妈再也不用担心我找不到修改函数的入口啦!


今天给大家演示一下,如何利用豆包MarsCode 的 #Workspace 功能,快速上手修改开源项目!

在修改之前,需要首先了解并熟悉当前项目的模块和架构,我们可以使用豆包 MarsCode #Workspace 功能,总结并输出模块的架构:


接着,我们可以把 #Workspace 输出的 mermaid 内容,贴到 mermaid 在线展示网站(https://sourl.cn/RUYGf4)上,就得到了 screenshot-to-code 项目的模块架构:

了解熟悉完架构后,接着使用 #Workspace 让豆包MarsCode 告诉我们如何本地启动项目:

根据它给的方式,我们在终端内一步步执行,即可成功启动项目。

下一步,我们需要增加调用 Gemini 模型,但模块代码这么多,该怎么快速找到需要修改哪些地方呢?

这时我们再次使出 #Workspace 利器,可以看到豆包MarsCode 精准地给出了我们要修改的地方:




接下来,让我们按照豆包MarsCode 给的教程,动手改起来。

首先,我们需要打开 frontend/src/lib/models.ts 增加 Gemini 模型的枚举。当输入 GEMINI 之后,豆包 MarsCode 自动为你推荐了补全代码,点击 Tab 键采纳即可。


紧接着,我们需要打开 frontend/src/App.tsx,找到修改入口:

我们先拿豆包 MarsCode 来解释一下这个函数的功能。

根据解释,该函数只是用于更新指令并生成代码,和我们要改的模型无关,我们忽略这个文件直接改后端代码。


后端部分主要修改 llm.py 和 generate_code.py 2 个文件。根据前面的架构介绍,llm.py 负责调用 LLM,generate_code.py 则是生成代码的路由处理。因此我们需要先改 llm.py,增加调用 Gemini。

在这里,我们直接复制上次的代码,修改 llm.py:


class Llm(Enum):    GPT_4_VISION = "gpt-4-vision-preview"    GPT_4_TURBO_2024_04_09 = "gpt-4-turbo-2024-04-09"    GPT_4O_2024_05_13 = "gpt-4o-2024-05-13"    CLAUDE_3_SONNET = "claude-3-sonnet-20240229"    CLAUDE_3_OPUS = "claude-3-opus-20240229"    CLAUDE_3_HAIKU = "claude-3-haiku-20240307"    CLAUDE_3_5_SONNET_2024_06_20 = "claude-3-5-sonnet-20240620"    //新增gemini    GEMINI_1_5_PRO_LATEST = "gemini-1.5-pro-latest"        async def stream_gemini_response(    messages: List[ChatCompletionMessageParam],    api_key: str,    callback: Callable[[str], Awaitable[None]],) -> str:  genai.configure(api_key=api_key)    generation_config = genai.GenerationConfig(    temperature = 0.0  )  model = genai.GenerativeModel(    model_name = "gemini-1.5-pro-latest",    generation_config = generation_config  )  contents = parse_openai_to_gemini_prompt(messages);    response = model.generate_content(    contents = contents,    #Support streaming    stream = True,   )     for chunk in response:    content = chunk.text or ""    await callback(content)
if not response: raise Exception("No HTML response found in AI response") else: return response.text;
def parse_openai_to_gemini_prompt(prompts): messages = [] for prompt in prompts: message = {} message['role'] = prompt['role'] if prompt['role'] == 'system': message['role'] = 'user' if prompt['role'] == 'assistant': message['role'] = 'model' message['parts'] = [] content = prompt['content'] if isinstance(content, list): for content in prompt['content']: part = {} if content['type'] == 'image_url': base64 = content['image_url']['url'] part['inline_data'] = { 'data': base64.split(",")[1], 'mime_type': base64.split(";")[0].split(":")[1] } elif content['type'] == 'text': part['text'] = content['text'] message['parts'].append(part) else: message['parts'] = [content] messages.append(message) return messages
复制代码


最后,我们再修改 generate_code.py,增加调用 Gemini 的函数:

if validated_input_mode == "video":                if not anthropic_api_key:                    await throw_error(                        "Video only works with Anthropic models. No Anthropic API key found. Please add the environment variable ANTHROPIC_API_KEY to backend/.env or in the settings dialog"                    )                    raise Exception("No Anthropic key")
completion = await stream_claude_response_native( system_prompt=VIDEO_PROMPT, messages=prompt_messages, # type: ignore api_key=anthropic_api_key, callback=lambda x: process_chunk(x), model=Llm.CLAUDE_3_OPUS, include_thinking=True, ) exact_llm_version = Llm.CLAUDE_3_OPUS elif ( code_generation_model == Llm.CLAUDE_3_SONNET or code_generation_model == Llm.CLAUDE_3_5_SONNET_2024_06_20 ): if not anthropic_api_key: await throw_error( "No Anthropic API key found. Please add the environment variable ANTHROPIC_API_KEY to backend/.env or in the settings dialog" ) raise Exception("No Anthropic key")
completion = await stream_claude_response( prompt_messages, # type: ignore api_key=anthropic_api_key, callback=lambda x: process_chunk(x), model=code_generation_model, ) exact_llm_version = code_generation_model # 增加调用gemini elif ( code_generation_model == Llm.GEMINI_1_5_PRO_LATEST ): if not GEMINI_API_KEY: await throw_error( "No GEMINI API key found. Please add the environment variable ANTHROPIC_API_KEY to backend/.env or in the settings dialog" ) raise Exception("No GEMINI key")
completion = await stream_gemini_response( prompt_messages, # type: ignore api_key=GEMINI_API_KEY, callback=lambda x: process_chunk(x), ) exact_llm_version = code_generation_model else: completion = await stream_openai_response( prompt_messages, # type: ignore api_key=openai_api_key, base_url=openai_base_url, callback=lambda x: process_chunk(x), model=code_generation_model, ) exact_llm_version = code_generation_model
复制代码

经过以上几个步骤的修改,我们就完成了代码修改部分,最后我们再安装 google-generativeai 库:

cd backendpoetry add google-generativeai
复制代码


安装完库后,再次启动项目,就可以愉快地使用 Gemini 来生成代码啦,大家赶快去试试使用 #Workspace 吧!

用户头像

还未添加个人签名 2024-08-27 加入

用 AI 激发创造

评论

发布
暂无评论
5分钟上手!用 #Workspace 为 screenshot-to-code 项目增加新功能_程序员_豆包MarsCode_InfoQ写作社区