python
复制代码
import os
import sys
sys.path.append(os.getcwd())
import base64
import pyautogui
from datetime import datetime
from langchain_openai import ChatOpenAI
from langchain_core.messages import SystemMessage, HumanMessage, AIMessage
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from share_config.screenshot_base64 import screenshot_to_base64
from share_config.api_key import ollama_vlm, ollama_base_url
# 初始化VLM模型
vlm = ChatOpenAI(model=ollama_vlm, base_url=ollama_base_url, api_key="empty")
# 创建支持多模态的核心模板结构 ✨
prompt = ChatPromptTemplate.from_messages([
SystemMessage(content=os.environ.get("system_prompt", "你是一个专业的图像分析助手,请仔细分析图片内容")),
MessagesPlaceholder(variable_name="context"),
MessagesPlaceholder(variable_name="input") # 支持多模态消息列表 🎯
])
def retrieve_context(query: str) -> list:
"""检索相关上下文(简化版本,可以后续扩展)"""
# 目前返回空上下文,可以根据需要扩展
return []
def create_multimodal_message(text: str, image_base64: str) -> HumanMessage:
"""创建多模态消息 - 标准格式 📷"""
return HumanMessage(
content=[
{"type": "text", "text": text},
{
"type": "image_url",
"image_url": {"url": f"data:image/jpeg;base64,{image_base64}"}
}
]
)
def analyze_screenshot_with_context(prompt_text: str, context: list = None) -> str:
"""使用结构化prompt分析截图 🚀"""
try:
if context is None:
context = []
# 截图
encoded_image = screenshot_to_base64()
# 构建多模态消息 💡
multimodal_message = create_multimodal_message(prompt_text, encoded_image)
# 使用结构化prompt调用VLM
chain = prompt | vlm
response = chain.invoke({
"context": context,
"input": [multimodal_message] # 作为消息列表传递 📋
})
print(f"VLM分析结果: {response.content}")
return response.content
except Exception as e:
print(f"分析失败: {e}")
return f"分析过程中出现错误: {str(e)}"
def analyze_screenshot(prompt: str = "请分析这张图片的内容") -> str:
"""分析截图内容(保持向后兼容)"""
context = retrieve_context(prompt)
return analyze_screenshot_with_context(prompt, context)
def chat_turn(user_input: str = None) -> str:
"""Flask服务调用接口 - 分析当前屏幕截图
Args:
user_input (str, optional): 用户的分析提示词,如果为None则使用默认提示
Returns:
str: VLM分析结果
"""
# 如果没有提供提示词,使用默认提示
if not user_input:
user_input = "请详细分析这张图片的内容,包括你能看到的所有文字、界面元素和重要信息"
# 使用结构化prompt进行分析
context = retrieve_context(user_input)
result = analyze_screenshot_with_context(user_input, context)
print(f"✅ 分析完成")
return result
def chat_loop():
"""主对话循环"""
print("🤖 VLM截图分析助手启动!")
print("📸 程序会自动截图并分析当前屏幕内容")
print("💬 输入自定义提示词来指导分析")
print("👋 输入 'quit', 'exit', 或 'bye' 退出")
print("-" * 50)
while True:
# 获取用户输入的分析提示词
user_prompt = input("\n🎯 请输入分析提示词 (直接回车使用默认提示): ").strip()
if user_prompt.lower() in ['quit', 'exit', 'bye']:
print("👋 再见!")
break
# 如果用户没有输入提示词,使用默认提示
if not user_prompt:
user_prompt = "请详细分析这张图片的内容,包括你能看到的所有文字、界面元素和重要信息"
print("📸 正在截图并分析...")
# 使用结构化prompt分析截图
context = retrieve_context(user_prompt)
result = analyze_screenshot_with_context(user_prompt, context)
print(f"\n💡 分析结果:")
print("-" * 30)
print(result)
print("-" * 30)
if __name__ == "__main__":
chat_loop()