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..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(), + 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(), + 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..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()) + 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 @@ -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") @@ -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 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"