anthropic.py 2.2 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768
  1. import anthropic
  2. from anthropic import Anthropic
  3. from anthropic.resources.completions import Completions
  4. from anthropic.types import completion_create_params, Completion
  5. from anthropic._types import NOT_GIVEN, NotGiven, Headers, Query, Body
  6. from _pytest.monkeypatch import MonkeyPatch
  7. from typing import List, Union, Literal, Any, Generator
  8. from time import sleep
  9. import pytest
  10. import os
  11. MOCK = os.getenv('MOCK_SWITCH', 'false') == 'true'
  12. class MockAnthropicClass(object):
  13. @staticmethod
  14. def mocked_anthropic_chat_create_sync(model: str) -> Completion:
  15. return Completion(
  16. completion='hello, I\'m a chatbot from anthropic',
  17. model=model,
  18. stop_reason='stop_sequence'
  19. )
  20. @staticmethod
  21. def mocked_anthropic_chat_create_stream(model: str) -> Generator[Completion, None, None]:
  22. full_response_text = "hello, I'm a chatbot from anthropic"
  23. for i in range(0, len(full_response_text) + 1):
  24. sleep(0.1)
  25. if i == len(full_response_text):
  26. yield Completion(
  27. completion='',
  28. model=model,
  29. stop_reason='stop_sequence'
  30. )
  31. else:
  32. yield Completion(
  33. completion=full_response_text[i],
  34. model=model,
  35. stop_reason=''
  36. )
  37. def mocked_anthropic(self: Completions, *,
  38. max_tokens_to_sample: int,
  39. model: Union[str, Literal["claude-2.1", "claude-instant-1"]],
  40. prompt: str,
  41. stream: Literal[True],
  42. **kwargs: Any
  43. ) -> Union[Completion, Generator[Completion, None, None]]:
  44. if len(self._client.api_key) < 18:
  45. raise anthropic.AuthenticationError('Invalid API key')
  46. if stream:
  47. return MockAnthropicClass.mocked_anthropic_chat_create_stream(model=model)
  48. else:
  49. return MockAnthropicClass.mocked_anthropic_chat_create_sync(model=model)
  50. @pytest.fixture
  51. def setup_anthropic_mock(request, monkeypatch: MonkeyPatch):
  52. if MOCK:
  53. monkeypatch.setattr(Completions, 'create', MockAnthropicClass.mocked_anthropic)
  54. yield
  55. if MOCK:
  56. monkeypatch.undo()