Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

added use of extra_headers with the python client #28

Open
wants to merge 2 commits into
base: main
Choose a base branch
from
Open

added use of extra_headers with the python client #28

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
802b345
added use of extra_headers with the python client
vedpatwardhan Jul 22, 2024
132a9fd
added check to limit headers that can be used
vedpatwardhan Jul 22, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
21 changes: 21 additions & 0 deletions unify/clients.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@
from unify.exceptions import BadRequestError, UnifyError, status_error_map
from unify.utils import ( # noqa:WPS450
_available_dynamic_modes,
_available_extra_headers,
_validate_api_key,
_validate_endpoint,
)
Expand Down Expand Up @@ -130,6 +131,7 @@ def generate( # noqa: WPS234, WPS211
temperature: Optional[float] = 1.0,
stop: Optional[List[str]] = None,
stream: bool = False,
extra_headers: Optional[Dict[str, str]] = None,
) -> Union[Generator[str, None, None], str]: # noqa: DAR101, DAR201, DAR401
"""Generate content using the Unify API.

Expand Down Expand Up @@ -157,6 +159,8 @@ def generate( # noqa: WPS234, WPS211
If False, generates content as a single response.
Defaults to False.

extra_headers: Send extra headers

Returns:
Union[Generator[str, None, None], str]: If stream is True,
returns a generator yielding chunks of content.
Expand All @@ -175,20 +179,23 @@ def generate( # noqa: WPS234, WPS211
else:
raise UnifyError("You must provider either the user_prompt or messages!")

extra_headers = {key: extra_headers[key] for key in _available_extra_headers}
if stream:
return self._generate_stream(
contents,
self._endpoint,
max_tokens=max_tokens,
temperature=temperature,
stop=stop,
extra_headers=extra_headers,
)
return self._generate_non_stream(
contents,
self._endpoint,
max_tokens=max_tokens,
temperature=temperature,
stop=stop,
extra_headers=extra_headers,
)

def get_credit_balance(self) -> float:
Expand Down Expand Up @@ -224,6 +231,7 @@ def _generate_stream(
max_tokens: Optional[int] = 1024,
temperature: Optional[float] = 1.0,
stop: Optional[List[str]] = None,
extra_headers: Optional[Dict[str, str]] = None,
) -> Generator[str, None, None]:
try:
chat_completion = self.client.chat.completions.create(
Expand All @@ -233,6 +241,7 @@ def _generate_stream(
temperature=temperature,
stop=stop,
stream=True,
extra_headers=extra_headers,
extra_body={"signature": "package"},
)
for chunk in chat_completion:
Expand All @@ -250,6 +259,7 @@ def _generate_non_stream(
max_tokens: Optional[int] = 1024,
temperature: Optional[float] = 1.0,
stop: Optional[List[str]] = None,
extra_headers: Optional[Dict[str, str]] = None,
) -> str:
try:
chat_completion = self.client.chat.completions.create(
Expand All @@ -259,6 +269,7 @@ def _generate_non_stream(
temperature=temperature,
stop=stop,
stream=False,
extra_headers=extra_headers,
extra_body={"signature": "package"},
)
if "router" not in endpoint:
Expand Down Expand Up @@ -422,6 +433,7 @@ async def generate( # noqa: WPS234, WPS211
temperature: Optional[float] = 1.0,
stop: Optional[List[str]] = None,
stream: bool = False,
extra_headers: Optional[Dict[str, str]] = None,
) -> Union[AsyncGenerator[str, None], str]: # noqa: DAR101, DAR201, DAR401
"""Generate content asynchronously using the Unify API.

Expand Down Expand Up @@ -449,6 +461,8 @@ async def generate( # noqa: WPS234, WPS211
If False, generates content as a single response.
Defaults to False.

extra_headers: Send extra headers

Returns:
Union[AsyncGenerator[str, None], List[str]]: If stream is True,
returns an asynchronous generator yielding chunks of content.
Expand All @@ -468,20 +482,23 @@ async def generate( # noqa: WPS234, WPS211
else:
raise UnifyError("You must provide either the user_prompt or messages!")

extra_headers = {key: extra_headers[key] for key in _available_extra_headers}
if stream:
return self._generate_stream(
contents,
self._endpoint,
max_tokens=max_tokens,
stop=stop,
temperature=temperature,
extra_headers=extra_headers,
)
return await self._generate_non_stream(
contents,
self._endpoint,
max_tokens=max_tokens,
stop=stop,
temperature=temperature,
extra_headers=extra_headers,
)

async def _generate_stream(
Expand All @@ -491,6 +508,7 @@ async def _generate_stream(
max_tokens: Optional[int] = None,
temperature: Optional[float] = 1.0,
stop: Optional[List[str]] = None,
extra_headers: Optional[Dict[str, str]] = None,
) -> AsyncGenerator[str, None]:
try:
async_stream = await self.client.chat.completions.create(
Expand All @@ -500,6 +518,7 @@ async def _generate_stream(
temperature=temperature,
stop=stop,
stream=True,
extra_headers=extra_headers,
extra_body={"signature": "package"},
)
async for chunk in async_stream: # type: ignore[union-attr]
Expand All @@ -515,6 +534,7 @@ async def _generate_non_stream(
max_tokens: Optional[int] = None,
temperature: Optional[float] = 1.0,
stop: Optional[List[str]] = None,
extra_headers: Optional[Dict[str, str]] = None,
) -> str:
try:
async_response = await self.client.chat.completions.create(
Expand All @@ -524,6 +544,7 @@ async def _generate_non_stream(
temperature=temperature,
stop=stop,
stream=False,
extra_headers=extra_headers,
extra_body={"signature": "package"},
)
self.set_provider(async_response.model.split("@")[-1]) # type: ignore
Expand Down
4 changes: 4 additions & 0 deletions unify/utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,10 @@
"tks-per-sec",
]

_available_extra_headers = [
"anthropic-beta"
]

_base_url = "https://api.unify.ai/v0"


Expand Down
Loading