modify the LLM utils.

Author: chengpeng-wang

src/llmtool/LLM_utils.py

@@ -1,8 +1,8 @@
-# Imports
 from openai import *
 from pathlib import Path
 from typing import Tuple
 import google.generativeai as genai
+import anthropic
 import signal
 import sys
 import tiktoken
@@ -18,8 +18,6 @@
 import boto3
 from utility.errors import RALLMAPIError, RAValueError
 from utility.logger import Logger
-
-
 class LLM:
     """
     An online inference model using different LLMs:
@@ -59,7 +57,7 @@ def infer(
         elif "o3-mini" in self.online_model_name or "o4-mini" in self.online_model_name:
             output = self.infer_with_On_mini_model(message)
         elif "claude" in self.online_model_name:
-            output = self.infer_with_claude(message)
+            output = self.infer_with_claude_key(message)
         elif "deepseek" in self.online_model_name:
             output = self.infer_with_deepseek_model(message)
         else:
@@ -223,63 +221,7 @@ def call_api():
 
         return ""
 
-    # def infer_with_claude(self, message):
-    #     """Infer using the Claude model via AWS Bedrock"""
-    #     if "3.5" in self.online_model_name:
-    #         model_id = "anthropic.claude-3-5-sonnet-20241022-v2:0"
-    #     if "3.7" in self.online_model_name:
-    #         model_id = "us.anthropic.claude-3-7-sonnet-20250219-v1:0"
-
-    #     model_input = [
-    #         {
-    #             "role": "assistant",
-    #             "content": self.systemRole,
-    #         },
-    #         {"role": "user", "content": message},
-    #     ]
-
-    #     body = json.dumps(
-    #         {
-    #             "messages": model_input,
-    #             "max_tokens": 4000,
-    #             "anthropic_version": "bedrock-2023-05-31",
-    #             "temperature": self.temperature,
-    #             "top_k": 50,
-    #         }
-    #     )
-
-    #     def call_api():
-    #         client = boto3.client(
-    #             "bedrock-runtime",
-    #             region_name="us-west-2",
-    #             config=Config(read_timeout=100),
-    #         )
-
-    #         response = (
-    #             client.invoke_model(
-    #                 modelId=model_id, contentType="application/json", body=body
-    #             )["body"]
-    #             .read()
-    #             .decode("utf-8")
-    #         )
-
-    #         response = json.loads(response)
-    #         return response["content"][0]["text"]
-
-    #     tryCnt = 0
-    #     while tryCnt < 5:
-    #         tryCnt += 1
-    #         try:
-    #             output = self.run_with_timeout(call_api, timeout=100)
-    #             if output:
-    #                 return output
-    #         except Exception as e:
-    #             self.logger.print_log(f"API error: {str(e)}")
-    #         time.sleep(2)
-
-    #     return ""
-
-    def infer_with_claude(self, message):
+    def infer_with_claude_aws_bedrock(self, message):
         """Infer using the Claude model via AWS Bedrock"""
         timeout = 500
         model_input = [
@@ -355,3 +297,78 @@ def call_api():
             time.sleep(2)
 
         return ""
+
+    def infer_with_claude_key(self, message):
+        """Infer using the Claude model via API key, with thinking mode for 3.7"""
+        api_key = os.environ.get("ANTHROPIC_API_KEY")
+        if not api_key:
+            raise RALLMAPIError(
+                "Please set the ANTHROPIC_API_KEY environment variable to use Claude models."
+            )
+
+        # Prepare messages - Claude prefers user messages over assistant system messages
+        model_input = [{"role": "user", "content": f"{self.systemRole}\n\n{message}"}]
+
+        def call_api():
+            client = anthropic.Anthropic(api_key=api_key)
+
+            # Determine model and settings based on version
+            if "3.7" in self.online_model_name:
+                # Claude 3.7 with thinking mode enabled
+                model_name = "claude-3-7-sonnet-20250219"
+                api_params = {
+                    "model": model_name,
+                    "messages": model_input,
+                    "max_tokens": self.max_output_length,
+                    "temperature": self.temperature,
+                    # "thinking": {
+                    #     "type": "enabled",
+                    #     "budget_tokens": 10000
+                    # },
+                }
+            else:
+                # Claude 3.5 standard mode
+                model_name = "claude-3-5-sonnet-20241022"
+                api_params = {
+                    "model": model_name,
+                    "messages": model_input,
+                    "max_tokens": self.max_output_length,
+                    "temperature": self.temperature,
+                    # No thinking parameter for 3.5
+                }
+
+            # Make the API call
+            response = client.messages.create(**api_params)
+
+            # Extract response text based on model type
+            if (
+                "3.7" in self.online_model_name
+                and hasattr(response, "content")
+                and len(response.content) > 1
+            ):
+                # For Claude 3.7 with thinking mode, get the final response (skip thinking content)
+                return response.content[-1].text
+            else:
+                # For Claude 3.5 or any standard response
+                return response.content[0].text
+
+        tryCnt = 0
+        max_retries = 5
+        while tryCnt < max_retries:
+            tryCnt += 1
+            try:
+                output = self.run_with_timeout(call_api, timeout=100)
+                if output:
+                    self.logger.print_log(
+                        f"Claude API call successful with {self.online_model_name}"
+                    )
+                    return output
+            except Exception as e:
+                self.logger.print_log(
+                    f"Claude API error (attempt {tryCnt}/{max_retries}): {e}"
+                )
+                if tryCnt == max_retries:
+                    self.logger.print_log("Max retries reached for Claude API")
+            time.sleep(2)
+
+        return ""