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

pygeai/assistant/mappers.py

@@ -0,0 +1,153 @@
+import json
+from typing import Dict, Any, Optional, List
+
+from pygeai.core.responses import NewAssistantResponse, ChatResponse, ProviderResponse, \
+    UsageDetails, TokenDetails, Choice, ChatMessage
+from pygeai.core.base.mappers import ModelMapper
+from pygeai.core.models import Assistant
+
+
+class AssistantResponseMapper:
+
+    @classmethod
+    def map_to_assistant_response(cls, data: dict) -> Assistant:
+        assistant = ModelMapper.map_to_assistant(data)
+
+        return assistant
+
+    @classmethod
+    def map_to_assistant_created_response(cls, data: dict) -> NewAssistantResponse:
+        response = NewAssistantResponse()
+        if 'assistantId' in data:
+            response.assistant = ModelMapper.map_to_assistant(data)
+        else:
+            response.project = ModelMapper.map_to_project(data)
+
+        return response
+
+    @classmethod
+    def map_to_chat_request_response(cls, data: Dict[str, Any]) -> ChatResponse:
+        """
+        Maps a dictionary to a `ChatResponse` object.
+
+        :param data: dict - The dictionary containing chat response details.
+        :return: ChatResponse - The mapped `ChatResponse` object.
+        """
+        provider_response = data.get("providerResponse")
+        return ChatResponse(
+            progress=data.get("progress"),
+            provider_name=data.get("providerName"),
+            provider_response=cls.map_to_provider_response(provider_response) if provider_response else None,
+            request_id=data.get("requestId"),
+            status=data.get("status"),
+            success=data.get("success"),
+            text=data.get("text")
+        )
+
+    @classmethod
+    def map_to_provider_response(cls, provider_response: Any) -> ProviderResponse:
+        """
+        Parses the provider response, which may be a JSON string.
+
+        :param provider_response: Any - The raw provider response.
+        :return: ProviderResponse - The parsed `ProviderResponse` object.
+        """
+        if isinstance(provider_response, str):
+            provider_response = json.loads(provider_response)
+
+        choices = provider_response.get("choices")
+        usage_data = provider_response.get("usage")
+        return ProviderResponse(
+            created=provider_response.get("created"),
+            usage=cls._parse_usage(usage_data) if usage_data else None,
+            model=provider_response.get("model"),
+            service_tier=provider_response.get("service_tier"),
+            id=provider_response.get("id"),
+            system_fingerprint=provider_response.get("system_fingerprint"),
+            choices=cls._parse_choices(choices) if choices else [],
+            object=provider_response.get("object")
+        )
+
+    @classmethod
+    def _parse_usage(cls, usage_data: Dict[str, Any]) -> UsageDetails:
+        """
+        Parses usage details from a dictionary.
+
+        :param usage_data: dict - The usage details data.
+        :return: UsageDetails - The parsed `UsageDetails` object.
+        """
+        completion_tokens_details_data = usage_data.get("completion_tokens_details")
+        prompt_tokens_details_data = usage_data.get("prompt_tokens_details")
+        return UsageDetails(
+            completion_tokens=usage_data.get("completion_tokens", 0),
+            prompt_tokens=usage_data.get("prompt_tokens", 0),
+            total_cost=usage_data.get("total_cost", 0.0),
+            completion_tokens_details=cls._parse_completion_tokens_details(completion_tokens_details_data) if completion_tokens_details_data else None,
+            prompt_tokens_details=cls._parse_prompt_tokens_details(prompt_tokens_details_data) if prompt_tokens_details_data else None,
+            total_tokens=usage_data.get("total_tokens", 0),
+            currency=usage_data.get("currency", "USD"),
+            completion_cost=usage_data.get("completion_cost", 0.0),
+            prompt_cost=usage_data.get("prompt_cost", 0.0)
+        )
+
+    @classmethod
+    def _parse_completion_tokens_details(cls, details: Dict[str, Any]) -> Optional[TokenDetails]:
+        """
+        Parses completion tokens details.
+
+        :param details: dict - The completion tokens details.
+        :return: CompletionTokensDetails - The parsed details object.
+        """
+        if not details:
+            return None
+        return TokenDetails(
+            reasoning_tokens=details.get("reasoning_tokens")
+        )
+
+    @classmethod
+    def _parse_prompt_tokens_details(cls, details: Dict[str, Any]) -> Optional[TokenDetails]:
+        """
+        Parses prompt tokens details.
+
+        :param details: dict - The prompt tokens details.
+        :return: PromptTokensDetails - The parsed details object.
+        """
+        if not details:
+            return None
+        return TokenDetails(
+            cached_tokens=details.get("cached_tokens")
+        )
+
+    @classmethod
+    def _parse_choices(cls, choices: list) -> List[Choice]:
+        """
+        Parses choices from a list.
+
+        :param choices: list - The choices data.
+        :return: List[Choice] - A list of parsed `Choice` objects.
+        """
+        return [
+            Choice(
+                finish_reason=choice.get("finish_reason"),
+                index=choice.get("index"),
+                message=cls._parse_message(choice.get("message")) if 'message' in choice else None
+            )
+            for choice in choices
+        ]
+
+    @classmethod
+    def _parse_message(cls, message_data: Dict[str, Any]) -> ChatMessage:
+        """
+        Parses a chat message.
+
+        :param message_data: dict - The chat message data.
+        :return: ChatMessage - The parsed `ChatMessage` object.
+        """
+        return ChatMessage(
+            role=message_data.get("role"),
+            function_call=message_data.get("function_call"),
+            refusal=message_data.get("refusal"),
+            tool_calls=message_data.get("tool_calls"),
+            content=message_data.get("content")
+        )
+

pygeai/assistant/rag/clients.py

@@ -1,11 +1,19 @@
 import json
+import logging
+from pathlib import Path
+from typing import Any
 
+from pygeai import logger
 from pygeai.assistant.rag.endpoints import GET_ASSISTANTS_FROM_PROJECT_V1, GET_ASSISTANT_V1, CREATE_ASSISTANT_V1, \
     UPDATE_ASSISTANT_V1, DELETE_ASSISTANT_V1, GET_DOCUMENTS_V1, DELETE_ALL_DOCUMENTS_V1, RETRIEVE_DOCUMENT_V1, \
     UPLOAD_DOCUMENT_V1, DELETE_DOCUMENT_V1, EXECUTE_QUERY_V1
 from pygeai.core.base.clients import BaseClient
 import urllib.parse
 
+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 RAGAssistantClient(BaseClient):
 
@@ -14,13 +22,14 @@ class RAGAssistantClient(BaseClient):
 
     def get_assistants_from_project(self) -> dict:
         response = self.api_service.get(endpoint=GET_ASSISTANTS_FROM_PROJECT_V1)
-        result = json.loads(response.content)
-        return result
+        validate_status_code(response)
+        return parse_json_response(response, "get assistants from project")
 
     def get_assistant_data(self, name: str) -> dict:
         endpoint = GET_ASSISTANT_V1.format(name=name)
         response = self.api_service.get(endpoint=endpoint)
-        return json.loads(response.content)
+        validate_status_code(response)
+        return parse_json_response(response, "get assistant data for name", name=name)
 
     def create_assistant(
             self,
@@ -38,7 +47,7 @@ class RAGAssistantClient(BaseClient):
         :param description: str, optional - A description of the RAG assistant's purpose or functionality.
         :param template: str, optional - Name of an existing RAG template to base the configuration on. Defaults to None.
         :param search_options: dict, optional - A dictionary containing search configuration options.
-            - "historyCount": int - Number of historical interactions to include in the search context.
+            - "history_count": int - Number of historical interactions to include in the search context.
             - "llm": A dictionary with LLM configuration:
                 - "cache": bool - Whether to enable caching for the LLM.
                 - "frequencyPenalty": float - Frequency penalty parameter for LLM responses.
@@ -69,11 +78,11 @@ class RAGAssistantClient(BaseClient):
         :param index_options: dict, optional - A dictionary containing index configuration options.
             - "chunks": Configuration for document chunking:
                 - "chunkOverlap": int - Overlap size between chunks in the main document.
-                - "chunkSize": int - Size of each chunk in the main document.
+                - "chunk_size": int - Size of each chunk in the main document.
             - "useParentDocument": bool - Whether to enable parent-child relationships.
-            - "childDocument": Configuration for child documents:
+            - "child_document": Configuration for child documents:
                 - "childK": float - Parameter for child document processing.
-                - "chunkSize": float - Chunk size for child documents.
+                - "chunk_size": float - Chunk size for child documents.
                 - "chunkOverlap": float - Overlap size between child document chunks.
         :param welcome_data: dict, optional - A dictionary containing welcome data for the assistant.
             - "title": str - Title of the welcome message.
@@ -100,7 +109,8 @@ class RAGAssistantClient(BaseClient):
             endpoint=CREATE_ASSISTANT_V1,
             data=data
         )
-        return json.loads(response.content)
+        validate_status_code(response)
+        return parse_json_response(response, "create assistant with name", name=name)
 
     def update_assistant(
             self,
@@ -133,45 +143,146 @@ class RAGAssistantClient(BaseClient):
             endpoint=endpoint,
             data=data
         )
-        return json.loads(response.content)
+        validate_status_code(response)
+        return parse_json_response(response, "update assistant with name", name=name)
 
     def delete_assistant(self, name: str) -> dict:
         safe_name = self.get_url_safe_name(name)
         endpoint = DELETE_ASSISTANT_V1.format(name=safe_name)
         response = self.api_service.delete(endpoint=endpoint)
-        return json.loads(response.content)
+        validate_status_code(response)
+        return parse_json_response(response, "delete assistant with name", name=name)
 
-    def get_documents(self, name: str) -> dict:
+    def get_documents(
+            self,
+            name: str,
+            skip: int = 0,
+            count: int = 10
+    ) -> dict:
         safe_name = self.get_url_safe_name(name)
         endpoint = GET_DOCUMENTS_V1.format(name=safe_name)
-        response = self.api_service.get(endpoint=endpoint)
-        return json.loads(response.content)
+        response = self.api_service.get(
+            endpoint=endpoint,
+            params={
+                "skip": skip,
+                "count": count
+            }
+        )
+        validate_status_code(response)
+        return parse_json_response(response, "get documents for assistant", name=name)
 
     def delete_all_documents(self, name: str) -> dict:
         safe_name = self.get_url_safe_name(name)
         endpoint = DELETE_ALL_DOCUMENTS_V1.format(name=safe_name)
         response = self.api_service.delete(endpoint=endpoint)
-        return json.loads(response.content)
+        validate_status_code(response)
+        return parse_json_response(response, "delete all documents for assistant", name=name)
 
     def retrieve_document(self, name: str, document_id: str) -> dict:
         safe_name = self.get_url_safe_name(name)
         endpoint = RETRIEVE_DOCUMENT_V1.format(name=safe_name, id=document_id)
         response = self.api_service.get(endpoint=endpoint)
-        return json.loads(response.content)
+        validate_status_code(response)
+        return parse_json_response(response, "retrieve document", document_id=document_id)
 
-    def upload_document(self, name: str, document: dict) -> dict:
+    def upload_document(
+            self,
+            name: str,
+            file_path: str,
+            upload_type: str = "multipart",
+            metadata: dict = None,
+            content_type: str = "application/pdf"
+    ) -> dict:
         safe_name = self.get_url_safe_name(name)
         endpoint = UPLOAD_DOCUMENT_V1.format(name=safe_name)
-        response = self.api_service.post(endpoint=endpoint, json=document)
-        return json.loads(response.content)
+
+        if upload_type == "binary":
+            response = self._upload_binary_document(endpoint, file_path, content_type)
+        elif upload_type == "multipart":
+            response = self._upload_multipart_document(endpoint, file_path, content_type, metadata)
+        else:
+            raise ValueError("Invalid upload_type. Use 'binary' or 'multipart'.")
+
+        validate_status_code(response)
+        return parse_json_response(response, "upload document for assistant", name=name)
+
+    def _upload_binary_document(self, endpoint: str, file_path: str, content_type: str):
+        """
+        Uploads a binary file to the specified API endpoint.
+
+        :param endpoint: str - The API endpoint where the file will be uploaded.
+        :param file_path: str - The path to the file to be uploaded.
+        :param content_type: str - The MIME type of the file (e.g., "application/pdf", "text/plain").
+        :return: Response object from the API request.
+        """
+        headers = dict()
+        headers["filename"] = Path(file_path).name
+        headers["Content-Type"] = content_type
+
+        with open(file_path, "rb") as file:
+            response = self.api_service.post_file_binary(
+                endpoint=endpoint,
+                headers=headers,
+                file=file
+            )
+        return response
+
+    def _upload_multipart_document(self, endpoint: str, file_path: str, content_type: str, metadata: Any):
+        """
+        Uploads a file using multipart/form-data encoding, optionally including metadata.
+
+        :param endpoint: str - The API endpoint where the file will be uploaded.
+        :param file_path: str - The path to the file to be uploaded.
+        :param content_type: str - The MIME type of the file (e.g., "application/pdf", "text/plain").
+        :param metadata: Any - Metadata to be included in the request. Can be a dictionary (sent as JSON) or a file path.
+
+        :raises ValueError: If metadata is neither a dictionary nor a valid file path.
+        :return: Response object from the API request, or None if an error occurs.
+        """
+        files_payload = []
+        file_name = Path(file_path).name
+        file_obj = open(file_path, "rb")
+        files_payload.append(
+            ("file", (file_name, file_obj, content_type))
+        )
+        if metadata:
+            if isinstance(metadata, dict):
+                files_payload.append(
+                    ("metadata", (None, json.dumps(metadata), "application/json"))
+                )
+            elif isinstance(metadata, str) and Path(metadata).is_file():
+                metadata_path = Path(metadata)
+                files_payload.append(
+                    ("metadata", metadata_path.open("rb"))
+                )
+            else:
+                raise ValueError("Metadata should be a dictionary or a valid file path.")
+        try:
+            response = self.api_service.post_files_multipart(
+                endpoint=endpoint,
+                files=files_payload
+            )
+        except Exception as e:
+            logging.error(f"Unable to receive answer from POST request: {e}")
+        else:
+            return response
+        finally:
+            for _, (file_name, file_obj, _) in files_payload:
+                if hasattr(file_obj, "close"):
+                    file_obj.close()
 
     def delete_document(self, name: str, document_id: str) -> dict:
         safe_name = self.get_url_safe_name(name)
         endpoint = DELETE_DOCUMENT_V1.format(name=safe_name, id=document_id)
         response = self.api_service.delete(endpoint=endpoint)
-        return json.loads(response.content)
+        validate_status_code(response)
+        return parse_json_response(response, "delete document", document_id=document_id)
 
     def execute_query(self, query: dict) -> dict:
-        response = self.api_service.post(endpoint=EXECUTE_QUERY_V1, json=query)
-        return json.loads(response.content)
+        response = self.api_service.post(
+            endpoint=EXECUTE_QUERY_V1,
+            data=query
+        )
+        validate_status_code(response)
+        return parse_json_response(response, "execute query")
 

pygeai/assistant/rag/mappers.py

@@ -0,0 +1,228 @@
+from typing import Optional, List
+
+from pygeai.assistant.rag.responses import DocumentListResponse
+from pygeai.core.base.mappers import ModelMapper
+from pygeai.core.models import ChatVariable, \
+    ChatVariableList
+from pygeai.assistant.rag.models import Search, RetrieverOptions, EmbeddingsOptions, IngestionOptions, \
+    SearchOptions, ChildOptions, ChildDocumentOptions, ChunkOptions, IndexOptions, RAGAssistant, Document, \
+    DocumentMetadata
+
+
+class RAGAssistantMapper(ModelMapper):
+
+    @classmethod
+    def map_to_rag_assistant(cls, data: dict) -> RAGAssistant:
+        welcome_data_data = data.get('welcomeData')
+        welcome_data = cls.map_to_welcome_data(welcome_data_data) if welcome_data_data else None
+
+        search_options_data = data.get('searchOptions')
+        search_options = cls.map_to_search_options(search_options_data) if search_options_data else None
+
+        llm_settings_data = data.get('llmSettings') or search_options_data.get('llm')
+        llm_settings = cls.map_to_llm_settings(llm_settings_data) if llm_settings_data else None
+
+        index_options_data = data.get('indexOptions')
+        index_options = cls.map_to_index_options(index_options_data) if index_options_data else None
+
+        return RAGAssistant(
+            name=data.get("name"),
+            type="rag",
+            status=data.get("status"),
+            description=data.get("description"),
+            prompt=data.get("prompt"),
+            template=data.get('template'),
+            search_options=search_options,
+            index_options=index_options,
+            welcome_data=welcome_data,
+            llm_settings=llm_settings
+        )
+
+    @classmethod
+    def map_to_search_options(cls, data: dict) -> SearchOptions:
+        """
+        Maps a dictionary to a `SearchOptions` object.
+
+        :param data: dict - The dictionary containing search options details.
+        :return: SearchOptions - The mapped `SearchOptions` object.
+        """
+        llm_data = data.get("llm")
+        search_data = data.get("search")
+        retriever_data = data.get("retriever")
+        embeddings_data = data.get('embeddings')
+        ingestion_data = data.get('ingestion')
+        variables_data = data.get('variables')
+        return SearchOptions(
+            history_count=data.get("historyCount"),
+            llm=ModelMapper.map_to_llm_settings(llm_data) if llm_data else None,
+            search=cls.map_to_search(search_data) if search_data else None,
+            retriever=cls.map_to_retriever_options(retriever_data) if retriever_data else None,
+            chain=data.get('chain'),
+            embeddings=cls.map_to_embeddings(embeddings_data) if embeddings_data else None,
+            ingestion=cls.map_to_ingestion(ingestion_data) if ingestion_data else None,
+            options=data.get('options'),
+            rerank=data.get('rerank'),
+            variables=cls.map_to_variable_list(variables_data) if variables_data else None,
+            vector_store=data.get('vector_store'),
+        )
+
+    @classmethod
+    def map_to_search(cls, data: dict) -> Search:
+        return Search(
+            k=data.get("k"),
+            type=data.get("type", "similarity"),
+            fetch_k=data.get("fetchK"),
+            lambda_=data.get("lambda"),
+            prompt=data.get("prompt"),
+            return_source_documents=data.get("returnSourceDocuments"),
+            score_threshold=data.get("scoreThreshold"),
+            template=data.get("template")
+        )
+
+    @classmethod
+    def map_to_retriever_options(cls, data: dict) -> RetrieverOptions:
+        return RetrieverOptions(
+            type=data.get("type"),
+            search_type=data.get("searchType", "similarity"),
+            step=data.get("step", "all"),
+            prompt=data.get("prompt")
+        )
+
+    @classmethod
+    def map_to_index_options(cls, data: dict) -> IndexOptions:
+        """
+        Maps a dictionary to an `IndexOptions` object.
+
+        :param data: dict - The dictionary containing index options details.
+        :return: IndexOptions - The mapped `IndexOptions` object.
+        """
+        return IndexOptions(
+            chunks=cls.map_to_chunk_options(data.get("chunks")),
+            use_parent_document=data.get("useParentDocument", False),
+            child_document=cls.map_to_child_document_options(data.get("childDocument"))
+            if data.get("useParentDocument") else None
+        )
+
+    @classmethod
+    def map_to_chunk_options(cls, data: dict) -> ChunkOptions:
+        return ChunkOptions(
+            chunk_overlap=data.get("chunkOverlap"),
+            chunk_size=data.get("chunkSize")
+        )
+
+    @classmethod
+    def map_to_child_document_options(cls, data: dict) -> ChildDocumentOptions:
+        return ChildDocumentOptions(
+            child_k=data.get("childK"),
+            child=cls.map_to_child_options(data.get("child"))
+        )
+
+    @classmethod
+    def map_to_child_options(cls, data: dict) -> ChildOptions:
+        return ChildOptions(
+            chunk_size=data.get("chunkSize"),
+            chunk_overlap=data.get("chunkOverlap"),
+            content_processing=data.get("contentProcessing", "")
+        )
+
+    @classmethod
+    def map_to_embeddings(cls, data: dict) -> Optional[EmbeddingsOptions]:
+        """
+        Maps a dictionary to an `EmbeddingsOptions` object.
+
+        :param data: dict - The dictionary containing embeddings options details.
+        :return: EmbeddingsOptions - The mapped `EmbeddingsOptions` object.
+        """
+        return EmbeddingsOptions(
+            dimensions=data.get('dimensions'),
+            model_name=data.get('modelName'),
+            provider=data.get('provider'),
+            use_proxy=data.get('useProxy')
+        )
+
+    @classmethod
+    def map_to_ingestion(cls, data: dict) -> Optional[IngestionOptions]:
+        """
+        Maps a dictionary to an `IngestionOptions` object.
+
+        :param data: dict - The dictionary containing ingestion options details.
+        :return: IngestionOptions - The mapped `IngestionOptions` object.
+        """
+        return IngestionOptions(
+            geai_options=data.get('geaiOptions'),
+            llama_parse_options=data.get('llamaParseOptions'),
+            provider=data.get('provider')
+        )
+
+    @classmethod
+    def map_to_variable_list(cls, data: list) -> Optional[ChatVariableList]:
+        """
+        Maps a list of dictionaries to a `ChatVariableList` object.
+
+        :param data: list - The list containing chat variable details.
+        :return: ChatVariableList - The mapped `ChatVariableList` object.
+        """
+        variables = list()
+        for variable in data:
+            for k, v in variable.items():
+                variables.append(
+                    ChatVariable(
+                        key=k,
+                        value=v
+                    )
+                )
+
+        return ChatVariableList(variables=variables)
+
+    @classmethod
+    def map_to_document_list_response(cls, data: dict) -> DocumentListResponse:
+        documents = cls.map_to_document_list(data)
+        count = data.get('count')
+
+        return DocumentListResponse(
+            count=count,
+            documents=documents
+        )
+
+    @classmethod
+    def map_to_document_list(cls, data: dict) -> list[Document]:
+        document_list = list()
+        documents = data.get("documents")
+        if documents is not None and any(documents):
+            for document_data in documents:
+                document = cls.map_to_document(document_data)
+                document_list.append(document)
+
+        return document_list
+
+    @classmethod
+    def map_to_document(cls, data: dict) -> Document:
+        metadata = cls.map_to_document_metadata_list(data)
+
+        return Document(
+            id=data.get('id'),
+            chunks=data.get('chunks'),
+            extension=data.get('extension'),
+            index_status=data.get('indexStatus'),
+            metadata=metadata,
+            timestamp=data.get('timestamp'),
+            url=data.get('url'),
+        )
+
+    @classmethod
+    def map_to_document_metadata_list(cls, data: dict) -> List[DocumentMetadata]:
+        metadata_list = list()
+        metadata = data.get('metadata')
+        if metadata is not None and any(metadata):
+            for metadata_data in metadata:
+                new_metadata = cls.map_to_document_metadata(metadata_data)
+                metadata_list.append(new_metadata)
+
+        return metadata_list
+
+    @classmethod
+    def map_to_document_metadata(cls, data: dict) -> DocumentMetadata:
+        return DocumentMetadata(
+            key=data.get('key'),
+            value=data.get('value')
+        )