langchain+vlm示例

Search | DeepWiki

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()