add anthropic messages support and stream (#118)

* add anthropic messages support and stream
add examples

* update anthropic version dependency

* add log10.load Anthropic client

* use claude-instant-1.2 in anthropic completions example

examples/logging/anthropic_completion.py

@@ -9,7 +9,7 @@
 client = anthropic.Anthropic(api_key=os.environ["ANTHROPIC_API_KEY"])
 
 response = client.completions.create(
- model="claude-1",
+ model="claude-instant-1.2",
  prompt=f"\n\nHuman:Write the names of all Star Wars movies and spinoffs along with the time periods in which they were set?{anthropic.AI_PROMPT}",
  temperature=0,
  max_tokens_to_sample=1024,

examples/logging/anthropic_messages.py

@@ -0,0 +1,18 @@
+import anthropic
+
+from log10.load import log10
+
+
+log10(anthropic)
+
+client = anthropic.Anthropic()
+
+message = client.messages.create(
+ model="claude-3-opus-20240229",
+ max_tokens=1000,
+ temperature=0.0,
+ system="Respond only in Yoda-speak.",
+ messages=[{"role": "user", "content": "How are you today?"}],
+)
+
+print(message.content[0].text)

examples/logging/anthropic_messages_stream.py

@@ -0,0 +1,26 @@
+import anthropic
+from anthropic import Anthropic
+
+from log10.load import log10
+
+
+log10(anthropic)
+
+
+client = Anthropic()
+
+stream = client.messages.create(
+ model="claude-3-opus-20240229",
+ messages=[
+ {
+ "role": "user",
+ "content": "Tell a 50 words joke.",
+ }
+ ],
+ max_tokens=128,
+ temperature=0.9,
+ stream=True,
+)
+for event in stream:
+ if event.type == "content_block_delta":
+ print(event.delta.text, end="", flush=True)

examples/logging/steaming.py → examples/logging/openai_chat_stream.py

@@ -6,7 +6,7 @@
 
 response = client.chat.completions.create(
  model="gpt-3.5-turbo",
- messages=[{"role": "user", "content": "Count to 200"}],
+ messages=[{"role": "user", "content": "Count to 20"}],
  temperature=0,
  stream=True,
 )

log10/anthropic.py

@@ -43,7 +43,7 @@ def chat(self, messages: List[Message], hparams: dict = None) -> ChatCompletion:
  completion = self.client.completions.create(**chat_request)
  content = completion.completion
 
- response = Anthropic.prepare_response(chat_request["prompt"], completion, "chat")
+ response = Anthropic.prepare_response(completion, chat_request["prompt"])
 
  self.log_end(
  completion_id,
@@ -83,7 +83,7 @@ def text(self, prompt: str, hparams: dict = None) -> TextCompletion:
  completion = self.client.completions.create(**text_request)
  text = completion.completion
 
- response = Anthropic.prepare_response(text_request["prompt"], completion, "text")
+ response = Anthropic.prepare_response(completion, text_request["prompt"])
 
  self.log_end(
  completion_id,
@@ -132,30 +132,35 @@ def create_tokens_usage(prompt: str, completion: str):
  }
 
  @staticmethod
- def prepare_response(prompt: str, completion: anthropic.types.Completion, type: str = "text") -> dict:
- if completion.stop_reason == "stop_sequence":
+ def prepare_response(
+ response: anthropic.types.Completion | anthropic.types.Message, input_prompt: str = ""
+ ) -> dict:
+ if response.stop_reason in ["stop_sequence", "end_turn"]:
  reason = "stop"
- elif completion.stop_reason == "max_tokens":
+ elif response.stop_reason == "max_tokens":
  reason = "length"
 
- tokens_usage = Anthropic.create_tokens_usage(prompt, completion.completion)
- response = {
- "id": completion.model_extra["log_id"],
+ ret_response = {
+ "id": response.id,
  "object": "completion",
- "model": completion.model,
+ "model": response.model,
  "choices": [
  {
  "index": 0,
  "finish_reason": reason,
  }
  ],
- "usage": tokens_usage,
  }
- if type == "text":
- response["choices"][0]["text"] = completion.completion
- elif type == "chat":
- response["choices"][0]["message"] = {
- "role": "assistant",
- "content": completion.completion,
+ if isinstance(response, anthropic.types.Message):
+ tokens_usage = {
+ "prompt_tokens": response.usage.input_tokens,
+ "completion_tokens": response.usage.output_tokens,
+ "total_tokens": response.usage.input_tokens + response.usage.output_tokens,
  }
- return response
+ ret_response["choices"][0]["message"] = {"role": response.role, "content": response.content[0].text}
+ elif isinstance(response, anthropic.types.Completion):
+ tokens_usage = Anthropic.create_tokens_usage(input_prompt, response.completion)
+ ret_response["choices"][0]["text"] = response.completion
+ ret_response["usage"] = tokens_usage
+
+ return ret_response

log10/load.py

@@ -214,9 +214,14 @@ async def log_async(completion_url, func, **kwargs):
  if "messages" in kwargs:
  kwargs["messages"] = flatten_messages(kwargs["messages"])
 
+ if "anthropic" in func.__module__:
+ if "system" in kwargs:
+ kwargs["messages"].insert(0, {"role": "system", "content": kwargs["system"]})
+
  log_row = {
  # do we want to also store args?
  "status": "started",
+ "kind": "chat" if "chat" in func.__module__ or "messages" in func.__module__ else "completion",
  "orig_module": func.__module__,
  "orig_qualname": func.__qualname__,
  "request": json.dumps(kwargs),
@@ -384,6 +389,69 @@ def __next__(self):
  raise se
 
 
+class AnthropicStreamingResponseWrapper:
+ """
+ Wraps a streaming response object to log the final result and duration to log10.
+ """
+
+ def __init__(self, completion_url, completionID, response, partial_log_row):
+ self.completionID = completionID
+ self.completion_url = completion_url
+ self.partial_log_row = partial_log_row
+ self.response = response
+ self.final_result = ""
+ self.start_time = time.perf_counter()
+ self.message_id = None
+ self.model = None
+ self.finish_reason = None
+ self.input_tokens = 0
+ self.output_tokens = 0
+
+ def __iter__(self):
+ return self
+
+ def __next__(self):
+ chunk = next(self.response)
+ if chunk.type == "message_start":
+ self.model = chunk.message.model
+ self.message_id = chunk.message.id
+ self.input_tokens = chunk.message.usage.input_tokens
+ elif chunk.type == "message_delta":
+ self.finish_reason = chunk.delta.stop_reason
+ self.output_tokens = chunk.usage.output_tokens
+ elif chunk.type == "content_block_delta":
+ self.final_result += chunk.delta.text
+ elif chunk.type == "message_stop":
+ response = {
+ "id": self.message_id,
+ "object": "chat",
+ "model": self.model,
+ "choices": [
+ {
+ "index": 0,
+ "finish_reason": self.finish_reason,
+ "message": {
+ "role": "assistant",
+ "content": self.final_result,
+ },
+ }
+ ],
+ "usage": {
+ "prompt_tokens": self.input_tokens,
+ "completion_tokens": self.output_tokens,
+ "total_tokens": self.input_tokens + self.output_tokens,
+ },
+ }
+ self.partial_log_row["response"] = json.dumps(response)
+ self.partial_log_row["duration"] = int((time.perf_counter() - self.start_time) * 1000)
+
+ res = post_request(self.completion_url + "/" + self.completionID, self.partial_log_row)
+ if res.status_code != 200:
+ logger.error(f"Failed to insert in log10: {self.partial_log_row} with error {res.text}. Skipping")
+
+ return chunk
+
+
 def flatten_messages(messages):
  flat_messages = []
  for message in messages:
@@ -491,30 +559,50 @@ def wrapper(*args, **kwargs):
  response = output
  # Adjust the Anthropic output to match OAI completion output
  if "anthropic" in type(output).__module__:
+ if type(output).__name__ == "Stream":
+ kind = "chat"
+ return AnthropicStreamingResponseWrapper(
+ completion_url=completion_url,
+ completionID=completionID,
+ response=response,
+ partial_log_row={
+ "response": response,
+ "status": "finished",
+ "stacktrace": json.dumps(stacktrace),
+ "kind": kind,
+ "orig_module": func.__module__,
+ "orig_qualname": func.__qualname__,
+ "request": json.dumps(kwargs),
+ "session_id": sessionID,
+ "tags": global_tags,
+ },
+ )
  from log10.anthropic import Anthropic
 
- response = Anthropic.prepare_response(kwargs["prompt"], output, "text")
- kind = "completion"
- elif type(output).__name__ == "Stream":
- kind = "chat" # Should be "stream", but we don't have that kind yet.
- return StreamingResponseWrapper(
- completion_url=completion_url,
- completionID=completionID,
- response=response,
- partial_log_row={
- "response": response,
- "status": "finished",
- "stacktrace": json.dumps(stacktrace),
- "kind": kind,
- "orig_module": func.__module__,
- "orig_qualname": func.__qualname__,
- "request": json.dumps(kwargs),
- "session_id": sessionID,
- "tags": global_tags,
- },
- )
+ kind = "chat" if response.type == "message" else "completion"
+ response = Anthropic.prepare_response(output, input_prompt=kwargs.get("prompt", ""))
 
+ if "system" in kwargs:
+ kwargs["messages"].insert(0, {"role": "system", "content": kwargs["system"]})
  else:
+ if type(output).__name__ == "Stream":
+ kind = "chat"
+ return StreamingResponseWrapper(
+ completion_url=completion_url,
+ completionID=completionID,
+ response=response,
+ partial_log_row={
+ "response": response,
+ "status": "finished",
+ "stacktrace": json.dumps(stacktrace),
+ "kind": kind,
+ "orig_module": func.__module__,
+ "orig_qualname": func.__qualname__,
+ "request": json.dumps(kwargs),
+ "session_id": sessionID,
+ "tags": global_tags,
+ },
+ )
  response = output
  kind = "chat" if output.object == "chat.completion" else "completion"
 
@@ -692,6 +780,11 @@ def log10(module, DEBUG_=False, USE_ASYNC_=True):
  attr = module.resources.completions.Completions
  method = getattr(attr, "create")
  setattr(attr, "create", intercepting_decorator(method))
+
+ # anthropic Messages completion
+ attr = module.resources.messages.Messages
+ method = getattr(attr, "create")
+ setattr(attr, "create", intercepting_decorator(method))
  elif module.__name__ == "openai":
  openai_version = parse(version("openai"))
  global OPENAI_V1
@@ -784,3 +877,36 @@ def __init__(self, *args, **kwargs):
  if not getattr(openai, "_log10_patched", False):
  log10(openai)
  openai._log10_patched = True
+
+
+try:
+ import anthropic
+except ImportError:
+ logger.warning("Anthropic not found. Skipping defining log10.load.Anthropic client.")
+else:
+ from anthropic import Anthropic
+
+ class Anthropic(Anthropic):
+ """
+ Example:
+ >>> from log10.load import Anthropic
+ >>> client = Anthropic(tags=["test", "load_anthropic"])
+ >>> message = client.messages.create(
+ ... model="claude-3-haiku-20240307",
+ ... max_tokens=100,
+ ... temperature=0.9,
+ ... system="Respond only in Yoda-speak.",
+ ... messages=[{"role": "user", "content": "How are you today?"}],
+ ... )
+ >>> print(message.content[0].text)
+ """
+
+ def __init__(self, *args, **kwargs):
+ if "tags" in kwargs:
+ global global_tags
+ global_tags = kwargs.pop("tags")
+ super().__init__(*args, **kwargs)
+
+ if not getattr(anthropic, "_log10_patched", False):
+ log10(anthropic)
+ anthropic._log10_patched = True

pyproject.toml

@@ -45,7 +45,7 @@ google-search-results = "^2.4.2"
 wikipedia = "^1.4.0"
 faker = "^19.2.0"
 backoff = "^2.2.1"
-anthropic = "^0.3.11"
+anthropic = "<1"
 mosaicml-cli = "^0.5.30"
 together = "^0.2.7"
 magentic = "^0.17.0"