Add Anthropic Model

README.md

@@ -120,9 +120,12 @@ Please checkout the [README.md in the `rest/agent` directory](rest/agent/README.
 
 ## Getting started with TraceRoot
 
-### TraceRoot Cloud (recommended)
+### TraceRoot Cloud (Recommended)
 
-The fastest and most reliable way to get started with TraceRoot is signing up to [TraceRoot Cloud](https://auth.traceroot.ai/).
+The fastest and most reliable way to get started with TraceRoot is signing up
+for free to [TraceRoot Cloud](https://auth.traceroot.ai/) for a 7 day trial.
+You will have 150k trace + logs storage with 30d retentions, 1.5M LLM tokens,
+and AI agent with chat mode.
 
 ### Self-hosting the open-source deploy (Advanced)
 

docker/public/Dockerfile

@@ -87,6 +87,7 @@ ENV DB_CONNECTION_TOKENS_COLLECTION=connection_tokens
 ENV DB_TRACEROOT_TOKENS_COLLECTION=traceroot_tokens
 ENV DB_SUBSCRIPTIONS_COLLECTION=user_subscriptions
 ENV OPENAI_API_KEY=""
+ENV ANTHROPIC_API_KEY=""
 ENV DB_PASSWORD=""
 ENV DB_USER_NAME=traceroot
 

pyproject.toml

@@ -52,6 +52,7 @@ dependencies = [
     "python-multipart==0.0.20",
     "requests==2.32.4",
     "openai==1.97.1",
+    "anthropic==0.29.0",
     "pymongo==4.13.2",
     "boto3==1.39.11",
     "numpy==2.3.1",
@@ -90,6 +91,7 @@ all = [
     "python-multipart==0.0.20",
     "requests==2.32.4",
     "openai==1.97.1",
+    "anthropic==0.29.0",
     "pymongo==4.13.2",
     "boto3==1.39.11",
     "numpy==2.3.1",

rest/agent/agent.py

@@ -28,6 +28,7 @@ async def chat(
         tree: SpanNode,
         chat_history: list[dict] | None = None,
         openai_token: str | None = None,
+        anthropic_token: str | None = None,
         github_token: str | None = None,
         github_file_tasks: set[tuple[str, str, str, str]] | None = None,
         is_github_issue: bool = False,

rest/agent/chat.py

@@ -2,6 +2,8 @@
 import os
 from datetime import datetime, timezone
 
+import httpx
+from anthropic import AsyncAnthropic
 from openai import AsyncOpenAI
 
 try:
@@ -28,16 +30,25 @@
 class Chat:
 
     def __init__(self):
-        api_key = os.getenv("OPENAI_API_KEY")
-        if api_key is None:
+        openai_api_key = os.getenv("OPENAI_API_KEY")
+        anthropic_api_key = os.getenv("ANTHROPIC_API_KEY")
+
+        if openai_api_key is None and anthropic_api_key is None:
             # This means that is using the local mode
             # and user needs to provide the token within
             # the integrate section at first
-            api_key = "fake_openai_api_key"
+            openai_api_key = "fake_openai_api_key"
+            anthropic_api_key = "fake_anthropic_api_key"
             self.local_mode = True
         else:
             self.local_mode = False
-        self.chat_client = AsyncOpenAI(api_key=api_key)
+        self.openai_client = AsyncOpenAI(api_key=openai_api_key)
+        # Explicitly create an httpx client to avoid the internal proxy issue.
+        http_client = httpx.AsyncClient()
+        self.anthropic_client = AsyncAnthropic(
+            api_key=anthropic_api_key,
+            http_client=http_client,
+        )
         self.system_prompt = (
             "You are a helpful TraceRoot.AI assistant that is the best "
             "assistant for debugging with logs, traces, metrics and source "
@@ -81,6 +92,7 @@ async def chat(
         tree: SpanNode,
         chat_history: list[dict] | None = None,
         openai_token: str | None = None,
+        anthropic_token: str | None = None,
     ) -> ChatbotResponse:
         """
         Args:
@@ -99,9 +111,20 @@ async def chat(
         if model == ChatModel.AUTO:
             model = ChatModel.GPT_4O
 
-        # Use local client to avoid race conditions in concurrent calls
+        # Determine if the model is from Anthropic or OpenAI
+        is_anthropic_model = "claude" in model.value
+
+        # Initialize clients, using user-provided tokens if available
+        openai_client = self.openai_client
         if openai_token is not None:
-            client = AsyncOpenAI(api_key=openai_token)
+            openai_client = AsyncOpenAI(api_key=openai_token)
+        anthropic_client = self.anthropic_client
+        if anthropic_token is not None:
+            http_client = httpx.AsyncClient()
+            anthropic_client = AsyncAnthropic(
+                api_key=anthropic_token,
+                http_client=http_client,
+            )
         else:
             client = self.chat_client
 
@@ -196,10 +219,25 @@ async def chat(
                     "status": ActionStatus.PENDING.value,
                 })
 
-        responses: list[ChatOutput] = await asyncio.gather(*[
-            self.chat_with_context_chunks(messages, model, client)
-            for messages in all_messages
-        ])
+        if is_anthropic_model:
+            chat_coros = [
+                self.chat_with_context_chunks_anthropic(
+                    messages, model, anthropic_client, self.system_prompt)
+                for messages in all_messages
+            ]
+        else:
+            for msg_list in all_messages:
+                msg_list.insert(0, {
+                    "role": "system",
+                    "content": self.system_prompt
+                })
+            chat_coros = [
+                self.chat_with_context_chunks_openai(messages, model,
+                                                     openai_client)
+                for messages in all_messages
+            ]
+
+        responses: list[ChatOutput] = await asyncio.gather(*chat_coros)
 
         response_time = datetime.now().astimezone(timezone.utc)
         if len(responses) == 1:
@@ -216,8 +254,8 @@ async def chat(
             response = await chunk_summarize(
                 response_answers=response_answers,
                 response_references=response_references,
-                client=client,
-                model=model,
+                client=openai_client,
+                model=ChatModel.GPT_4O,
             )
             response_content = response.answer
             response_references = response.reference
@@ -243,22 +281,69 @@ async def chat(
             chat_id=chat_id,
         )
 
-    async def chat_with_context_chunks(
+    async def chat_with_context_chunks_openai(
         self,
         messages: list[dict[str, str]],
         model: ChatModel,
         chat_client: AsyncOpenAI,
     ) -> ChatOutput:
-        r"""Chat with context chunks.
-        """
+        r"""Chat with context chunks using an OpenAI model."""
+        # NOTE: `chat_client.responses.parse` seems to be a custom wrapper or
+        # part of a library like `instructor` for structured output.
         response = await chat_client.responses.parse(
-            model=model,
+            model=model.value,
             input=messages,
             text_format=ChatOutput,
             temperature=0.8,
         )
         return response.output[0].content[0].parsed
 
+    async def chat_with_context_chunks_anthropic(
+        self,
+        messages: list[dict[str, str]],
+        model: ChatModel,
+        chat_client: AsyncAnthropic,
+        system_prompt: str,
+    ) -> ChatOutput:
+        r"""Chat with context chunks using an Anthropic model."""
+        try:
+            # Use Anthropic's tool-use feature for structured output
+            response = await chat_client.messages.create(
+                model=model.value,
+                system=system_prompt,
+                messages=messages,
+                max_tokens=4096,
+                temperature=0.8,
+                tools=[{
+                    "name": "provide_answer",
+                    "description": "Answer with references.",
+                    "input_schema": ChatOutput.model_json_schema(),
+                }],
+                tool_choice={
+                    "type": "tool",
+                    "name": "provide_answer"
+                },
+            )
+
+            tool_call = next(
+                (block
+                 for block in response.content if block.type == "tool_use"),
+                None)
+            if tool_call and tool_call.name == "provide_answer":
+                return ChatOutput(**tool_call.input)
+            else:
+                # Fallback if the model fails to use the tool
+                text_content = "".join(block.text for block in response.content
+                                       if block.type == "text")
+                return ChatOutput(
+                    answer=f"Unstructured model response: {text_content}",
+                    reference=[])
+        except Exception as e:
+            print(f"Error calling Anthropic API: {e}")
+            return ChatOutput(
+                answer=f"An error occurred with the Anthropic API: {str(e)}",
+                reference=[])
+
     def get_context_messages(self, context: str) -> list[str]:
         r"""Get the context message.
         """