Handle Anthropic API errors gracefully

Fixes #71, #38

gptcli/anthropic.py

@@ -2,7 +2,12 @@
 from typing import Iterator, List
 import anthropic
 
-from gptcli.completion import CompletionProvider, Message
+from gptcli.completion import (
+ CompletionProvider,
+ Message,
+ CompletionError,
+ BadRequestError,
+)
 
 api_key = os.environ.get("ANTHROPIC_API_KEY")
 
@@ -53,13 +58,18 @@ def complete(
  kwargs["messages"] = messages
 
  client = get_client()
- if stream:
- with client.messages.stream(**kwargs) as stream:
- for text in stream.text_stream:
- yield text
- else:
- response = client.messages.create(**kwargs, stream=False)
- yield "".join(c.text for c in response.content)
+ try:
+ if stream:
+ with client.messages.stream(**kwargs) as completion:
+ for text in completion.text_stream:
+ yield text
+ else:
+ response = client.messages.create(**kwargs, stream=False)
+ yield "".join(c.text for c in response.content)
+ except anthropic.BadRequestError as e:
+ raise BadRequestError(e.message) from e
+ except anthropic.APIError as e:
+ raise CompletionError(e.message) from e
 
 
 def num_tokens_from_messages_anthropic(messages: List[Message], model: str) -> int:

gptcli/completion.py

@@ -19,3 +19,11 @@ def complete(
  self, messages: List[Message], args: dict, stream: bool = False
  ) -> Iterator[str]:
  pass
+
+
+class CompletionError(Exception):
+ pass
+
+
+class BadRequestError(CompletionError):
+ pass

gptcli/openai.py

@@ -5,7 +5,12 @@
 
 import tiktoken
 
-from gptcli.completion import CompletionProvider, Message
+from gptcli.completion import (
+ CompletionProvider,
+ Message,
+ CompletionError,
+ BadRequestError,
+)
 
 
 class OpenAICompletionProvider(CompletionProvider):
@@ -21,28 +26,33 @@ def complete(
  if "top_p" in args:
  kwargs["top_p"] = args["top_p"]
 
- if stream:
- response_iter = self.client.chat.completions.create(
- messages=cast(List[ChatCompletionMessageParam], messages),
- stream=True,
- model=args["model"],
- **kwargs,
- )
+ try:
+ if stream:
+ response_iter = self.client.chat.completions.create(
+ messages=cast(List[ChatCompletionMessageParam], messages),
+ stream=True,
+ model=args["model"],
+ **kwargs,
+ )
 
- for response in response_iter:
+ for response in response_iter:
+ next_choice = response.choices[0]
+ if next_choice.finish_reason is None and next_choice.delta.content:
+ yield next_choice.delta.content
+ else:
+ response = self.client.chat.completions.create(
+ messages=cast(List[ChatCompletionMessageParam], messages),
+ model=args["model"],
+ stream=False,
+ **kwargs,
+ )
  next_choice = response.choices[0]
- if next_choice.finish_reason is None and next_choice.delta.content:
- yield next_choice.delta.content
- else:
- response = self.client.chat.completions.create(
- messages=cast(List[ChatCompletionMessageParam], messages),
- model=args["model"],
- stream=False,
- **kwargs,
- )
- next_choice = response.choices[0]
- if next_choice.message.content:
- yield next_choice.message.content
+ if next_choice.message.content:
+ yield next_choice.message.content
+ except openai.BadRequestError as e:
+ raise BadRequestError(e.message) from e
+ except openai.APIError as e:
+ raise CompletionError(e.message) from e
 
 
 def num_tokens_from_messages_openai(messages: List[Message], model: str) -> int:

gptcli/session.py

@@ -1,8 +1,7 @@
 from abc import abstractmethod
 from typing_extensions import TypeGuard
 from gptcli.assistant import Assistant
-from gptcli.completion import Message, ModelOverrides
-from openai import BadRequestError, OpenAIError
+from gptcli.completion import Message, ModelOverrides, CompletionError, BadRequestError
 from typing import Any, Dict, List, Tuple
 
 
@@ -115,7 +114,7 @@ def _respond(self, args: ModelOverrides) -> bool:
  except BadRequestError as e:
  self.listener.on_error(e)
  return False
- except OpenAIError as e:
+ except CompletionError as e:
  self.listener.on_error(e)
  return True