Add web search capability with Tavily (per-bot control)

.github/workflows/lint-test.yml

@@ -0,0 +1,66 @@
+name: Lint and Test
+
+on:
+  push:
+    branches: [main]
+  pull_request:
+    branches: [main]
+
+jobs:
+  frontend:
+    runs-on: ubuntu-latest
+
+    steps:
+      - uses: actions/checkout@v4
+
+      - name: Setup Node.js
+        uses: actions/setup-node@v4
+        with:
+          node-version: '20'
+          cache: 'npm'
+          cache-dependency-path: front/package-lock.json
+
+      - name: Install dependencies
+        run: |
+          cd front
+          npm ci
+          npm install
+
+      - name: Run lint
+        run: |
+          cd front
+          npm run lint
+
+      - name: Run tests
+        run: |
+          cd front
+          npm test -- --passWithNoTests --coverage=false --watchAll=false
+
+  backend:
+    runs-on: ubuntu-latest
+
+    steps:
+      - uses: actions/checkout@v4
+
+      - name: Set up Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: '3.12'
+          cache: 'pip'
+          cache-dependency-path: back/requirements.txt
+
+      - name: Install dependencies
+        run: |
+          cd back
+          pip install -r requirements.txt
+          pip install ruff
+
+      - name: Run ruff
+        run: |
+          cd back
+          ruff check bots/ server/
+
+      - name: Run tests
+        run: |
+          cd back
+          pytest

back/.gitignore

@@ -3,3 +3,4 @@ __pycache__
 db.sqlite3
 test.http
 *.pem
+venv/

back/bots/admin.py

@@ -1,5 +1,4 @@
 from django.contrib import admin
-from django.apps import apps
 from .models import Chat, Message, Profile, Bot, UserAccount, UsageLimitHit, AiModel, Device, RevenueCatWebhookEvent
 from django.contrib.auth.models import User
 from django.contrib.auth.admin import UserAdmin as BaseUserAdmin

back/bots/apps.py

@@ -6,4 +6,4 @@ class BotsConfig(AppConfig):
     name = 'bots'
 
     def ready(self):
-        import bots.signals
+        import bots.signals  # noqa: F401

back/bots/migrations/0033_add_enable_web_search.py

@@ -0,0 +1,18 @@
+# Generated by Django 5.1.4 on 2026-04-02
+
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ('bots', '0032_usagelimithit_modified_at'),
+    ]
+
+    operations = [
+        migrations.AddField(
+            model_name='bot',
+            name='enable_web_search',
+            field=models.BooleanField(default=False),
+        ),
+    ]

back/bots/models/bot.py

@@ -23,6 +23,7 @@ class Bot(models.Model):
     response_length = models.IntegerField(default=200)
     restrict_language = models.BooleanField(default=True)
     restrict_adult_topics = models.BooleanField(default=True)
+    enable_web_search = models.BooleanField(default=False)
 
     deleted_at = models.DateTimeField(null=True, blank=True)
     created_at = models.DateTimeField(auto_now_add=True)

back/bots/models/chat.py

@@ -2,15 +2,33 @@
 from django.db import models
 import uuid
 from langchain_aws import ChatBedrock
-from langchain.schema import HumanMessage, SystemMessage, AIMessage
-import requests
+from langchain_core.messages import HumanMessage, SystemMessage, AIMessage
+from langchain_core.tools import tool
+from langchain_core.callbacks.base import BaseCallbackHandler
+from langchain.agents import create_agent
+from tavily import TavilyClient
+import logging
 import base64
 import boto3
 
 from .profile import Profile
 from .bot import Bot
 from .ai_model import AiModel
 
+logger = logging.getLogger(__name__)
+
+class TokenTracker(BaseCallbackHandler):
+    def __init__(self):
+        self.input_tokens = 0
+        self.output_tokens = 0
+
+    def on_llm_end(self, response, *, run_id=None, parent_run_id=None, **kwargs):
+        usage_metadata = getattr(response, 'usage_metadata', None)
+        if usage_metadata:
+            self.input_tokens += usage_metadata.get('input_tokens', 0)
+            self.output_tokens += usage_metadata.get('output_tokens', 0)
+
+
 S3_CLIENT = boto3.client('s3')
 S3_BUCKET = settings.AWS_STORAGE_BUCKET_NAME
 
@@ -45,7 +63,7 @@ def __init__(self, *args, **kwargs):
         self.ai = None
 
     def __str__(self):
-        return self.title if self.user == None else self.user.email + ' - ' + self.title
+        return self.title if self.user is None else self.user.email + ' - ' + self.title
 
     def use_default_model(self, ai=None):
         try:
@@ -63,15 +81,129 @@ def get_response(self, ai=None):
 
         message_list, contains_image = self.get_input()
 
-        # Check if any messages have image_filename and if the model supports images
         if contains_image and self.bot and 'image' not in self.bot.ai_model.supported_input_modalities:
             self.use_default_model(ai)
 
         if self.user.user_account.over_limit():
             return "You have exceeded your daily limit. Please try again tomorrow or upgrade your subscription."
-        response = self.ai.invoke(
-            message_list
-        )
+
+        # Use agent with tool calling if web search is enabled
+        if self.bot and self.bot.enable_web_search and settings.TAVILY_API_KEY:
+            logger.info(f"Web search enabled for bot {self.bot.name}")
+            tavily_client = TavilyClient(api_key=settings.TAVILY_API_KEY)
+
+            @tool
+            def web_search(query: str) -> str:
+                """Search the web for current information. Use this when you need up-to-date information or facts that may not be in your training data."""
+                logger.info(f"🔍 WEB_SEARCH_TOOL_INVOKED: query='{query}'")
+                try:
+                    results = tavily_client.search(query=query)
+                    num_results = len(results.get('results', []))
+                    logger.info(f"🔍 WEB_SEARCH_SUCCESS: returned {num_results} results")
+                    # Format results as a readable string for the model
+                    if results.get('results'):
+                        formatted = "\n".join([
+                            f"- {r.get('title', 'No title')}: {r.get('content', '')[:200]}"
+                            for r in results['results'][:3]
+                        ])
+                        logger.debug(f"🔍 WEB_SEARCH_FORMATTED_RESULTS:\n{formatted}")
+                        return formatted
+                    else:
+                        logger.info("🔍 WEB_SEARCH_NO_RESULTS: empty result set")
+                        return "No results found."
+                except Exception as e:
+                    logger.error(f"🔍 WEB_SEARCH_ERROR: {str(e)}")
+                    return f"Error during search: {str(e)}"
+
+
+            # Create chat model
+            chat_model = ChatBedrock(model_id=self.ai.model_id)
+            tools = [web_search]
+
+            # Create modern agent with tool calling support
+            # This is the recommended approach per LangChain docs
+            agent = create_agent(
+                model=chat_model,
+                tools=tools,
+                system_prompt=self.get_system_message(),
+                debug=settings.DEBUG
+            )
+
+            # Extract text input from message_list for agent
+            agent_input = self._extract_agent_input(message_list)
+
+            logger.info(f"Invoking agent with input: {agent_input[:100]}...")
+            logger.info("🤖 AGENT_INVOKE_START: web_search tool available")
+
+            # Invoke agent - the CompiledStateGraph handles tool loop internally
+            response = agent.invoke({"messages": [HumanMessage(content=agent_input)]})
+
+            logger.info("🤖 AGENT_INVOKE_COMPLETE: got response")
+
+            # Extract response text from the agent result
+            # The response is a dict with 'messages' key containing final messages
+            response_text = ""
+            usage_metadata = {"input_tokens": 0, "output_tokens": 0}
+
+            if isinstance(response, dict) and "messages" in response:
+                for msg in reversed(response["messages"]):
+                    if isinstance(msg, AIMessage):
+                        response_text = msg.content
+                        # Extract token usage from the message metadata
+                        if hasattr(msg, 'usage_metadata') and msg.usage_metadata:
+                            usage_metadata = msg.usage_metadata
+                        break
+            elif isinstance(response, dict) and "output" in response:
+                response_text = response["output"]
+            else:
+                response_text = str(response)
+
+            message_order = self.messages.count()
+
+            input_tokens = usage_metadata.get('input_tokens', 0)
+            output_tokens = usage_metadata.get('output_tokens', 0)
+
+            self.messages.create(
+                text=response_text, 
+                role='assistant', 
+                order=message_order,
+                input_tokens=input_tokens,
+                output_tokens=output_tokens
+            )
+            self.input_tokens += input_tokens
+            self.output_tokens += output_tokens
+            self.save()
+            return response_text
+
+        # Standard response without web search
+        return self.get_response_standard(message_list, ai)
+
+    def _extract_agent_input(self, message_list):
+        """Extract text input from message_list for the agent.
+
+        Handles both simple text and multimodal content.
+        Assumes message_list has system message at index 0 and user message at end.
+        """
+        # Find the last non-system message (should be the user's query)
+        user_input = ""
+        for msg in reversed(message_list):
+            if isinstance(msg, HumanMessage):
+                if isinstance(msg.content, list):
+                    # Multimodal content - extract text
+                    for item in msg.content:
+                        if isinstance(item, dict) and item.get('type') == 'text':
+                            user_input = item.get('text', '')
+                            break
+                else:
+                    # Simple text content
+                    user_input = msg.content
+                break
+
+        return user_input if user_input else "Please help me."
+
+    def get_response_standard(self, message_list, ai=None):
+        """Handle response without web search."""
+        response = self.ai.invoke(message_list)
 
         response_text = response.content
         usage_metadata = response.usage_metadata

back/bots/models/usage_limit_hit.py

@@ -1,8 +1,6 @@
-from datetime import datetime, time
 from django.db import models
 from .user_account import UserAccount
-from django.utils import timezone
-import pytz
+
 
 class UsageLimitHit(models.Model):
     user_account = models.ForeignKey(UserAccount, 

back/bots/serializers/bot_serializer.py

@@ -19,6 +19,7 @@ class Meta:
             'response_length',
             'restrict_adult_topics',
             'restrict_language',
+            'enable_web_search',
             'created_at',
             'modified_at',
             'deleted_at',

back/bots/serializers/chat_serializer.py

@@ -1,5 +1,5 @@
 from rest_framework import serializers
-from bots.models import Chat, Profile, Bot
+from bots.models import Chat
 from .message_serializer import MessageSerializer
 from .profile_serializer import ProfileIdSerializer
 from .bot_serializer import BotSerializer