Building an AI Agent to interact with Databases

def _generate_documentation(
    table_name: str, columns_data: List[dict], config: ChatCompletionConfig
) -> _TableDocumentation:
    try:
        # quick fix: for clickhouse do not have this type of infrommation, as it does not enforec foregins relationships
        new_cols = []
        for col in columns_data:
            _col = {k: v for k, v in col.items() if k != "foreign_key_table_docs"}
            if col.get("foreign_key_reference"):
                _col["foreign_key_info"] = col.get("foreign_key_reference")
            new_cols.append(_col)
    except:
        new_cols = columns_data

    prompt = f"""
        I need comprehensive documentation for a database table named '{table_name}' with the following columns:
        
        {json.dumps(new_cols, indent=2)}
        
        For each table, provide:
        1. A business summary explaining the purpose and role of the table
        2. Possible usage scenarios for the table
        
        For each column, provide:
        3. A detailed summary of what the column represents
        4. A business summary explaining the significance of this data
        5. Possible usage scenarios for this column
        6. A list of relevant tags (e.g., identifier, text, timestamp, foreign key, etc.)
    """

    messages = [
        {
            "role": "system",
            "content": "You are a data documentation expert.",
        },
        {"role": "user", "content": prompt},
    ]
    try:
        response = chat_completion_from_config(
            messages, config, response_format=_TableDocumentation
        )
        return _TableDocumentation(**json.loads(response.choices[0].message.content))
    except Exception as e:
        raise e

class ContextRetrieverTool(Tool):
    name = "context_retiver"
    description = """this tool provides you similarity search on provided natural langague (question) against vector database to obtain most similar related tables and column names and their bussiness and usage documentation. and then joins the retireved 
    similar table names and columns and query the database to obatin its schema, as in datatypes, constraints. etc."""
    inputs = {
        "query": {
            "type": "string",
            "description": "a natural question to fetch most similar tables and columns with",
        }
    }
    output_type = "string"

    def forward(self, query: str):
        return json.dumps(
            get_context(
                query=query, index=index, internal_db=internal_db, source=source_db
            )
        )

class DatabaseConfig(BaseModel):
    engine: str = Field(description="The engine of the database")

class PostgresDatabaseConfig(DatabaseConfig):
    user: str = Field(description="The user of the database")
    password: str = Field(description="The password of the database")
    host: str = Field(description="The host of the database")
    port: int = Field(description="The port of the database")
    dbname: str = Field(description="The name of the database")
    schema_query_path: str = Field(description="The path to the schema query file")

class DuckDBDatabaseConfig(DatabaseConfig):
    path: str = Field(description="The path to the DuckDB database file")
    ddl_query_path: str = Field(description="The path to the ddl query file")

class ClickhouseDatabaseConfig(DatabaseConfig):
    ...

def generate_sql_from_context(query: str, context: List[dict], database : str, config: ChatCompletionConfig):
        # quick fix for numpy array, should be fixed in the source
        ...
        prompt = f"""
        You are a SQL expert tasked with generating a SQL query based on a natural language question.

        CONTEXT:
        {context}

        DATABASE: {database}

        INSTRUCTIONS:
        1. Generate a SQL query that answers the following question: "{query}"
        2. Use only tables and columns that are mentioned in the CONTEXT section.
        3. Make sure your query is syntactically correct and follows best practices for {database}.
        4. Include appropriate JOINs when querying across multiple tables.
        5. Add comments to explain your reasoning for complex parts of the query.
        6. If the question cannot be answered with the available tables/columns, explain why.

        Your response should be structured as follows:

        SQL:
        [Your SQL query here]

        EXPLANATION:
        [Your explanation of how the query works and any assumptions made]

        RELEVANT SOURCES FROM CONTEXT:
        [Your derived infromation from the context provided]
        """
        ...

def create_agent(
    task: str = "query",
    engine=ENGINE,
    planning_interval: int = 5,
    additional_authorized_imports: list = DEFULAT_AUTH_IMPORTS,
    tools: list = DEFAULT_TOOLS
):
    if task == "query":
        custom_prompt_templates = yaml.safe_load(
            importlib.resources.files("core.prompts").joinpath("query_agent.yaml").read_text()
        )
    if task == "finder":
        custom_prompt_templates = yaml.safe_load(
            importlib.resources.files("core.prompts").joinpath("finder_agent.yaml").read_text()
        )

    agent = CodeAgent(
        model=engine,
        tools=tools,
        additional_authorized_imports=additional_authorized_imports,
        planning_interval=planning_interval,
        prompt_templates=custom_prompt_templates,
    )
    return agent