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

Adding support for google in pl_client.run #223

Open
wants to merge 2 commits into
base: master
Choose a base branch
from
+204 −5
Open

Adding support for google in pl_client.run #223

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
47fad8f
Adding support for mistral in pl_client.run
raymond-chii Jan 17, 2025
9300670
Apply linting fixes
raymond-chii Jan 17, 2025
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
2 changes: 1 addition & 1 deletion promptlayer/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
from .promptlayer import AsyncPromptLayer, PromptLayer

__version__ = "1.0.36"
__version__ = "1.0.37"
__all__ = ["PromptLayer", "AsyncPromptLayer", "__version__"]
8 changes: 6 additions & 2 deletions promptlayer/promptlayer.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -173,7 +173,9 @@ def _run_internal(
group_id,
pl_run_span_id,
metadata=metadata,
request_response=response.model_dump(),
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can we figure out a way to keep this as it is but call model_dump from the corresponding LLM providers functions that do support it?

request_response=response.model_dump()
if hasattr(response, "model_dump")
else response,
)

return {
Expand Down Expand Up @@ -603,7 +605,9 @@ async def _run_internal(
group_id,
pl_run_span_id,
metadata=metadata,
request_response=response.model_dump(),
request_response=response.model_dump()
if hasattr(response, "model_dump")
else response,
)

return {
Expand Down
18 changes: 18 additions & 0 deletions promptlayer/promptlayer_mixins.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,8 @@
aanthropic_stream_completion,
aanthropic_stream_message,
aazure_openai_request,
agoogle_request,
agoogle_stream_chat,
amistral_request,
amistral_stream_chat,
anthropic_request,
Expand All @@ -24,6 +26,8 @@
aopenai_stream_chat,
aopenai_stream_completion,
azure_openai_request,
google_request,
google_stream_chat,
mistral_request,
mistral_stream_chat,
openai_request,
Expand Down Expand Up @@ -72,6 +76,12 @@
"stream_function": None,
},
},
"google": {
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@raymond-chii why no completion support?

"chat": {
"function_name": "google.convo.send_message",
"stream_function": google_stream_chat,
},
},
}


Expand All @@ -80,6 +90,7 @@
"anthropic": anthropic_request,
"openai.azure": azure_openai_request,
"mistral": mistral_request,
"google": google_request,
}

AMAP_PROVIDER_TO_FUNCTION_NAME = {
Expand Down Expand Up @@ -123,6 +134,12 @@
"stream_function": None,
},
},
"google": {
"chat": {
"function_name": "google.convo.send_message",
"stream_function": agoogle_stream_chat,
},
},
}


Expand All @@ -131,6 +148,7 @@
"anthropic": aanthropic_request,
"openai.azure": aazure_openai_request,
"mistral": amistral_request,
"google": agoogle_request,
}


Expand Down
179 changes: 178 additions & 1 deletion promptlayer/utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1428,7 +1428,11 @@ def stream_response(
data["raw_response"] = result
yield data
request_response = map_results(results)
response = after_stream(request_response=request_response.model_dump())
response = after_stream(
request_response=request_response.model_dump()
if hasattr(request_response, "model_dump")
else request_response
)
data["request_id"] = response.get("request_id")
data["prompt_blueprint"] = response.get("prompt_blueprint")
yield data
Expand Down Expand Up @@ -1457,6 +1461,8 @@ async def async_generator_from_list(lst):
request_response = await map_results(async_generator_from_list(results))
after_stream_response = await after_stream(
request_response=request_response.model_dump()
if hasattr(request_response, "model_dump")
else request_response
)
data["request_id"] = after_stream_response.get("request_id")
data["prompt_blueprint"] = after_stream_response.get("prompt_blueprint")
Expand Down Expand Up @@ -1795,3 +1801,174 @@ async def amistral_stream_chat(generator: AsyncIterable[Any]) -> Any:
response.choices[0].message.content = content
response.choices[0].message.tool_calls = tool_calls
return response


def google_request(
prompt_blueprint: GetPromptTemplateResponse,
**kwargs,
):
import google.generativeai as genai

genai.configure(api_key=os.environ.get("GOOGLE_API_KEY"))

generation_config = genai.GenerationConfig(
candidate_count=kwargs.pop("candidateCount", 1),
max_output_tokens=kwargs.pop("maxOutputTokens", 256),
temperature=kwargs.pop("temperature", 0),
top_p=kwargs.pop("topP", 1),
top_k=kwargs.pop("topK", 0),
)

model_name = kwargs["model"]
message = kwargs["history"][0]["parts"][0]["text"]
system_instruction = kwargs.pop("system_instruction", None)
model = genai.GenerativeModel(model_name, system_instruction=system_instruction)

stream = kwargs.pop("stream", False)

response = model.generate_content(
message, generation_config=generation_config, stream=stream
)
if stream:
return response
return response.to_dict()


async def agoogle_request(
prompt_blueprint: GetPromptTemplateResponse,
**kwargs,
):
import google.generativeai as genai

genai.configure(api_key=os.environ.get("GOOGLE_API_KEY"))

generation_config = genai.GenerationConfig(
candidate_count=kwargs.pop("candidateCount", 1),
max_output_tokens=kwargs.pop("maxOutputTokens", 256),
temperature=kwargs.pop("temperature", 0),
top_p=kwargs.pop("topP", 1),
top_k=kwargs.pop("topK", 0),
)

model_name = kwargs["model"]
message = kwargs["history"][0]["parts"][0]["text"]
system_instruction = kwargs.pop("system_instruction", None)
model = genai.GenerativeModel(model_name, system_instruction=system_instruction)

stream = kwargs.pop("stream", False)
response = await model.generate_content_async(
message, generation_config=generation_config, stream=stream
)

if stream:
return response
return response.to_dict()


def google_stream_chat(results: list):
# Get the last chunk to access final state
last_result = results[-1] # Remove .result access

# Combine all content from the stream
content = ""
for chunk in results:
if hasattr(chunk, "candidates"):
for candidate in chunk.candidates:
if hasattr(candidate, "content"):
for part in candidate.content.parts:
content += part.text

# Create response in Google's format
response = {
"candidates": [
{
"content": {"parts": [{"text": content}], "role": "model"},
"finish_reason": 1,
"safety_ratings": [],
"token_count": getattr(
last_result.usage_metadata
if hasattr(last_result, "usage_metadata")
else None,
"candidates_token_count",
0,
),
"grounding_attributions": [],
}
],
"usage_metadata": {
"prompt_token_count": getattr(
last_result.usage_metadata
if hasattr(last_result, "usage_metadata")
else None,
"prompt_token_count",
0,
),
"candidates_token_count": getattr(
last_result.usage_metadata
if hasattr(last_result, "usage_metadata")
else None,
"candidates_token_count",
0,
),
"total_token_count": getattr(
last_result.usage_metadata
if hasattr(last_result, "usage_metadata")
else None,
"total_token_count",
0,
),
"cached_content_token_count": 0,
},
}

return response


async def agoogle_stream_chat(generator: AsyncIterable[Any]) -> Any:
last_result = None
content = ""

async for chunk in generator:
last_result = chunk
if hasattr(chunk, "candidates"):
for candidate in chunk.candidates:
if hasattr(candidate, "content"):
for part in candidate.content.parts:
content += part.text

# Create response in Google's format using the final state
response = {
"candidates": [
{
"content": {"parts": [{"text": content}], "role": "model"},
"finish_reason": 1,
"safety_ratings": [],
"token_count": getattr(
last_result.usage_metadata if last_result else None,
"candidates_token_count",
0,
),
"grounding_attributions": [],
}
],
"usage_metadata": {
"prompt_token_count": getattr(
last_result.usage_metadata if last_result else None,
"prompt_token_count",
0,
),
"candidates_token_count": getattr(
last_result.usage_metadata if last_result else None,
"candidates_token_count",
0,
),
"total_token_count": getattr(
last_result.usage_metadata if last_result else None,
"total_token_count",
0,
),
"cached_content_token_count": 0,
},
}

return response
2 changes: 1 addition & 1 deletion pyproject.toml
View file
Open in desktop
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@raymond-chii can you add google generativeai as a dev dependency as well?

Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
[tool.poetry]
name = "promptlayer"
version = "1.0.36"
version = "1.0.37"
description = "PromptLayer is a platform for prompt engineering and tracks your LLM requests."
authors = ["Magniv <[email protected]>"]
license = "Apache-2.0"
Expand Down
Loading