main.py

"""
Main Entry Point for KB manager.

This script provides a command-line interface (CLI) using Typer for managing knowledge bases, running benchmarks, and generating reports with MindsDB.
"""

import typer
from typing import List, Optional
from core.file_ops import FileOperations
from core.kb_ops import KnowledgeBaseOperations, KnowledgeBaseConfig, OpenAIEmbeddingModel, OllamaEmbeddingModel, ReRankingModel
from core.agent_ops import AgentOperations, ModelConfig
import pandas as pd

app = typer.Typer(help="CLI tool for managing MindsDB knowledge bases and benchmarking.")

def _get_kb_ops() -> KnowledgeBaseOperations:
    """Helper to get KnowledgeBaseOperations instance."""
    return KnowledgeBaseOperations()

def _get_file_ops() -> FileOperations:
    """Helper to get FileOperations instance."""
    return FileOperations()

def _get_agent_ops() -> AgentOperations:
    """Helper to get AgentOperations instance."""
    return AgentOperations()

@app.command()
def init_kb(
    name: str = typer.Option(..., help="Name of the knowledge base"),
    model_name: str = typer.Option("text-embedding-ada-002", help="Embedding model name"),
    api_key: Optional[str] = typer.Option(None, help="API key for the embedding model provider (required for OpenAI)"),
    embedding_provider: str = typer.Option("openai", help="Provider for the embedding model (openai or ollama)"),
    base_url: Optional[str] = typer.Option(None, help="Base URL for Ollama embedding model, if applicable"),
    reranker_model_name: Optional[str] = typer.Option(None, help="Reranker model name"),
    reranker_api_key: Optional[str] = typer.Option(None, help="API key for the reranker model provider"),
    reranker_provider: Optional[str] = typer.Option("openai", help="Provider for the reranker model (openai or gemini)"),
    reranker_base_url: Optional[str] = typer.Option(None, help="Base URL for reranker model, if applicable"),
    id_column: str = typer.Option("id", help="Column name for document IDs"),
    metadata_columns: List[str] = typer.Option([], help="List of metadata column names"),
    content_columns: List[str] = typer.Option(["content"], help="List of content column names"),
    project_name: str = typer.Option("mindsdb", help="Project name in MindsDB")
):
    """Initialize a new knowledge base in MindsDB."""
    kb_ops = _get_kb_ops()
    
    if embedding_provider == "openai":
        if not api_key:
            typer.echo("API key is required for OpenAI provider.")
            raise typer.Exit(code=1)
        embedding_model = OpenAIEmbeddingModel(provider="openai", model_name=model_name, api_key=api_key, base_url=base_url)
    elif embedding_provider == "ollama":
        if not base_url:
            typer.echo("Base URL is required for Ollama provider.")
            raise typer.Exit(code=1)
        embedding_model = OllamaEmbeddingModel(provider="ollama", model_name=model_name, base_url=base_url)
    else:
        typer.echo(f"Unsupported embedding provider: {embedding_provider}. Use 'openai' or 'ollama'.")
        raise typer.Exit(code=1)
    
    reranking_model = None
    if reranker_model_name and reranker_api_key:
        if reranker_provider == "openai":
            reranking_model = ReRankingModel(provider="openai", model_name=reranker_model_name, api_key=reranker_api_key, base_url=reranker_base_url)
        elif reranker_provider == "gemini":
            reranking_model = ReRankingModel(provider="gemini", model_name=reranker_model_name, api_key=reranker_api_key, base_url=reranker_base_url)
        else:
            typer.echo(f"Unsupported reranker provider: {reranker_provider}. Use 'openai' or 'gemini'.")
            raise typer.Exit(code=1)
    
    config = KnowledgeBaseConfig(
        name=name,
        embedding_model=embedding_model,
        reranking_model=reranking_model,
        metadata_columns=metadata_columns,
        content_columns=content_columns,
        id_column=id_column
    )
    
    typer.echo(f"Creating knowledge base {name}...")
    kb_ops.create_knowledge_base(config, project_name=project_name)
    typer.echo(f"Knowledge base {name} created successfully.")

@app.command()
def ingest_kb(
    kb_name: str = typer.Option(..., help="Name of the knowledge base"),
    file_path: str = typer.Option(..., help="Path to the CSV file to upload and ingest"),
    project_name: str = typer.Option("mindsdb", help="Project name in MindsDB")
):
    """Ingest data into a knowledge base from a CSV file."""
    kb_ops = _get_kb_ops()
    file_ops = _get_file_ops()
    
    # Derive table name from file path
    table_name = file_path.split('/')[-1].replace('.csv', '')
    
    typer.echo(f"Uploading file {file_path} to MindsDB as {table_name}...")
    df = pd.read_csv(file_path)
    df.insert(0, "id", df.index)
    file_ops.upload_file(table_name, df, replace=True)
    typer.echo(f"File uploaded as {table_name}.")
    
    typer.echo(f"Ingesting data from {table_name} into knowledge base {kb_name}...")
    kb_ops.ingest_data(kb_name, f"files.{table_name}", project_name=project_name)
    typer.echo(f"Data ingested into {kb_name} successfully.")
    
    typer.echo(f"Cleaning up: Deleting file {table_name} from MindsDB...")
    file_ops.delete_file(table_name)
    typer.echo(f"File {table_name} deleted from MindsDB.")

@app.command()
def query_kb(
    kb_name: str = typer.Option(..., help="Name of the knowledge base"),
    query: str = typer.Option(..., help="Query text to search in the knowledge base"),
    relevance: float = typer.Option(0.2, help="Minimum relevance score for results"),
    metadata_filter: Optional[str] = typer.Option(None, help="Optional metadata filter condition, e.g., 'category = \"tech\"'"),
    project_name: str = typer.Option("mindsdb", help="Project name in MindsDB")
):
    """Query a knowledge base with a specific text and optional metadata filter."""
    kb_ops = _get_kb_ops()
    
    typer.echo(f"Querying knowledge base {kb_name} with: {query}" + (f" and filter: {metadata_filter}" if metadata_filter else ""))
    conditions = f"AND {metadata_filter}" if metadata_filter else None
    results = kb_ops.query_knowledge_base(kb_name, query, relevance=relevance, conditions=conditions, project_name=project_name)
    if not results.empty:
        typer.echo("Results:")
        for _, row in results.iterrows():
            typer.echo(f"- {row.get('chunk_content', 'No content available')}")
    else:
        typer.echo("No results found.")

@app.command()
def delete_kb(
    kb_name: str = typer.Option(..., help="Name of the knowledge base to delete"),
    project_name: str = typer.Option("mindsdb", help="Project name in MindsDB")
):
    """Delete a knowledge base from MindsDB."""
    kb_ops = _get_kb_ops()
    
    typer.echo(f"Deleting knowledge base {kb_name}...")
    kb_ops.delete_knowledge_base(kb_name, project_name=project_name)
    typer.echo(f"Knowledge base {kb_name} deleted successfully.")

@app.command()
def summarize_kb(
    kb_name: str = typer.Option(..., help="Name of the knowledge base to summarize"),
    query: str = typer.Option(..., help="Query text to summarize the knowledge base contents"),
    project_name: str = typer.Option("mindsdb", help="Project name in MindsDB")
):
    """Summarize the contents of a knowledge base."""
    kb_ops = _get_kb_ops()
    
    typer.echo(f"Summarizing knowledge base {kb_name}...")
    
    summary = kb_ops.summarize_knowledge_base(kb_name, content=query, project_name=project_name)

    typer.echo("Summary:")
    typer.echo(summary)


@app.command()
def create_agent(
    agent_name: str = typer.Option(..., help="Name of the agent to create"),
    model_name: str = typer.Option(..., help="Model name for the agent"),
    provider: str = typer.Option("openai", help="Provider for the model (openai, anthropic, google)"),
    api_key: str = typer.Option(..., help="API key for the model provider"),
    system_prompt: str = typer.Option("You are a helpful AI assistant.", help="System prompt for the agent"),
    knowledge_bases: List[str] = typer.Option(..., help="List of knowledge base names to associate with the agent"),
    project_name: str = typer.Option("mindsdb", help="Project name in MindsDB")
):
    """Create an AI agent with specified model and associated knowledge bases."""
    agent_ops = _get_agent_ops()
    
    if provider not in ["openai", "anthropic", "google"]:
        typer.echo(f"Unsupported provider: {provider}. Use 'openai', 'anthropic', or 'google'.")
        raise typer.Exit(code=1)
    
    # Cast provider to satisfy static type checking
    model_config = ModelConfig(
        provider=provider,  # type: ignore
        model_name=model_name,
        api_key=api_key,
        system_prompt=system_prompt
    )
    
    typer.echo(f"Creating agent {agent_name} with model {model_name}...")
    agent_ops.create_agent(agent_name, model_config, knowledge_bases=knowledge_bases, project_name=project_name)
    typer.echo(f"Agent {agent_name} created successfully.")

@app.command()
def chat_agent(
    agent_name: str = typer.Option(..., help="Name of the agent to chat with"),
    question: str = typer.Option(..., help="Question to ask the agent"),
    project_name: str = typer.Option("mindsdb", help="Project name in MindsDB")
):
    """Chat with an AI agent by asking a question."""
    agent_ops = _get_agent_ops()
    
    typer.echo(f"Asking agent {agent_name}: {question}")
    answer = agent_ops.query_agent(agent_name, question, project_name=project_name)
    typer.echo(f"Answer: {answer}")

@app.command()
def delete_agent(
    agent_name: str = typer.Option(..., help="Name of the agent to delete"),
    project_name: str = typer.Option("mindsdb", help="Project name in MindsDB")
):
    """Delete an AI agent from MindsDB."""
    agent_ops = _get_agent_ops()
    
    typer.echo(f"Deleting agent {agent_name}...")
    agent_ops.delete_agent(agent_name, project_name=project_name)
    typer.echo(f"Agent {agent_name} deleted successfully.")

if __name__ == "__main__":
    app()