A custom Langchain Tool for converting a table to a Weaviate Schema

TableToWeaviateClassTool.py

from langchain.agents import AgentType, initialize_agent
from langchain.chat_models import ChatOpenAI
from langchain.tools import BaseTool
from typing import Optional, Type
from langchain.callbacks.manager import (
    AsyncCallbackManagerForToolRun,
    CallbackManagerForToolRun,
)

class TableToWeaviateClassTool(BaseTool):
    name = "TableToWeaviateClass"
    description = "Converts a table into a Weaviate class schema."

    def table_to_weaviate_class(self, query: str) -> str:
        # Assuming the table is passed as a string representation of a list of dictionaries
        table = eval(query)

        class_schema = {
            "class": "YourClassName",
            "properties": []
        }

        for column in table:
            class_schema["properties"].append({
                "name": column["name"],
                "dataType": [column["type"]]
            })

        return str(class_schema)

    def _run(self, query: str, run_manager: Optional[CallbackManagerForToolRun] = None) -> str:
        return self.table_to_weaviate_class(query)

    async def _arun(self, query: str, run_manager: Optional[AsyncCallbackManagerForToolRun] = None) -> str:
        raise NotImplementedError("This tool does not support async yet.")

tools = [TableToWeaviateClassTool()]

agent = initialize_agent(
    tools,
    ChatOpenAI(temperature=0),
    agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
    verbose=True
)

# Sample run for the agent
print(agent.run('[{"name": "column1", "type": "string"}, {"name": "column2", "type": "int"}]'))