Merge pull request #3 from danilaplee/ollama

feat: add ollama

Author: danilaplee

.github/workflows/docker.yml

@@ -37,7 +37,7 @@ jobs:
           id: meta
           uses: docker/metadata-action@9ec57ed1fcdbf14dcef7dfbe97b2010124a938b7
           with:
-            images: browseruse/browser-use
+            images: browseruser/browser-use
 
         - name: Build and push Docker image
           id: push
@@ -49,5 +49,5 @@ jobs:
             push: true
             tags: ${{ steps.meta.outputs.tags }}
             labels: ${{ steps.meta.outputs.labels }}
-            cache-from: type=registry,ref=browseruse/browser-use:buildcache
-            cache-to: type=registry,ref=browseruse/browser-use:buildcache,mode=max
+            cache-from: type=registry,ref=browseruser/browser-user:buildcache
+            cache-to: type=registry,ref=browseruser/browser-user:buildcache,mode=max

.gitignore

@@ -2,7 +2,7 @@
 __pycache__/
 *.py[cod]
 *$py.class
-
+ollama
 # C extensions
 *.so
 

Dockerfile

@@ -141,6 +141,8 @@ RUN playwright install-deps
 
 RUN apt-get install xauth -y
 
+RUN pip install --no-cache-dir langchain-ollama
+
 # ensure correct permissions for /tmp/.X11-unix to prevent Xvfb from issuing warnings
 RUN mkdir -p /tmp/.X11-unix && chmod 1777 /tmp/.X11-unix
 
@@ -161,7 +163,7 @@ RUN chown -R appuser:appuser /app
 USER appuser
 
 # Expose port
-EXPOSE 8000
+EXPOSE 9000
 
 # Command to start the application
 CMD ["uvicorn", "server:app", "--host", "0.0.0.0", "--port", "8000"]

api.py

@@ -130,6 +130,7 @@ def calculate_max_tasks():
 
 # Function to execute a task
 async def execute_task(task_id: int, task: str, config: Dict[str, Any], db: Session):
+    result = None
     try:
         # Update status to running
         db_task = db.query(Task).filter(Task.id == task_id).first()
@@ -148,7 +149,11 @@ async def execute_task(task_id: int, task: str, config: Dict[str, Any], db: Sess
         if db_task:
             db_task.status = "completed"
             db_task.result = json.dumps({
-                "videopath":result.videopath
+                "videopath":result.videopath,
+                "result":result.result,
+                "task":result.task,
+                "steps_executed":result.steps_executed,
+                "success":result.success
             })
             db_task.completed_at = datetime.utcnow()
             db.commit()
@@ -158,9 +163,12 @@ async def execute_task(task_id: int, task: str, config: Dict[str, Any], db: Sess
         if db_task:
             db_task.status = "failed"
             db_task.error = str(e)
-            db_task.result = json.dumps({
-                "videopath":result.videopath
-            })
+            if result != None : 
+                db_task.result = json.dumps({
+                    "videopath":result.videopath
+                })
+            else : 
+                db_task.result = json.dumps({})
             db_task.completed_at = datetime.utcnow()
             db.commit()
         await send_error_to_webhook(str(e), "execute_task", task_id)

browser.py

@@ -57,6 +57,9 @@ async def execute_task(self, task: str, config: Dict[str, Any]) -> AgentResponse
             # Initialize browser
             browser = Browser(config=browser_config)
 
+            tool_calling_method = "auto"
+            if "deepseek-r1" in llm_config.model_name:
+                tool_calling_method = "json_mode"
             # Initialize and run agent
             agent = Agent(
                 task=task, 
@@ -65,7 +68,8 @@ async def execute_task(self, task: str, config: Dict[str, Any]) -> AgentResponse
                 max_failures=config.get("max_failures", 5),
                 use_vision=config.get("use_vision", True),
                 memory_interval=config.get("memory_interval", 10),
-                planner_interval=config.get("planner_interval", 1)
+                planner_interval=config.get("planner_interval", 1),
+                tool_calling_method=tool_calling_method
             )
 
             result = await agent.run(max_steps=config.get("max_steps", 5))

compose.yml

@@ -17,8 +17,10 @@ services:
   browser:
     depends_on:
       - db
-    image: browseruse/browser-use:pr-2
+    image: browseruse/browser-use:pr-3
     environment:
+      # - OLLAMA_HOST=http://host.docker.internal:11434
+      - OLLAMA_HOST=${OLLAMA_HOST}
       - ERROR_WEBHOOK_URL=http://localhost:3000
       - NOTIFY_WEBHOOK_URL=http://localhost:3000
       - METRICS_WEBHOOK_URL=http://localhost:3000
@@ -35,6 +37,7 @@ services:
       - CHROME_PERSISTENT_SESSION=true
       - RESOLUTION_WIDTH=1920
       - RESOLUTION_HEIGHT=1080
+      - APP_PORT=9000
     ports:
       - "9000:8000"
     volumes:

server.py

@@ -75,7 +75,10 @@ async def run_agent(
 
         # Initialize browser
         browser = Browser(config=browser_config)
-        
+        tool_calling_method = "auto"
+        if "deepseek-r1" in request.llm_config.model_name:
+            tool_calling_method = "json_mode"
+
         # Initialize and run agent
         agent = Agent(
             task=request.task, 
@@ -85,7 +88,8 @@ async def run_agent(
             generate_gif=request.generate_gif,
             max_failures=request.max_failures,
             memory_interval=request.memory_interval,
-            planner_interval=request.planner_interval
+            planner_interval=request.planner_interval,
+            tool_calling_method=tool_calling_method
         )
 
         result = await agent.run(max_steps=request.max_steps)

settings.py

@@ -6,9 +6,12 @@
 from dotenv import load_dotenv
 from pydantic import BaseModel, Field
 from langchain_openai import ChatOpenAI, AzureChatOpenAI
+from langchain_ollama import ChatOllama
 from pydantic import SecretStr
 from fastapi import HTTPException
 from logging_config import setup_logging, log_info, log_error, log_debug, log_warning
+from langchain_core.messages import BaseMessage, AIMessage
+from langchain_core.runnables import RunnableConfig
 
 # Logging configuration
 logger = logging.getLogger('browser-use.settings')
@@ -26,24 +29,21 @@ class ModelConfig(BaseModel):
     api_key: Optional[str] = Field(None, description="API key for the provider (if needed)")
     azure_endpoint: Optional[str] = Field(None, description="Endpoint for Azure OpenAI (if provider=azure)")
     azure_api_version: Optional[str] = Field(None, description="Azure OpenAI API version (if provider=azure)")
-    temperature: float = Field(0.0, description="Generation temperature (0.0 to 1.0)")
+    temperature: float = Field(0.5, description="Generation temperature (0.0 to 1.0)")
+    base_url: Optional[str] = Field(None, description="api base url")
 
 
-# Configurações do banco de dados
 SQLALCHEMY_DATABASE_URL = os.getenv("DATABASE_URL", "sqlite:///./browser_use.db")
 engine = create_engine(SQLALCHEMY_DATABASE_URL, connect_args={"check_same_thread": False})
 SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
 
-# Configurações da API
 API_HOST = os.getenv("API_HOST", "0.0.0.0")
-API_PORT = int(os.getenv("API_PORT", "8000"))
+API_PORT = int(os.getenv("API_PORT", "9000"))
 API_DEBUG = os.getenv("API_DEBUG", "False").lower() == "true"
 
-# Configurações do OpenAI
 OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
 OPENAI_MODEL = os.getenv("OPENAI_MODEL", "gpt-3.5-turbo")
 
-# Configurações do navegador
 BROWSER_HEADLESS = os.getenv("BROWSER_HEADLESS", "True").lower() == "true"
 BROWSER_TIMEOUT = int(os.getenv("BROWSER_TIMEOUT", "30000")) 
 
@@ -95,6 +95,19 @@ def get_llm(model_config: ModelConfig):
                 azure_endpoint=model_config.azure_endpoint or os.getenv("AZURE_OPENAI_ENDPOINT", ""),
                 api_version=model_config.azure_api_version or "2024-10-21"
             )
+        elif provider == "ollama":
+            if "deepseek-r1" in model_config.model_name :
+                log_info(logger, "initializing special provider for ollama deepseek-r1")
+                return DeepSeekR1ChatOllama(
+                    model=model_config.model_name,
+                    temperature=model_config.temperature,
+                    # num_ctx=32000,
+                    base_url=os.getenv("OLLAMA_HOST")
+                )
+            else: 
+                return ChatOllama(
+                    model=model_config.model_name
+                ) 
         else:
             raise ValueError(f"Unsupported provider: {provider}")
     except Exception as e:
@@ -104,3 +117,62 @@ def get_llm(model_config: ModelConfig):
             "error": str(e)
         }, exc_info=True)
         raise HTTPException(status_code=500, detail=f"Error initializing LLM: {str(e)}")
+    
+class DeepSeekR1ChatOllama(ChatOllama):
+    """Custom chat model for DeepSeek-R1."""
+
+    def invoke(
+        self,
+        input: List[BaseMessage],
+        config: Optional[RunnableConfig] = None,
+        **kwargs: Any,
+    ) -> AIMessage:
+        """Invoke the chat model with DeepSeek-R1 specific processing."""
+        org_ai_message = super().invoke(input, config, **kwargs)
+        org_content = org_ai_message.content
+
+        # Extract reasoning content and main content
+        org_content = str(org_ai_message.content)
+        if "</think>" in org_content:
+            parts = org_content.split("</think>")
+            reasoning_content = parts[0].replace("<think>", "").strip()
+            content = parts[1].strip()
+
+            # Remove JSON Response tag if present
+            if "**JSON Response:**" in content:
+                content = content.split("**JSON Response:**")[-1].strip()
+
+            # Create AIMessage with extra attributes
+            message = AIMessage(content=content)
+            setattr(message, "reasoning_content", reasoning_content)
+            return message
+
+        return AIMessage(content=org_ai_message.content)
+
+    async def ainvoke(
+        self,
+        input: List[BaseMessage],
+        config: Optional[RunnableConfig] = None,
+        **kwargs: Any,
+    ) -> AIMessage:
+        """Async invoke the chat model with DeepSeek-R1 specific processing."""
+        org_ai_message = await super().ainvoke(input, config, **kwargs)
+        org_content = org_ai_message.content
+
+        # Extract reasoning content and main content
+        org_content = str(org_ai_message.content)
+        if "</think>" in org_content:
+            parts = org_content.split("</think>")
+            reasoning_content = parts[0].replace("<think>", "").strip()
+            content = parts[1].strip()
+
+            # Remove JSON Response tag if present
+            if "**JSON Response:**" in content:
+                content = content.split("**JSON Response:**")[-1].strip()
+
+            # Create AIMessage with extra attributes
+            message = AIMessage(content=content)
+            setattr(message, "reasoning_content", reasoning_content)
+            return message
+
+        return AIMessage(content=org_ai_message.content)