From e54c1fda0b3c44008a021cec46113a09f7466db1 Mon Sep 17 00:00:00 2001 From: warmshao Date: Sat, 4 Jan 2025 10:58:48 +0800 Subject: [PATCH 1/2] add ollama --- requirements.txt | 3 ++- src/utils/utils.py | 10 ++++++++-- tests/test_browser_use.py | 10 ++++++++-- tests/test_llm_api.py | 12 ++++++++++-- webui.py | 3 ++- 5 files changed, 30 insertions(+), 8 deletions(-) diff --git a/requirements.txt b/requirements.txt index eb339d6d..14719094 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,4 +1,5 @@ browser-use langchain-google-genai pyperclip -gradio \ No newline at end of file +gradio +langchain-ollama \ No newline at end of file diff --git a/src/utils/utils.py b/src/utils/utils.py index cc3b9e46..6fbbd6c5 100644 --- a/src/utils/utils.py +++ b/src/utils/utils.py @@ -11,6 +11,7 @@ from langchain_openai import ChatOpenAI, AzureChatOpenAI from langchain_anthropic import ChatAnthropic from langchain_google_genai import ChatGoogleGenerativeAI +from langchain_ollama import ChatOllama def get_llm_model(provider: str, **kwargs): @@ -39,7 +40,7 @@ def get_llm_model(provider: str, **kwargs): ) elif provider == 'openai': if not kwargs.get("base_url", ""): - base_url = "https://api.openai.com/v1" + base_url = os.getenv("OPENAI_ENDPOINT", "https://api.openai.com/v1") else: base_url = kwargs.get("base_url") @@ -66,7 +67,7 @@ def get_llm_model(provider: str, **kwargs): api_key = kwargs.get("api_key") return ChatOpenAI( - model=kwargs.get("model_name", 'gpt-4o'), + model=kwargs.get("model_name", 'deepseek-chat'), temperature=kwargs.get("temperature", 0.0), base_url=base_url, api_key=api_key @@ -81,6 +82,11 @@ def get_llm_model(provider: str, **kwargs): temperature=kwargs.get("temperature", 0.0), google_api_key=api_key, ) + elif provider == 'ollama': + return ChatOllama( + model=kwargs.get("model_name", 'qwen2.5:7b'), + temperature=kwargs.get("temperature", 0.0), + ) elif provider == "azure_openai": if not kwargs.get("base_url", ""): base_url = os.getenv("AZURE_OPENAI_ENDPOINT", "") diff --git a/tests/test_browser_use.py b/tests/test_browser_use.py index cc6c11ed..84ed23a9 100644 --- a/tests/test_browser_use.py +++ b/tests/test_browser_use.py @@ -105,9 +105,15 @@ async def test_browser_use_custom(): # api_key=os.getenv("GOOGLE_API_KEY", "") # ) + # llm = utils.get_llm_model( + # provider="deepseek", + # model_name="deepseek-chat", + # temperature=0.8 + # ) + llm = utils.get_llm_model( - provider="deepseek", - model_name="deepseek-chat", + provider="ollama", + model_name="qwen2.5:7b", temperature=0.8 ) diff --git a/tests/test_llm_api.py b/tests/test_llm_api.py index 03d57537..9e2a1d6d 100644 --- a/tests/test_llm_api.py +++ b/tests/test_llm_api.py @@ -106,7 +106,6 @@ def test_deepseek_model(): base_url=os.getenv("DEEPSEEK_ENDPOINT", ""), api_key=os.getenv("DEEPSEEK_API_KEY", "") ) - pdb.set_trace() message = HumanMessage( content=[ {"type": "text", "text": "who are you?"} @@ -116,8 +115,17 @@ def test_deepseek_model(): print(ai_msg.content) +def test_ollama_model(): + from langchain_ollama import ChatOllama + + llm = ChatOllama(model="qwen2.5:7b") + ai_msg = llm.invoke("Sing a ballad of LangChain.") + print(ai_msg.content) + + if __name__ == '__main__': # test_openai_model() # test_gemini_model() # test_azure_openai_model() - test_deepseek_model() + # test_deepseek_model() + test_ollama_model() diff --git a/webui.py b/webui.py index f44bc143..b5741153 100644 --- a/webui.py +++ b/webui.py @@ -255,7 +255,8 @@ def main(): use_vision = gr.Checkbox(label="use vision", value=True) with gr.Row(): llm_provider = gr.Dropdown( - ["anthropic", "openai", "gemini", "azure_openai", "deepseek"], label="LLM Provider", value="gemini" + ["anthropic", "openai", "gemini", "azure_openai", "deepseek", "ollama"], label="LLM Provider", + value="gemini" ) llm_model_name = gr.Textbox(label="LLM Model Name", value="gemini-2.0-flash-exp") llm_temperature = gr.Number(label="LLM Temperature", value=1.0) From 3287b6e00af88b51999d583444b646d590e7b16b Mon Sep 17 00:00:00 2001 From: warmshao Date: Sat, 4 Jan 2025 10:59:22 +0800 Subject: [PATCH 2/2] add ollama to readme --- README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/README.md b/README.md index ecb9bfeb..6b40a9f3 100644 --- a/README.md +++ b/README.md @@ -6,7 +6,7 @@ This project builds upon the foundation of the [browser-use](https://github.com/ 1. **A Brand New WebUI:** We offer a comprehensive web interface that supports a wide range of `browser-use` functionalities. This UI is designed to be user-friendly and enables easy interaction with the browser agent. -2. **Expanded LLM Support:** We've integrated support for various Large Language Models (LLMs), including: Gemini, OpenAI, Azure OpenAI, Anthropic, DeepSeek etc. And we plan to add support for even more models in the future. +2. **Expanded LLM Support:** We've integrated support for various Large Language Models (LLMs), including: Gemini, OpenAI, Azure OpenAI, Anthropic, DeepSeek, Ollama etc. And we plan to add support for even more models in the future. 3. **Custom Browser Support:** You can use your own browser with our tool, eliminating the need to re-login to sites or deal with other authentication challenges. This feature also supports high-definition screen recording.