From 47fad8fd275189744aed9ebce4a3f5057c2aa12e Mon Sep 17 00:00:00 2001 From: raymond-chii Date: Thu, 16 Jan 2025 22:33:55 -0500 Subject: [PATCH 1/2] Adding support for mistral in pl_client.run --- promptlayer/__init__.py | 2 +- promptlayer/promptlayer.py | 4 +- promptlayer/promptlayer_mixins.py | 18 +++ promptlayer/utils.py | 180 +++++++++++++++++++++++++++++- pyproject.toml | 2 +- 5 files changed, 200 insertions(+), 6 deletions(-) diff --git a/promptlayer/__init__.py b/promptlayer/__init__.py index c6300cb..100d50c 100644 --- a/promptlayer/__init__.py +++ b/promptlayer/__init__.py @@ -1,4 +1,4 @@ from .promptlayer import AsyncPromptLayer, PromptLayer -__version__ = "1.0.36" +__version__ = "1.0.37" __all__ = ["PromptLayer", "AsyncPromptLayer", "__version__"] diff --git a/promptlayer/promptlayer.py b/promptlayer/promptlayer.py index 98e5640..b0126f5 100644 --- a/promptlayer/promptlayer.py +++ b/promptlayer/promptlayer.py @@ -173,7 +173,7 @@ 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 { @@ -603,7 +603,7 @@ 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 { diff --git a/promptlayer/promptlayer_mixins.py b/promptlayer/promptlayer_mixins.py index 4bee7c6..2b88343 100644 --- a/promptlayer/promptlayer_mixins.py +++ b/promptlayer/promptlayer_mixins.py @@ -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, @@ -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, @@ -72,6 +76,12 @@ "stream_function": None, }, }, + "google": { + "chat": { + "function_name": "google.convo.send_message", + "stream_function": google_stream_chat, + }, + }, } @@ -80,6 +90,7 @@ "anthropic": anthropic_request, "openai.azure": azure_openai_request, "mistral": mistral_request, + "google": google_request, } AMAP_PROVIDER_TO_FUNCTION_NAME = { @@ -123,6 +134,12 @@ "stream_function": None, }, }, + "google": { + "chat": { + "function_name": "google.convo.send_message", + "stream_function": agoogle_stream_chat, + }, + }, } @@ -131,6 +148,7 @@ "anthropic": aanthropic_request, "openai.azure": aazure_openai_request, "mistral": amistral_request, + "google": agoogle_request, } diff --git a/promptlayer/utils.py b/promptlayer/utils.py index f45c040..e10d2b2 100644 --- a/promptlayer/utils.py +++ b/promptlayer/utils.py @@ -1428,7 +1428,7 @@ 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 @@ -1456,7 +1456,7 @@ 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() + 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") @@ -1795,3 +1795,179 @@ 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 diff --git a/pyproject.toml b/pyproject.toml index 825465b..02db60b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -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 "] license = "Apache-2.0" From 93006705223831c3cef1303efb62a874ee7d3a64 Mon Sep 17 00:00:00 2001 From: raymond-chii Date: Thu, 16 Jan 2025 22:51:39 -0500 Subject: [PATCH 2/2] Apply linting fixes --- promptlayer/promptlayer.py | 8 +- promptlayer/utils.py | 155 +++++++++++++++++++------------------ 2 files changed, 84 insertions(+), 79 deletions(-) diff --git a/promptlayer/promptlayer.py b/promptlayer/promptlayer.py index b0126f5..d0f1970 100644 --- a/promptlayer/promptlayer.py +++ b/promptlayer/promptlayer.py @@ -173,7 +173,9 @@ def _run_internal( group_id, pl_run_span_id, metadata=metadata, - request_response=response.model_dump() if hasattr(response, 'model_dump') else response, + request_response=response.model_dump() + if hasattr(response, "model_dump") + else response, ) return { @@ -603,7 +605,9 @@ async def _run_internal( group_id, pl_run_span_id, metadata=metadata, - request_response=response.model_dump() if hasattr(response, 'model_dump') else response, + request_response=response.model_dump() + if hasattr(response, "model_dump") + else response, ) return { diff --git a/promptlayer/utils.py b/promptlayer/utils.py index e10d2b2..4ceb3f6 100644 --- a/promptlayer/utils.py +++ b/promptlayer/utils.py @@ -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() if hasattr(request_response, 'model_dump') else request_response) + 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 @@ -1456,7 +1460,9 @@ 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 + 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") @@ -1796,11 +1802,13 @@ async def amistral_stream_chat(generator: AsyncIterable[Any]) -> Any: 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( @@ -1808,23 +1816,18 @@ def google_request( max_output_tokens=kwargs.pop("maxOutputTokens", 256), temperature=kwargs.pop("temperature", 0), top_p=kwargs.pop("topP", 1), - top_k=kwargs.pop("topK", 0) + top_k=kwargs.pop("topK", 0), ) model_name = kwargs["model"] - message = kwargs['history'][0]['parts'][0]['text'] + message = kwargs["history"][0]["parts"][0]["text"] system_instruction = kwargs.pop("system_instruction", None) - model = genai.GenerativeModel( - model_name, - system_instruction=system_instruction - ) + 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 + message, generation_config=generation_config, stream=stream ) if stream: return response @@ -1836,6 +1839,7 @@ async def agoogle_request( **kwargs, ): import google.generativeai as genai + genai.configure(api_key=os.environ.get("GOOGLE_API_KEY")) generation_config = genai.GenerationConfig( @@ -1843,24 +1847,19 @@ async def agoogle_request( max_output_tokens=kwargs.pop("maxOutputTokens", 256), temperature=kwargs.pop("temperature", 0), top_p=kwargs.pop("topP", 1), - top_k=kwargs.pop("topK", 0) + top_k=kwargs.pop("topK", 0), ) model_name = kwargs["model"] - message = kwargs['history'][0]['parts'][0]['text'] + message = kwargs["history"][0]["parts"][0]["text"] system_instruction = kwargs.pop("system_instruction", None) - model = genai.GenerativeModel( - model_name, - system_instruction=system_instruction - ) + 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 + message, generation_config=generation_config, stream=stream ) - + if stream: return response return response.to_dict() @@ -1869,52 +1868,57 @@ async def agoogle_request( 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'): + if hasattr(chunk, "candidates"): for candidate in chunk.candidates: - if hasattr(candidate, 'content'): + 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": [] - }], + "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 + 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 + 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 + last_result.usage_metadata + if hasattr(last_result, "usage_metadata") + else None, + "total_token_count", + 0, ), - "cached_content_token_count": 0 - } + "cached_content_token_count": 0, + }, } return response @@ -1926,48 +1930,45 @@ async def agoogle_stream_chat(generator: AsyncIterable[Any]) -> Any: async for chunk in generator: last_result = chunk - if hasattr(chunk, 'candidates'): + if hasattr(chunk, "candidates"): for candidate in chunk.candidates: - if hasattr(candidate, 'content'): + 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": [] - }], + "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 + "prompt_token_count", + 0, ), "candidates_token_count": getattr( last_result.usage_metadata if last_result else None, - 'candidates_token_count', - 0 + "candidates_token_count", + 0, ), "total_token_count": getattr( last_result.usage_metadata if last_result else None, - 'total_token_count', - 0 + "total_token_count", + 0, ), - "cached_content_token_count": 0 - } + "cached_content_token_count": 0, + }, } return response