Merge branch 'pr/491' into development

agent.py

@@ -343,19 +343,28 @@ class Agent:
                         # create log message right away, more responsive
                         self.loop_data.params_temporary["log_item_generating"] = (
                             self.context.log.log(
-                                type="agent", heading=f"{self.agent_name}: Thinking..."
+                                type="agent", heading=f"{self.agent_name}: Generating..."
                             )
                         )
 
+                        async def reasoning_callback(chunk: str, full: str):
+                            if chunk == full:
+                                printer.print("Reasoning: ")  # start of reasoning
+                            printer.stream(chunk)
+                            await self.handle_reasoning_stream(full)
+
                         async def stream_callback(chunk: str, full: str):
                             # output the agent response stream
-                            if chunk:
-                                printer.stream(chunk)
-                                await self.handle_response_stream(full)
-
-                        agent_response = await self.call_chat_model(
-                            prompt, callback=stream_callback
-                        )  # type: ignore
+                            if chunk == full:
+                                printer.print("Response: ")  # start of response
+                            printer.stream(chunk)
+                            await self.handle_response_stream(full)
+
+                        agent_response, _reasoning = await self.call_chat_model(
+                            messages=prompt,
+                            response_callback=stream_callback,
+                            reasoning_callback=reasoning_callback,
+                        )
 
                         await self.handle_intervention(agent_response)
 
@@ -409,7 +418,7 @@ class Agent:
                 # call monologue_end extensions
                 await self.call_extensions("monologue_end", loop_data=self.loop_data)  # type: ignore
 
-    async def prepare_prompt(self, loop_data: LoopData) -> ChatPromptTemplate:
+    async def prepare_prompt(self, loop_data: LoopData) -> list[BaseMessage]:
         self.context.log.set_progress("Building prompt")
 
         # call extensions before setting prompts
@@ -422,12 +431,10 @@ class Agent:
         # and allow extensions to edit them
         await self.call_extensions("message_loop_prompts_after", loop_data=loop_data)
 
-        # extras (memory etc.)
-        # extras: list[history.OutputMessage] = []
-        # for extra in loop_data.extras_persistent.values():
-        #     extras += history.Message(False, content=extra).output()
-        # for extra in loop_data.extras_temporary.values():
-        #     extras += history.Message(False, content=extra).output()
+        # concatenate system prompt
+        system_text = "\n\n".join(loop_data.system)
+
+        # join extras
         extras = history.Message(
             False,
             content=self.read_prompt(
@@ -444,28 +451,23 @@ class Agent:
             loop_data.history_output + extras
         )
 
-        # build chain from system prompt, message history and model
-        system_text = "\n\n".join(loop_data.system)
-        prompt = ChatPromptTemplate.from_messages(
-            [
-                SystemMessage(content=system_text),
-                *history_langchain,
-                # AIMessage(content="JSON:"), # force the LLM to start with json
-            ]
-        )
+        # build full prompt from system prompt, message history and extrS
+        full_prompt: list[BaseMessage] = [
+            SystemMessage(content=system_text),
+            *history_langchain,
+        ]
+        full_text = ChatPromptTemplate.from_messages(full_prompt).format()
 
         # store as last context window content
         self.set_data(
             Agent.DATA_NAME_CTX_WINDOW,
             {
-                "text": prompt.format(),
-                "tokens": self.history.get_tokens()
-                + tokens.approximate_tokens(system_text)
-                + tokens.approximate_tokens(history.output_text(extras)),
+                "text": full_text,
+                "tokens": tokens.approximate_tokens(full_text),
             },
         )
 
-        return prompt
+        return full_prompt
 
     def handle_critical_exception(self, exception: Exception):
         if isinstance(exception, HandledException):
@@ -586,24 +588,21 @@ class Agent:
         return self.history.output_text(human_label="user", ai_label="assistant")
 
     def get_chat_model(self):
-        return models.get_model(
-            models.ModelType.CHAT,
+        return models.get_chat_model(
             self.config.chat_model.provider,
             self.config.chat_model.name,
             **self.config.chat_model.kwargs,
         )
 
     def get_utility_model(self):
-        return models.get_model(
-            models.ModelType.CHAT,
+        return models.get_chat_model(
             self.config.utility_model.provider,
             self.config.utility_model.name,
             **self.config.utility_model.kwargs,
         )
 
     def get_embedding_model(self):
-        return models.get_model(
-            models.ModelType.EMBEDDING,
+        return models.get_embedding_model(
             self.config.embeddings_model.provider,
             self.config.embeddings_model.name,
             **self.config.embeddings_model.kwargs,
@@ -616,36 +615,37 @@ class Agent:
         callback: Callable[[str], Awaitable[None]] | None = None,
         background: bool = False,
     ):
-        prompt = ChatPromptTemplate.from_messages(
-            [SystemMessage(content=system), HumanMessage(content=message)]
-        )
-
-        response = ""
-
-        # model class
         model = self.get_utility_model()
 
         # rate limiter
         limiter = await self.rate_limiter(
-            self.config.utility_model, prompt.format(), background
+            self.config.utility_model, f"SYSTEM: {system}\nUSER: {message}", background
         )
 
-        async for chunk in (prompt | model).astream({}):
-            await self.handle_intervention()  # wait for intervention and handle it, if paused
-
-            content = models.parse_chunk(chunk)
-            limiter.add(output=tokens.approximate_tokens(content))
-            response += content
+        # add output tokens to rate limiter in tokens callback
+        async def tokens_callback(delta: str, tokens: int):
+            await self.handle_intervention()
+            limiter.add(output=tokens)
 
+        # propagate stream to callback if set
+        async def stream_callback(chunk: str, total: str):
             if callback:
-                await callback(content)
+                await callback(chunk)
+
+        response, _reasoning = await model.unified_call(
+            system_message=system,
+            user_message=message,
+            response_callback=stream_callback,
+            tokens_callback=tokens_callback,
+        )
 
         return response
 
     async def call_chat_model(
         self,
-        prompt: ChatPromptTemplate,
-        callback: Callable[[str, str], Awaitable[None]] | None = None,
+        messages: list[BaseMessage],
+        response_callback: Callable[[str, str], Awaitable[None]] | None = None,
+        reasoning_callback: Callable[[str, str], Awaitable[None]] | None = None,
     ):
         response = ""
 
@@ -653,19 +653,24 @@ class Agent:
         model = self.get_chat_model()
 
         # rate limiter
-        limiter = await self.rate_limiter(self.config.chat_model, prompt.format())
-
-        async for chunk in (prompt | model).astream({}):
-            await self.handle_intervention()  # wait for intervention and handle it, if paused
-
-            content = models.parse_chunk(chunk)
-            limiter.add(output=tokens.approximate_tokens(content))
-            response += content
+        limiter = await self.rate_limiter(
+            self.config.chat_model, ChatPromptTemplate.from_messages(messages).format()
+        )
 
-            if callback:
-                await callback(content, response)
+        # add output tokens to rate limiter in tokens callback
+        async def tokens_callback(delta: str, tokens: int):
+            await self.handle_intervention()
+            limiter.add(output=tokens)
+
+        # call model
+        response, reasoning = await model.unified_call(
+            messages=messages,
+            reasoning_callback=reasoning_callback,
+            response_callback=response_callback,
+            tokens_callback=tokens_callback,
+        )
 
-        return response
+        return response, reasoning
 
     async def rate_limiter(
         self, model_config: ModelConfig, input: str, background: bool = False
@@ -786,6 +791,13 @@ class Agent:
                 content=f"{self.agent_name}: Message misformat, no valid tool request found.",
             )
 
+    async def handle_reasoning_stream(self, stream: str):
+        await self.call_extensions(
+            "reasoning_stream",
+            loop_data=self.loop_data,
+            text=stream,
+        )
+
     async def handle_response_stream(self, stream: str):
         try:
             if len(stream) < 25:

initialize.py

@@ -7,6 +7,24 @@ from python.helpers.print_style import PrintStyle
 def initialize_agent():
     current_settings = settings.get_settings()
 
+    def _normalize_model_kwargs(kwargs: dict) -> dict:
+        # convert string values that represent valid Python numbers to numeric types
+        result = {}
+        for key, value in kwargs.items():
+            if isinstance(value, str):
+                # try to convert string to number if it's a valid Python number
+                try:
+                    # try int first, then float
+                    result[key] = int(value)
+                except ValueError:
+                    try:
+                        result[key] = float(value)
+                    except ValueError:
+                        result[key] = value
+            else:
+                result[key] = value
+        return result
+
     # chat model from user settings
     chat_llm = ModelConfig(
         provider=models.ModelProvider[current_settings["chat_model_provider"]],
@@ -16,7 +34,7 @@ def initialize_agent():
         limit_requests=current_settings["chat_model_rl_requests"],
         limit_input=current_settings["chat_model_rl_input"],
         limit_output=current_settings["chat_model_rl_output"],
-        kwargs=current_settings["chat_model_kwargs"],
+        kwargs=_normalize_model_kwargs(current_settings["chat_model_kwargs"]),
     )
 
     # utility model from user settings
@@ -27,21 +45,21 @@ def initialize_agent():
         limit_requests=current_settings["util_model_rl_requests"],
         limit_input=current_settings["util_model_rl_input"],
         limit_output=current_settings["util_model_rl_output"],
-        kwargs=current_settings["util_model_kwargs"],
+        kwargs=_normalize_model_kwargs(current_settings["util_model_kwargs"]),
     )
     # embedding model from user settings
     embedding_llm = ModelConfig(
         provider=models.ModelProvider[current_settings["embed_model_provider"]],
         name=current_settings["embed_model_name"],
         limit_requests=current_settings["embed_model_rl_requests"],
-        kwargs=current_settings["embed_model_kwargs"],
+        kwargs=_normalize_model_kwargs(current_settings["embed_model_kwargs"]),
     )
     # browser model from user settings
     browser_llm = ModelConfig(
         provider=models.ModelProvider[current_settings["browser_model_provider"]],
         name=current_settings["browser_model_name"],
         vision=current_settings["browser_model_vision"],
-        kwargs=current_settings["browser_model_kwargs"],
+        kwargs=_normalize_model_kwargs(current_settings["browser_model_kwargs"]),
     )
     # agent configuration
     config = AgentConfig(

models.py

@@ -1,38 +1,37 @@
 from enum import Enum
 import os
-from typing import Any
-from langchain_openai import (
-    ChatOpenAI,
-    OpenAI,
-    OpenAIEmbeddings,
-    AzureChatOpenAI,
-    AzureOpenAIEmbeddings,
-    AzureOpenAI,
+from typing import (
+    Any,
+    Awaitable,
+    Callable,
+    List,
+    Optional,
+    Iterator,
+    AsyncIterator,
+    Tuple,
+    TypedDict,
 )
-from langchain_community.llms.ollama import Ollama
-from langchain_ollama import ChatOllama
-from langchain_community.embeddings import OllamaEmbeddings
-from langchain_anthropic import ChatAnthropic
-from langchain_groq import ChatGroq
-from langchain_huggingface import (
-    HuggingFaceEmbeddings,
-    ChatHuggingFace,
-    HuggingFaceEndpoint,
-)
-from langchain_google_genai import (
-    ChatGoogleGenerativeAI,
-    HarmBlockThreshold,
-    HarmCategory,
-    embeddings as google_embeddings,
-)
-from langchain_mistralai import ChatMistralAI
 
-# from pydantic.v1.types import SecretStr
-from python.helpers import dotenv, runtime
+from litellm import completion, acompletion, embedding
+from python.helpers import dotenv
 from python.helpers.dotenv import load_dotenv
 from python.helpers.rate_limiter import RateLimiter
+from python.helpers.tokens import approximate_tokens
+
+from langchain_core.language_models.chat_models import SimpleChatModel
+from langchain_core.outputs.chat_generation import ChatGenerationChunk
+from langchain_core.callbacks.manager import (
+    CallbackManagerForLLMRun,
+    AsyncCallbackManagerForLLMRun,
+)
+from langchain_core.messages import (
+    BaseMessage,
+    AIMessageChunk,
+    HumanMessage,
+    SystemMessage,
+)
+from langchain.embeddings.base import Embeddings
 
-# environment variables
 load_dotenv()
 
 
@@ -52,40 +51,71 @@ class ModelProvider(Enum):
     MISTRALAI = "Mistral AI"
     OLLAMA = "Ollama"
     OPENAI = "OpenAI"
-    OPENAI_AZURE = "OpenAI Azure"
+    AZURE = "OpenAI Azure"
     OPENROUTER = "OpenRouter"
     SAMBANOVA = "Sambanova"
     OTHER = "Other"
 
 
+class ChatChunk(TypedDict):
+    """Simplified response chunk for chat models."""
+
+    response_delta: str
+    reasoning_delta: str
+
+
 rate_limiters: dict[str, RateLimiter] = {}
 
 
-# Utility function to get API keys from environment variables
-def get_api_key(service):
+def configure_litellm_environment():
+    env_mappings = {
+        "API_KEY_OPENAI": "OPENAI_API_KEY",
+        "API_KEY_ANTHROPIC": "ANTHROPIC_API_KEY",
+        "API_KEY_GROQ": "GROQ_API_KEY",
+        "API_KEY_GOOGLE": "GOOGLE_API_KEY",
+        "API_KEY_MISTRAL": "MISTRAL_API_KEY",
+        "API_KEY_OLLAMA": "OLLAMA_API_KEY",
+        "API_KEY_HUGGINGFACE": "HUGGINGFACE_API_KEY",
+        "API_KEY_OPENAI_AZURE": "AZURE_API_KEY",
+        "API_KEY_DEEPSEEK": "DEEPSEEK_API_KEY",
+        "API_KEY_SAMBANOVA": "SAMBANOVA_API_KEY",
+    }
+    base_url_mappings = {
+        "OPENAI_BASE_URL": "OPENAI_API_BASE",
+        "ANTHROPIC_BASE_URL": "ANTHROPIC_API_BASE",
+        "GROQ_BASE_URL": "GROQ_API_BASE",
+        "GOOGLE_BASE_URL": "GOOGLE_API_BASE",
+        "MISTRAL_BASE_URL": "MISTRAL_API_BASE",
+        "OLLAMA_BASE_URL": "OLLAMA_API_BASE",
+        "HUGGINGFACE_BASE_URL": "HUGGINGFACE_API_BASE",
+        "AZURE_BASE_URL": "AZURE_API_BASE",
+        "DEEPSEEK_BASE_URL": "DEEPSEEK_API_BASE",
+        "SAMBANOVA_BASE_URL": "SAMBANOVA_API_BASE",
+    }
+    for a0, llm in env_mappings.items():
+        val = dotenv.get_dotenv_value(a0)
+        if val and not os.getenv(llm):
+            os.environ[llm] = val
+    for a0_base, llm_base in base_url_mappings.items():
+        val = dotenv.get_dotenv_value(a0_base)
+        if val and not os.getenv(llm_base):
+            os.environ[llm_base] = val
+
+
+def get_api_key(service: str) -> str:
     return (
         dotenv.get_dotenv_value(f"API_KEY_{service.upper()}")
         or dotenv.get_dotenv_value(f"{service.upper()}_API_KEY")
-        or dotenv.get_dotenv_value(
-            f"{service.upper()}_API_TOKEN"
-        )  # Added for CHUTES_API_TOKEN
+        or dotenv.get_dotenv_value(f"{service.upper()}_API_TOKEN")
         or "None"
     )
 
 
-def get_model(type: ModelType, provider: ModelProvider, name: str, **kwargs):
-    fnc_name = f"get_{provider.name.lower()}_{type.name.lower()}"  # function name of model getter
-    model = globals()[fnc_name](name, **kwargs)  # call function by name
-    return model
-
-
 def get_rate_limiter(
     provider: ModelProvider, name: str, requests: int, input: int, output: int
 ) -> RateLimiter:
-    # get or create
     key = f"{provider.name}\\{name}"
     rate_limiters[key] = limiter = rate_limiters.get(key, RateLimiter(seconds=60))
-    # always update
     limiter.limits["requests"] = requests or 0
     limiter.limits["input"] = input or 0
     limiter.limits["output"] = output or 0
@@ -102,332 +132,385 @@ def parse_chunk(chunk: Any):
     return content
 
 
-# Ollama models
-def get_ollama_base_url():
-    return (
-        dotenv.get_dotenv_value("OLLAMA_BASE_URL")
-        or f"http://{runtime.get_local_url()}:11434"
-    )
-
-
-def get_ollama_chat(
-    model_name: str,
-    base_url=None,
-    num_ctx=8192,
-    **kwargs,
-):
-    if not base_url:
-        base_url = get_ollama_base_url()
-    return ChatOllama(
-        model=model_name,
-        base_url=base_url,
-        num_ctx=num_ctx,
-        **kwargs,
+def _parse_chunk(chunk: Any) -> ChatChunk:
+    delta = chunk["choices"][0].get("delta", {})
+    message = chunk["choices"][0].get("model_extra", {}).get("message", {})
+    response_delta = (
+        delta.get("content", "")
+        if isinstance(delta, dict)
+        else getattr(delta, "content", "")
+    ) or (
+        message.get("content", "") 
+        if isinstance(message, dict) 
+        else getattr(message, "content", "")
     )
-
-
-def get_ollama_embedding(
-    model_name: str,
-    base_url=None,
-    num_ctx=8192,
-    **kwargs,
-):
-    if not base_url:
-        base_url = get_ollama_base_url()
-    return OllamaEmbeddings(
-        model=model_name, base_url=base_url, num_ctx=num_ctx, **kwargs
+    reasoning_delta = (
+        delta.get("reasoning_content", "")
+        if isinstance(delta, dict)
+        else getattr(delta, "reasoning_content", "")
     )
-
-
-# HuggingFace models
-def get_huggingface_chat(
-    model_name: str,
-    api_key=None,
-    **kwargs,
-):
-    # different naming convention here
-    if not api_key:
-        api_key = get_api_key("huggingface") or os.environ["HUGGINGFACEHUB_API_TOKEN"]
-
-    # Initialize the HuggingFaceEndpoint with the specified model and parameters
-    llm = HuggingFaceEndpoint(
-        repo_id=model_name,
-        task="text-generation",
-        do_sample=True,
-        **kwargs,
-    )
-
-    # Initialize the ChatHuggingFace with the configured llm
-    return ChatHuggingFace(llm=llm)
-
-
-def get_huggingface_embedding(model_name: str, **kwargs):
-    return HuggingFaceEmbeddings(model_name=model_name, **kwargs)
-
-
-# LM Studio and other OpenAI compatible interfaces
-def get_lmstudio_base_url():
-    return (
-        dotenv.get_dotenv_value("LM_STUDIO_BASE_URL")
-        or f"http://{runtime.get_local_url()}:1234/v1"
-    )
-
-
-def get_lmstudio_chat(
-    model_name: str,
-    base_url=None,
-    **kwargs,
-):
-    if not base_url:
-        base_url = get_lmstudio_base_url()
-    return ChatOpenAI(model_name=model_name, base_url=base_url, api_key="none", **kwargs)  # type: ignore
-
-
-def get_lmstudio_embedding(
-    model_name: str,
-    base_url=None,
-    **kwargs,
-):
-    if not base_url:
-        base_url = get_lmstudio_base_url()
-    return OpenAIEmbeddings(model=model_name, api_key="none", base_url=base_url, check_embedding_ctx_length=False, **kwargs)  # type: ignore
-
-
-# Anthropic models
-def get_anthropic_chat(
-    model_name: str,
-    api_key=None,
-    base_url=None,
-    **kwargs,
-):
-    if not api_key:
-        api_key = get_api_key("anthropic")
-    if not base_url:
-        base_url = (
-            dotenv.get_dotenv_value("ANTHROPIC_BASE_URL") or "https://api.anthropic.com"
+    return ChatChunk(reasoning_delta=reasoning_delta, response_delta=response_delta)
+
+
+class LiteLLMChatWrapper(SimpleChatModel):
+    model_name: str
+    provider: str
+    kwargs: dict = {}
+
+    def __init__(self, model: str, provider: str, **kwargs: Any):
+        model_value = f"{provider}/{model}"
+        super().__init__(model_name=model_value, provider=provider, kwargs=kwargs)  # type: ignore
+
+    @property
+    def _llm_type(self) -> str:
+        return "litellm-chat"
+
+    def _convert_messages(self, messages: List[BaseMessage]) -> List[dict]:
+        result = []
+        # Map LangChain message types to LiteLLM roles
+        role_mapping = {
+            "human": "user",
+            "ai": "assistant",
+            "system": "system",
+            "tool": "tool",
+        }
+        for m in messages:
+            role = role_mapping.get(m.type, m.type)
+            message_dict = {"role": role, "content": m.content}
+
+            # Handle tool calls for AI messages
+            tool_calls = getattr(m, "tool_calls", None)
+            if tool_calls:
+                # Convert LangChain tool calls to LiteLLM format
+                new_tool_calls = []
+                for tool_call in tool_calls:
+                    # Ensure arguments is a JSON string
+                    args = tool_call["args"]
+                    if isinstance(args, dict):
+                        import json
+
+                        args_str = json.dumps(args)
+                    else:
+                        args_str = str(args)
+
+                    new_tool_calls.append(
+                        {
+                            "id": tool_call.get("id", ""),
+                            "type": "function",
+                            "function": {
+                                "name": tool_call["name"],
+                                "arguments": args_str,
+                            },
+                        }
+                    )
+                message_dict["tool_calls"] = new_tool_calls
+
+            # Handle tool call ID for ToolMessage
+            tool_call_id = getattr(m, "tool_call_id", None)
+            if tool_call_id:
+                message_dict["tool_call_id"] = tool_call_id
+
+            result.append(message_dict)
+        return result
+
+    def _call(
+        self,
+        messages: List[BaseMessage],
+        stop: Optional[List[str]] = None,
+        run_manager: Optional[CallbackManagerForLLMRun] = None,
+        **kwargs: Any,
+    ) -> str:
+        msgs = self._convert_messages(messages)
+        resp = completion(
+            model=self.model_name, messages=msgs, stop=stop, **{**self.kwargs, **kwargs}
         )
-    return ChatAnthropic(model_name=model_name, api_key=api_key, base_url=base_url, **kwargs)  # type: ignore
-
-
-# right now anthropic does not have embedding models, but that might change
-def get_anthropic_embedding(
-    model_name: str,
-    api_key=None,
-    **kwargs,
-):
-    if not api_key:
-        api_key = get_api_key("anthropic")
-    return OpenAIEmbeddings(model=model_name, api_key=api_key, **kwargs)  # type: ignore
-
-
-# OpenAI models
-def get_openai_chat(
-    model_name: str,
-    api_key=None,
-    **kwargs,
-):
-    if not api_key:
-        api_key = get_api_key("openai")
-    return ChatOpenAI(model_name=model_name, api_key=api_key, **kwargs)  # type: ignore
-
-
-def get_openai_embedding(model_name: str, api_key=None, **kwargs):
-    if not api_key:
-        api_key = get_api_key("openai")
-    return OpenAIEmbeddings(model=model_name, api_key=api_key, **kwargs)  # type: ignore
-
-
-def get_openai_azure_chat(
-    deployment_name: str,
-    api_key=None,
-    azure_endpoint=None,
-    **kwargs,
-):
-    if not api_key:
-        api_key = get_api_key("openai_azure")
-    if not azure_endpoint:
-        azure_endpoint = dotenv.get_dotenv_value("OPENAI_AZURE_ENDPOINT")
-    return AzureChatOpenAI(deployment_name=deployment_name, api_key=api_key, azure_endpoint=azure_endpoint, **kwargs)  # type: ignore
-
-
-def get_openai_azure_embedding(
-    deployment_name: str,
-    api_key=None,
-    azure_endpoint=None,
-    **kwargs,
-):
-    if not api_key:
-        api_key = get_api_key("openai_azure")
-    if not azure_endpoint:
-        azure_endpoint = dotenv.get_dotenv_value("OPENAI_AZURE_ENDPOINT")
-    return AzureOpenAIEmbeddings(deployment_name=deployment_name, api_key=api_key, azure_endpoint=azure_endpoint, **kwargs)  # type: ignore
-
-
-# Google models
-def get_google_chat(
-    model_name: str,
-    api_key=None,
-    **kwargs,
-):
-    if not api_key:
-        api_key = get_api_key("google")
-    return ChatGoogleGenerativeAI(model=model_name, google_api_key=api_key, safety_settings={HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT: HarmBlockThreshold.BLOCK_NONE}, **kwargs)  # type: ignore
-
-
-def get_google_embedding(
-    model_name: str,
-    api_key=None,
-    **kwargs,
-):
-    if not api_key:
-        api_key = get_api_key("google")
-    return google_embeddings.GoogleGenerativeAIEmbeddings(model=model_name, google_api_key=api_key, **kwargs)  # type: ignore
-
-
-# Mistral models
-def get_mistralai_chat(
-    model_name: str,
-    api_key=None,
-    **kwargs,
-):
-    if not api_key:
-        api_key = get_api_key("mistral")
-    return ChatMistralAI(model=model_name, api_key=api_key, **kwargs)  # type: ignore
-
-
-# Groq models
-def get_groq_chat(
-    model_name: str,
-    api_key=None,
-    **kwargs,
-):
-    if not api_key:
-        api_key = get_api_key("groq")
-    return ChatGroq(model_name=model_name, api_key=api_key, **kwargs)  # type: ignore
-
-
-# DeepSeek models
-def get_deepseek_chat(
-    model_name: str,
-    api_key=None,
-    base_url=None,
-    **kwargs,
-):
-    if not api_key:
-        api_key = get_api_key("deepseek")
-    if not base_url:
-        base_url = (
-            dotenv.get_dotenv_value("DEEPSEEK_BASE_URL") or "https://api.deepseek.com"
+        parsed = _parse_chunk(resp)
+        return parsed["response_delta"]
+
+    def _stream(
+        self,
+        messages: List[BaseMessage],
+        stop: Optional[List[str]] = None,
+        run_manager: Optional[CallbackManagerForLLMRun] = None,
+        **kwargs: Any,
+    ) -> Iterator[ChatGenerationChunk]:
+        msgs = self._convert_messages(messages)
+        for chunk in completion(
+            model=self.model_name,
+            messages=msgs,
+            stream=True,
+            stop=stop,
+            **{**self.kwargs, **kwargs},
+        ):
+            parsed = _parse_chunk(chunk)
+            # Only yield chunks with non-None content
+            if parsed["response_delta"]:
+                yield ChatGenerationChunk(
+                    message=AIMessageChunk(content=parsed["response_delta"])
+                )
+
+    async def _astream(
+        self,
+        messages: List[BaseMessage],
+        stop: Optional[List[str]] = None,
+        run_manager: Optional[AsyncCallbackManagerForLLMRun] = None,
+        **kwargs: Any,
+    ) -> AsyncIterator[ChatGenerationChunk]:
+        msgs = self._convert_messages(messages)
+        response = await acompletion(
+            model=self.model_name,
+            messages=msgs,
+            stream=True,
+            stop=stop,
+            **{**self.kwargs, **kwargs},
+        )
+        async for chunk in response:  # type: ignore
+            parsed = _parse_chunk(chunk)
+            # Only yield chunks with non-None content
+            if parsed["response_delta"]:
+                yield ChatGenerationChunk(
+                    message=AIMessageChunk(content=parsed["response_delta"])
+                )
+
+    async def unified_call(
+        self,
+        system_message="",
+        user_message="",
+        messages: List[BaseMessage] = [],
+        response_callback: Callable[[str, str], Awaitable[None]] | None = None,
+        reasoning_callback: Callable[[str, str], Awaitable[None]] | None = None,
+        tokens_callback: Callable[[str, int], Awaitable[None]] | None = None,
+        **kwargs: Any,
+    ) -> Tuple[str, str]:
+        # construct messages
+        if system_message:
+            messages.insert(0, SystemMessage(content=system_message))
+        if user_message:
+            messages.append(HumanMessage(content=user_message))
+
+        # convert to litellm format
+        msgs_conv = self._convert_messages(messages)
+
+        # call model
+        _completion = await acompletion(
+            model=self.model_name,
+            messages=msgs_conv,
+            stream=True,
+            **{**self.kwargs, **kwargs},
         )
 
-    return ChatOpenAI(api_key=api_key, model=model_name, base_url=base_url, **kwargs)  # type: ignore
-
-
-# OpenRouter models
-def get_openrouter_chat(
-    model_name: str,
-    api_key=None,
-    base_url=None,
-    **kwargs,
-):
-    if not api_key:
-        api_key = get_api_key("openrouter")
-    if not base_url:
-        base_url = (
-            dotenv.get_dotenv_value("OPEN_ROUTER_BASE_URL")
-            or "https://openrouter.ai/api/v1"
+        # results
+        reasoning = ""
+        response = ""
+
+        # iterate over chunks
+        async for chunk in _completion:  # type: ignore
+            parsed = _parse_chunk(chunk)
+            # collect reasoning delta and call callbacks
+            if parsed["reasoning_delta"]:
+                reasoning += parsed["reasoning_delta"]
+                if reasoning_callback:
+                    await reasoning_callback(parsed["reasoning_delta"], reasoning)
+                if tokens_callback:
+                    await tokens_callback(
+                        parsed["reasoning_delta"],
+                        approximate_tokens(parsed["reasoning_delta"]),
+                    )
+            # collect response delta and call callbacks
+            if parsed["response_delta"]:
+                response += parsed["response_delta"]
+                if response_callback:
+                    await response_callback(parsed["response_delta"], response)
+                if tokens_callback:
+                    await tokens_callback(
+                        parsed["response_delta"],
+                        approximate_tokens(parsed["response_delta"]),
+                    )
+
+        # return complete results
+        return response, reasoning
+
+
+class BrowserCompatibleChatWrapper(LiteLLMChatWrapper):
+    """
+    A wrapper for browser agent that can filter/sanitize messages
+    before sending them to the LLM.
+    """
+
+    def _call(
+        self,
+        messages: List[BaseMessage],
+        stop: Optional[List[str]] = None,
+        run_manager: Optional[CallbackManagerForLLMRun] = None,
+        **kwargs: Any,
+    ) -> str:
+        # In the future, message filtering logic can be added here.
+        result = super()._call(messages, stop, run_manager, **kwargs)
+        return result
+
+    async def _astream(
+        self,
+        messages: List[BaseMessage],
+        stop: Optional[List[str]] = None,
+        run_manager: Optional[AsyncCallbackManagerForLLMRun] = None,
+        **kwargs: Any,
+    ) -> AsyncIterator[ChatGenerationChunk]:
+        # In the future, message filtering logic can be added here.
+        async for chunk in super()._astream(messages, stop, run_manager, **kwargs):
+            yield chunk
+
+
+class LiteLLMEmbeddingWrapper(Embeddings):
+    model_name: str
+    kwargs: dict = {}
+
+    def __init__(self, model: str, provider: str, **kwargs: Any):
+        self.model_name = f"{provider}/{model}" if provider != "openai" else model
+        self.kwargs = kwargs
+
+    def embed_documents(self, texts: List[str]) -> List[List[float]]:
+        resp = embedding(model=self.model_name, input=texts, **self.kwargs)
+        return [
+            item.get("embedding") if isinstance(item, dict) else item.embedding
+            for item in resp.data
+        ]
+
+    def embed_query(self, text: str) -> List[float]:
+        resp = embedding(model=self.model_name, input=[text], **self.kwargs)
+        item = resp.data[0]
+        return item.get("embedding") if isinstance(item, dict) else item.embedding
+
+
+class LocalSentenceTransformerWrapper(Embeddings):
+    """Local wrapper for sentence-transformers models to avoid HuggingFace API calls"""
+
+    def __init__(self, model_name: str, **kwargs: Any):
+        try:
+            from sentence_transformers import SentenceTransformer
+        except ImportError:
+            raise ImportError(
+                "sentence-transformers library is required for local embeddings. Install with: pip install sentence-transformers"
+            )
+
+        # Remove the "sentence-transformers/" prefix if present
+        if model_name.startswith("sentence-transformers/"):
+            model_name = model_name[len("sentence-transformers/") :]
+
+        self.model = SentenceTransformer(model_name, **kwargs)
+        self.model_name = model_name
+
+    def embed_documents(self, texts: List[str]) -> List[List[float]]:
+        embeddings = self.model.encode(texts, convert_to_tensor=False)
+        return embeddings.tolist() if hasattr(embeddings, "tolist") else embeddings
+
+    def embed_query(self, text: str) -> List[float]:
+        embedding = self.model.encode([text], convert_to_tensor=False)
+        result = (
+            embedding[0].tolist() if hasattr(embedding[0], "tolist") else embedding[0]
         )
-    return ChatOpenAI(
-        api_key=api_key, # type: ignore
-        model=model_name,
-        base_url=base_url,
-        stream_usage=True,
-        model_kwargs={
-            "extra_headers": {
-                "HTTP-Referer": "https://agent-zero.ai",
-                "X-Title": "Agent Zero",
-            }
-        },
-        **kwargs,
-    )
+        return result
 
 
-def get_openrouter_embedding(
-    model_name: str,
-    api_key=None,
-    base_url=None,
-    **kwargs,
+def _get_litellm_chat(
+    cls: type = LiteLLMChatWrapper,
+    model_name: str = "",
+    provider_name: str = "",
+    **kwargs: Any,
 ):
-    if not api_key:
-        api_key = get_api_key("openrouter")
-    if not base_url:
-        base_url = (
-            dotenv.get_dotenv_value("OPEN_ROUTER_BASE_URL")
-            or "https://openrouter.ai/api/v1"
-        )
-    return OpenAIEmbeddings(model=model_name, api_key=api_key, base_url=base_url, **kwargs)  # type: ignore
+    provider_name = provider_name.lower()
+
+    configure_litellm_environment()
+    # Use original provider name for API key lookup, fallback to mapped provider name
+    api_key = kwargs.pop("api_key", None) or get_api_key(provider_name)
+
+    # litellm will pick up base_url from env. We just need to control the api_key.
+    base_url = dotenv.get_dotenv_value(f"{provider_name.upper()}_BASE_URL")
+
+    # If a base_url is set, ensure api_key is not passed to litellm
+    if base_url:
+        if "api_key" in kwargs:
+            del kwargs["api_key"]
+    # Only pass API key if no base_url is set and key is not a placeholder
+    elif api_key and api_key not in ("None", "NA"):
+        kwargs["api_key"] = api_key
+
+    # for openrouter add app reference
+    if provider_name == "openrouter":
+        kwargs["extra_headers"] = {
+            "HTTP-Referer": "https://agent-zero.ai",
+            "X-Title": "Agent Zero",
+        }
+
+    return cls(model=model_name, provider=provider_name, **kwargs)
+
+
+def get_litellm_embedding(model_name: str, provider: str, **kwargs: Any):
+    # Check if this is a local sentence-transformers model
+    if provider == "huggingface" and model_name.startswith("sentence-transformers/"):
+        # Use local sentence-transformers instead of LiteLLM for local models
+        return LocalSentenceTransformerWrapper(model_name=model_name, **kwargs)
+
+    configure_litellm_environment()
+    # Use original provider name for API key lookup, fallback to mapped provider name
+    api_key = kwargs.pop("api_key", None) or get_api_key(provider)
+
+    # litellm will pick up base_url from env. We just need to control the api_key.
+    base_url = dotenv.get_dotenv_value(f"{provider.upper()}_BASE_URL")
+
+    # If a base_url is set, ensure api_key is not passed to litellm
+    if base_url:
+        if "api_key" in kwargs:
+            del kwargs["api_key"]
+    # Only pass API key if no base_url is set and key is not a placeholder
+    elif api_key and api_key not in ("None", "NA"):
+        kwargs["api_key"] = api_key
+
+    return LiteLLMEmbeddingWrapper(model=model_name, provider=provider, **kwargs)
+
+
+def get_model(type: ModelType, provider: ModelProvider, name: str, **kwargs: Any):
+    provider_name = provider.name.lower()
+    kwargs = _normalize_chat_kwargs(kwargs)
+    if type == ModelType.CHAT:
+        return _get_litellm_chat(LiteLLMChatWrapper, name, provider_name, **kwargs)
+    elif type == ModelType.EMBEDDING:
+        return get_litellm_embedding(name, provider_name, **kwargs)
+    else:
+        raise ValueError(f"Unsupported model type: {type}")
 
 
-# Sambanova models
-def get_sambanova_chat(
-    model_name: str,
-    api_key=None,
-    base_url=None,
-    max_tokens=1024,
-    **kwargs,
-):
-    if not api_key:
-        api_key = get_api_key("sambanova")
-    if not base_url:
-        base_url = (
-            dotenv.get_dotenv_value("SAMBANOVA_BASE_URL")
-            or "https://fast-api.snova.ai/v1"
-        )
-    return ChatOpenAI(api_key=api_key, model=model_name, base_url=base_url, max_tokens=max_tokens, **kwargs)  # type: ignore
+def get_chat_model(
+    provider: ModelProvider, name: str, **kwargs: Any
+) -> LiteLLMChatWrapper:
+    provider_name = provider.name.lower()
+    kwargs = _normalize_chat_kwargs(kwargs)
+    model = _get_litellm_chat(LiteLLMChatWrapper, name, provider_name, **kwargs)
+    return model
 
 
-# right now sambanova does not have embedding models, but that might change
-def get_sambanova_embedding(
-    model_name: str,
-    api_key=None,
-    base_url=None,
-    **kwargs,
-):
-    if not api_key:
-        api_key = get_api_key("sambanova")
-    if not base_url:
-        base_url = (
-            dotenv.get_dotenv_value("SAMBANOVA_BASE_URL")
-            or "https://fast-api.snova.ai/v1"
-        )
-    return OpenAIEmbeddings(model=model_name, api_key=api_key, base_url=base_url, **kwargs)  # type: ignore
+def get_browser_model(
+    provider: ModelProvider, name: str, **kwargs: Any
+) -> BrowserCompatibleChatWrapper:
+    provider_name = provider.name.lower()
+    kwargs = _normalize_chat_kwargs(kwargs)
+    model = _get_litellm_chat(
+        BrowserCompatibleChatWrapper, name, provider_name, **kwargs
+    )
+    return model
 
 
-# Other OpenAI compatible models
-def get_other_chat(
-    model_name: str,
-    api_key=None,
-    base_url=None,
-    **kwargs,
-):
-    return ChatOpenAI(api_key=api_key, model=model_name, base_url=base_url, **kwargs)  # type: ignore
+def get_embedding_model(
+    provider: ModelProvider, name: str, **kwargs: Any
+) -> LiteLLMEmbeddingWrapper | LocalSentenceTransformerWrapper:
+    provider_name = provider.name.lower()
+    kwargs = _normalize_embedding_kwargs(kwargs)
+    model = get_litellm_embedding(name, provider_name, **kwargs)
+    return model
 
 
-def get_other_embedding(model_name: str, api_key=None, base_url=None, **kwargs):
-    return OpenAIEmbeddings(model=model_name, api_key=api_key, base_url=base_url, **kwargs)  # type: ignore
+def _normalize_chat_kwargs(kwargs: Any) -> Any:
+    return kwargs
 
 
-# Chutes models
-def get_chutes_chat(
-    model_name: str,
-    api_key=None,
-    base_url=None,
-    **kwargs,
-):
-    if not api_key:
-        api_key = get_api_key("chutes")
-    if not base_url:
-        base_url = (
-            dotenv.get_dotenv_value("CHUTES_BASE_URL") or "https://llm.chutes.ai/v1"
-        )
-    return ChatOpenAI(api_key=api_key, model=model_name, base_url=base_url, **kwargs)  # type: ignore
+def _normalize_embedding_kwargs(kwargs: Any) -> Any:
+    return kwargs

preload.py

@@ -22,7 +22,11 @@ async def preload():
         async def preload_embedding():
             if set["embed_model_provider"] == models.ModelProvider.HUGGINGFACE.name:
                 try:
-                    emb_mod = models.get_huggingface_embedding(set["embed_model_name"])
+                    # Use the new LiteLLM-based model system
+                    emb_mod = models.get_embedding_model(
+                        models.ModelProvider.HUGGINGFACE,
+                        set["embed_model_name"]
+                    )
                     emb_txt = await emb_mod.aembed_query("test")
                     return emb_txt
                 except Exception as e:

prompts/agent0/agent.system.tool.response.md

@@ -3,6 +3,7 @@ final answer to user
 ends task processing use only when done or no task active
 put result in text arg
 always use markdown formatting headers bold text lists
+full message is automatically markdown do not wrap ~~~markdown
 use emojis as icons improve readability
 prefer using tables
 focus nice structured output key selling point

prompts/default/agent.system.tool.call_sub.md

@@ -1,63 +1,26 @@
 ### call_subordinate
 
 you can use subordinates for subtasks
-subordinates can be specialized roles
-message field: always describe task details goal overview important details for new subordinate
+subordinates can be scientist coder engineer etc
+message field: always describe role, task details goal overview for new subordinate
 delegate specific subtasks not entire task
 reset arg usage:
   "true": spawn new subordinate
-  "false": continue current conversation
-prompt_profile defines subordinate specialization
-
-#### if you are superior
-- identify new tasks which your main task's completion depends upon
-- break down your main task into subtasks if possible. If the task can not be split execute it yourself
-- only let saubtasks and new depended upon tasks of your main task be handled by subordinates
-- never forward your entire task to a subordinate to avoid endless delegation loops
-
-#### if you are subordinate:
-- superior is {{agent_name}} minus 1
-- execute the task you were assigned
-- delegate further if asked
-- break down tasks and delegate if necessary
-- do not delegate tasks you can accomplish yourself without refining them
-- only subtasks of your current main task are allowed to be delegated. Never delegate your entire task ro prevent endless loops.
-
-#### Arguments:
-- message (string): always describe task details goal overview important details for new subordinate
-- reset (boolean): true: spawn new subordinate, false: continue current conversation
-- prompt_profile (string): defines specialization, only available prompt profiles below, can omit when reset false
-
-##### Prompt Profiles available
-{{prompt_profiles}}
-
-#### example usage
-~~~json
-{
-    "thoughts": [
-        "This task is challenging and requires a data analyst",
-        "The research_agent profile supports data analysis",
-    ],
-    "headline": "Delegating coding fix to subordinate agent",
-    "tool_name": "call_subordinate",
-    "tool_args": {
-        "message": "...",
-        "reset": "true",
-        "prompt_profile": "research_agent",
-    }
-}
-~~~
+  "false": continue existing subordinate
+if superior, orchestrate
+respond to existing subordinates using call_subordinate tool with reset false
 
+example usage
 ~~~json
 {
     "thoughts": [
-        "The response is missing...",
-        "I will ask a subordinate to add...",
+        "The result seems to be ok but...",
+        "I will ask a coder subordinate to fix...",
     ],
     "tool_name": "call_subordinate",
     "tool_args": {
         "message": "...",
-        "reset": "false",
+        "reset": "true"
     }
 }
-~~~
+~~~

python/extensions/reasoning_stream/_10_log_from_stream.py

@@ -0,0 +1,29 @@
+from python.helpers import persist_chat, tokens
+from python.helpers.extension import Extension
+from agent import LoopData
+import asyncio
+from python.helpers.log import LogItem
+from python.helpers import log
+import math
+
+
+class LogFromStream(Extension):
+
+    async def execute(self, loop_data: LoopData = LoopData(), text: str = "", **kwargs):
+
+        # thought length indicator
+        length = math.ceil(len(text) / 10) * 10
+        heading = f"{self.agent.agent_name}: Reasoning ({length})..."
+
+        # create log message and store it in loop data temporary params
+        if "log_item_generating" not in loop_data.params_temporary:
+            loop_data.params_temporary["log_item_generating"] = (
+                self.agent.context.log.log(
+                    type="agent",
+                    heading=heading,
+                )
+            )
+
+        # update log message
+        log_item = loop_data.params_temporary["log_item_generating"]
+        log_item.update(heading=heading, reasoning=text)

python/extensions/response_stream/_10_log_from_stream.py

@@ -4,6 +4,7 @@ from agent import LoopData
 import asyncio
 from python.helpers.log import LogItem
 from python.helpers import log
+import math
 
 
 class LogFromStream(Extension):
@@ -13,20 +14,38 @@ class LogFromStream(Extension):
         loop_data: LoopData = LoopData(),
         text: str = "",
         parsed: dict = {},
-        **kwargs
+        **kwargs,
     ):
 
-        heading = f"{self.agent.agent_name}: Thinking..."
+        heading = f"{self.agent.agent_name}: Generating..."
         if "headline" in parsed:
             heading = f"{self.agent.agent_name}: {parsed['headline']}"
+        elif "thoughts" in parsed:
+            # thought length indicator
+            thoughts = "\n".join(parsed["thoughts"])
+            length = math.ceil(len(thoughts) / 10) * 10
+            heading = f"{self.agent.agent_name}: Thinking ({length})..."
+
+        if "tool_name" in parsed:
+            heading += f" ({parsed['tool_name']})"
 
         # create log message and store it in loop data temporary params
         if "log_item_generating" not in loop_data.params_temporary:
-            loop_data.params_temporary["log_item_generating"] = self.agent.context.log.log(
-                type="agent",
-                heading=heading,
+            loop_data.params_temporary["log_item_generating"] = (
+                self.agent.context.log.log(
+                    type="agent",
+                    heading=heading,
+                )
             )
 
         # update log message
         log_item = loop_data.params_temporary["log_item_generating"]
-        log_item.update(heading=heading, content=text, kvps=parsed)
+
+        # keep reasoning from previous logs in kvps
+        kvps = {}
+        if log_item.kvps is not None and "reasoning" in log_item.kvps:
+            kvps["reasoning"] = log_item.kvps["reasoning"]
+        kvps.update(parsed)
+
+        # update the log item
+        log_item.update(heading=heading, content=text, kvps=kvps)

python/helpers/document_query.py

@@ -42,6 +42,7 @@ from langchain.text_splitter import RecursiveCharacterTextSplitter
 
 DEFAULT_SEARCH_THRESHOLD = 0.5
 
+
 class DocumentQueryStore:
     """
     FAISS Store for document query results.
@@ -85,7 +86,7 @@ class DocumentQueryStore:
             Normalized URI
         """
         # Convert to lowercase
-        normalized = uri.strip() # uri.lower()
+        normalized = uri.strip()  # uri.lower()
 
         # Parse the URL to get scheme
         parsed = urlparse(normalized)
@@ -368,7 +369,9 @@ class DocumentQueryStore:
 
 class DocumentQueryHelper:
 
-    def __init__(self, agent: Agent, progress_callback: Callable[[str], None] | None = None):
+    def __init__(
+        self, agent: Agent, progress_callback: Callable[[str], None] | None = None
+    ):
         self.agent = agent
         self.store = DocumentQueryStore.get(agent)
         self.progress_callback = progress_callback or (lambda x: None)
@@ -414,30 +417,34 @@ class DocumentQueryHelper:
             content = f"!!! No content found for document: {document_uri} matching queries: {json.dumps(questions)}"
             return False, content
 
-        self.progress_callback(f"Processing {len(questions)} questions in context of {len(selected_chunks)} chunks")
+        self.progress_callback(
+            f"Processing {len(questions)} questions in context of {len(selected_chunks)} chunks"
+        )
 
         questions_str = "\n".join([f" *  {question}" for question in questions])
-        content = "\n\n----\n\n".join([chunk.page_content for chunk in selected_chunks.values()])
+        content = "\n\n----\n\n".join(
+            [chunk.page_content for chunk in selected_chunks.values()]
+        )
 
         qa_system_message = self.agent.parse_prompt(
             "fw.document_query.system_prompt.md"
         )
         qa_user_message = f"# Document:\n{content}\n\n# Queries:\n{questions_str}"
 
-        ai_response = await self.agent.call_chat_model(
-            prompt=ChatPromptTemplate.from_messages(
-                [
-                    SystemMessage(content=qa_system_message),
-                    HumanMessage(content=qa_user_message),
-                ]
-            )
+        ai_response, _reasoning = await self.agent.call_chat_model(
+            messages=[
+                SystemMessage(content=qa_system_message),
+                HumanMessage(content=qa_user_message),
+            ]
         )
 
         self.progress_callback(f"Q&A process completed")
 
         return True, str(ai_response)
 
-    async def document_get_content(self, document_uri: str, add_to_db: bool = False) -> str:
+    async def document_get_content(
+        self, document_uri: str, add_to_db: bool = False
+    ) -> str:
         self.progress_callback(f"Fetching document content")
         url = urlparse(document_uri)
         scheme = url.scheme or "file"
@@ -518,7 +525,9 @@ class DocumentQueryHelper:
                 )
             if add_to_db:
                 self.progress_callback(f"Indexing document")
-                success, ids = await self.store.add_document(document_content, document_uri_norm)
+                success, ids = await self.store.add_document(
+                    document_content, document_uri_norm
+                )
                 if not success:
                     self.progress_callback(f"Failed to index document")
                     raise ValueError(

python/helpers/history.py

@@ -534,10 +534,17 @@ def _merge_outputs(a: MessageContent, b: MessageContent) -> MessageContent:
     if isinstance(a, str) and isinstance(b, str):
         return a + "\n" + b
 
-    if not isinstance(a, list):
-        a = [a]
-    if not isinstance(b, list):
-        b = [b]
+    def make_list(obj: MessageContent) -> list[MessageContent]:
+        if isinstance(obj, list):
+            return obj  # type: ignore
+        if isinstance(obj, dict):
+            return [obj]
+        if isinstance(obj, str):
+            return [{"type": "text", "text": obj}]
+        return [obj]
+
+    a = make_list(a)
+    b = make_list(b)
 
     return cast(MessageContent, a + b)
 

python/helpers/memory.py

@@ -117,8 +117,7 @@ class Memory:
             os.makedirs(em_dir, exist_ok=True)
             store = LocalFileStore(em_dir)
 
-        embeddings_model = models.get_model(
-            models.ModelType.EMBEDDING,
+        embeddings_model = models.get_embedding_model(
             model_config.provider,
             model_config.name,
             **model_config.kwargs,

python/tools/browser_agent.py

@@ -1,5 +1,4 @@
 import asyncio
-import json
 import time
 from typing import Optional
 from agent import Agent, InterventionException
@@ -57,6 +56,8 @@ class State:
                 screen={"width": 1024, "height": 2048},
                 viewport={"width": 1024, "height": 2048},
                 args=["--headless=new"],
+                # Use a unique user data directory to avoid conflicts
+                user_data_dir=str(Path.home() / ".config" / "browseruse" / "profiles" / f"agent_{self.agent.context.id}"),
             )
         )
 
@@ -118,25 +119,28 @@ class State:
             )
             return result
 
-        model = models.get_model(
-            type=models.ModelType.CHAT,
+
+        model = models.get_browser_model(
             provider=self.agent.config.browser_model.provider,
             name=self.agent.config.browser_model.name,
             **self.agent.config.browser_model.kwargs,
         )
 
-        self.use_agent = browser_use.Agent(
-            task=task,
-            browser_session=self.browser_session,
-            llm=model,
-            use_vision=self.agent.config.browser_model.vision,
-            extend_system_message=self.agent.read_prompt(
-                "prompts/browser_agent.system.md"
-            ),
-            controller=controller,
-            enable_memory=False,  # Disable memory to avoid state conflicts
-            # available_file_paths=[],
-        )
+        try:
+            self.use_agent = browser_use.Agent(
+                task=task,
+                browser_session=self.browser_session,
+                llm=model,
+                use_vision=self.agent.config.browser_model.vision,
+                extend_system_message=self.agent.read_prompt(
+                    "prompts/browser_agent.system.md"
+                ),
+                controller=controller,
+                enable_memory=False,  # Disable memory to avoid state conflicts
+                # available_file_paths=[],
+            )
+        except Exception as e:
+            raise Exception(f"Browser agent initialization failed. This might be due to model compatibility issues. Error: {e}") from e
 
         self.iter_no = get_iter_no(self.agent)
 

requirements.txt

@@ -10,15 +10,7 @@ flask-basicauth==0.2.0
 flaredantic==0.1.4
 GitPython==3.1.43
 inputimeout==1.0.4
-langchain-anthropic==0.3.3
-langchain-community==0.3.19
-langchain-google-genai==2.1.2
-langchain-groq==0.2.2
-langchain-huggingface==0.1.2
-langchain-mistralai==0.2.4
-langchain-ollama==0.3.0
-langchain-openai==0.3.11
-langchain-unstructured[all-docs]==0.1.6
+langchain-core==0.3.49
 openai-whisper==20240930
 lxml_html_clean==0.3.1
 markdown==3.7
@@ -35,6 +27,8 @@ unstructured[all-docs]==0.16.23
 unstructured-client==0.31.0
 webcolors==24.6.0
 nest-asyncio==1.6.0
+crontab==1.0.1
+litellm==1.72.4
 markdownify==1.1.0
 pymupdf==1.25.3
 pytesseract==0.3.13

test.py

@@ -0,0 +1,74 @@
+import asyncio
+from os import sep
+from langchain_core.messages import HumanMessage, SystemMessage
+from langchain_core.prompts import ChatPromptTemplate
+import models
+from python.helpers import dotenv
+
+
+async def test():
+
+    dotenv.load_dotenv()
+
+    # model_name = "moonshotai/kimi-dev-72b:free"
+    # model_name = "qwen/qwq-32b"
+    # model_name = "qwen/qwen3-32b"
+    # model_name = "anthropic/claude-3.7-sonnet:thinking"
+    model_name = "openai/gpt-4.1-nano"
+    system = ""
+    message = "hello"
+
+    model = models.get_chat_model(models.ModelProvider.OPENROUTER, model_name)
+
+    async def response_callback(chunk: str, full: str):
+        if chunk == full:
+            print("\n")
+            print("Response:")
+        print(chunk, end="", flush=True)
+
+    async def reasoning_callback(chunk: str, full: str):
+        if chunk == full:
+            print("\n")
+            print("Reasoning:")
+        print(chunk, end="", flush=True)
+
+    response, reasoning = await model.unified_call(
+        system_message=system,
+        user_message=message,
+        response_callback=response_callback,
+        reasoning_callback=reasoning_callback,
+    )
+
+    print("\n")
+    print("Final:")
+    print("Reasoning:", reasoning)
+    print("Response:", response)
+
+
+async def test2():
+    
+    dotenv.load_dotenv()
+
+    import initialize
+    config = initialize.initialize_agent()
+
+    model = models.get_browser_model(
+        provider=config.browser_model.provider,
+        name=config.browser_model.name,
+        **config.browser_model.kwargs,
+    )
+
+    response, reasoning = await model.unified_call(
+        system_message="",
+        user_message="hi",
+    )
+
+    print("\n")
+    print("Final:")
+    print("Reasoning:", reasoning)
+    print("Response:", response)
+
+
+if __name__ == "__main__":
+    # asyncio.run(test())
+    asyncio.run(test2())

webui/index.css

@@ -1572,6 +1572,11 @@ input:checked + .slider:before {
   display: auto;
 }
 
+.msg-thoughts .kvps-val {
+  max-height: 20em;
+  overflow: auto;
+}
+
 .msg-content {
   margin-bottom: 0;
 }

webui/js/messages.js

@@ -537,7 +537,7 @@ function drawKvps(container, kvps, latex) {
     for (let [key, value] of Object.entries(kvps)) {
       const row = table.insertRow();
       row.classList.add("kvps-row");
-      if (key === "thoughts" || key === "reflection")
+      if (key === "thoughts" || key === "reasoning") // TODO: find a better way to determine special class assignment
         row.classList.add("msg-thoughts");
 
       const th = row.insertCell();
@@ -545,6 +545,9 @@ function drawKvps(container, kvps, latex) {
       th.classList.add("kvps-key");
 
       const td = row.insertCell();
+      const tdiv = document.createElement("div");
+      tdiv.classList.add("kvps-val");
+      td.appendChild(tdiv);
 
       if (Array.isArray(value)) {
         for (const item of value) {
@@ -562,7 +565,7 @@ function drawKvps(container, kvps, latex) {
           imgElement.classList.add("kvps-img");
           imgElement.src = value.replace("img://", "/image_get?path=");
           imgElement.alt = "Image Attachment";
-          td.appendChild(imgElement);
+          tdiv.appendChild(imgElement);
 
           // Add click handler and cursor change
           imgElement.style.cursor = "pointer";
@@ -570,15 +573,14 @@ function drawKvps(container, kvps, latex) {
             openImageModal(imgElement.src, 1000);
           });
 
-          td.appendChild(imgElement);
         } else {
           const pre = document.createElement("pre");
-          pre.classList.add("kvps-val");
+          // pre.classList.add("kvps-val");
           //   if (row.classList.contains("msg-thoughts")) {
           const span = document.createElement("span");
           span.innerHTML = convertHTML(value);
           pre.appendChild(span);
-          td.appendChild(pre);
+          tdiv.appendChild(pre);
           addCopyButtonToElement(row);
 
           // Add click handler