Add reasoning (extended thinking) for claude 3.7

backend/app/agents/tools/internet_search.py

@@ -1,9 +1,15 @@
+import logging
+
 from app.agents.tools.agent_tool import AgentTool
-from app.repositories.models.custom_bot import BotModel
+from app.repositories.models.custom_bot import BotModel, InternetToolModel
 from app.routes.schemas.conversation import type_model_name
 from duckduckgo_search import DDGS
+from firecrawl.firecrawl import FirecrawlApp
 from pydantic import BaseModel, Field, root_validator
 
+logger = logging.getLogger(__name__)
+logger.setLevel(logging.INFO)
+
 
 class InternetSearchInput(BaseModel):
     query: str = Field(description="The query to search for on the internet.")
@@ -33,38 +39,133 @@ def validate_country(cls, values):
         return values
 
 
-def internet_search(
-    tool_input: InternetSearchInput, bot: BotModel | None, model: type_model_name | None
-) -> list:
-    query = tool_input.query
-    time_limit = tool_input.time_limit
-    country = tool_input.country
-
+def _search_with_duckduckgo(query: str, time_limit: str, country: str) -> list:
     REGION = country
     SAFE_SEARCH = "moderate"
     MAX_RESULTS = 20
     BACKEND = "api"
+    logger.info(
+        f"Executing DuckDuckGo search with query: {query}, region: {REGION}, time_limit: {time_limit}"
+    )
     with DDGS() as ddgs:
-        return [
-            {
-                "content": result["body"],
-                "source_name": result["title"],
-                "source_link": result["href"],
-            }
-            for result in ddgs.text(
+        results = list(
+            ddgs.text(
                 keywords=query,
                 region=REGION,
                 safesearch=SAFE_SEARCH,
                 timelimit=time_limit,
                 max_results=MAX_RESULTS,
                 backend=BACKEND,
             )
+        )
+        logger.info(f"DuckDuckGo search completed. Found {len(results)} results")
+        return [
+            {
+                "content": result["body"],
+                "source_name": result["title"],
+                "source_link": result["href"],
+            }
+            for result in results
+        ]
+
+
+def _search_with_firecrawl(
+    query: str, api_key: str, country: str, max_results: int = 10
+) -> list:
+    logger.info(f"Searching with Firecrawl. Query: {query}, Max Results: {max_results}")
+
+    try:
+        app = FirecrawlApp(api_key=api_key)
+
+        # Search using Firecrawl
+        # SearchParams: /~https://github.com/mendableai/firecrawl/blob/main/apps/python-sdk/firecrawl/firecrawl.py#L24
+        results = app.search(
+            query,
+            {
+                "limit": max_results,
+                "lang": country,
+                "scrapeOptions": {"formats": ["markdown"], "onlyMainContent": True},
+            },
+        )
+
+        if not results:
+            logger.warning("No results found")
+            return []
+        logger.info(f"results of firecrawl: {results}")
+
+        # Format search results
+        search_results = [
+            {
+                "content": data.get("markdown", {}),
+                "source_name": data.get("title", ""),
+                "source_link": data.get("metadata", {}).get("sourceURL", ""),
+            }
+            for data in results.get("data", [])
+            if isinstance(data, dict)
         ]
 
+        logger.info(f"Found {len(search_results)} results from Firecrawl")
+        return search_results
+
+    except Exception as e:
+        logger.error(f"Error searching with Firecrawl: {e}")
+        raise e
+
+
+def _internet_search(
+    tool_input: InternetSearchInput, bot: BotModel | None, model: type_model_name | None
+) -> list:
+    query = tool_input.query
+    time_limit = tool_input.time_limit
+    country = tool_input.country
+
+    logger.info(
+        f"Internet search request - Query: {query}, Time Limit: {time_limit}, Country: {country}"
+    )
+
+    if bot is None:
+        logger.warning("Bot is None, defaulting to DuckDuckGo search")
+        return _search_with_duckduckgo(query, time_limit, country)
+
+    # Find internet search tool
+    internet_tool = next(
+        (tool for tool in bot.agent.tools if isinstance(tool, InternetToolModel)),
+        None,
+    )
+
+    # If no internet tool found or search engine is duckduckgo, use DuckDuckGo
+    if not internet_tool or internet_tool.search_engine == "duckduckgo":
+        logger.info("No internet tool found or search engine is DuckDuckGo")
+        return _search_with_duckduckgo(query, time_limit, country)
+
+    # Handle Firecrawl search
+    if internet_tool.search_engine == "firecrawl":
+        if not internet_tool.firecrawl_config:
+            raise ValueError("Firecrawl configuration is not set in the bot.")
+
+        try:
+            api_key = internet_tool.firecrawl_config.api_key
+            if not api_key:
+                raise ValueError("Firecrawl API key is empty")
+
+            return _search_with_firecrawl(
+                query=query,
+                api_key=api_key,
+                country=country,
+                max_results=internet_tool.firecrawl_config.max_results,
+            )
+        except Exception as e:
+            logger.error(f"Error with Firecrawl search: {e}")
+            raise e
+
+    # Fallback to DuckDuckGo for any unexpected cases
+    logger.warning("Unexpected search engine configuration, falling back to DuckDuckGo")
+    return _search_with_duckduckgo(query, time_limit, country)
+
 
 internet_search_tool = AgentTool(
     name="internet_search",
     description="Search the internet for information.",
     args_schema=InternetSearchInput,
-    function=internet_search,
+    function=_internet_search,
 )

backend/app/bedrock.py

@@ -19,7 +19,7 @@
     from mypy_boto3_bedrock_runtime.type_defs import (
         ContentBlockTypeDef,
         ConverseResponseTypeDef,
-        ConverseStreamRequestRequestTypeDef,
+        ConverseStreamRequestTypeDef,
         GuardrailConverseContentBlockTypeDef,
         InferenceConfigurationTypeDef,
         MessageTypeDef,
@@ -105,7 +105,8 @@ def compose_args_for_converse_api(
     grounding_source: GuardrailConverseContentBlockTypeDef | None = None,
     tools: dict[str, AgentTool] | None = None,
     stream: bool = True,
-) -> ConverseStreamRequestRequestTypeDef:
+    enable_reasoning: bool = False,
+) -> ConverseStreamRequestTypeDef:
     def process_content(c: ContentModel, role: str) -> list[ContentBlockTypeDef]:
         if c.content_type == "text":
             if (
@@ -142,6 +143,7 @@ def process_content(c: ContentModel, role: str) -> list[ContentBlockTypeDef]:
     inference_config: InferenceConfigurationTypeDef
     additional_model_request_fields: dict[str, Any]
     system_prompts: list[SystemContentBlockTypeDef]
+
     if is_nova_model(model):
         # Special handling for Nova models
         inference_config, additional_model_request_fields = _prepare_nova_model_params(
@@ -159,35 +161,76 @@ def process_content(c: ContentModel, role: str) -> list[ContentBlockTypeDef]:
 
     else:
         # Standard handling for non-Nova models
-        inference_config = {
-            "maxTokens": (
+        if enable_reasoning:
+            budget_tokens = (
+                generation_params.reasoning_params.budget_tokens
+                if generation_params and generation_params.reasoning_params
+                else DEFAULT_GENERATION_CONFIG["reasoning_params"]["budget_tokens"]  # type: ignore
+            )
+            max_tokens = (
                 generation_params.max_tokens
                 if generation_params
                 else DEFAULT_GENERATION_CONFIG["max_tokens"]
-            ),
-            "temperature": (
-                generation_params.temperature
-                if generation_params
-                else DEFAULT_GENERATION_CONFIG["temperature"]
-            ),
-            "topP": (
-                generation_params.top_p
-                if generation_params
-                else DEFAULT_GENERATION_CONFIG["top_p"]
-            ),
-            "stopSequences": (
-                generation_params.stop_sequences
-                if generation_params
-                else DEFAULT_GENERATION_CONFIG.get("stop_sequences", [])
-            ),
-        }
-        additional_model_request_fields = {
-            "top_k": (
-                generation_params.top_k
-                if generation_params
-                else DEFAULT_GENERATION_CONFIG["top_k"]
             )
-        }
+
+            if max_tokens <= budget_tokens:
+                logger.warning(
+                    f"max_tokens ({max_tokens}) must be greater than budget_tokens ({budget_tokens}). "
+                    f"Setting max_tokens to {budget_tokens + 1024}"
+                )
+                max_tokens = budget_tokens + 1024
+
+            inference_config = {
+                "maxTokens": max_tokens,
+                "temperature": 1.0,  # Force temperature to 1.0 when reasoning is enabled
+                "topP": (
+                    generation_params.top_p
+                    if generation_params
+                    else DEFAULT_GENERATION_CONFIG["top_p"]
+                ),
+                "stopSequences": (
+                    generation_params.stop_sequences
+                    if generation_params
+                    else DEFAULT_GENERATION_CONFIG.get("stop_sequences", [])
+                ),
+            }
+            additional_model_request_fields = {
+                # top_k cannot be used with reasoning
+                "thinking": {
+                    "type": "enabled",
+                    "budget_tokens": budget_tokens,
+                },
+            }
+        else:
+            inference_config = {
+                "maxTokens": (
+                    generation_params.max_tokens
+                    if generation_params
+                    else DEFAULT_GENERATION_CONFIG["max_tokens"]
+                ),
+                "temperature": (
+                    generation_params.temperature
+                    if generation_params
+                    else DEFAULT_GENERATION_CONFIG["temperature"]
+                ),
+                "topP": (
+                    generation_params.top_p
+                    if generation_params
+                    else DEFAULT_GENERATION_CONFIG["top_p"]
+                ),
+                "stopSequences": (
+                    generation_params.stop_sequences
+                    if generation_params
+                    else DEFAULT_GENERATION_CONFIG.get("stop_sequences", [])
+                ),
+            }
+            additional_model_request_fields = {
+                "top_k": (
+                    generation_params.top_k
+                    if generation_params
+                    else DEFAULT_GENERATION_CONFIG["top_k"]
+                ),
+            }
         system_prompts = [
             {
                 "text": instruction,
@@ -197,7 +240,7 @@ def process_content(c: ContentModel, role: str) -> list[ContentBlockTypeDef]:
         ]
 
     # Construct the base arguments
-    args: ConverseStreamRequestRequestTypeDef = {
+    args: ConverseStreamRequestTypeDef = {
         "inferenceConfig": inference_config,
         "modelId": get_model_id(model),
         "messages": arg_messages,
@@ -230,7 +273,7 @@ def process_content(c: ContentModel, role: str) -> list[ContentBlockTypeDef]:
 
 
 def call_converse_api(
-    args: ConverseStreamRequestRequestTypeDef,
+    args: ConverseStreamRequestTypeDef,
 ) -> ConverseResponseTypeDef:
     client = get_bedrock_runtime_client()
 

backend/app/bot_remove.py

@@ -9,6 +9,7 @@
     find_usage_plan_by_id,
 )
 from app.repositories.common import RecordNotFoundError, decompose_bot_id
+from app.utils import delete_api_key_from_secret_manager
 
 DOCUMENT_BUCKET = os.environ.get("DOCUMENT_BUCKET", "documents")
 BEDROCK_REGION = os.environ.get("BEDROCK_REGION", "us-east-1")
@@ -75,6 +76,7 @@ def handler(event: dict, context: Any) -> None:
 
     delete_from_s3(user_id, bot_id)
     delete_custom_bot_stack_by_bot_id(bot_id)
+    delete_api_key_from_secret_manager(user_id, bot_id, "firecrawl")
 
     # Check if api published stack exists
     try:

backend/app/config.py

@@ -1,4 +1,4 @@
-from typing_extensions import TypedDict
+from typing_extensions import NotRequired, TypedDict
 
 
 class GenerationParams(TypedDict):
@@ -7,6 +7,7 @@ class GenerationParams(TypedDict):
     top_p: float
     temperature: float
     stop_sequences: list[str]
+    reasoning_params: NotRequired[dict[str, int]]
 
 
 class EmbeddingConfig(TypedDict):
@@ -20,11 +21,15 @@ class EmbeddingConfig(TypedDict):
 # Adjust the values according to your application.
 # See: https://docs.anthropic.com/claude/reference/complete_post
 DEFAULT_GENERATION_CONFIG: GenerationParams = {
-    "max_tokens": 2000,
+    # Minimum (Haiku) is 4096
+    # Ref: https://docs.anthropic.com/en/docs/about-claude/models/all-models#model-comparison
+    "max_tokens": 4096,
     "top_k": 250,
     "top_p": 0.999,
-    "temperature": 0.6,
+    "temperature": 1.0,
     "stop_sequences": ["Human: ", "Assistant: "],
+    # Budget tokens must NOT exceeds max_tokens
+    "reasoning_params": {"budget_tokens": 1024},
 }
 
 # Ref: https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-mistral.html#model-parameters-mistral-request-response