refactor(anthropic): AnthropicLLM to use Messages API (#32290)

re: #32189

Author: mdrxy

libs/partners/anthropic/langchain_anthropic/llms.py

@@ -3,11 +3,7 @@
 import re
 import warnings
 from collections.abc import AsyncIterator, Iterator, Mapping
-from typing import (
-    Any,
-    Callable,
-    Optional,
-)
+from typing import Any, Callable, Optional
 
 import anthropic
 from langchain_core._api.deprecation import deprecated
@@ -19,25 +15,19 @@
 from langchain_core.language_models.llms import LLM
 from langchain_core.outputs import GenerationChunk
 from langchain_core.prompt_values import PromptValue
-from langchain_core.utils import (
-    get_pydantic_field_names,
-)
-from langchain_core.utils.utils import (
-    _build_model_kwargs,
-    from_env,
-    secret_from_env,
-)
+from langchain_core.utils import get_pydantic_field_names
+from langchain_core.utils.utils import _build_model_kwargs, from_env, secret_from_env
 from pydantic import ConfigDict, Field, SecretStr, model_validator
 from typing_extensions import Self
 
 
 class _AnthropicCommon(BaseLanguageModel):
     client: Any = None  #: :meta private:
     async_client: Any = None  #: :meta private:
-    model: str = Field(default="claude-2", alias="model_name")
+    model: str = Field(default="claude-3-5-sonnet-latest", alias="model_name")
     """Model name to use."""
 
-    max_tokens_to_sample: int = Field(default=1024, alias="max_tokens")
+    max_tokens: int = Field(default=1024, alias="max_tokens_to_sample")
     """Denotes the number of tokens to predict per generation."""
 
     temperature: Optional[float] = None
@@ -104,15 +94,16 @@ def validate_environment(self) -> Self:
             timeout=self.default_request_timeout,
             max_retries=self.max_retries,
         )
-        self.HUMAN_PROMPT = anthropic.HUMAN_PROMPT
-        self.AI_PROMPT = anthropic.AI_PROMPT
+        # Keep for backward compatibility but not used in Messages API
+        self.HUMAN_PROMPT = getattr(anthropic, "HUMAN_PROMPT", None)
+        self.AI_PROMPT = getattr(anthropic, "AI_PROMPT", None)
         return self
 
     @property
     def _default_params(self) -> Mapping[str, Any]:
         """Get the default parameters for calling Anthropic API."""
         d = {
-            "max_tokens_to_sample": self.max_tokens_to_sample,
+            "max_tokens": self.max_tokens,
             "model": self.model,
         }
         if self.temperature is not None:
@@ -129,16 +120,8 @@ def _identifying_params(self) -> Mapping[str, Any]:
         return {**self._default_params}
 
     def _get_anthropic_stop(self, stop: Optional[list[str]] = None) -> list[str]:
-        if not self.HUMAN_PROMPT or not self.AI_PROMPT:
-            msg = "Please ensure the anthropic package is loaded"
-            raise NameError(msg)
-
         if stop is None:
             stop = []
-
-        # Never want model to invent new turns of Human / Assistant dialog.
-        stop.extend([self.HUMAN_PROMPT])
-
         return stop
 
 
@@ -192,7 +175,7 @@ def _identifying_params(self) -> dict[str, Any]:
         """Get the identifying parameters."""
         return {
             "model": self.model,
-            "max_tokens": self.max_tokens_to_sample,
+            "max_tokens": self.max_tokens,
             "temperature": self.temperature,
             "top_k": self.top_k,
             "top_p": self.top_p,
@@ -211,27 +194,51 @@ def _get_ls_params(
         params = super()._get_ls_params(stop=stop, **kwargs)
         identifying_params = self._identifying_params
         if max_tokens := kwargs.get(
-            "max_tokens_to_sample",
+            "max_tokens",
             identifying_params.get("max_tokens"),
         ):
             params["ls_max_tokens"] = max_tokens
         return params
 
-    def _wrap_prompt(self, prompt: str) -> str:
-        if not self.HUMAN_PROMPT or not self.AI_PROMPT:
-            msg = "Please ensure the anthropic package is loaded"
-            raise NameError(msg)
-
-        if prompt.startswith(self.HUMAN_PROMPT):
-            return prompt  # Already wrapped.
-
-        # Guard against common errors in specifying wrong number of newlines.
-        corrected_prompt, n_subs = re.subn(r"^\n*Human:", self.HUMAN_PROMPT, prompt)
-        if n_subs == 1:
-            return corrected_prompt
-
-        # As a last resort, wrap the prompt ourselves to emulate instruct-style.
-        return f"{self.HUMAN_PROMPT} {prompt}{self.AI_PROMPT} Sure, here you go:\n"
+    def _format_messages(self, prompt: str) -> list[dict[str, str]]:
+        """Convert prompt to Messages API format."""
+        messages = []
+
+        # Handle legacy prompts that might have HUMAN_PROMPT/AI_PROMPT markers
+        if self.HUMAN_PROMPT and self.HUMAN_PROMPT in prompt:
+            # Split on human/assistant turns
+            parts = prompt.split(self.HUMAN_PROMPT)
+
+            for _, part in enumerate(parts):
+                if not part.strip():
+                    continue
+
+                if self.AI_PROMPT and self.AI_PROMPT in part:
+                    # Split human and assistant parts
+                    human_part, assistant_part = part.split(self.AI_PROMPT, 1)
+                    if human_part.strip():
+                        messages.append({"role": "user", "content": human_part.strip()})
+                    if assistant_part.strip():
+                        messages.append(
+                            {"role": "assistant", "content": assistant_part.strip()}
+                        )
+                else:
+                    # Just human content
+                    if part.strip():
+                        messages.append({"role": "user", "content": part.strip()})
+        else:
+            # Handle modern format or plain text
+            # Clean prompt for Messages API
+            content = re.sub(r"^\n*Human:\s*", "", prompt)
+            content = re.sub(r"\n*Assistant:\s*.*$", "", content)
+            if content.strip():
+                messages.append({"role": "user", "content": content.strip()})
+
+        # Ensure we have at least one message
+        if not messages:
+            messages = [{"role": "user", "content": prompt.strip() or "Hello"}]
+
+        return messages
 
     def _call(
         self,
@@ -272,15 +279,19 @@ def _call(
 
         stop = self._get_anthropic_stop(stop)
         params = {**self._default_params, **kwargs}
-        response = self.client.completions.create(
-            prompt=self._wrap_prompt(prompt),
-            stop_sequences=stop,
+
+        # Remove parameters not supported by Messages API
+        params = {k: v for k, v in params.items() if k != "max_tokens_to_sample"}
+
+        response = self.client.messages.create(
+            messages=self._format_messages(prompt),
+            stop_sequences=stop if stop else None,
             **params,
         )
-        return response.completion
+        return response.content[0].text
 
     def convert_prompt(self, prompt: PromptValue) -> str:
-        return self._wrap_prompt(prompt.to_string())
+        return prompt.to_string()
 
     async def _acall(
         self,
@@ -304,12 +315,15 @@ async def _acall(
         stop = self._get_anthropic_stop(stop)
         params = {**self._default_params, **kwargs}
 
-        response = await self.async_client.completions.create(
-            prompt=self._wrap_prompt(prompt),
-            stop_sequences=stop,
+        # Remove parameters not supported by Messages API
+        params = {k: v for k, v in params.items() if k != "max_tokens_to_sample"}
+
+        response = await self.async_client.messages.create(
+            messages=self._format_messages(prompt),
+            stop_sequences=stop if stop else None,
             **params,
         )
-        return response.completion
+        return response.content[0].text
 
     def _stream(
         self,
@@ -343,17 +357,20 @@ def _stream(
         stop = self._get_anthropic_stop(stop)
         params = {**self._default_params, **kwargs}
 
-        for token in self.client.completions.create(
-            prompt=self._wrap_prompt(prompt),
-            stop_sequences=stop,
-            stream=True,
-            **params,
-        ):
-            chunk = GenerationChunk(text=token.completion)
+        # Remove parameters not supported by Messages API
+        params = {k: v for k, v in params.items() if k != "max_tokens_to_sample"}
 
-            if run_manager:
-                run_manager.on_llm_new_token(chunk.text, chunk=chunk)
-            yield chunk
+        with self.client.messages.stream(
+            messages=self._format_messages(prompt),
+            stop_sequences=stop if stop else None,
+            **params,
+        ) as stream:
+            for event in stream:
+                if event.type == "content_block_delta" and hasattr(event.delta, "text"):
+                    chunk = GenerationChunk(text=event.delta.text)
+                    if run_manager:
+                        run_manager.on_llm_new_token(chunk.text, chunk=chunk)
+                    yield chunk
 
     async def _astream(
         self,
@@ -386,17 +403,20 @@ async def _astream(
         stop = self._get_anthropic_stop(stop)
         params = {**self._default_params, **kwargs}
 
-        async for token in await self.async_client.completions.create(
-            prompt=self._wrap_prompt(prompt),
-            stop_sequences=stop,
-            stream=True,
-            **params,
-        ):
-            chunk = GenerationChunk(text=token.completion)
+        # Remove parameters not supported by Messages API
+        params = {k: v for k, v in params.items() if k != "max_tokens_to_sample"}
 
-            if run_manager:
-                await run_manager.on_llm_new_token(chunk.text, chunk=chunk)
-            yield chunk
+        async with self.async_client.messages.stream(
+            messages=self._format_messages(prompt),
+            stop_sequences=stop if stop else None,
+            **params,
+        ) as stream:
+            async for event in stream:
+                if event.type == "content_block_delta" and hasattr(event.delta, "text"):
+                    chunk = GenerationChunk(text=event.delta.text)
+                    if run_manager:
+                        await run_manager.on_llm_new_token(chunk.text, chunk=chunk)
+                    yield chunk
 
     def get_num_tokens(self, text: str) -> int:
         """Calculate number of tokens."""