aroniscunt/Browser_Web_UI_Automation

1import os
2import pdb
3from dataclasses import dataclass
4
5from dotenv import load_dotenv
6from langchain_core.messages import HumanMessage, SystemMessage
7from langchain_ollama import ChatOllama
8
9load_dotenv()
10
11import sys
12
13sys.path.append(".")
14
15
16@dataclass
17class LLMConfig:
18    provider: str
19    model_name: str
20    temperature: float = 0.8
21    base_url: str = None
22    api_key: str = None
23
24
25def create_message_content(text, image_path=None):
26    content = [{"type": "text", "text": text}]
27    image_format = "png" if image_path and image_path.endswith(".png") else "jpeg"
28    if image_path:
29        from src.utils import utils
30        image_data = utils.encode_image(image_path)
31        content.append({
32            "type": "image_url",
33            "image_url": {"url": f"data:image/{image_format};base64,{image_data}"}
34        })
35    return content
36
37
38def get_env_value(key, provider):
39    env_mappings = {
40        "openai": {"api_key": "OPENAI_API_KEY", "base_url": "OPENAI_ENDPOINT"},
41        "azure_openai": {"api_key": "AZURE_OPENAI_API_KEY", "base_url": "AZURE_OPENAI_ENDPOINT"},
42        "google": {"api_key": "GOOGLE_API_KEY"},
43        "deepseek": {"api_key": "DEEPSEEK_API_KEY", "base_url": "DEEPSEEK_ENDPOINT"},
44        "mistral": {"api_key": "MISTRAL_API_KEY", "base_url": "MISTRAL_ENDPOINT"},
45        "alibaba": {"api_key": "ALIBABA_API_KEY", "base_url": "ALIBABA_ENDPOINT"},
46        "moonshot": {"api_key": "MOONSHOT_API_KEY", "base_url": "MOONSHOT_ENDPOINT"},
47        "ibm": {"api_key": "IBM_API_KEY", "base_url": "IBM_ENDPOINT"}
48    }
49
50    if provider in env_mappings and key in env_mappings[provider]:
51        return os.getenv(env_mappings[provider][key], "")
52    return ""
53
54
55def test_llm(config, query, image_path=None, system_message=None):
56    from src.utils import utils, llm_provider
57
58    # Special handling for Ollama-based models
59    if config.provider == "ollama":
60        if "deepseek-r1" in config.model_name:
61            from src.utils.llm_provider import DeepSeekR1ChatOllama
62            llm = DeepSeekR1ChatOllama(model=config.model_name)
63        else:
64            llm = ChatOllama(model=config.model_name)
65
66        ai_msg = llm.invoke(query)
67        print(ai_msg.content)
68        if "deepseek-r1" in config.model_name:
69            pdb.set_trace()
70        return
71
72    # For other providers, use the standard configuration
73    llm = llm_provider.get_llm_model(
74        provider=config.provider,
75        model_name=config.model_name,
76        temperature=config.temperature,
77        base_url=config.base_url or get_env_value("base_url", config.provider),
78        api_key=config.api_key or get_env_value("api_key", config.provider)
79    )
80
81    # Prepare messages for non-Ollama models
82    messages = []
83    if system_message:
84        messages.append(SystemMessage(content=create_message_content(system_message)))
85    messages.append(HumanMessage(content=create_message_content(query, image_path)))
86    ai_msg = llm.invoke(messages)
87
88    # Handle different response types
89    if hasattr(ai_msg, "reasoning_content"):
90        print(ai_msg.reasoning_content)
91    print(ai_msg.content)
92
93def test_openai_model():
94    config = LLMConfig(provider="openai", model_name="gpt-4o")
95    test_llm(config, "Describe this image", "assets/examples/test.png")
96
97
98def test_google_model():
99    # Enable your API key first if you haven't: https://ai.google.dev/palm_docs/oauth_quickstart
100    config = LLMConfig(provider="google", model_name="gemini-2.0-flash-exp")
101    test_llm(config, "Describe this image", "assets/examples/test.png")
102
103
104def test_azure_openai_model():
105    config = LLMConfig(provider="azure_openai", model_name="gpt-4o")
106    test_llm(config, "Describe this image", "assets/examples/test.png")
107
108
109def test_deepseek_model():
110    config = LLMConfig(provider="deepseek", model_name="deepseek-chat")
111    test_llm(config, "Who are you?")
112
113
114def test_deepseek_r1_model():
115    config = LLMConfig(provider="deepseek", model_name="deepseek-reasoner")
116    test_llm(config, "Which is greater, 9.11 or 9.8?", system_message="You are a helpful AI assistant.")
117
118
119def test_ollama_model():
120    config = LLMConfig(provider="ollama", model_name="qwen2.5:7b")
121    test_llm(config, "Sing a ballad of LangChain.")
122
123
124def test_deepseek_r1_ollama_model():
125    config = LLMConfig(provider="ollama", model_name="deepseek-r1:14b")
126    test_llm(config, "How many 'r's are in the word 'strawberry'?")
127
128
129def test_mistral_model():
130    config = LLMConfig(provider="mistral", model_name="pixtral-large-latest")
131    test_llm(config, "Describe this image", "assets/examples/test.png")
132
133
134def test_moonshot_model():
135    config = LLMConfig(provider="moonshot", model_name="moonshot-v1-32k-vision-preview")
136    test_llm(config, "Describe this image", "assets/examples/test.png")
137
138
139def test_ibm_model():
140    config = LLMConfig(provider="ibm", model_name="meta-llama/llama-4-maverick-17b-128e-instruct-fp8")
141    test_llm(config, "Describe this image", "assets/examples/test.png")
142
143
144def test_qwen_model():
145    config = LLMConfig(provider="alibaba", model_name="qwen-vl-max")
146    test_llm(config, "How many 'r's are in the word 'strawberry'?")
147
148
149if __name__ == "__main__":
150    # test_openai_model()
151    # test_google_model()
152    test_azure_openai_model()
153    # test_deepseek_model()
154    # test_ollama_model()
155    # test_deepseek_r1_model()
156    # test_deepseek_r1_ollama_model()
157    # test_mistral_model()
158    # test_ibm_model()
159    # test_qwen_model()
160