weaviate_utils.py

import weaviate
from weaviate.embedded import EmbeddedOptions
from weaviate import Client

def initialize_weaviate_client():
    return weaviate.Client(embedded_options=EmbeddedOptions())

def class_exists(client, class_name):
    try:
        client.schema.get_class(class_name)
        return True
    except:
        return False

def map_dtype_to_weaviate(dtype):
    if "int" in str(dtype):
        return "int"
    elif "float" in str(dtype):
        return "number"
    elif "bool" in str(dtype):
        return "boolean"
    else:
        return "string"

def create_new_class_schema(client, class_name, class_description):
    class_schema = {
        "class": class_name,
        "description": class_description,
        "properties": []
    }
    try:
        client.schema.create({"classes": [class_schema]})
        st.success(f"Class {class_name} created successfully!")
    except Exception as e:
        st.error(f"Error creating class: {e}")

def ingest_data_to_weaviate(client, csv_file, selected_class):
    # Convert CSV to DataFrame
    data = csv_file.read().decode("utf-8")
    dataframe = pd.read_csv(StringIO(data))

    # Check if columns match the selected class schema
    class_schema = get_class_schema(client, selected_class)
    if class_schema:
        schema_columns = [prop["name"] for prop in class_schema["properties"]]
        if set(dataframe.columns) == set(schema_columns):
            data = dataframe.to_dict(orient="records")
            client.data_object.create(data, selected_class)
            st.success("Data ingested successfully!")
        else:
            st.error("The columns in the uploaded CSV do not match the schema of the selected class.")

def get_class_schema(client, class_name):
    try:
        return client.schema.get_class(class_name)
    except weaviate.exceptions.SchemaValidationException:
        return None