feat: chat model VISION support for user and tool msg attachments

Rafael Uzarowski committed Mar 2, 2025 at 18:00 UTC 22eef83cc6628e13938a851beece8e7e8c6af19e
8 files changed +105 -22
agent.py
+14 -7
@@ -10,8 +10,9 @@ import models
10 from langchain_core.prompt_values import ChatPromptValue
11 from python.helpers import extract_tools, rate_limiter, files, errors, history, tokens
12 from python.helpers.print_style import PrintStyle
13 -from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
14 -from langchain_core.messages import HumanMessage, SystemMessage, AIMessage
13 +from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder, HumanMessagePromptTemplate, StringPromptTemplate
14 +from langchain_core.prompts.image import ImagePromptTemplate
15 +from langchain_core.messages import HumanMessage, SystemMessage, AIMessage, BaseMessage
16 from langchain_core.language_models.chat_models import BaseChatModel
17 from langchain_core.language_models.llms import BaseLLM
18 from langchain_core.embeddings import Embeddings
@@ -19,6 +20,7 @@ import python.helpers.log as Log
20 from python.helpers.dirty_json import DirtyJson
21 from python.helpers.defer import DeferredTask
22 from typing import Callable
23 +from python.helpers.history import OutputMessage
24
25
26 class AgentContext:
@@ -187,7 +189,7 @@ class AgentConfig:
189 @dataclass
190 class UserMessage:
191 message: str
190 - attachments: list[str]
192 + attachments: list[str] = field(default_factory=list[str])
193
194
195 class LoopData:
@@ -357,10 +359,14 @@ class Agent:
359 loop_data.extras_temporary.clear()
360
361 # combine history and extras
360 - history_combined = history.group_outputs_abab(loop_data.history_output + extras)
362 + history_combined: list[OutputMessage] = history.group_outputs_abab(loop_data.history_output + extras)
363
364 # convert history to LLM format
363 - history_langchain = history.output_langchain(history_combined)
365 + history_langchain: list[BaseMessage] = history.output_langchain(history_combined)
366 +
367 + PrintStyle(font_color="grey", background_color="black", bold=True, padding=True).print(
368 + f"History Langchain: {history_langchain}"
369 + )
370
371 # build chain from system prompt, message history and model
372 prompt = ChatPromptTemplate.from_messages(
@@ -479,9 +485,10 @@ class Agent:
485 content = self.parse_prompt("fw.warning.md", message=message)
486 return self.hist_add_message(False, content=content)
487
482 - async def hist_add_tool_result(self, tool_name: str, tool_result: str):
488 + async def hist_add_tool_result(self, tool_name: str, tool_result: str, attachments: list[str] = []):
489 + attachments_str = json.dumps(attachments).replace("\n", "")
490 content = self.parse_prompt(
484 - "fw.tool_result.md", tool_name=tool_name, tool_result=tool_result
491 + "fw.tool_result.md", tool_name=tool_name, tool_result=tool_result, attachments=attachments_str
492 )
493 return self.hist_add_message(False, content=content)
494
initialize.py
+1
@@ -13,6 +13,7 @@ def initialize():
13 provider=models.ModelProvider[current_settings["chat_model_provider"]],
14 name=current_settings["chat_model_name"],
15 ctx_length=current_settings["chat_model_ctx_length"],
16 + vision=current_settings["chat_model_vision"],
17 limit_requests=current_settings["chat_model_rl_requests"],
18 limit_input=current_settings["chat_model_rl_input"],
19 limit_output=current_settings["chat_model_rl_output"],
prompts/default/agent.system.tools_vision.md new
+3
@@ -0,0 +1,3 @@
1 +## "Multimodal (Vision) Agent Tools" available:
2 +
3 +None yet. In future, this section will contain vision-only tools
prompts/default/fw.tool_result.md
+3 -2
@@ -1,6 +1,7 @@
1 ~~~json
2 {
3 "tool_name": {{tool_name}},
4 - "tool_result": {{tool_result}}
4 + "tool_result": {{tool_result}},
5 + "attachments": {{attachments}}
6 }
6 -~~~
\ No newline at end of file
7 +~~~
python/extensions/system_prompt/_10_system_prompt.py
+8 -2
@@ -12,11 +12,17 @@ class SystemPrompt(Extension):
12 system_prompt.append(main)
13 system_prompt.append(tools)
14
15 +
16 def get_main_prompt(agent: Agent):
17 return get_prompt("agent.system.main.md", agent)
18
19 +
20 def get_tools_prompt(agent: Agent):
19 - return get_prompt("agent.system.tools.md", agent)
21 + prompt = get_prompt("agent.system.tools.md", agent)
22 + if agent.config.chat_model.vision:
23 + prompt += '\n' + get_prompt("agent.system.tools_vision.md", agent)
24 + return prompt
25 +
26
27 def get_prompt(file: str, agent: Agent):
28 # variables for system prompts
@@ -26,4 +32,4 @@ def get_prompt(file: str, agent: Agent):
32 "date_time": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
33 "agent_name": agent.agent_name,
34 }
29 - return agent.read_prompt(file, **vars)
\ No newline at end of file
35 + return agent.read_prompt(file, **vars)
python/helpers/history.py
+54 -4
@@ -3,11 +3,14 @@ import asyncio
3 from collections import OrderedDict
4 import json
5 import math
6 +import os
7 from typing import Coroutine, Literal, TypedDict, cast
8 from python.helpers import messages, tokens, settings, call_llm
9 from enum import Enum
9 -from langchain_core.messages import HumanMessage, SystemMessage, AIMessage
10 -
10 +from langchain_core.messages import HumanMessage, SystemMessage, AIMessage, BaseMessage
11 +from python.helpers.print_style import PrintStyle
12 +from langchain_core.prompts import HumanMessagePromptTemplate
13 +from typing import Any
14 BULK_MERGE_COUNT = 3
15 TOPICS_KEEP_COUNT = 3
16 CURRENT_TOPIC_RATIO = 0.5
@@ -425,13 +428,60 @@ def group_outputs_abab(outputs: list[OutputMessage]) -> list[OutputMessage]:
428 return result
429
430
428 -def output_langchain(messages: list[OutputMessage]):
431 +def output_langchain(messages: list[OutputMessage]) -> list[BaseMessage]:
432 result = []
433 for m in messages:
434 if m["ai"]:
435 result.append(AIMessage(content=serialize_content(m["content"])))
436 else:
434 - result.append(HumanMessage(content=serialize_content(m["content"])))
437 + contents = m["content"]
438 +
439 + # sometimes content is a list sometimes not
440 + if not isinstance(contents, list):
441 + contents = [contents]
442 +
443 + PrintStyle(font_color="grey", background_color="black", bold=True, padding=True).print(
444 + f"Contents: {json.dumps(contents, indent=2)}"
445 + )
446 +
447 + template: list[dict[str, str]] = [] # type: ignore
448 + message = ""
449 + images = {}
450 + for _, content in enumerate(contents):
451 + if message:
452 + # the first message is the user message, then the memory and solutions
453 + message += "\n\n--- Memory & Solutions Section: ---\n\n"
454 +
455 + message += serialize_content(content)
456 +
457 + if isinstance(content, dict) and "attachments" in content:
458 + attachments: list[str] = cast(list[str], content["attachments"])
459 + for attachment in attachments:
460 + if not os.path.exists(str(attachment)):
461 + continue
462 + if attachment not in images:
463 + import base64
464 + from mimetypes import guess_type
465 + mime_type, _ = guess_type(str(attachment))
466 + if mime_type.startswith("image/"):
467 + # Read and encode the image file
468 + with open(str(attachment), "rb") as image_file:
469 + base64_encoded_data = base64.b64encode(image_file.read()).decode('utf-8')
470 + # Construct the data URL
471 + images[attachment] = f"data:{mime_type};base64,{base64_encoded_data}"
472 +
473 + if message:
474 + template.append({"type": "text", "text": message})
475 + if images:
476 + for _, image in images.items():
477 + template.append({"type": "image_url", "image_url": image})
478 + if template:
479 + # only jinja2 is safe for json, both mustache({{...}}) and f-string({...}) are not
480 + result.append(HumanMessagePromptTemplate.from_template(template=template, partial_variables={}, template_format="jinja2")) # type: ignore
481 +
482 + PrintStyle(font_color="grey", background_color="black", bold=True, padding=True).print(
483 + f"Result: {result}"
484 + )
485 return result
486
487
python/helpers/settings.py
+13
@@ -16,6 +16,7 @@ class Settings(TypedDict):
16 chat_model_kwargs: dict[str, str]
17 chat_model_ctx_length: int
18 chat_model_ctx_history: float
19 + chat_model_vision: bool
20 chat_model_rl_requests: int
21 chat_model_rl_input: int
22 chat_model_rl_output: int
@@ -149,6 +150,17 @@ def convert_out(settings: Settings) -> SettingsOutput:
150 }
151 )
152
153 +
154 + chat_model_fields.append(
155 + {
156 + "id": "chat_model_vision",
157 + "title": "Supports Vision",
158 + "description": "Models capable of Vision can for example natively see the content of image attachments.",
159 + "type": "switch",
160 + "value": settings["chat_model_vision"],
161 + }
162 + )
163 +
164 chat_model_fields.append(
165 {
166 "id": "chat_model_rl_requests",
@@ -777,6 +789,7 @@ def get_default_settings() -> Settings:
789 chat_model_kwargs={ "temperature": "0" },
790 chat_model_ctx_length=120000,
791 chat_model_ctx_history=0.7,
792 + chat_model_vision=False,
793 chat_model_rl_requests=0,
794 chat_model_rl_input=0,
795 chat_model_rl_output=0,
python/helpers/tool.py
+9 -7
@@ -1,14 +1,16 @@
1 from abc import abstractmethod
2 -from dataclasses import dataclass
2 +from dataclasses import dataclass, field
3 from agent import Agent
4 from python.helpers.print_style import PrintStyle
5 -from python.helpers import messages
5 +
6
7 @dataclass
8 class Response:
9 message:str
10 - break_loop:bool
11 -
10 + break_loop: bool
11 + attachments: list[str] = field(default_factory=list[str])
12 +
13 +
14 class Tool:
15
16 def __init__(self, agent: Agent, name: str, args: dict[str,str], message: str, **kwargs) -> None:
@@ -29,10 +31,10 @@ class Tool:
31 PrintStyle(font_color="#85C1E9", bold=True).stream(self.nice_key(key)+": ")
32 PrintStyle(font_color="#85C1E9", padding=isinstance(value,str) and "\n" in value).stream(value)
33 PrintStyle().print()
32 -
34 +
35 async def after_execution(self, response: Response, **kwargs):
36 text = response.message.strip()
35 - await self.agent.hist_add_tool_result(self.name, text)
37 + await self.agent.hist_add_tool_result(self.name, text, response.attachments)
38 PrintStyle(font_color="#1B4F72", background_color="white", padding=True, bold=True).print(f"{self.agent.agent_name}: Response from tool '{self.name}'")
39 PrintStyle(font_color="#85C1E9").print(response.message)
40 self.log.update(content=response.message)
@@ -44,4 +46,4 @@ class Tool:
46 words = key.split('_')
47 words = [words[0].capitalize()] + [word.lower() for word in words[1:]]
48 result = ' '.join(words)
47 - return result
\ No newline at end of file
49 + return result