diff --git a/.env.docker b/.env.docker index 89c409ea..a65b5423 100644 --- a/.env.docker +++ b/.env.docker @@ -27,7 +27,7 @@ REDIS_PORT="6379" ELASTICSEARCH_URL=http://index:9200 # Chatbot -GOOGLE_APPLICATION_CREDENTIALS= +CHATBOT_CREDENTIALS= BILLING_PROJECT_ID= QUERY_PROJECT_ID= MODEL_URI= diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 8981cb01..8241de2c 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -36,4 +36,5 @@ repos: name: yamlfix types: [yaml] language: system - entry: yamlfix --exclude "charts/**/*" . + entry: yamlfix + exclude: ^charts/ diff --git a/backend/apps/chatbot/agent/tools.py b/backend/apps/chatbot/agent/tools.py index d4e6dbcc..4cc64e69 100644 --- a/backend/apps/chatbot/agent/tools.py +++ b/backend/apps/chatbot/agent/tools.py @@ -1,8 +1,7 @@ # -*- coding: utf-8 -*- import json -import os from collections.abc import Callable -from functools import cache, wraps +from functools import wraps from typing import Any, Literal, Self import httpx @@ -11,6 +10,8 @@ from langchain_core.tools import BaseTool, tool from pydantic import BaseModel, model_validator +from backend.apps.chatbot.utils.gcloud import get_bigquery_client + # HTTPX Default Timeout TIMEOUT = 5.0 @@ -234,19 +235,6 @@ def wrapper(*args, **kwargs) -> Any: return decorator(_func) -@cache -def get_bigquery_client() -> bq.Client: - """Return a cached BigQuery client. - - The client is initialized once using the project ID from the - `QUERY_PROJECT_ID` environment variable and reused on subsequent calls. - - Returns: - bigquery.Client: A cached, authenticated BigQuery client. - """ - return bq.Client(project=os.environ["QUERY_PROJECT_ID"]) - - @tool @handle_tool_errors def search_datasets(query: str) -> str: diff --git a/backend/apps/chatbot/apps.py b/backend/apps/chatbot/apps.py index ce6b6312..8ce9bcff 100644 --- a/backend/apps/chatbot/apps.py +++ b/backend/apps/chatbot/apps.py @@ -9,7 +9,8 @@ class ChatbotConfig(AppConfig): verbose_name = "Chatbot" def ready(self): - # Enable logs from the chatbot package + import backend.apps.chatbot.checks # noqa: F401 import chatbot + # Enable logs from the chatbot package logger.enable(chatbot.__name__) diff --git a/backend/apps/chatbot/checks.py b/backend/apps/chatbot/checks.py new file mode 100644 index 00000000..91b28bed --- /dev/null +++ b/backend/apps/chatbot/checks.py @@ -0,0 +1,66 @@ +# -*- coding: utf-8 -*- +import os + +from django.core.checks import Warning, register + + +@register() +def check_gcloud_env_vars(app_configs, **kwargs): + """Validate Google Cloud environment variables (warnings only).""" + warnings = [] + + sa_file = os.getenv("CHATBOT_CREDENTIALS") + if not sa_file: + warnings.append( + Warning( + "CHATBOT_CREDENTIALS not set - chatbot will be disabled", + hint="Set CHATBOT_CREDENTIALS=/path/to/service-account.json\n", + id="chatbot.W001", + ) + ) + elif not os.path.exists(sa_file): + warnings.append( + Warning( + f"Service account file {sa_file} not found - chatbot will be disabled", + hint="Ensure the file exists at the specified path\n", + id="chatbot.W002", + ) + ) + + if not os.getenv("QUERY_PROJECT_ID"): + warnings.append( + Warning( + "QUERY_PROJECT_ID not set - chatbot will be disabled", + hint="Set QUERY_PROJECT_ID=your-gcp-project-id\n", + id="chatbot.W003", + ) + ) + + if not os.getenv("MODEL_URI"): + warnings.append( + Warning( + "MODEL_URI not set - chatbot will be disabled", + hint="Set a valid model uri like 'google_vertexai:gemini-2.5-flash'\n", + id="chatbot.W004", + ) + ) + + if not os.getenv("LANGCHAIN_TRACING_V2"): + warnings.append( + Warning( + "LANGCHAIN_TRACING_V2 not set - tracing will be disabled", + hint="Set LANGCHAIN_TRACING_V2=true\n", + id="chatbot.W005", + ) + ) + + if not os.getenv("LANGCHAIN_API_KEY"): + warnings.append( + Warning( + "LANGCHAIN_API_KEY not set - tracing will be disabled", + hint="Set LANGCHAIN_API_KEY=your-langsmith-api-key\n", + id="chatbot.W006", + ) + ) + + return warnings diff --git a/backend/apps/chatbot/utils/gcloud.py b/backend/apps/chatbot/utils/gcloud.py new file mode 100644 index 00000000..341f2a41 --- /dev/null +++ b/backend/apps/chatbot/utils/gcloud.py @@ -0,0 +1,47 @@ +# -*- coding: utf-8 -*- +import os +from functools import cache + +from google.cloud import bigquery as bq +from google.oauth2.service_account import Credentials + + +@cache +def get_chatbot_credentials() -> Credentials: + """Return cached Google Cloud service account credentials.""" + sa_file = os.getenv("CHATBOT_CREDENTIALS") + + if not sa_file: + raise ValueError( + "CHATBOT_CREDENTIALS environment variable must be set. " + "Please provide the path to your service account JSON file." + ) + + if not os.path.exists(sa_file): + raise FileNotFoundError(f"Service account file not found: {sa_file}") + + return Credentials.from_service_account_file(sa_file) + + +@cache +def get_bigquery_client() -> bq.Client: + """Return a cached BigQuery client. + + The client is initialized once using the project ID from the + `QUERY_PROJECT_ID` environment variable and reused on subsequent calls. + + Returns: + bigquery.Client: A cached, authenticated BigQuery client. + """ + project = os.getenv("QUERY_PROJECT_ID") + + if not project: + raise ValueError( + "QUERY_PROJECT_ID environment variable must be set. " + "Please provide the ID of your BigQuery project." + ) + + return bq.Client( + project=project, + credentials=get_chatbot_credentials(), + ) diff --git a/backend/apps/chatbot/views.py b/backend/apps/chatbot/views.py index a31b041e..683141bc 100644 --- a/backend/apps/chatbot/views.py +++ b/backend/apps/chatbot/views.py @@ -38,6 +38,7 @@ ThreadSerializer, UserMessageSerializer, ) +from backend.apps.chatbot.utils.gcloud import get_chatbot_credentials from backend.apps.chatbot.utils.stream import EventData, StreamEvent, process_chunk ModelSerializer = TypeVar("ModelSerializer", bound=Serializer) @@ -336,7 +337,9 @@ def _get_sql_agent() -> Generator[ReActAgent]: conn = f"postgresql://{db_user}:{db_password}@{db_host}:{db_port}/{db_name}" - model = init_chat_model(MODEL_URI, temperature=0) + credentials = get_chatbot_credentials() + + model = init_chat_model(MODEL_URI, temperature=0, credentials=credentials) def start_hook(state: State): messages = state["messages"] diff --git a/charts/basedosdados-api/templates/deployment.yaml b/charts/basedosdados-api/templates/deployment.yaml index f40c079f..c4a05a7a 100644 --- a/charts/basedosdados-api/templates/deployment.yaml +++ b/charts/basedosdados-api/templates/deployment.yaml @@ -90,3 +90,5 @@ spec: items: - key: GCP_SA path: gcp-sa.json + - key: CHATBOT_SA + path: chatbot-sa.json diff --git a/docker-compose.override.yaml b/docker-compose.override.yaml index 0de2f9b8..d2f472b7 100644 --- a/docker-compose.override.yaml +++ b/docker-compose.override.yaml @@ -63,7 +63,7 @@ services: database: condition: service_healthy healthcheck: - test: [CMD, curl, -f, http://localhost/healthcheck/] + test: [CMD, curl, -f, http://localhost:8000/healthcheck/] interval: 1m timeout: 30s retries: 5