Knowledge base API

# This sample code is for reference only. Do not use it directly in a production environment.
import os

from alibabacloud_bailian20231229 import models as bailian_20231229_models
from alibabacloud_bailian20231229.client import Client as bailian20231229Client
from alibabacloud_tea_openapi import models as open_api_models
from alibabacloud_tea_util import models as util_models
from alibabacloud_tea_util.client import Client as UtilClient

def check_environment_variables():
    """Checks for required environment variables and reports any that are missing."""
    required_vars = {
        'ALIBABA_CLOUD_ACCESS_KEY_ID': 'Your Alibaba Cloud AccessKey ID',
        'ALIBABA_CLOUD_ACCESS_KEY_SECRET': 'Your Alibaba Cloud AccessKey secret',
        'WORKSPACE_ID': 'Your Model Studio workspace ID'
    }
    missing_vars = []
    for var, description in required_vars.items():
        if not os.environ.get(var):
            missing_vars.append(var)
            print(f"Error: Set the {var} environment variable ({description}).")

    return len(missing_vars) == 0

def create_client() -> bailian20231229Client:
    """
    Creates and configures a client.

    Returns:
        bailian20231229Client: The configured client.
    """
    config = open_api_models.Config(
        access_key_id=os.environ.get('ALIBABA_CLOUD_ACCESS_KEY_ID'),
        access_key_secret=os.environ.get('ALIBABA_CLOUD_ACCESS_KEY_SECRET')
    )
        # This is a public cloud endpoint. You can change it as needed.
    config.endpoint = 'bailian.ap-southeast-1.aliyuncs.com'
    return bailian20231229Client(config)

def retrieve_index(client, workspace_id, index_id, query):
    """
    Retrieves information from a specified knowledge base.

    Args:
        client (bailian20231229Client): The client.
        workspace_id (str): The workspace ID.
        index_id (str): The knowledge base ID.
        query (str): The search query.

    Returns:
        The response from the Model Studio service.
    """
    headers = {}
    retrieve_request = bailian_20231229_models.RetrieveRequest(
        index_id=index_id,
        query=query
    )
    runtime = util_models.RuntimeOptions()
    return client.retrieve_with_options(workspace_id, retrieve_request, headers, runtime)

def main():
    """
    Uses the Model Studio service to retrieve from a knowledge base.

    Returns:
        str or None: The retrieved text segments if the call is successful. Otherwise, None is returned.
    """
    if not check_environment_variables():
        print("Environment variable validation failed.")
        return
    try:
        print("Step 1: Create a client.")
        client = create_client()
        print("Step 2: Retrieve from the knowledge base.")
        index_id = input("Enter the knowledge base ID: ")  # This is the Data.Id value returned by the CreateIndex operation. You can also obtain it from the knowledge base page in the Model Studio console.
        query = input("Enter the search query: ")
        workspace_id = os.environ.get('WORKSPACE_ID')
        resp = retrieve_index(client, workspace_id, index_id, query)
        result = UtilClient.to_jsonstring(resp.body)
        print(result)
    except Exception as e:
        print(f"An error occurred: {e}")
        return None

if __name__ == '__main__':
    main()

# This sample code is for reference only. Do not use it in a production environment.
import hashlib
import os
import time

import requests
from alibabacloud_bailian20231229 import models as bailian_20231229_models
from alibabacloud_bailian20231229.client import Client as bailian20231229Client
from alibabacloud_tea_openapi import models as open_api_models
from alibabacloud_tea_util import models as util_models

def check_environment_variables():
    """Checks for required environment variables and prompts for any that are missing."""
    required_vars = {
        'ALIBABA_CLOUD_ACCESS_KEY_ID': 'Alibaba Cloud AccessKey ID',
        'ALIBABA_CLOUD_ACCESS_KEY_SECRET': 'Alibaba Cloud AccessKey secret',
        'WORKSPACE_ID': 'Model Studio workspace ID'
    }
    missing_vars = []
    for var, description in required_vars.items():
        if not os.environ.get(var):
            missing_vars.append(var)
            print(f"Error: Please set the {var} environment variable ({description}).")

    return len(missing_vars) == 0

# Create a client.
def create_client() -> bailian20231229Client:
    """
    Creates and configures a client.

    Returns:
        bailian20231229Client: The configured client.
    """
    config = open_api_models.Config(
        access_key_id=os.environ.get('ALIBABA_CLOUD_ACCESS_KEY_ID'),
        access_key_secret=os.environ.get('ALIBABA_CLOUD_ACCESS_KEY_SECRET')
    )
    # This is a public endpoint. You can change it as needed.
    config.endpoint = 'bailian.ap-southeast-1.aliyuncs.com'
    return bailian20231229Client(config)

def calculate_md5(file_path: str) -> str:
    """
    Calculates the MD5 hash of a file.

    Args:
        file_path (str): The local path of the file.

    Returns:
        str: The MD5 hash of the file.
    """
    md5_hash = hashlib.md5()

    # Read the file in binary mode.
    with open(file_path, "rb") as f:
        # Read the file in chunks to avoid high memory usage for large files.
        for chunk in iter(lambda: f.read(4096), b""):
            md5_hash.update(chunk)

    return md5_hash.hexdigest()

def get_file_size(file_path: str) -> int:
    """
    Gets the size of a file in bytes.
    Args:
        file_path (str): The local path of the file.
    Returns:
        int: The file size in bytes.
    """
    return os.path.getsize(file_path)

# Apply for a file upload lease.
def apply_lease(client, category_id, file_name, file_md5, file_size, workspace_id):
    """
    Applies for a file upload lease from the Model Studio service.

    Args:
        client (bailian20231229Client): The client.
        category_id (str): The category ID.
        file_name (str): The name of the file.
        file_md5 (str): The MD5 hash of the file.
        file_size (int): The file size in bytes.
        workspace_id (str): The workspace ID.

    Returns:
        The response from the Model Studio service.
    """
    headers = {}
    request = bailian_20231229_models.ApplyFileUploadLeaseRequest(
        file_name=file_name,
        md_5=file_md5,
        size_in_bytes=file_size,
    )
    runtime = util_models.RuntimeOptions()
    return client.apply_file_upload_lease_with_options(category_id, workspace_id, request, headers, runtime)

# Upload the file to temporary storage.
def upload_file(pre_signed_url, headers, file_path):
    """
    Uploads a file to the Model Studio service.
    Args:
        pre_signed_url (str): The URL from the upload lease.
        headers (dict): The request headers for the upload.
        file_path (str): The local path of the file.
    """
    with open(file_path, 'rb') as f:
        file_content = f.read()
    upload_headers = {
        "X-bailian-extra": headers["X-bailian-extra"],
        "Content-Type": headers["Content-Type"]
    }
    response = requests.put(pre_signed_url, data=file_content, headers=upload_headers)
    response.raise_for_status()

# Add the file to a category.
def add_file(client: bailian20231229Client, lease_id: str, parser: str, category_id: str, workspace_id: str):
    """
    Adds a file to a specified category in the Model Studio service.

    Args:
        client (bailian20231229Client): The client.
        lease_id (str): The lease ID.
        parser (str): The parser to use for the file.
        category_id (str): The category ID.
        workspace_id (str): The workspace ID.

    Returns:
        The response from the Model Studio service.
    """
    headers = {}
    request = bailian_20231229_models.AddFileRequest(
        lease_id=lease_id,
        parser=parser,
        category_id=category_id,
    )
    runtime = util_models.RuntimeOptions()
    return client.add_file_with_options(workspace_id, request, headers, runtime)

# Query the parsing status of the file.
def describe_file(client, workspace_id, file_id):
    """
    Queries the parsing status of a file.

    Args:
        client (bailian20231229Client): The client.
        workspace_id (str): The workspace ID.
        file_id (str): The file ID.

    Returns:
        The response from the Model Studio service.
    """
    headers = {}
    runtime = util_models.RuntimeOptions()
    return client.describe_file_with_options(workspace_id, file_id, headers, runtime)

# Submit a job to append documents.
def submit_index_add_documents_job(client, workspace_id, index_id, file_id, source_type):
    """
    Appends a parsed file to a knowledge base.

    Args:
        client (bailian20231229Client): The client.
        workspace_id (str): The workspace ID.
        index_id (str): The knowledge base ID.
        file_id (str): The file ID.
        source_type(str): The data type.

    Returns:
        The response from the Model Studio service.
    """
    headers = {}
    submit_index_add_documents_job_request = bailian_20231229_models.SubmitIndexAddDocumentsJobRequest(
        index_id=index_id,
        document_ids=[file_id],
        source_type=source_type
    )
    runtime = util_models.RuntimeOptions()
    return client.submit_index_add_documents_job_with_options(workspace_id, submit_index_add_documents_job_request,
                                                              headers, runtime)

# Wait for the document appending job to complete.
def get_index_job_status(client, workspace_id, job_id, index_id):
    """
    Queries the status of an index job.

    Args:
        client (bailian20231229Client): The client.
        workspace_id (str): The workspace ID.
        index_id (str): The knowledge base ID.
        job_id (str): The job ID.

    Returns:
        The response from the Model Studio service.
    """
    headers = {}
    get_index_job_status_request = bailian_20231229_models.GetIndexJobStatusRequest(
        index_id=index_id,
        job_id=job_id
    )
    runtime = util_models.RuntimeOptions()
    return client.get_index_job_status_with_options(workspace_id, get_index_job_status_request, headers, runtime)

# Delete the old file.
def delete_index_document(client, workspace_id, index_id, file_id):
    """
    Permanently deletes one or more files from a specified knowledge base.

    Args:
        client (bailian20231229Client): The client.
        workspace_id (str): The workspace ID.
        index_id (str): The knowledge base ID.
        file_id (str): The file ID.

    Returns:
        The response from the Model Studio service.
    """
    headers = {}
    delete_index_document_request = bailian_20231229_models.DeleteIndexDocumentRequest(
        index_id=index_id,
        document_ids=[file_id]
    )
    runtime = util_models.RuntimeOptions()
    return client.delete_index_document_with_options(workspace_id, delete_index_document_request, headers, runtime)

def update_knowledge_base(
        file_path: str,
        workspace_id: str,
        index_id: str,
        old_file_id: str
):
    """
    Updates a knowledge base in Model Studio.
    Args:
        file_path (str): The local path of the updated file.
        workspace_id (str): The workspace ID.
        index_id (str): The ID of the knowledge base to update.
        old_file_id (str): The FileId of the file to be replaced.
    Returns:
        str or None: The knowledge base ID if successful, otherwise None.
    """
    # Set default values.
    category_id = 'default'
    parser = 'DASHSCOPE_DOCMIND'
    source_type = 'DATA_CENTER_FILE'
    try:
        # Step 1: Create a client.
        print("Step 1: Create a client.")
        client = create_client()
        # Step 2: Prepare file information.
        print("Step 2: Prepare file information.")
        file_name = os.path.basename(file_path)
        file_md5 = calculate_md5(file_path)
        file_size = get_file_size(file_path)
        # Step 3: Apply for an upload lease.
        print("Step 3: Apply for an upload lease from Model Studio.")
        lease_response = apply_lease(client, category_id, file_name, file_md5, file_size, workspace_id)
        lease_id = lease_response.body.data.file_upload_lease_id
        upload_url = lease_response.body.data.param.url
        upload_headers = lease_response.body.data.param.headers
        # Step 4: Upload the file to temporary storage.
        print("Step 4: Upload the file to temporary storage.")
        upload_file(upload_url, upload_headers, file_path)
        # Step 5: Add the file to the category.
        print("Step 5: Add the file to the category.")
        add_response = add_file(client, lease_id, parser, category_id, workspace_id)
        file_id = add_response.body.data.file_id
        # Step 6: Check the file status.
        print("Step 6: Check the file status in Model Studio.")
        while True:
            describe_response = describe_file(client, workspace_id, file_id)
            status = describe_response.body.data.status
            print(f"Current file status: {status}")
            if status == 'INIT':
                print("File is waiting to be parsed. Please wait...")
            elif status == 'PARSING':
                print("File is being parsed. Please wait...")
            elif status == 'PARSE_SUCCESS':
                print("File parsing completed successfully.")
                break
            else:
                print(f"Unknown file status: {status}. Please contact technical support.")
                return None
            time.sleep(5)
        # Step 7: Submit a job to append documents.
        print("Step 7: Submit a job to append documents.")
        index_add_response = submit_index_add_documents_job(client, workspace_id, index_id, file_id, source_type)
        job_id = index_add_response.body.data.id
        # Step 8: Wait for the indexing job to complete.
        print("Step 8: Wait for the indexing job to complete.")
        while True:
            get_index_job_status_response = get_index_job_status(client, workspace_id, job_id, index_id)
            status = get_index_job_status_response.body.data.status
            print(f"Current index job status: {status}")
            if status == 'COMPLETED':
                break
            time.sleep(5)
        # Step 9: Delete the old file.
        print("Step 9: Delete the old file.")
        delete_index_document(client, workspace_id, index_id, old_file_id)
        print("Model Studio knowledge base updated successfully.")
        return index_id
    except Exception as e:
        print(f"An error occurred: {e}")
        return None

def main():
    if not check_environment_variables():
        print("Environment variable check failed.")
        return
    file_path = input("Enter the local path of the updated file (e.g., on a Linux system: /path/to/Model_Studio_Phone_Series_Introduction.docx):")
    index_id = input("Enter the ID of the knowledge base to update: ")  # This is the Data.Id returned by the CreateIndex operation. You can also obtain it from the knowledge base page in the Model Studio console.
    old_file_id = input("Enter the FileId of the file to be replaced: ")  # This is the FileId returned by the AddFile operation. You can also obtain it by clicking the ID icon next to the file name on the application data page in the Model Studio console.
    workspace_id = os.environ.get('WORKSPACE_ID')
    update_knowledge_base(file_path, workspace_id, index_id, old_file_id)

if __name__ == '__main__':
    main()

# This sample code is for reference only. Do not use it in a production environment.
import os

from alibabacloud_bailian20231229 import models as bailian_20231229_models
from alibabacloud_bailian20231229.client import Client as bailian20231229Client
from alibabacloud_tea_openapi import models as open_api_models
from alibabacloud_tea_util import models as util_models
from alibabacloud_tea_util.client import Client as UtilClient

def check_environment_variables():
    """Verifies that the required environment variables are set."""
    required_vars = {
        'ALIBABA_CLOUD_ACCESS_KEY_ID': 'Alibaba Cloud AccessKey ID',
        'ALIBABA_CLOUD_ACCESS_KEY_SECRET': 'Alibaba Cloud AccessKey secret',
        'WORKSPACE_ID': 'Alibaba Cloud Model Studio workspace ID'
    }
    missing_vars = []
    for var, description in required_vars.items():
        if not os.environ.get(var):
            missing_vars.append(var)
            print(f"Error: Please set the {var} environment variable ({description})")

    return len(missing_vars) == 0

# Create a client.
def create_client() -> bailian20231229Client:
    """
    Creates and configures a client.

    Returns:
        bailian20231229Client: The configured client.
    """
    config = open_api_models.Config(
        access_key_id=os.environ.get('ALIBABA_CLOUD_ACCESS_KEY_ID'),
        access_key_secret=os.environ.get('ALIBABA_CLOUD_ACCESS_KEY_SECRET')
    )
    # The following endpoint is for the public cloud in the Singapore region. Change it as needed.
    config.endpoint = 'bailian.ap-southeast-1.aliyuncs.com'
    return bailian20231229Client(config)

# List knowledge bases.
def list_indices(client, workspace_id):
    """
    Lists the knowledge bases in a specified workspace.

    Args:
        client (bailian20231229Client): The client.
        workspace_id (str): The workspace ID.

    Returns:
        The response from the Alibaba Cloud Model Studio service.
    """
    headers = {}
    list_indices_request = bailian_20231229_models.ListIndicesRequest()
    runtime = util_models.RuntimeOptions()
    return client.list_indices_with_options(workspace_id, list_indices_request, headers, runtime)

# Delete a knowledge base.
def delete_index(client, workspace_id, index_id):
    """
    Permanently deletes the specified knowledge base.

    Args:
        client (bailian20231229Client): The client.
        workspace_id (str): The workspace ID.
        index_id (str): The knowledge base ID.

    Returns:
        The response from the Alibaba Cloud Model Studio service.
    """
    headers = {}
    delete_index_request = bailian_20231229_models.DeleteIndexRequest(
        index_id=index_id
    )
    runtime = util_models.RuntimeOptions()
    return client.delete_index_with_options(workspace_id, delete_index_request, headers, runtime)

def main():
    if not check_environment_variables():
        print("Environment variable check failed.")
        return
    try:
        start_option = input(
            "Select an operation to perform:\n1. List knowledge bases\n2. Delete a knowledge base\nEnter an option (1 or 2):")
        if start_option == '1':
            # List knowledge bases.
            print("\nListing knowledge bases...")
            workspace_id = os.environ.get('WORKSPACE_ID')
            client = create_client()
            list_indices_response = list_indices(client, workspace_id)
            print(UtilClient.to_jsonstring(list_indices_response.body.data))
        elif start_option == '2':
            print("\nDeleting a knowledge base...")
            workspace_id = os.environ.get('WORKSPACE_ID')
            index_id = input("Enter the knowledge base ID: ")  # This is the `Data.Id` value from the `CreateIndex` response. You can also find it on the knowledge base page in the Model Studio console.
            # Confirm before deletion.
            while True:
                confirm = input(f"Are you sure you want to permanently delete the knowledge base {index_id}? (y/n): ").strip().lower()
                if confirm == 'y':
                    break
                elif confirm == 'n':
                    print("Deletion canceled.")
                    return
                else:
                    print("Invalid input. Please enter y or n.")
            client = create_client()
            resp = delete_index(client, workspace_id, index_id)
            if resp.body.status == 200:
                print(f"Knowledge base {index_id} deleted successfully.")
            else:
                err_info = UtilClient.to_jsonstring(resp.body)
                print(f"An error occurred: {err_info}")
        else:
            print("Invalid option. Exiting program.")
            return
    except Exception as e:
        print(f"An error occurred: {e}")
        return

if __name__ == '__main__':
    main()