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

pygeai/assistant/rag/models.py

@@ -0,0 +1,396 @@
+from datetime import datetime
+from typing import Optional, Literal, Dict, List
+
+from pydantic import Field, field_validator
+
+from pygeai.core import CustomBaseModel
+from pygeai.core.models import ChatVariableList, Assistant, LlmSettings
+
+
+class DocumentMetadata(CustomBaseModel):
+    """
+    Represents metadata for a document in key-value format.
+    """
+    key: str = Field(..., alias="key", description="The key of the metadata")
+    value: str = Field(..., alias="value", description="The value of the metadata")
+
+    def to_dict(self):
+        return self.model_dump(by_alias=True, exclude_none=True)
+
+    def __str__(self):
+        return str(self.to_dict())
+
+
+class Document(CustomBaseModel):
+    """
+    Represents a document with associated metadata and properties.
+    """
+    id: str = Field(..., alias="id", description="Unique identifier of the document")
+    chunks: str = Field(..., alias="chunks", description="Content chunks of the document")
+    name: Optional[str] = Field(None, alias="name", description="Name of the document")
+    extension: str = Field(..., alias="extension", description="File extension of the document")
+    index_status: str = Field(..., alias="indexStatus", description="Indexing status of the document")
+    metadata: Optional[List[DocumentMetadata]] = Field([], alias="metadata", description="List of metadata associated with the document")
+    timestamp: Optional[datetime] = Field(None, alias="timestamp", description="Timestamp of the document")
+    url: str = Field(..., alias="url", description="URL of the document")
+
+    @field_validator("metadata", mode="before")
+    @classmethod
+    def normalize_metadata(cls, value):
+        if isinstance(value, list):
+            return [DocumentMetadata.model_validate(item) if isinstance(item, dict) else item for item in value]
+        return value
+
+    def to_dict(self):
+        result = {
+            "id": self.id,
+            "chunks": self.chunks,
+            "name": self.name,
+            "extension": self.extension,
+            "indexStatus": self.index_status,
+            "metadata": [meta.to_dict() for meta in self.metadata] if self.metadata else None,
+            "timestamp": self.timestamp,
+            "url": self.url
+        }
+        return {k: v for k, v in result.items() if v is not None}
+
+    def __str__(self):
+        return str(self.to_dict())
+
+
+class Search(CustomBaseModel):
+    """
+    Represents search configuration for querying documents.
+    """
+    k: int = Field(..., alias="k", description="Number of results to return")
+    type: Literal["similarity", "mmr"] = Field("similarity", alias="type", description="Type of search algorithm to use")
+    fetch_k: Optional[float] = Field(None, alias="fetchK", description="Number of documents to fetch for MMR search")
+    lambda_: Optional[float] = Field(None, alias="lambda", description="Lambda parameter for MMR search")
+    prompt: str = Field(..., alias="prompt", description="Search prompt or query")
+    return_source_documents: bool = Field(..., alias="returnSourceDocuments", description="Whether to return source documents")
+    score_threshold: float = Field(..., alias="scoreThreshold", description="Score threshold for search results")
+    template: str = Field(..., alias="template", description="Template for formatting search results")
+
+    def to_dict(self):
+        result = {
+            "k": self.k,
+            "type": self.type,
+            "fetchK": self.fetch_k,
+            "lambda": self.lambda_,
+            "prompt": self.prompt,
+            "returnSourceDocuments": self.return_source_documents,
+            "scoreThreshold": self.score_threshold,
+            "template": self.template
+        }
+        return {k: v for k, v in result.items() if v is not None}
+
+    def __str__(self):
+        return str(self.to_dict())
+
+
+class RetrieverOptions(CustomBaseModel):
+    """
+    Represents options for the retriever used in search.
+    """
+    type: Literal["vectorStore", "multiQuery", "selfQuery", "hyde", "contextualCompression"] = Field(..., alias="type", description="Type of retriever to use")
+    search_type: Optional[str] = Field("similarity", alias="searchType", description="Search type for Azure AISearch")
+    step: Optional[Literal["all", "documents"]] = Field("all", alias="step", description="Step for retrieval process")
+    prompt: Optional[str] = Field(None, alias="prompt", description="Custom prompt for retriever, if applicable")
+
+    def to_dict(self):
+        result = {
+            "type": self.type,
+            "searchType": self.search_type,
+            "step": self.step,
+            "prompt": self.prompt
+        }
+        return {k: v for k, v in result.items() if v is not None}
+
+    def __str__(self):
+        return str(self.to_dict())
+
+
+class ChainOptions(CustomBaseModel):
+    """
+    Represents options for the processing chain.
+    """
+    type: str = Field(..., alias="type", description="Type of processing chain")
+
+    def to_dict(self):
+        return {"type": self.type}
+
+    def __str__(self):
+        return str(self.to_dict())
+
+
+class EmbeddingsOptions(CustomBaseModel):
+    """
+    Represents configuration for embeddings.
+    """
+    dimensions: int = Field(..., alias="dimensions", description="Number of dimensions for embeddings")
+    model_name: str = Field(..., alias="modelName", description="Name of the embedding model")
+    provider: str = Field(..., alias="provider", description="Provider of the embedding model")
+    use_proxy: Optional[bool] = Field(False, alias="useProxy", description="Whether to use a proxy for embedding requests")
+
+    def to_dict(self):
+        result = {
+            "dimensions": self.dimensions,
+            "modelName": self.model_name,
+            "provider": self.provider,
+            "useProxy": self.use_proxy
+        }
+        return {k: v for k, v in result.items() if v is not None}
+
+    def __str__(self):
+        return str(self.to_dict())
+
+
+class IngestionOptions(CustomBaseModel):
+    """
+    Represents configuration for document ingestion.
+    """
+    geai_options: Dict = Field(..., alias="geaiOptions", description="GEAI-specific ingestion options")
+    llama_parse_options: Dict = Field(..., alias="llamaParseOptions", description="LlamaParse-specific ingestion options")
+    provider: str = Field(..., alias="provider", description="Provider for ingestion")
+
+    def to_dict(self):
+        result = {
+            "geaiOptions": self.geai_options,
+            "llamaParseOptions": self.llama_parse_options,
+            "provider": self.provider
+        }
+        return {k: v for k, v in result.items() if v is not None}
+
+    def __str__(self):
+        return str(self.to_dict())
+
+
+class SearchOptions(CustomBaseModel):
+    """
+    Represents comprehensive search options for RAG.
+    """
+    history_count: int = Field(..., alias="historyCount", description="Number of history items to consider")
+    llm: LlmSettings = Field(..., alias="llm", description="LLM settings for search")
+    search: Search = Field(..., alias="search", description="Search configuration")
+    retriever: RetrieverOptions = Field(..., alias="retriever", description="Retriever configuration")
+    chain: Optional[ChainOptions] = Field(None, alias="chain", description="Chain configuration")
+    embeddings: Optional[EmbeddingsOptions] = Field(None, alias="embeddings", description="Embeddings configuration")
+    ingestion: Optional[IngestionOptions] = Field(None, alias="ingestion", description="Ingestion configuration")
+    options: Optional[Dict] = Field(None, alias="options", description="Additional options")
+    rerank: Optional[Dict] = Field(None, alias="rerank", description="Rerank configuration")
+    variables: Optional[ChatVariableList] = Field(None, alias="variables", description="Variables for search")
+    vector_store: Optional[Dict] = Field(None, alias="vectorStore", description="Vector store configuration")
+
+    @field_validator("search", mode="before")
+    @classmethod
+    def normalize_search(cls, value):
+        if isinstance(value, dict):
+            return Search.model_validate(value)
+        return value
+
+    @field_validator("retriever", mode="before")
+    @classmethod
+    def normalize_retriever(cls, value):
+        if isinstance(value, dict):
+            return RetrieverOptions.model_validate(value)
+        return value
+
+    @field_validator("chain", mode="before")
+    @classmethod
+    def normalize_chain(cls, value):
+        if isinstance(value, dict):
+            return ChainOptions.model_validate(value)
+        return value
+
+    @field_validator("embeddings", mode="before")
+    @classmethod
+    def normalize_embeddings(cls, value):
+        if isinstance(value, dict):
+            return EmbeddingsOptions.model_validate(value)
+        return value
+
+    @field_validator("ingestion", mode="before")
+    @classmethod
+    def normalize_ingestion(cls, value):
+        if isinstance(value, dict):
+            return IngestionOptions.model_validate(value)
+        return value
+
+    def to_dict(self):
+        result = {
+            "historyCount": self.history_count,
+            "llm": self.llm.to_dict() if self.llm else None,
+            "search": self.search.to_dict() if self.search else None,
+            "retriever": self.retriever.to_dict() if self.retriever else None,
+            "chain": self.chain.to_dict() if self.chain else None,
+            "embeddings": self.embeddings.to_dict() if self.embeddings else None,
+            "ingestion": self.ingestion.to_dict() if self.ingestion else None,
+            "options": self.options,
+            "rerank": self.rerank,
+            "variables": self.variables.to_list() if self.variables else None,
+            "vectorStore": self.vector_store
+        }
+        return {k: v for k, v in result.items() if v is not None}
+
+    def __str__(self):
+        return str(self.to_dict())
+
+
+class ChildOptions(CustomBaseModel):
+    """
+    Represents options for child document processing.
+    """
+    chunk_size: float = Field(..., alias="chunkSize", description="Size of chunks for child documents")
+    chunk_overlap: float = Field(..., alias="chunkOverlap", description="Overlap between chunks for child documents")
+    content_processing: Optional[Literal["", "clean"]] = Field("", alias="contentProcessing", description="Content processing mode for child documents")
+
+    def to_dict(self):
+        result = {
+            "chunkSize": self.chunk_size,
+            "chunkOverlap": self.chunk_overlap,
+            "contentProcessing": self.content_processing
+        }
+        return {k: v for k, v in result.items() if v is not None}
+
+    def __str__(self):
+        return str(self.to_dict())
+
+
+class ChildDocumentOptions(CustomBaseModel):
+    """
+    Represents options for child document retrieval.
+    """
+    child_k: float = Field(..., alias="childK", description="Number of child documents to retrieve")
+    child: ChildOptions = Field(..., alias="child", description="Child document processing options")
+
+    @field_validator("child", mode="before")
+    @classmethod
+    def normalize_child(cls, value):
+        if isinstance(value, dict):
+            return ChildOptions.model_validate(value)
+        return value
+
+    def to_dict(self):
+        result = {
+            "childK": self.child_k,
+            "child": self.child.to_dict() if self.child else None
+        }
+        return {k: v for k, v in result.items() if v is not None}
+
+    def __str__(self):
+        return str(self.to_dict())
+
+
+class ChunkOptions(CustomBaseModel):
+    """
+    Represents chunking options for document processing.
+    """
+    chunk_overlap: int = Field(..., alias="chunkOverlap", description="Overlap between chunks")
+    chunk_size: int = Field(..., alias="chunkSize", description="Size of each chunk")
+
+    def to_dict(self):
+        result = {
+            "chunkOverlap": self.chunk_overlap,
+            "chunkSize": self.chunk_size
+        }
+        return {k: v for k, v in result.items() if v is not None}
+
+    def __str__(self):
+        return str(self.to_dict())
+
+
+class IndexOptions(CustomBaseModel):
+    """
+    Represents indexing options for documents.
+    """
+    chunks: ChunkOptions = Field(..., alias="chunks", description="Chunking configuration for indexing")
+    use_parent_document: Optional[bool] = Field(False, alias="useParentDocument", description="Whether to use parent document for indexing")
+    child_document: Optional[ChildDocumentOptions] = Field(None, alias="childDocument", description="Child document options if use_parent_document is True")
+
+    @field_validator("chunks", mode="before")
+    @classmethod
+    def normalize_chunks(cls, value):
+        if isinstance(value, dict):
+            return ChunkOptions.model_validate(value)
+        return value
+
+    @field_validator("child_document", mode="before")
+    @classmethod
+    def normalize_child_document(cls, value):
+        if isinstance(value, dict):
+            return ChildDocumentOptions.model_validate(value)
+        return value
+
+    def to_dict(self):
+        result = {
+            "chunks": self.chunks.to_dict() if self.chunks else None,
+            "useParentDocument": self.use_parent_document,
+            "childDocument": self.child_document.to_dict() if self.child_document else None
+        }
+        return {k: v for k, v in result.items() if v is not None}
+
+    def __str__(self):
+        return str(self.to_dict())
+
+
+class RAGAssistant(Assistant):
+    """
+    Represents a RAG (Retrieval-Augmented Generation) assistant configuration.
+    """
+    template: Optional[str] = Field(None, alias="template", description="Template for RAG assistant responses")
+    search_options: Optional[SearchOptions] = Field(None, alias="searchOptions", description="Search options for RAG")
+    index_options: Optional[IndexOptions] = Field(None, alias="indexOptions", description="Indexing options for RAG")
+
+    @field_validator("search_options", mode="before")
+    @classmethod
+    def normalize_search_options(cls, value):
+        if isinstance(value, dict):
+            return SearchOptions.model_validate(value)
+        return value
+
+    @field_validator("index_options", mode="before")
+    @classmethod
+    def normalize_index_options(cls, value):
+        if isinstance(value, dict):
+            return IndexOptions.model_validate(value)
+        return value
+
+    def to_dict(self):
+        assistant = super().to_dict()
+        result = {
+            "template": self.template,
+            "searchOptions": self.search_options.to_dict() if self.search_options else None,
+            "indexOptions": self.index_options.to_dict() if self.index_options else None
+        }
+        assistant.update({k: v for k, v in result.items() if v is not None})
+        return assistant
+
+    def __str__(self):
+        return str(self.to_dict())
+
+
+class UploadDocument(CustomBaseModel):
+    """
+    Represents a document to be uploaded.
+    """
+    path: str = Field(..., alias="path", description="Path to the document file")
+    upload_type: Literal["binary", "multipart"] = Field("multipart", alias="uploadType", description="Type of upload method")
+    metadata: Optional[dict] = Field(None, alias="metadata", description="Metadata associated with the document")
+    content_type: str = Field(..., alias="contentType", description="Content type of the document")
+
+    def to_dict(self):
+        result = {
+            "path": self.path,
+            "uploadType": self.upload_type,
+            "metadata": self.metadata,
+            "contentType": self.content_type
+        }
+        return {k: v for k, v in result.items() if v is not None}
+
+    def __str__(self):
+        return str(self.to_dict())
+
+
+class UploadType:
+    BINARY = "binary"
+    MULTIPART = "multipart"

pygeai/assistant/rag/responses.py

@@ -0,0 +1,10 @@
+from typing import List
+
+from pydantic.main import BaseModel
+
+from pygeai.assistant.rag.models import Document
+
+
+class DocumentListResponse(BaseModel):
+    count: int
+    documents: List[Document]

pygeai/auth/clients.py

@@ -0,0 +1,129 @@
+from pygeai import logger
+from pygeai.auth.endpoints import GET_USER_PROFILE_INFO, GET_OAUTH2_ACCESS_TOKEN, \
+    CREATE_PROJECT_API_TOKEN_V2, DELETE_PROJECT_API_TOKEN_V2, UPDATE_PROJECT_API_TOKEN_V2, GET_PROJECT_API_TOKEN_V2
+from pygeai.core.base.clients import BaseClient
+from pygeai.core.common.exceptions import InvalidAPIResponseException, APIResponseError
+from pygeai.core.utils.validators import validate_status_code
+from pygeai.core.utils.parsers import parse_json_response
+
+
+class AuthClient(BaseClient):
+
+    def get_oauth2_access_token(
+            self,
+            client_id: str,
+            username: str,
+            password: str,
+            scope: str = "gam_user_data gam_user_roles"
+    ) -> dict:
+        """
+        Retrieves the list of projects that the user is authorized to access within a specific organization.
+
+        :param client_id: str - The client identifier provided by Globant
+        :param username: str - Username for authentication
+        :param password: str - Password for authentication
+        :param scope: str - Space-separated list of requested scopes
+        :return: dict - The API response containing the list of authorized projects in JSON format.
+        """
+        response = self.api_service.get(
+            endpoint=GET_OAUTH2_ACCESS_TOKEN,
+            params={
+                "client_id": client_id,
+                "scope": scope,
+                "username": username,
+                "password": password
+            }
+        )
+        validate_status_code(response)
+        return parse_json_response(response, "obtain Oauth2 access token")
+
+    def get_user_profile_information(self, access_token: str) -> dict:
+        """
+        Get user profile information
+
+        :param access_token: str - Token obtained in the POST to /oauth/access_token
+        :return: dict - The API response containing the user profile information
+        """
+        self.api_service.token = access_token
+        response = self.api_service.get(endpoint=GET_USER_PROFILE_INFO)
+        validate_status_code(response)
+        return parse_json_response(response, "retrieve user profile information")
+
+    def create_project_api_token(
+            self,
+            project_id: str,
+            name: str,
+            description: str = None
+    ) -> dict:
+        """
+        Creates a new API token for a project.
+
+        :param project_id: str - The project identifier (required). Will be sent as header.
+        :param name: str - The name of the API token (required).
+        :param description: str - A description of the API token (optional).
+        :return: dict - The API response containing the created API token details in JSON format.
+        """
+        headers = {"project-id": project_id}
+        data = {"name": name}
+        if description:
+            data["description"] = description
+
+        response = self.api_service.post(
+            endpoint=CREATE_PROJECT_API_TOKEN_V2,
+            data=data,
+            headers=headers
+        )
+        validate_status_code(response)
+        return parse_json_response(response, "create project API token")
+
+    def delete_project_api_token(self, api_token_id: str) -> dict:
+        """
+        Revokes an API token by setting its status to "revoked".
+
+        :param api_token_id: str - The unique identifier of the API token to delete (required).
+        :return: dict - The API response confirming the deletion, in JSON format.
+        """
+        endpoint = DELETE_PROJECT_API_TOKEN_V2.format(id=api_token_id)
+        response = self.api_service.delete(endpoint=endpoint)
+        validate_status_code(response)
+        return parse_json_response(response, "delete project API token", api_token_id=api_token_id)
+
+    def update_project_api_token(
+            self,
+            api_token_id: str,
+            description: str = None,
+            status: str = None
+    ) -> dict:
+        """
+        Updates an existing API token's description and/or status.
+
+        :param api_token_id: str - The unique identifier of the API token to update (required).
+        :param description: str - A new description for the API token (optional).
+        :param status: str - The new status for the API token: 'active' or 'blocked' (optional).
+        :return: dict - The API response containing the update result messages in JSON format.
+        """
+        endpoint = UPDATE_PROJECT_API_TOKEN_V2.format(id=api_token_id)
+        data = {}
+        if description is not None:
+            data["description"] = description
+        if status is not None:
+            data["status"] = status
+
+        response = self.api_service.put(
+            endpoint=endpoint,
+            data=data
+        )
+        validate_status_code(response)
+        return parse_json_response(response, "update project API token", api_token_id=api_token_id)
+
+    def get_project_api_token(self, api_token_id: str) -> dict:
+        """
+        Retrieves data for a specific project API token.
+
+        :param api_token_id: str - The unique identifier of the API token (required).
+        :return: dict - The API response containing the API token details in JSON format.
+        """
+        endpoint = GET_PROJECT_API_TOKEN_V2.format(id=api_token_id)
+        response = self.api_service.get(endpoint=endpoint)
+        validate_status_code(response)
+        return parse_json_response(response, "get project API token", api_token_id=api_token_id)

pygeai/auth/endpoints.py

@@ -0,0 +1,6 @@
+GET_OAUTH2_ACCESS_TOKEN = "/oauth/access_token"      # POST -> Obtain an OAuth 2.0 access token
+GET_USER_PROFILE_INFO = "/openid/userinfo"          # GET -> Retrieve user profile information
+CREATE_PROJECT_API_TOKEN_V2 = "v2/projects/tokens"          # POST -> Create a new API token for a project
+DELETE_PROJECT_API_TOKEN_V2 = "v2/projects/tokens/{id}"     # DELETE -> Revoke an API token
+UPDATE_PROJECT_API_TOKEN_V2 = "v2/projects/tokens/{id}"     # PUT -> Update an API token
+GET_PROJECT_API_TOKEN_V2 = "v2/projects/tokens/{id}"        # GET -> Get data of a specific project API token