databao/api.py

from databao.caches.in_mem_cache import InMemCache
from databao.configs.llm import LLMConfig, LLMConfigDirectory
from databao.core import Agent, Cache, Executor, Visualizer
from databao.executors.lighthouse.executor import LighthouseExecutor
from databao.visualizers.vega_chat import VegaChatVisualizer


def new_agent(
    name: str | None = None,
    llm_config: LLMConfig | None = None,
    data_executor: Executor | None = None,
    visualizer: Visualizer | None = None,
    cache: Cache | None = None,
    rows_limit: int = 1000,
    stream_ask: bool = True,
    stream_plot: bool = False,
    lazy_threads: bool = False,
    auto_output_modality: bool = True,
) -> Agent:
    """This is an entry point for users to create a new agent.
    Agent can't be modified after it's created. Only new data sources can be added.
    """
    llm_config = llm_config if llm_config else LLMConfigDirectory.DEFAULT
    return Agent(
        llm_config,
        name=name or "default_agent",
        data_executor=data_executor or LighthouseExecutor(),
        visualizer=visualizer or VegaChatVisualizer(llm_config),
        cache=cache or InMemCache(),
        rows_limit=rows_limit,
        stream_ask=stream_ask,
        stream_plot=stream_plot,
        lazy_threads=lazy_threads,
        auto_output_modality=auto_output_modality,
    )