Add support for Anthropic/Claude and other LLM providers

.env

@@ -1,4 +1,16 @@
+# LLM Provider: "openai" (default), "anthropic", or "ollama"
+LLM_PROVIDER=openai
+
 # OpenAI Configuration
 OPENAI_API_KEY=
 OPENAI_API_VERSION=
-OPENAI_API_TYPE=
+OPENAI_API_TYPE=
+OPENAI_MODEL=gpt-5.2
+
+# Anthropic Configuration
+ANTHROPIC_API_KEY=
+ANTHROPIC_MODEL=claude-sonnet-4-5
+
+# Ollama Configuration (local models)
+OLLAMA_MODEL=llama3
+OLLAMA_BASE_URL=http://localhost:11434

pyproject.toml

@@ -39,6 +39,9 @@ dependencies = [
     "numpy>=1.26.4",
 ]
 
+[project.optional-dependencies]
+anthropic = ["langchain-anthropic~=0.3.12"]
+
 [project.urls]
 "Homepage" = "https://github.com/nasa-jpl/rosa"
 "Bug Tracker" = "https://github.com/nasa-jpl/rosa/issues"

src/rosa/rosa.py

@@ -12,24 +12,31 @@
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
 
-from typing import Any, AsyncIterable, Dict, Literal, Optional, Union
+from __future__ import annotations
 
-from langchain.agents import AgentExecutor
-from langchain.agents.format_scratchpad.openai_tools import (
-    format_to_openai_tool_messages,
-)
-from langchain.agents.output_parsers.openai_tools import OpenAIToolsAgentOutputParser
+import logging
+from contextlib import contextmanager
+from typing import TYPE_CHECKING, Any, AsyncIterable, Dict, Literal, Optional
+
+from langchain.agents import AgentExecutor, create_tool_calling_agent
 from langchain.prompts import MessagesPlaceholder
 from langchain_community.callbacks import get_openai_callback
+from langchain_core.language_models import BaseChatModel
 from langchain_core.messages import AIMessage, HumanMessage
 from langchain_core.prompts import ChatPromptTemplate
-from langchain_ollama import ChatOllama
 from langchain_openai import AzureChatOpenAI, ChatOpenAI
 
+if TYPE_CHECKING:
+    from langchain_anthropic import ChatAnthropic
+
 from .prompts import RobotSystemPrompts, system_prompts
 from .tools import ROSATools
 
-ChatModel = Union[ChatOpenAI, AzureChatOpenAI, ChatOllama]
+logger = logging.getLogger(__name__)
+
+# Runtime-safe type alias: accepts any BaseChatModel, covering OpenAI, Azure,
+# Anthropic, Ollama and any future langchain provider that implements tool calling.
+ChatModel = BaseChatModel
 
 
 class ROSA:
@@ -38,7 +45,9 @@ class ROSA:
 
     Args:
         ros_version (Literal[1, 2]): The version of ROS that the agent will interact with.
-        llm (Union[AzureChatOpenAI, ChatOpenAI, ChatOllama]): The language model to use for generating responses.
+        llm (BaseChatModel): Any langchain chat model that supports tool calling. Tested with
+            ChatOpenAI, AzureChatOpenAI, ChatOllama, and ChatAnthropic. Note: token usage
+            tracking only works with ChatOpenAI and AzureChatOpenAI.
         tools (Optional[list]): A list of additional LangChain tool functions to use with the agent.
         tool_packages (Optional[list]): A list of Python packages containing LangChain tool functions to use.
         prompts (Optional[RobotSystemPrompts]): Custom prompts to use with the agent.
@@ -95,7 +104,6 @@ def __init__(
             ros_version, packages=tool_packages, tools=tools, blacklist=self.__blacklist
         )
         self.__prompts = self._get_prompts(prompts)
-        self.__llm_with_tools = self.__llm.bind_tools(self.__tools.get_tools())
         self.__agent = self._get_agent()
         self.__executor = self._get_executor(verbose=verbose)
         self.__show_token_usage = show_token_usage if not streaming else False
@@ -131,7 +139,7 @@ def invoke(self, query: str) -> str:
             - Token usage is printed if the show_token_usage flag is set.
         """
         try:
-            with get_openai_callback() as cb:
+            with self._token_callback() as cb:
                 result = self.__executor.invoke(
                     {"input": query, "chat_history": self.__chat_history}
                 )
@@ -245,18 +253,15 @@ def _get_executor(self, verbose: bool) -> AgentExecutor:
         return executor
 
     def _get_agent(self):
-        """Create and return an agent for processing user inputs and generating responses."""
-        agent = (
-            {
-                "input": lambda x: x["input"],
-                "agent_scratchpad": lambda x: format_to_openai_tool_messages(
-                    x["intermediate_steps"]
-                ),
-                "chat_history": lambda x: x["chat_history"],
-            }
-            | self.__prompts
-            | self.__llm_with_tools
-            | OpenAIToolsAgentOutputParser()
+        """Create and return an agent for processing user inputs and generating responses.
+
+        Uses create_tool_calling_agent which is provider-agnostic and works with
+        any LLM that supports tool calling (OpenAI, Anthropic, Ollama, etc).
+        """
+        agent = create_tool_calling_agent(
+            llm=self.__llm,
+            tools=self.__tools.get_tools(),
+            prompt=self.__prompts,
         )
         return agent
 
@@ -296,12 +301,28 @@ def _get_prompts(
         )
         return template
 
+    @contextmanager
+    def _token_callback(self):
+        """Context manager for token usage tracking.
+
+        Uses the OpenAI callback when the LLM is an OpenAI-based model,
+        otherwise yields None so the rest of the flow is unaffected.
+        """
+        if isinstance(self.__llm, (ChatOpenAI, AzureChatOpenAI)):
+            with get_openai_callback() as cb:
+                yield cb
+        else:
+            if self.__show_token_usage:
+                logger.warning("Token usage tracking is only supported for OpenAI and Azure models.")
+            yield None
+
     def _print_usage(self, cb):
         """Print the token usage if show_token_usage is enabled."""
-        if cb and self.__show_token_usage:
-            print(f"[bold]Prompt Tokens:[/bold] {cb.prompt_tokens}")
-            print(f"[bold]Completion Tokens:[/bold] {cb.completion_tokens}")
-            print(f"[bold]Total Cost (USD):[/bold] ${cb.total_cost}")
+        if cb is None or not self.__show_token_usage:
+            return
+        print(f"[bold]Prompt Tokens:[/bold] {cb.prompt_tokens}")
+        print(f"[bold]Completion Tokens:[/bold] {cb.completion_tokens}")
+        print(f"[bold]Total Cost (USD):[/bold] ${cb.total_cost}")
 
     def _record_chat_history(self, query: str, response: str):
         """Record the chat history if accumulation is enabled."""

src/turtle_agent/scripts/llm.py

@@ -19,14 +19,55 @@
 
 
 def get_llm(streaming: bool = False):
-    """A helper function to get the LLM instance."""
+    """A helper function to get the LLM instance.
+
+    Supports OpenAI (default), Anthropic and Ollama models.
+    Set the LLM_PROVIDER env variable to switch between providers:
+      - "openai" (default): uses OPENAI_API_KEY
+      - "anthropic": uses ANTHROPIC_API_KEY
+      - "ollama": uses local Ollama instance
+    """
     dotenv.load_dotenv(dotenv.find_dotenv())
 
-    llm = ChatOpenAI(
-        api_key=get_env_variable("OPENAI_API_KEY"),
-        model="gpt-5.1",
-        streaming=streaming,
-    )
+    provider = os.getenv("LLM_PROVIDER", "openai").lower()
+
+    if provider == "openai":
+        llm = ChatOpenAI(
+            api_key=get_env_variable("OPENAI_API_KEY"),
+            model=get_env_variable("OPENAI_MODEL"),
+            streaming=streaming,
+        )
+    elif provider == "anthropic":
+        try:
+            from langchain_anthropic import ChatAnthropic
+        except ImportError:
+            raise ImportError(
+                "langchain-anthropic is required for Anthropic support. "
+                "Install it with: pip install langchain-anthropic"
+            )
+        llm = ChatAnthropic(
+            api_key=get_env_variable("ANTHROPIC_API_KEY"),
+            model=get_env_variable("ANTHROPIC_MODEL"),
+            streaming=streaming,
+        )
+    elif provider == "ollama":
+        try:
+            from langchain_ollama import ChatOllama
+        except ImportError:
+            raise ImportError(
+                "langchain-ollama is required for Ollama support. "
+                "Install it with: pip install langchain-ollama"
+            )
+        llm = ChatOllama(
+            model=os.getenv("OLLAMA_MODEL", "llama3"),
+            base_url=os.getenv("OLLAMA_BASE_URL", "http://localhost:11434"),
+            streaming=streaming,
+        )
+    else:
+        raise ValueError(
+            f"Unknown LLM provider: '{provider}'. "
+            "Supported providers are: 'openai', 'anthropic', 'ollama'."
+        )
 
     return llm