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