vanna_setup.py

import os
from dotenv import load_dotenv
from vanna import Agent
from vanna.core.registry import ToolRegistry
from vanna.core.user import UserResolver, User
from vanna.tools import RunSqlTool, VisualizeDataTool
from vanna.tools.agent_memory import SaveQuestionToolArgsTool, SearchSavedCorrectToolUsesTool
from vanna.integrations.sqlite import SqliteRunner
from vanna.integrations.local.agent_memory import DemoAgentMemory
from vanna.integrations.google import GeminiLlmService

load_dotenv()

class SimpleUserResolver(UserResolver):
    def resolve_user(self, **kwargs) -> User:
        return User(user_id="default_user")

def get_agent():
    
    llm = GeminiLlmService(os.getenv("GOOGLE_API_KEY"), "gemini-2.0-flash")

    runner = SqliteRunner("clinic.db")

    memory = DemoAgentMemory()

    tools = ToolRegistry([
        RunSqlTool(runner), 
        VisualizeDataTool(),
        SaveQuestionToolArgsTool(),
        SearchSavedCorrectToolUsesTool()
    ])

    return Agent(llm, tools, memory, SimpleUserResolver())