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

pygeai/cli/commands/rag.py

@@ -1,27 +1,28 @@
 import json
-import sys
+from pathlib import Path
 
 from pygeai.assistant.rag.clients import RAGAssistantClient
 from pygeai.cli.commands import Command, Option, ArgumentsEnum
 from pygeai.cli.commands.builders import build_help_text
-from pygeai.cli.commands.common import get_messages, get_welcome_data, get_search_options, get_boolean_value, \
+from pygeai.cli.commands.common import get_welcome_data, get_search_options, get_boolean_value, \
     get_welcome_data_feature_list, get_welcome_data_example_prompt, get_index_options
 from pygeai.cli.texts.help import RAG_ASSISTANT_HELP_TEXT
 from pygeai.core.common.exceptions import MissingRequirementException, WrongArgumentError
+from pygeai.core.utils.console import Console
 
 
-def show_assistant_help():
+def show_help():
     """
     Displays help text in stdout
     """
     help_text = build_help_text(rag_commands, RAG_ASSISTANT_HELP_TEXT)
-    sys.stdout.write(help_text)
+    Console.write_stdout(help_text)
 
 
 def get_assistants_from_project():
     client = RAGAssistantClient()
     result = client.get_assistants_from_project()
-    sys.stdout.write(f"RAG Assistants in project: \n{result}\n")
+    Console.write_stdout(f"RAG Assistants in project: \n{result}")
 
 
 def get_assistant_detail(option_list: list):
@@ -35,7 +36,7 @@ def get_assistant_detail(option_list: list):
 
     client = RAGAssistantClient()
     result = client.get_assistant_data(name)
-    sys.stdout.write(f"Assistant detail: \n{result}\n")
+    Console.write_stdout(f"Assistant detail: \n{result}")
 
 
 assistant_detail_options = [
@@ -264,7 +265,7 @@ def create_assistant(option_list: list):
         index_options=index_options,
         welcome_data=welcome_data
     )
-    sys.stdout.write(f"New RAG Assistant: \n{result}\n")
+    Console.write_stdout(f"New RAG Assistant: \n{result}")
 
 
 create_assistant_options = [
@@ -685,7 +686,7 @@ def update_assistant(option_list: list):
         search_options=search_options,
         welcome_data=welcome_data
     )
-    sys.stdout.write(f"Updated RAG Assistant: \n{result}\n")
+    Console.write_stdout(f"Updated RAG Assistant: \n{result}")
 
 
 update_assistant_options = [
@@ -894,12 +895,254 @@ update_assistant_options = [
 ]
 
 
+def delete_assistant(option_list: list):
+    name = None
+    for option_flag, option_arg in option_list:
+        if option_flag.name == "name":
+            name = option_arg
+
+    if not name:
+        raise MissingRequirementException("Cannot delete assistant detail without name")
+
+    client = RAGAssistantClient()
+    result = client.delete_assistant(name)
+    Console.write_stdout(f"Deleted assistant: \n{result}")
+
+
+delete_assistant_options = [
+    Option(
+        "name",
+        ["--name", "-n"],
+        "RAG assistant name (required)",
+        True
+    ),
+]
+
+
+def list_documents(option_list: list):
+    name = None
+    skip = 0
+    count = 10
+    for option_flag, option_arg in option_list:
+        if option_flag.name == "name":
+            name = option_arg
+        if option_flag.name == "skip":
+            skip = int(option_arg) if str(option_arg).isdigit() else 0
+        if option_flag.name == "count":
+            count = int(option_arg) if str(option_arg).isdigit() else 10
+
+    if not name:
+        raise MissingRequirementException("Cannot list documents without assistant name")
+
+    client = RAGAssistantClient()
+    result = client.get_documents(
+        name=name,
+        skip=skip,
+        count=count
+    )
+    Console.write_stdout(f"Assistant's documents: \n{result}")
+
+
+list_documents_options = [
+    Option(
+        "name",
+        ["--name", "-n"],
+        "RAG assistant name (required)",
+        True
+    ),
+    Option(
+        "skip",
+        ["--skip", "-s"],
+        "Number of documents to skip",
+        True
+    ),
+    Option(
+        "count",
+        ["--count", "-c"],
+        "Number of documents to return (defaults to 10)",
+        True
+    ),
+]
+
+
+def delete_all_documents(option_list: list):
+    name = None
+    for option_flag, option_arg in option_list:
+        if option_flag.name == "name":
+            name = option_arg
+
+    if not name:
+        raise MissingRequirementException("Cannot delete all documents without assistant name")
+
+    client = RAGAssistantClient()
+    result = client.delete_all_documents(
+        name=name,
+    )
+    Console.write_stdout(f"Deleted documents: \n{result}")
+
+
+delete_all_documents_options = [
+    Option(
+        "name",
+        ["--name", "-n"],
+        "RAG assistant name (required)",
+        True
+    ),
+]
+
+
+def get_document_data(option_list: list):
+    name = None
+    document_id = None
+    for option_flag, option_arg in option_list:
+        if option_flag.name == "name":
+            name = option_arg
+        if option_flag.name == "document_id":
+            document_id = option_arg
+
+    if not (name and document_id):
+        raise MissingRequirementException("Cannot retrieve document data without assistant name and id")
+
+    client = RAGAssistantClient()
+    result = client.retrieve_document(
+        name=name,
+        document_id=document_id
+    )
+    Console.write_stdout(f"Document detail: \n{result}")
+
+
+get_document_data_options = [
+    Option(
+        "name",
+        ["--name", "-n"],
+        "RAG assistant name (required)",
+        True
+    ),
+    Option(
+        "document_id",
+        ["--document-id", "--id"],
+        "Document id (required)",
+        True
+    ),
+]
+
+
+def upload_document(option_list: list):
+    name = None
+    file_path = None
+    upload_type = 'multipart'
+    metadata = {}
+    content_type = None
+    for option_flag, option_arg in option_list:
+        if option_flag.name == "name":
+            name = option_arg
+        if option_flag.name == "file_path":
+            file_path = option_arg
+        if option_flag.name == "upload_type":
+            upload_type = option_arg
+        if option_flag.name == "metadata":
+            metadata_str = str(option_arg)
+            if not Path(metadata_str).is_file():
+                try:
+                    metadata = json.loads(metadata_str)
+                except Exception as e:
+                    raise WrongArgumentError(f"Metadata should be either a valid dictionary or a file path.")
+            else:
+                metadata = metadata_str
+        if option_flag.name == "content_type":
+            content_type = option_arg
+
+    if not (name and file_path):
+        raise MissingRequirementException("Cannot upload document without assistant name and file name")
+
+    client = RAGAssistantClient()
+    result = client.upload_document(
+        name=name,
+        file_path=file_path,
+        upload_type=upload_type,
+        metadata=metadata,
+        content_type=content_type
+    )
+    Console.write_stdout(f"Uploaded: \n{result}")
+
+
+upload_document_options = [
+    Option(
+        "name",
+        ["--name", "-n"],
+        "RAG assistant name (required)",
+        True
+    ),
+    Option(
+        "file_path",
+        ["--file-path", "-f"],
+        "Path to document file (required)",
+        True
+    ),
+    Option(
+        "upload_type",
+        ["--upload-type", "-t"],
+        "Upload type. Available options: binary or multipart (multipart/form-data). Defaults to multipart",
+        True
+    ),
+    Option(
+        "metadata",
+        ["--metadata", "-m"],
+        "Document metadata (only available for multipart/form-data). Can be valid JSON or a path to metadata file.",
+        True
+    ),
+    Option(
+        "content_type",
+        ["--content-type", "--ct"],
+        "Document content type",
+        True
+    ),
+
+]
+
+
+def delete_document(option_list: list):
+    name = None
+    document_id = None
+    for option_flag, option_arg in option_list:
+        if option_flag.name == "name":
+            name = option_arg
+        if option_flag.name == "document_id":
+            document_id = option_arg
+
+    if not (name and document_id):
+        raise MissingRequirementException("Cannot delete document without assistant name and id")
+
+    client = RAGAssistantClient()
+    result = client.delete_document(
+        name=name,
+        document_id=document_id
+    )
+    Console.write_stdout(f"Deleted document: \n{result}")
+
+
+delete_document_options = [
+    Option(
+        "name",
+        ["--name", "-n"],
+        "RAG assistant name (required)",
+        True
+    ),
+    Option(
+        "document_id",
+        ["--document-id", "--id"],
+        "Document id (required)",
+        True
+    ),
+]
+
+
 rag_commands = [
     Command(
         "help",
         ["help", "h"],
         "Display help text",
-        show_assistant_help,
+        show_help,
         ArgumentsEnum.NOT_AVAILABLE,
         [],
         []
@@ -940,5 +1183,58 @@ rag_commands = [
         [],
         update_assistant_options
     ),
-
+    Command(
+        "delete_assistant",
+        ["delete-assistant"],
+        "Delete existing RAG assistant",
+        delete_assistant,
+        ArgumentsEnum.REQUIRED,
+        [],
+        delete_assistant_options
+    ),
+    Command(
+        "list_documents",
+        ["list-documents"],
+        "List documents for RAG assistant",
+        list_documents,
+        ArgumentsEnum.REQUIRED,
+        [],
+        list_documents_options
+    ),
+    Command(
+        "delete_all_documents",
+        ["delete-all-documents", "del-docs"],
+        "Delete all documents for RAG assistant",
+        delete_all_documents,
+        ArgumentsEnum.REQUIRED,
+        [],
+        delete_all_documents_options
+    ),
+    Command(
+        "get_document",
+        ["get-document", "get-doc"],
+        "Get document for RAG assistant",
+        get_document_data,
+        ArgumentsEnum.REQUIRED,
+        [],
+        get_document_data_options
+    ),
+    Command(
+        "upload_document",
+        ["upload-document", "up-doc"],
+        "Upload document for RAG assistant",
+        upload_document,
+        ArgumentsEnum.REQUIRED,
+        [],
+        upload_document_options
+    ),
+    Command(
+        "delete_document",
+        ["delete-document", "del-doc"],
+        "Delete document for RAG assistant by id",
+        delete_document,
+        ArgumentsEnum.REQUIRED,
+        [],
+        delete_document_options
+    )
 ]

pygeai/cli/commands/rerank.py

@@ -0,0 +1,108 @@
+import json
+from pygeai.cli.commands import Command, Option, ArgumentsEnum
+from pygeai.cli.commands.builders import build_help_text
+from pygeai.cli.texts.help import RERANK_HELP_TEXT
+from pygeai.core.common.exceptions import MissingRequirementException, WrongArgumentError
+from pygeai.core.rerank.clients import RerankClient
+from pygeai.core.utils.console import Console
+
+
+def show_help():
+    """
+    Displays help text in stdout
+    """
+    help_text = build_help_text(rerank_commands, RERANK_HELP_TEXT)
+    Console.write_stdout(help_text)
+
+
+def rerank_chunks(option_list: list):
+    query = None
+    model = None
+    documents = []
+    top_n = 3
+    for option_flag, option_arg in option_list:
+        if option_flag.name == "query":
+            query = option_arg
+        if option_flag.name == "model":
+            model = option_arg
+        if option_flag.name == "documents":
+            if "[" not in option_arg:
+                documents.append(option_arg)
+            else:
+                try:
+                    documents_json = json.loads(option_arg)
+                    if not isinstance(documents_json, list):
+                        raise ValueError
+
+                    documents = documents_json
+                except Exception as e:
+                    raise WrongArgumentError(
+                        "Documents must be a list of strings: '[\"text_chunk_1\", \"text_chunk_2\"]'. "
+                        "Each element in the list must be a string representing a text chunk.."
+                    )
+        if option_flag.name == "top_n":
+            top_n = option_arg
+
+    if not (model and query and documents):
+        raise MissingRequirementException("Cannot rerank chunks without model, query and documents")
+
+    client = RerankClient()
+    result = client.rerank_chunks(
+        query=query,
+        model=model,
+        documents=documents,
+        top_n=top_n
+    )
+    Console.write_stdout(f"Rerank details: \n{result}")
+
+
+rerank_chunks_options = [
+    Option(
+        "query",
+        ["--query", "-q"],
+        "string: Input query",
+        True
+    ),
+    Option(
+        "model",
+        ["--model", "-m"],
+        "string: provider/modelName reranker to use; supported values: cohere/rerank-v3.5, "
+        "awsbedrock/cohere.rerank-v3.5, awsbedrock/amazon.rerank-v1",
+        True
+    ),
+    Option(
+        "documents",
+        ["--documents", "--doc", "-d"],
+        "string or array: A list of text chunks",
+        True
+    ),
+    Option(
+        "top_n",
+        ["--top-n"],
+        "string: Count of best n results to return",
+        True
+    ),
+
+]
+
+
+rerank_commands = [
+    Command(
+        "help",
+        ["help", "h"],
+        "Display help text",
+        show_help,
+        ArgumentsEnum.NOT_AVAILABLE,
+        [],
+        []
+    ),
+    Command(
+        "rerank",
+        ["rerank-chunks", "chunks", "rc"],
+        "Rerank chunks based on a query",
+        rerank_chunks,
+        ArgumentsEnum.REQUIRED,
+        [],
+        rerank_chunks_options
+    ),
+]