pygeai 0.1.6__py3-none-any.whl → 0.6.0b15__py3-none-any.whl

pygeai/chat/clients.py

@@ -1,7 +1,14 @@
 import json
+from json import JSONDecodeError
+from pathlib import Path
+from typing import List, Dict, Optional, Union, Generator
 
-from pygeai.chat.endpoints import CHAT_V1, CHAT_COMPLETION_V1
+from pygeai import logger
+from pygeai.chat.endpoints import CHAT_V1, GENERATE_IMAGE_V1, EDIT_IMAGE_V1, RESPONSES_V1
 from pygeai.core.base.clients import BaseClient
+from pygeai.core.common.exceptions import InvalidAPIResponseException
+from pygeai.core.utils.validators import validate_status_code
+from pygeai.core.utils.parsers import parse_json_response
 
 
 class ChatClient(BaseClient):
@@ -10,69 +17,437 @@ class ChatClient(BaseClient):
         response = self.api_service.post(
             endpoint=CHAT_V1
         )
-        result = json.loads(response.content)
-        return result
+        validate_status_code(response)
+        return parse_json_response(response, "chat")
 
     def chat_completion(
             self,
             model: str,
-            messages: list,
+            messages: List[Dict[str, str]],
             stream: bool = False,
-            temperature: int = None,
-            max_tokens: int = None,
-            thread_id: str = None,
-            frequency_penalty: float = None,
-            presence_penalty: float = None
-    ) -> dict:
+            temperature: Optional[float] = None,
+            max_tokens: Optional[int] = None,
+            thread_id: Optional[str] = None,
+            frequency_penalty: Optional[float] = None,
+            presence_penalty: Optional[float] = None,
+            variables: Optional[List] = None,
+            top_p: Optional[float] = None,
+            stop: Optional[Union[str, List[str]]] = None,
+            response_format: Optional[Dict] = None,
+            tools: Optional[List[Dict]] = None,
+            tool_choice: Optional[Union[str, Dict]] = None,
+            logprobs: Optional[bool] = None,
+            top_logprobs: Optional[int] = None,
+            seed: Optional[int] = None,
+            stream_options: Optional[Dict] = None,
+            store: Optional[bool] = None,
+            metadata: Optional[Dict] = None,
+            user: Optional[str] = None,
+            reasoning_effort: Optional[str] = None
+    ) -> Union[dict, str, Generator[str, None, None]]:
         """
         Generates a chat completion response using the specified model and parameters.
 
         :param model: str - The model specification in the format
             "saia:<assistant_type>:<assistant_name>|<bot_id>". Determines the assistant type and associated configuration. (Required)
-        :param messages: list - A list of messages to include in the chat completion. Each message should be a dictionary with
+        :param messages: List[Dict[str, str]] - A list of messages to include in the chat completion. Each message should be a dictionary with
             the following structure:
                 {
-                    "role": "string",  # Possible values: "user", "system", or others supported by the model
-                    "content": "string"  # The content of the message
+                    "role": "string",
+                    "content": "string"
                 } (Required)
         :param stream: bool - Whether the response should be streamed. Possible values:
-            - False: Do not stream the response (default).
-            - True: Stream the response.
-        :param temperature: int - Controls the randomness of the response. Higher values (e.g., 2) produce more random responses,
+            - False: Returns the complete response as a dictionary or string (default).
+            - True: Returns a generator yielding content strings as they are received.
+        :param temperature: Optional[float] - Controls the randomness of the response. Higher values (e.g., 2.0) produce more random responses,
             while lower values (e.g., 0.2) produce more deterministic responses. (Optional)
-        :param max_tokens: int - The maximum number of tokens to generate in the response. (Optional)
-        :param thread_id: str - An optional UUID to identify the conversation thread. (Optional)
-        :param frequency_penalty: float - A value between -2.0 and 2.0. Positive values decrease the model's likelihood of
+        :param max_tokens: Optional[int] - The maximum number of tokens to generate in the response. (Optional)
+        :param thread_id: Optional[str] - An optional UUID to identify the conversation thread. (Optional)
+        :param frequency_penalty: Optional[float] - A value between -2.0 and 2.0. Positive values decrease the model's likelihood of
             repeating tokens based on their frequency in the text so far. (Optional)
-        :param presence_penalty: float - A value between -2.0 and 2.0. Positive values increase the model's likelihood of
+        :param presence_penalty: Optional[float] - A value between -2.0 and 2.0. Positive values increase the model's likelihood of
             discussing new topics by penalizing tokens that have already appeared in the text. (Optional)
-        :return: dict - The API response containing the chat completion result.
+        :param variables: Optional[List] - A list of additional variables. These must be defined at the time of creation, otherwise
+            their use will throw an error. (Optional)
+        :param top_p: Optional[float] - An alternative to temperature, nucleus sampling considers tokens with top_p probability mass. (Optional)
+        :param stop: Optional[Union[str, List[str]]] - Up to 4 sequences where the API stops generating further tokens. (Optional)
+        :param response_format: Optional[Dict] - Specifies the output format, e.g., JSON schema for structured outputs. (Optional)
+        :param tools: Optional[List[Dict]] - A list of tools (e.g., functions) the model may call. (Optional)
+        :param tool_choice: Optional[Union[str, Dict]] - Controls which tool is called (e.g., "none", "auto", or specific tool). (Optional)
+        :param logprobs: Optional[bool] - Whether to return log probabilities of output tokens. (Optional)
+        :param top_logprobs: Optional[int] - Number of most likely tokens to return with log probabilities (0-20). (Optional)
+        :param seed: Optional[int] - For deterministic sampling, in Beta. (Optional)
+        :param stream_options: Optional[Dict] - Options for streaming, e.g., include_usage. (Optional)
+        :param store: Optional[bool] - Whether to store the output for model distillation or evals. (Optional)
+        :param metadata: Optional[Dict] - Up to 16 key-value pairs to attach to the object. (Optional)
+        :param user: Optional[str] - A unique identifier for the end-user to monitor abuse. (Optional)
+        :param reasoning_effort: Optional[str] - Controls the depth of reasoning applied by supported models.
+            Possible values: "low", "medium", "high". Supported by OpenAI models from version 5,
+            Claude models from version 4.1, and Gemini models from version 2.0. (Optional)
+        :return: Union[dict, str, Generator[str, None, None]] - For non-streaming (stream=False), returns a dictionary containing the chat completion
+            result or a string if JSON decoding fails. For streaming (stream=True), returns a generator yielding content strings extracted from the
+            streaming response chunks.
         """
         data = {
             'model': model,
             'messages': messages,
             'stream': stream
         }
-        if temperature:
+        if temperature is not None:
             data['temperature'] = temperature
 
         if max_tokens:
-            # TODO -> Review deprecation of max_tokens in favor of max_completion_tokens
-            data['max_tokens'] = max_tokens
-
-        if thread_id:
-            data['threadId'] = thread_id
+            data['max_completion_tokens'] = max_tokens
 
-        if frequency_penalty:
+        if frequency_penalty is not None:
             data['frequency_penalty'] = frequency_penalty
 
-        if presence_penalty:
+        if presence_penalty is not None:
             data['presence_penalty'] = presence_penalty
 
+        if variables is not None and any(variables):
+            data['variables'] = variables
+
+        if top_p is not None:
+            data['top_p'] = top_p
+
+        if stop is not None:
+            data['stop'] = stop
+
+        if response_format is not None:
+            data['response_format'] = response_format
+
+        if tools is not None:
+            data['tools'] = tools
+
+        if tool_choice is not None:
+            data['tool_choice'] = tool_choice
+
+        if logprobs is not None:
+            data['logprobs'] = logprobs
+
+        if top_logprobs is not None:
+            data['top_logprobs'] = top_logprobs
+
+        if seed is not None:
+            data['seed'] = seed
+
+        if stream_options is not None:
+            data['stream_options'] = stream_options
+
+        if store is not None:
+            data['store'] = store
+
+        if metadata is not None:
+            data['metadata'] = metadata
+
+        if user is not None:
+            data['user'] = user
+
+        if reasoning_effort is not None:
+            data['reasoning_effort'] = reasoning_effort
+
+        headers = {}
+        if thread_id:
+            headers["saia-conversation-id"] = thread_id
+
+        logger.debug(f"Generating chat completion with data: {data}")
+
+        if stream:
+            response = self.api_service.stream_post(
+                endpoint=CHAT_V1,
+                data=data,
+                headers=headers
+            )
+            return self.stream_chat_generator(response)
+        else:
+            response = self.api_service.post(
+                endpoint=CHAT_V1,
+                data=data,
+                headers=headers
+            )
+            validate_status_code(response)
+
+            result = parse_json_response(response, "process chat request")
+
+            logger.debug(f"Chat completion result: {result}")
+
+            return result
+
+    def stream_chat_generator(self, response) -> Generator[str, None, None]:
+        """
+        Processes a streaming response and yields content strings.
+
+        :param response: The streaming response from the API.
+        :return: Generator[str, None, None] - Yields content strings extracted from streaming chunks.
+        """
+        try:
+            for line in response:
+                if line.startswith("data:"):
+                    chunk = line[5:].strip()
+                    if chunk == "[DONE]":
+                        break
+                    try:
+                        json_data = json.loads(chunk)
+                        if (
+                                json_data.get("choices")
+                                and len(json_data["choices"]) > 0
+                                and "delta" in json_data["choices"][0]
+                                and "content" in json_data["choices"][0]["delta"]
+                        ):
+                            content = json_data["choices"][0]["delta"]["content"]
+                            yield content
+                    except JSONDecodeError as e:
+                        continue
+        except Exception as e:
+            raise InvalidAPIResponseException(f"Unable to process streaming chat response: {e}")
+
+    def stream_response_generator(self, response) -> Generator[str, None, None]:
+        """
+        Processes a streaming response from the Responses API and yields content strings.
+
+        :param response: The streaming response from the API.
+        :return: Generator[str, None, None] - Yields content strings extracted from streaming chunks.
+        """
+        try:
+            for line in response:
+                if line.startswith("data:"):
+                    chunk = line[5:].strip()
+                    if chunk == "[DONE]":
+                        break
+                    try:
+                        json_data = json.loads(chunk)
+                        if (
+                                json_data.get("choices")
+                                and len(json_data["choices"]) > 0
+                                and "delta" in json_data["choices"][0]
+                                and "content" in json_data["choices"][0]["delta"]
+                        ):
+                            content = json_data["choices"][0]["delta"]["content"]
+                            yield content
+                    except JSONDecodeError as e:
+                        continue
+        except Exception as e:
+            raise InvalidAPIResponseException(f"Unable to process streaming response: {e}")
+
+    def generate_image(
+            self,
+            model: str,
+            prompt: str,
+            n: int,
+            quality: str,
+            size: str,
+            aspect_ratio: Optional[str] = None
+    ) -> dict:
+        """
+        Generates an image based on the provided parameters.
+
+        :param model: str - The model specification for image generation, e.g., "openai/gpt-image-1". (Required)
+        :param prompt: str - Description of the desired image. (Required)
+        :param n: int - Number of images to generate (1-10, depending on the model). (Required)
+        :param quality: str - Rendering quality, e.g., "high". (Required)
+        :param size: str - Image dimensions, e.g., "1024x1024". (Required)
+        :param aspect_ratio: Optional[str] - Relationship between image's width and height, e.g., "1:1", "9:16", "16:9", "3:4", "4:3". (Optional)
+        :return: dict - The API response containing the generated image data.
+        :raises InvalidAPIResponseException: If the API response cannot be processed.
+        """
+        data = {
+            'model': model,
+            'prompt': prompt,
+            'n': n,
+            'quality': quality,
+            'size': size
+        }
+
+        if aspect_ratio:
+            data['aspect_ratio'] = aspect_ratio
+
+        logger.debug(f"Generating image with data: {data}")
+
         response = self.api_service.post(
-            endpoint=CHAT_COMPLETION_V1,
+            endpoint=GENERATE_IMAGE_V1,
             data=data
         )
-        result = json.loads(response.content)
+
+        validate_status_code(response)
+
+        result = parse_json_response(response, "generate image")
+        logger.debug(f"Image generation result: {result}")
         return result
 
+    def edit_image(
+            self,
+            model: str,
+            prompt: str,
+            image: str,
+            size: str,
+            n: int = 1,
+            quality: Optional[str] = None
+    ) -> dict:
+        """
+        Edits an existing image based on the provided parameters.
+
+        :param model: str - The model specification for image editing, e.g., "openai/gpt-image-1". (Required)
+        :param prompt: str - Description of the desired edit, e.g., "remove background people". (Required)
+        :param image: str - URL of the image to be edited, e.g., "https://example.com/image.jpg". (Required)
+        :param size: str - Desired dimensions of the output image in pixels, e.g., "1024x1024". (Required)
+        :param n: int - Number of edited images to generate (1-10, depending on the model). Default is 1. (Optional)
+        :param quality: Optional[str] - Rendering quality, e.g., "high", "medium", "low". (Optional)
+        :return: dict - The API response containing the edited image data.
+        :raises InvalidAPIResponseException: If the API response cannot be processed.
+        """
+        data = {
+            'model': model,
+            'prompt': prompt,
+            'image': image,
+            'size': size,
+            'n': n
+        }
+
+        if quality:
+            data['quality'] = quality
+
+        logger.debug(f"Editing image with data: {data}")
+
+        response = self.api_service.post(
+            endpoint=EDIT_IMAGE_V1,
+            data=data
+        )
+
+        validate_status_code(response)
+
+        result = parse_json_response(response, "edit image")
+        logger.debug(f"Image editing result: {result}")
+        return result
+
+    def get_response(
+            self,
+            model: str,
+            input: str,
+            files: Optional[List[str]] = None,
+            tools: Optional[List[Dict]] = None,
+            tool_choice: Optional[Union[str, Dict]] = None,
+            temperature: Optional[float] = None,
+            max_output_tokens: Optional[int] = None,
+            top_p: Optional[float] = None,
+            metadata: Optional[Dict] = None,
+            user: Optional[str] = None,
+            instructions: Optional[str] = None,
+            reasoning: Optional[Dict] = None,
+            truncation: Optional[str] = None,
+            parallel_tool_calls: Optional[bool] = None,
+            store: Optional[bool] = None,
+            stream: bool = False
+    ) -> Union[dict, str, Generator[str, None, None]]:
+        """
+        Generates a response using the Responses API with support for images and PDF files.
+
+        :param model: str - The model specification, e.g., "openai/o1-pro". (Required)
+        :param input: str - The user input text. (Required)
+        :param files: Optional[List[str]] - List of file paths (images or PDFs) to include in the request. (Optional)
+        :param tools: Optional[List[Dict]] - A list of tools (e.g., functions) the model may call. (Optional)
+        :param tool_choice: Optional[Union[str, Dict]] - Controls which tool is called (e.g., "none", "auto", or specific tool). (Optional)
+        :param temperature: Optional[float] - Controls the randomness of the response. (Optional)
+        :param max_output_tokens: Optional[int] - The maximum number of tokens to generate in the response. (Optional)
+        :param top_p: Optional[float] - Nucleus sampling parameter. (Optional)
+        :param metadata: Optional[Dict] - Up to 16 key-value pairs to attach to the object. (Optional)
+        :param user: Optional[str] - A unique identifier for the end-user. (Optional)
+        :param instructions: Optional[str] - Additional instructions for the model. (Optional)
+        :param reasoning: Optional[Dict] - Reasoning configuration, e.g., {"effort": "medium"}. (Optional)
+        :param truncation: Optional[str] - Truncation strategy, e.g., "disabled". (Optional)
+        :param parallel_tool_calls: Optional[bool] - Whether to enable parallel tool calls. (Optional)
+        :param store: Optional[bool] - Whether to store the output. (Optional)
+        :param stream: bool - Whether the response should be streamed. Possible values:
+            - False: Returns the complete response as a dictionary or string (default).
+            - True: Returns a generator yielding content strings as they are received. (Optional)
+        :return: Union[dict, str, Generator[str, None, None]] - For non-streaming (stream=False), returns a dictionary containing the response
+            result or a string if JSON decoding fails. For streaming (stream=True), returns a generator yielding content strings extracted from the
+            streaming response chunks.
+        :raises InvalidAPIResponseException: If the API response cannot be processed.
+        """
+        data = {
+            'model': model,
+            'input': input,
+            'stream': stream
+        }
+
+        if temperature is not None:
+            data['temperature'] = temperature
+
+        if max_output_tokens is not None:
+            data['max_output_tokens'] = max_output_tokens
+
+        if top_p is not None:
+            data['top_p'] = top_p
+
+        if tools is not None:
+            data['tools'] = tools
+
+        if tool_choice is not None:
+            data['tool_choice'] = tool_choice
+
+        if metadata is not None:
+            data['metadata'] = metadata
+
+        if user is not None:
+            data['user'] = user
+
+        if instructions is not None:
+            data['instructions'] = instructions
+
+        if reasoning is not None:
+            data['reasoning'] = reasoning
+
+        if truncation is not None:
+            data['truncation'] = truncation
+
+        if parallel_tool_calls is not None:
+            data['parallel_tool_calls'] = parallel_tool_calls
+
+        if store is not None:
+            data['store'] = store
+
+        logger.debug(f"Generating response with data: {data}")
+
+        if files:
+            if stream:
+                raise InvalidAPIResponseException("Streaming is not supported when uploading files")
+            
+            file_handles = []
+            try:
+                files_dict = {}
+                for idx, file_path in enumerate(files):
+                    path = Path(file_path)
+                    if not path.is_file():
+                        raise FileNotFoundError(f"File not found: {file_path}")
+                    
+                    file_handle = path.open("rb")
+                    file_handles.append(file_handle)
+                    files_dict[f"file{idx}"] = file_handle
+
+                response = self.api_service.post_files_multipart(
+                    endpoint=RESPONSES_V1,
+                    data=data,
+                    files=files_dict
+                )
+            finally:
+                for fh in file_handles:
+                    fh.close()
+        else:
+            if stream:
+                response = self.api_service.stream_post(
+                    endpoint=RESPONSES_V1,
+                    data=data
+                )
+                return self.stream_response_generator(response)
+            else:
+                response = self.api_service.post(
+                    endpoint=RESPONSES_V1,
+                    data=data
+                )
+
+        validate_status_code(response)
+        result = parse_json_response(response, "get response")
+        logger.debug(f"Response result: {result}")
+
+        return result

pygeai/chat/endpoints.py

@@ -1,2 +1,5 @@
 CHAT_V1 = "/chat"                           # POST Chat endpoint
 CHAT_COMPLETION_V1 = "/chat/completion"     # POST Chat completion endpoint
+GENERATE_IMAGE_V1 = "/images"               # POST Generate image
+EDIT_IMAGE_V1 = "/images/edits"             # POST Edit image
+RESPONSES_V1 = "/responses"                 # POST Responses API endpoint

pygeai/chat/iris.py

@@ -0,0 +1,17 @@
+from pygeai.chat.clients import ChatClient
+from pygeai.chat.settings import LLM_SETTINGS
+
+
+class Iris:
+    def __init__(self):
+        self.client = ChatClient()
+
+    def stream_answer(self, messages):
+        result = self.client.chat_completion(
+            model="saia:agent:com.globant.iris",
+            messages=messages,
+            stream=True,
+            **LLM_SETTINGS
+        )
+        return result
+

pygeai/chat/managers.py

@@ -0,0 +1,64 @@
+from pygeai import logger
+from pygeai.core.common.exceptions import APIError
+from pygeai.core.handlers import ErrorHandler
+from pygeai.core.models import ChatMessageList, \
+    ChatVariableList, LlmSettings, ChatToolList, ToolChoice
+from pygeai.assistant.mappers import AssistantResponseMapper
+from pygeai.chat.clients import ChatClient
+from pygeai.core.responses import ChatResponse, ProviderResponse
+
+
+class ChatManager:
+
+    def __init__(self, api_key: str = None, base_url: str = None, alias: str = None):
+        self.__chat_client = ChatClient(api_key, base_url, alias)
+
+    def chat_completion(
+            self,
+            model: str,
+            messages: ChatMessageList,
+            llm_settings: LlmSettings,
+            thread_id: str = None,
+            variables: ChatVariableList = None,
+            tool_choice: ToolChoice = None,
+            tools: ChatToolList = None,
+    ) -> ProviderResponse:
+        """
+        Generates a chat completion response using the specified language model.
+
+        This method sends a chat completion request to the ChatClient with the provided
+        model, messages, and settings. It processes the response by mapping errors if
+        present or converting it into an assistant response format.
+
+        :param model: str - The identifier of the language model to use.
+        :param messages: ChatMessageList - A list of chat messages to provide context for the completion.
+        :param llm_settings: LlmSettings - Configuration settings for the language model,
+                             including temperature, max tokens, and penalties.
+        :param thread_id: str, optional - An optional identifier for the conversation thread.
+        :param variables: ChatVariableList, optional - Additional variables to include in the request.
+        :param tool_choice: ToolChoice, optional - Indicates which tool to call.
+        :param tools: ChatToolList, optional - Additional tools the model may call
+        :return: ProviderResponse - The structured chat response.
+        :raises APIError: If the API returns errors.
+        """
+        response_data = self.__chat_client.chat_completion(
+            model=model,
+            messages=messages.to_list(),
+            stream=False,
+            temperature=llm_settings.temperature,
+            max_tokens=llm_settings.max_tokens,
+            thread_id=thread_id,
+            frequency_penalty=llm_settings.frequency_penalty,
+            presence_penalty=llm_settings.presence_penalty,
+            variables=variables.to_list() if variables else None,
+            tool_choice=tool_choice.to_dict() if tool_choice else None,
+            tools=tools.to_list() if tools else None
+        )
+
+        if ErrorHandler.has_errors(response_data):
+            error = ErrorHandler.extract_error(response_data)
+            logger.error(f"Error received while chatting while sending chat request: {error}")
+            raise APIError(f"Error received while chatting while sending chat request: {error}")
+
+        result = AssistantResponseMapper.map_to_provider_response(response_data)
+        return result

pygeai/chat/session.py

@@ -0,0 +1,38 @@
+from pygeai import logger
+from pygeai.chat.clients import ChatClient
+
+from pygeai.chat.settings import LLM_SETTINGS
+from pygeai.core.common.exceptions import InvalidAPIResponseException
+
+
+class AgentChatSession:
+
+    def __init__(self, agent_name: str):
+        self.client = ChatClient()
+        self.agent_name = agent_name
+
+    def stream_answer(self, messages):
+        result = self.client.chat_completion(
+            model=f"saia:agent:{self.agent_name}",
+            messages=messages,
+            stream=True,
+            **LLM_SETTINGS
+        )
+        return result
+
+    def get_answer(self, messages):
+        answer = ""
+        try:
+            result = self.client.chat_completion(
+                model=f"saia:agent:{self.agent_name}",
+                messages=messages,
+                stream=False,
+                **LLM_SETTINGS
+            )
+            answer = result['choices'][0]['message']['content']
+        except Exception as e:
+            logger.error(f"Unable to communicate with specified agent {self.agent_name}: {e}")
+            raise InvalidAPIResponseException(f"Unable to communicate with specified agent {self.agent_name}: {e}")
+
+        return answer
+

pygeai/chat/settings.py

@@ -0,0 +1,6 @@
+LLM_SETTINGS = {
+    "temperature": 0.6,
+    "max_tokens": 800,
+    "frequency_penalty": 0.1,
+    "presence_penalty": 0.2
+}