symbolicai 1.4.0__py3-none-any.whl → 1.6.0__py3-none-any.whl

symai/backend/engines/search/engine_firecrawl.py

@@ -0,0 +1,333 @@
+import json
+import logging
+from copy import deepcopy
+from dataclasses import dataclass
+from typing import Any
+from urllib.parse import parse_qsl, urlencode, urlsplit, urlunsplit
+
+from firecrawl import Firecrawl
+from firecrawl.v2.types import ScrapeOptions
+
+from ....symbol import Result
+from ....utils import UserMessage
+from ...base import Engine
+from ...settings import SYMAI_CONFIG
+
+logging.getLogger("requests").setLevel(logging.ERROR)
+logging.getLogger("urllib3").setLevel(logging.ERROR)
+logging.getLogger("httpx").setLevel(logging.ERROR)
+
+TRACKING_KEYS = {
+    "utm_source",
+    "utm_medium",
+    "utm_campaign",
+    "utm_term",
+    "utm_content",
+}
+
+
+@dataclass
+class Citation:
+    id: int
+    title: str
+    url: str
+    start: int
+    end: int
+
+    def __hash__(self):
+        return hash((self.url,))
+
+
+class FirecrawlSearchResult(Result):
+    def __init__(
+        self, value: dict[str, Any] | Any, max_chars_per_result: int | None = None, **kwargs
+    ) -> None:
+        raw_dict = value.model_dump() if hasattr(value, "model_dump") else value
+        super().__init__(raw_dict, **kwargs)
+        self._citations: list[Citation] = []
+        self._max_chars_per_result = max_chars_per_result
+        try:
+            text, citations = self._build_text_and_citations(raw_dict)
+            self._value = text
+            self._citations = citations
+        except Exception as e:
+            self._value = None
+            UserMessage(f"Failed to parse Firecrawl search response: {e}", raise_with=ValueError)
+
+    def _build_text_and_citations(self, data: dict[str, Any]) -> tuple[str, list[Citation]]:
+        results = []
+        for source in ["web", "news", "images"]:
+            source_data = data.get(source) or []
+            results.extend(source_data)
+
+        if not results:
+            return "", []
+
+        parts = []
+        citations = []
+        cursor = 0
+
+        for idx, item in enumerate(results, 1):
+            # Handle both SearchResultWeb (url/title at top level) and Document (url/title in metadata)
+            metadata = item.get("metadata") or {}
+            url = item.get("url") or metadata.get("url") or metadata.get("source_url") or ""
+            title = item.get("title") or metadata.get("title") or ""
+
+            if not url:
+                continue
+
+            # Check if this is a scraped result (has markdown content)
+            markdown = item.get("markdown", "")
+            if markdown:
+                content = markdown
+                if self._max_chars_per_result and len(content) > self._max_chars_per_result:
+                    content = content[: self._max_chars_per_result] + "..."
+                result_text = f"{title}\n{url}\n{content}"
+            else:
+                description = (
+                    item.get("description")
+                    or item.get("snippet")
+                    or metadata.get("description")
+                    or ""
+                )
+                result_text = f"{title}\n{url}"
+                if description:
+                    if self._max_chars_per_result and len(description) > self._max_chars_per_result:
+                        description = description[: self._max_chars_per_result] + "..."
+                    result_text += f"\n{description}"
+
+            if parts:
+                parts.append("\n\n")
+                cursor += 2
+
+            parts.append(result_text)
+            cursor += len(result_text)
+
+            marker = f"[{idx}]"
+            start = cursor
+            parts.append(marker)
+            cursor += len(marker)
+
+            citations.append(Citation(id=idx, title=title, url=url, start=start, end=cursor))
+
+        text = "".join(parts)
+        return text, citations
+
+    def __str__(self) -> str:
+        if isinstance(self._value, str) and self._value:
+            return self._value
+        try:
+            return json.dumps(self.raw, indent=2)
+        except TypeError:
+            return str(self.raw)
+
+    def _repr_html_(self) -> str:
+        if isinstance(self._value, str) and self._value:
+            return f"<pre>{self._value}</pre>"
+        try:
+            return f"<pre>{json.dumps(self.raw, indent=2)}</pre>"
+        except Exception:
+            return f"<pre>{self.raw!s}</pre>"
+
+    def get_citations(self) -> list[Citation]:
+        return self._citations
+
+
+class FirecrawlExtractResult(Result):
+    """Result wrapper for Firecrawl scrape API responses."""
+
+    def __init__(self, value: Any, **kwargs) -> None:
+        raw_dict = value.model_dump() if hasattr(value, "model_dump") else value
+        super().__init__(raw_dict, **kwargs)
+        try:
+            self._value = self._extract_content(raw_dict)
+        except Exception as e:
+            self._value = None
+            UserMessage(f"Failed to parse Firecrawl scrape response: {e}", raise_with=ValueError)
+
+    def _extract_content(self, data: dict[str, Any]) -> str:
+        content = data.get("markdown") or data.get("html") or data.get("raw_html")
+        if content:
+            return str(content)
+        json_data = data.get("json")
+        if json_data:
+            return json.dumps(json_data, indent=2)
+        return ""
+
+    def __str__(self) -> str:
+        try:
+            return str(self._value or "")
+        except Exception:
+            return ""
+
+    def _repr_html_(self) -> str:
+        try:
+            return f"<pre>{self._value or ''}</pre>"
+        except Exception:
+            return "<pre></pre>"
+
+
+class FirecrawlEngine(Engine):
+    def __init__(self, api_key: str | None = None):
+        super().__init__()
+        self.config = deepcopy(SYMAI_CONFIG)
+        self.api_key = api_key or self.config.get("SEARCH_ENGINE_API_KEY")
+        self.model = self.config.get("SEARCH_ENGINE_MODEL")
+        self.name = self.__class__.__name__
+
+        if not self.api_key:
+            UserMessage(
+                "Firecrawl API key not found. Set SEARCH_ENGINE_API_KEY in config or environment.",
+                raise_with=ValueError,
+            )
+
+        try:
+            self.client = Firecrawl(api_key=self.api_key)
+        except Exception as e:
+            UserMessage(f"Failed to initialize Firecrawl client: {e}", raise_with=ValueError)
+
+    def id(self) -> str:
+        if (
+            self.config.get("SEARCH_ENGINE_API_KEY")
+            and str(self.config.get("SEARCH_ENGINE_MODEL", "")).lower() == "firecrawl"
+        ):
+            return "search"
+        return super().id()
+
+    def command(self, *args, **kwargs):
+        super().command(*args, **kwargs)
+        if "SEARCH_ENGINE_API_KEY" in kwargs:
+            self.api_key = kwargs["SEARCH_ENGINE_API_KEY"]
+        if "SEARCH_ENGINE_MODEL" in kwargs:
+            self.model = kwargs["SEARCH_ENGINE_MODEL"]
+
+    def _normalize_url(self, url: str) -> str:
+        parts = urlsplit(url)
+        filtered_query = [
+            (k, v)
+            for k, v in parse_qsl(parts.query, keep_blank_values=True)
+            if k not in TRACKING_KEYS and not k.lower().startswith("utm_")
+        ]
+        query = urlencode(filtered_query, doseq=True)
+        return urlunsplit((parts.scheme, parts.netloc, parts.path, query, parts.fragment))
+
+    def _search(self, query: str, kwargs: dict[str, Any]):
+        if not query:
+            UserMessage(
+                "FirecrawlEngine._search requires a non-empty query.", raise_with=ValueError
+            )
+
+        max_chars_per_result = kwargs.get("max_chars_per_result")
+
+        # Build search kwargs
+        search_kwargs = {}
+        if "limit" in kwargs:
+            search_kwargs["limit"] = kwargs["limit"]
+        if "location" in kwargs:
+            search_kwargs["location"] = kwargs["location"]
+        if "tbs" in kwargs:
+            search_kwargs["tbs"] = kwargs["tbs"]
+        if "sources" in kwargs:
+            search_kwargs["sources"] = kwargs["sources"]
+        if "categories" in kwargs:
+            search_kwargs["categories"] = kwargs["categories"]
+        if "timeout" in kwargs:
+            search_kwargs["timeout"] = kwargs["timeout"]
+
+        # Build scrape options for search results content
+        scrape_opts = {}
+        if "formats" in kwargs:
+            scrape_opts["formats"] = kwargs["formats"]
+        if "proxy" in kwargs:
+            scrape_opts["proxy"] = kwargs["proxy"]
+        if "only_main_content" in kwargs:
+            scrape_opts["only_main_content"] = kwargs["only_main_content"]
+        if "scrape_location" in kwargs:
+            scrape_opts["location"] = kwargs["scrape_location"]
+        if "include_tags" in kwargs:
+            scrape_opts["include_tags"] = kwargs["include_tags"]
+        if "exclude_tags" in kwargs:
+            scrape_opts["exclude_tags"] = kwargs["exclude_tags"]
+
+        if scrape_opts:
+            search_kwargs["scrape_options"] = ScrapeOptions(**scrape_opts)
+
+        try:
+            result = self.client.search(query, **search_kwargs)
+        except Exception as e:
+            UserMessage(f"Failed to call Firecrawl Search API: {e}", raise_with=ValueError)
+
+        raw = result.model_dump() if hasattr(result, "model_dump") else result
+        return [FirecrawlSearchResult(result, max_chars_per_result=max_chars_per_result)], {
+            "raw_output": raw
+        }
+
+    def _extract(self, url: str, kwargs: dict[str, Any]):
+        normalized_url = self._normalize_url(url)
+
+        # Build scrape kwargs
+        scrape_kwargs = {"formats": kwargs.get("formats", ["markdown"])}
+        if "only_main_content" in kwargs:
+            scrape_kwargs["only_main_content"] = kwargs["only_main_content"]
+        if "timeout" in kwargs:
+            scrape_kwargs["timeout"] = kwargs["timeout"]
+        if "proxy" in kwargs:
+            scrape_kwargs["proxy"] = kwargs["proxy"]
+        if "location" in kwargs:
+            scrape_kwargs["location"] = kwargs["location"]
+        if "max_age" in kwargs:
+            scrape_kwargs["max_age"] = kwargs["max_age"]
+        if "store_in_cache" in kwargs:
+            scrape_kwargs["store_in_cache"] = kwargs["store_in_cache"]
+        if "actions" in kwargs:
+            scrape_kwargs["actions"] = kwargs["actions"]
+        if "headers" in kwargs:
+            scrape_kwargs["headers"] = kwargs["headers"]
+        if "include_tags" in kwargs:
+            scrape_kwargs["include_tags"] = kwargs["include_tags"]
+        if "exclude_tags" in kwargs:
+            scrape_kwargs["exclude_tags"] = kwargs["exclude_tags"]
+        if "wait_for" in kwargs:
+            scrape_kwargs["wait_for"] = kwargs["wait_for"]
+        if "mobile" in kwargs:
+            scrape_kwargs["mobile"] = kwargs["mobile"]
+
+        try:
+            result = self.client.scrape(normalized_url, **scrape_kwargs)
+        except Exception as e:
+            UserMessage(f"Failed to call Firecrawl Scrape API: {e}", raise_with=ValueError)
+
+        raw = result.model_dump() if hasattr(result, "model_dump") else result
+        return [FirecrawlExtractResult(result)], {"raw_output": raw, "final_url": normalized_url}
+
+    def forward(self, argument):
+        kwargs = argument.kwargs
+        url = argument.prop.url or kwargs.get("url")
+        if url:
+            return self._extract(str(url), kwargs)
+
+        raw_query = argument.prop.prepared_input
+        if raw_query is None:
+            raw_query = argument.prop.query
+
+        query = str(raw_query or "").strip() if raw_query else ""
+        if not query:
+            UserMessage(
+                "FirecrawlEngine.forward requires at least one non-empty query or url.",
+                raise_with=ValueError,
+            )
+
+        return self._search(query, kwargs)
+
+    def prepare(self, argument):
+        url = argument.kwargs.get("url") or argument.prop.url
+        if url:
+            argument.prop.prepared_input = str(url)
+            return
+
+        query = argument.prop.query
+        if isinstance(query, list):
+            argument.prop.prepared_input = " ".join(str(q) for q in query if q)
+            return
+
+        argument.prop.prepared_input = str(query or "").strip()

symai/backend/engines/search/engine_parallel.py

@@ -66,7 +66,7 @@ class Citation:
         return hash((self.url,))
 
 
-class SearchResult(Result):
+class ParallelSearchResult(Result):
     def __init__(self, value: dict[str, Any] | Any, **kwargs) -> None:
         super().__init__(value, **kwargs)
         if isinstance(value, dict) and value.get("error"):
@@ -286,7 +286,7 @@ class SearchResult(Result):
         return self._citations
 
 
-class ExtractResult(Result):
+class ParallelExtractResult(Result):
     """Result wrapper for Parallel Extract API responses."""
 
     def __init__(self, value: dict[str, Any] | Any, **kwargs) -> None:
@@ -485,7 +485,7 @@ class ParallelEngine(Engine):
             )
         except Exception as e:
             UserMessage(f"Failed to call Parallel Search API: {e}", raise_with=ValueError)
-        return [SearchResult(result)], {"raw_output": result}
+        return [ParallelSearchResult(result)], {"raw_output": result}
 
     def _task(self, queries: list[str], kwargs: dict[str, Any]):
         processor_name = self._coerce_processor(kwargs.get("processor"))
@@ -521,7 +521,7 @@ class ParallelEngine(Engine):
         result = self._fetch_task_result(run.run_id, timeout=timeout, api_timeout=api_timeout)
 
         payload = self._task_result_to_search_payload(result)
-        return [SearchResult(payload)], {
+        return [ParallelSearchResult(payload)], {
             "raw_output": result,
             "task_output": payload.get("task_output"),
             "task_output_type": payload.get("task_output_type"),
@@ -699,7 +699,7 @@ class ParallelEngine(Engine):
             )
         except Exception as e:
             UserMessage(f"Failed to call Parallel Extract API: {e}", raise_with=ValueError)
-        return [ExtractResult(result)], {"raw_output": result, "final_url": url}
+        return [ParallelExtractResult(result)], {"raw_output": result, "final_url": url}
 
     def forward(self, argument):
         kwargs = argument.kwargs

symai/backend/mixin/__init__.py

@@ -11,6 +11,8 @@ from .groq import SUPPORTED_REASONING_MODELS as GROQ_REASONING_MODELS
 from .openai import SUPPORTED_CHAT_MODELS as OPENAI_CHAT_MODELS
 from .openai import SUPPORTED_REASONING_MODELS as OPENAI_REASONING_MODELS
 from .openai import SUPPORTED_RESPONSES_MODELS as OPENAI_RESPONSES_MODELS
+from .openrouter import SUPPORTED_CHAT_MODELS as OPENROUTER_CHAT_MODELS
+from .openrouter import SUPPORTED_REASONING_MODELS as OPENROUTER_REASONING_MODELS
 
 __all__ = [
     "ANTHROPIC_CHAT_MODELS",
@@ -26,4 +28,6 @@ __all__ = [
     "OPENAI_CHAT_MODELS",
     "OPENAI_REASONING_MODELS",
     "OPENAI_RESPONSES_MODELS",
+    "OPENROUTER_CHAT_MODELS",
+    "OPENROUTER_REASONING_MODELS",
 ]

symai/backend/mixin/openrouter.py

@@ -0,0 +1,2 @@
+SUPPORTED_CHAT_MODELS = ["openrouter:moonshotai/kimi-k2.5"]
+SUPPORTED_REASONING_MODELS = []

symai/components.py

@@ -1229,6 +1229,7 @@ class MetadataTracker(Expression):
             and frame.f_code.co_name == "forward"
             and "self" in frame.f_locals
             and isinstance(frame.f_locals["self"], Engine)
+            and arg is not None  # Ensure arg is not None to avoid unpacking error on exceptions
         ):
             _, metadata = arg  # arg contains return value on 'return' event
             engine_name = frame.f_locals["self"].__class__.__name__
@@ -1350,6 +1351,116 @@ class MetadataTracker(Expression):
                     token_details[(engine_name, model_name)]["completion_breakdown"][
                         "reasoning_tokens"
                     ] += 0
+                elif engine_name in ("ClaudeXChatEngine", "ClaudeXReasoningEngine"):
+                    raw_output = metadata["raw_output"]
+                    usage = self._extract_claude_usage(raw_output)
+                    if usage is None:
+                        # Skip if we can't extract usage (shouldn't happen normally)
+                        logger.warning(f"Could not extract usage from {engine_name} response.")
+                        token_details[(engine_name, model_name)]["usage"]["total_calls"] += 1
+                        token_details[(engine_name, model_name)]["prompt_breakdown"][
+                            "cached_tokens"
+                        ] += 0
+                        token_details[(engine_name, model_name)]["completion_breakdown"][
+                            "reasoning_tokens"
+                        ] += 0
+                        continue
+                    input_tokens = getattr(usage, "input_tokens", 0) or 0
+                    output_tokens = getattr(usage, "output_tokens", 0) or 0
+                    token_details[(engine_name, model_name)]["usage"]["prompt_tokens"] += (
+                        input_tokens
+                    )
+                    token_details[(engine_name, model_name)]["usage"]["completion_tokens"] += (
+                        output_tokens
+                    )
+                    # Calculate total tokens
+                    total = input_tokens + output_tokens
+                    token_details[(engine_name, model_name)]["usage"]["total_tokens"] += total
+                    token_details[(engine_name, model_name)]["usage"]["total_calls"] += 1
+                    # Track cache tokens if available
+                    cache_creation = getattr(usage, "cache_creation_input_tokens", 0) or 0
+                    cache_read = getattr(usage, "cache_read_input_tokens", 0) or 0
+                    token_details[(engine_name, model_name)]["prompt_breakdown"][
+                        "cache_creation_tokens"
+                    ] += cache_creation
+                    token_details[(engine_name, model_name)]["prompt_breakdown"][
+                        "cache_read_tokens"
+                    ] += cache_read
+                    # For backward compatibility, also track as cached_tokens
+                    token_details[(engine_name, model_name)]["prompt_breakdown"][
+                        "cached_tokens"
+                    ] += cache_read
+                    # Track reasoning/thinking tokens for ClaudeXReasoningEngine
+                    if engine_name == "ClaudeXReasoningEngine":
+                        thinking_output = metadata.get("thinking", "")
+                        # Store thinking content if available
+                        if thinking_output:
+                            if "thinking_content" not in token_details[(engine_name, model_name)]:
+                                token_details[(engine_name, model_name)]["thinking_content"] = []
+                            token_details[(engine_name, model_name)]["thinking_content"].append(
+                                thinking_output
+                            )
+                    # Note: Anthropic doesn't break down reasoning tokens separately in usage,
+                    # but extended thinking is included in output_tokens
+                    token_details[(engine_name, model_name)]["completion_breakdown"][
+                        "reasoning_tokens"
+                    ] += 0
+                elif engine_name == "GeminiXReasoningEngine":
+                    usage = metadata["raw_output"].usage_metadata
+                    token_details[(engine_name, model_name)]["usage"]["prompt_tokens"] += (
+                        usage.prompt_token_count
+                    )
+                    token_details[(engine_name, model_name)]["usage"]["completion_tokens"] += (
+                        usage.candidates_token_count
+                    )
+                    token_details[(engine_name, model_name)]["usage"]["total_tokens"] += (
+                        usage.total_token_count
+                    )
+                    token_details[(engine_name, model_name)]["usage"]["total_calls"] += 1
+                    # Track cache tokens if available
+                    cache_read = getattr(usage, "cached_content_token_count", 0) or 0
+                    token_details[(engine_name, model_name)]["prompt_breakdown"][
+                        "cached_tokens"
+                    ] += cache_read
+                    # Track thinking content if available
+                    thinking_output = metadata.get("thinking", "")
+                    if thinking_output:
+                        if "thinking_content" not in token_details[(engine_name, model_name)]:
+                            token_details[(engine_name, model_name)]["thinking_content"] = []
+                        token_details[(engine_name, model_name)]["thinking_content"].append(
+                            thinking_output
+                        )
+                    # Note: Gemini reasoning tokens are part of candidates_token_count
+                    token_details[(engine_name, model_name)]["completion_breakdown"][
+                        "reasoning_tokens"
+                    ] += 0
+                elif engine_name == "DeepSeekXReasoningEngine":
+                    usage = metadata["raw_output"].usage
+                    token_details[(engine_name, model_name)]["usage"]["completion_tokens"] += (
+                        usage.completion_tokens
+                    )
+                    token_details[(engine_name, model_name)]["usage"]["prompt_tokens"] += (
+                        usage.prompt_tokens
+                    )
+                    token_details[(engine_name, model_name)]["usage"]["total_tokens"] += (
+                        usage.total_tokens
+                    )
+                    token_details[(engine_name, model_name)]["usage"]["total_calls"] += 1
+                    # Track thinking content if available
+                    thinking_output = metadata.get("thinking", "")
+                    if thinking_output:
+                        if "thinking_content" not in token_details[(engine_name, model_name)]:
+                            token_details[(engine_name, model_name)]["thinking_content"] = []
+                        token_details[(engine_name, model_name)]["thinking_content"].append(
+                            thinking_output
+                        )
+                    # Note: DeepSeek reasoning tokens might be in completion_tokens_details
+                    reasoning_tokens = 0
+                    if hasattr(usage, "completion_tokens_details") and usage.completion_tokens_details:
+                        reasoning_tokens = getattr(usage.completion_tokens_details, "reasoning_tokens", 0) or 0
+                    token_details[(engine_name, model_name)]["completion_breakdown"][
+                        "reasoning_tokens"
+                    ] += reasoning_tokens
                 else:
                     logger.warning(f"Tracking {engine_name} is not supported.")
                     continue
@@ -1361,8 +1472,60 @@ class MetadataTracker(Expression):
         # Convert to normal dict
         return {**token_details}
 
+    def _extract_claude_usage(self, raw_output):
+        """Extract usage information from Claude response (handles both streaming and non-streaming).
+
+        For non-streaming responses, raw_output is a Message object with a .usage attribute.
+        For streaming responses, raw_output is a list of stream events. Usage info is in:
+        - RawMessageStartEvent.message.usage (input_tokens)
+        - RawMessageDeltaEvent.usage (output_tokens)
+        """
+        # Non-streaming: raw_output is a Message with .usage
+        if hasattr(raw_output, "usage"):
+            return raw_output.usage
+
+        # Streaming: raw_output is a list of events
+        if isinstance(raw_output, list):
+            # Accumulate usage from stream events
+            input_tokens = 0
+            output_tokens = 0
+            cache_creation = 0
+            cache_read = 0
+
+            for event in raw_output:
+                # RawMessageStartEvent contains initial usage with input_tokens
+                if hasattr(event, "message") and hasattr(event.message, "usage"):
+                    msg_usage = event.message.usage
+                    input_tokens += getattr(msg_usage, "input_tokens", 0) or 0
+                    cache_creation += getattr(msg_usage, "cache_creation_input_tokens", 0) or 0
+                    cache_read += getattr(msg_usage, "cache_read_input_tokens", 0) or 0
+                # RawMessageDeltaEvent contains usage with output_tokens
+                elif hasattr(event, "usage") and event.usage is not None:
+                    evt_usage = event.usage
+                    output_tokens += getattr(evt_usage, "output_tokens", 0) or 0
+
+            # Create a simple object-like dict to hold usage (using Box for attribute access)
+            return Box({
+                "input_tokens": input_tokens,
+                "output_tokens": output_tokens,
+                "cache_creation_input_tokens": cache_creation,
+                "cache_read_input_tokens": cache_read,
+            })
+
+        return None
+
     def _can_accumulate_engine(self, engine_name: str) -> bool:
-        supported_engines = ("GPTXChatEngine", "GPTXReasoningEngine", "GPTXSearchEngine")
+        supported_engines = (
+            "GPTXChatEngine",
+            "GPTXReasoningEngine",
+            "GPTXSearchEngine",
+            "ClaudeXChatEngine",
+            "ClaudeXReasoningEngine",
+            "GeminiXReasoningEngine",
+            "DeepSeekXReasoningEngine",
+            "GroqEngine",
+            "CerebrasEngine",
+        )
         return engine_name in supported_engines
 
     def _track_parallel_usage_items(self, token_details, engine_name, metadata):
@@ -1388,21 +1551,48 @@ class MetadataTracker(Expression):
 
         metadata_raw_output = metadata["raw_output"]
         accumulated_raw_output = accumulated["raw_output"]
-        if not hasattr(metadata_raw_output, "usage") or not hasattr(
-            accumulated_raw_output, "usage"
-        ):
-            return
 
-        current_usage = metadata_raw_output.usage
-        accumulated_usage = accumulated_raw_output.usage
+        # Handle both OpenAI/Anthropic-style (usage) and Gemini-style (usage_metadata)
+        current_usage = getattr(metadata_raw_output, "usage", None) or getattr(
+            metadata_raw_output, "usage_metadata", None
+        )
+        accumulated_usage = getattr(accumulated_raw_output, "usage", None) or getattr(
+            accumulated_raw_output, "usage_metadata", None
+        )
+
+        if not current_usage or not accumulated_usage:
+            return
 
-        for attr in ["completion_tokens", "prompt_tokens", "total_tokens"]:
+        # Handle both OpenAI-style (completion_tokens, prompt_tokens),
+        # Anthropic-style (output_tokens, input_tokens),
+        # and Gemini-style (candidates_token_count, prompt_token_count) fields
+        token_attrs = [
+            "completion_tokens",
+            "prompt_tokens",
+            "total_tokens",
+            "input_tokens",
+            "output_tokens",
+            "candidates_token_count",
+            "prompt_token_count",
+            "total_token_count",
+        ]
+        for attr in token_attrs:
             if hasattr(current_usage, attr) and hasattr(accumulated_usage, attr):
-                setattr(
-                    accumulated_usage,
-                    attr,
-                    getattr(accumulated_usage, attr) + getattr(current_usage, attr),
-                )
+                current_val = getattr(current_usage, attr) or 0
+                accumulated_val = getattr(accumulated_usage, attr) or 0
+                setattr(accumulated_usage, attr, accumulated_val + current_val)
+
+        # Handle Anthropic cache tokens and Gemini cached tokens
+        cache_attrs = [
+            "cache_creation_input_tokens",
+            "cache_read_input_tokens",
+            "cached_content_token_count",
+        ]
+        for attr in cache_attrs:
+            if hasattr(current_usage, attr) and hasattr(accumulated_usage, attr):
+                current_val = getattr(current_usage, attr) or 0
+                accumulated_val = getattr(accumulated_usage, attr) or 0
+                setattr(accumulated_usage, attr, accumulated_val + current_val)
 
         for detail_attr in ["completion_tokens_details", "prompt_tokens_details"]:
             if not hasattr(current_usage, detail_attr) or not hasattr(
@@ -1508,12 +1698,18 @@ class DynamicEngine(Expression):
         """Create an engine instance based on the model name."""
         # Deferred to avoid components <-> neurosymbolic engine circular imports.
         from .backend.engines.neurosymbolic import ENGINE_MAPPING  # noqa
-        from .backend.engines.neurosymbolic.engine_cerebras import CerebrasEngine  # noqa
+        from .backend.engines.search import SEARCH_ENGINE_MAPPING  # noqa
 
         try:
+            # Check neurosymbolic engines first
             engine_class = ENGINE_MAPPING.get(self.model)
-            if engine_class is None and self.model.startswith("cerebras:"):
-                engine_class = CerebrasEngine
+
+            # Check search engines
+            if engine_class is None:
+                engine_class = SEARCH_ENGINE_MAPPING.get(self.model)
+                if engine_class is not None:
+                    return engine_class(api_key=self.api_key)
+
             if engine_class is None:
                 UserMessage(f"Unsupported model '{self.model}'", raise_with=ValueError)
             return engine_class(api_key=self.api_key, model=self.model)