feat: add Ollama and AWS Bedrock AI assistants

plugins/ai_assistant_plugin/__init__.py

@@ -5,3 +5,15 @@
 # from lib.ai_assistant.assistants.openai_assistant import OpenAIAssistant
 #
 # ALL_PLUGIN_AI_ASSISTANTS = [OpenAIAssistant()]
+
+# Example to add Ollama assistant (requires OLLAMA_BASE_URL env var)
+#
+# from lib.ai_assistant.assistants.ollama_assistant import OllamaAIAssistant
+#
+# ALL_PLUGIN_AI_ASSISTANTS = [OllamaAIAssistant()]
+
+# Example to add AWS Bedrock assistant (requires boto3 credentials)
+#
+# from lib.ai_assistant.assistants.bedrock_assistant import BedrockAIAssistant
+#
+# ALL_PLUGIN_AI_ASSISTANTS = [BedrockAIAssistant()]

querybook/server/lib/ai_assistant/assistants/bedrock_assistant.py

@@ -0,0 +1,63 @@
+from langchain_aws import ChatBedrock
+
+from lib.ai_assistant.base_ai_assistant import BaseAIAssistant
+from lib.logger import get_logger
+
+LOG = get_logger(__file__)
+
+BEDROCK_MODEL_CONTEXT_WINDOW_SIZE = {
+    "amazon.nova-pro-v1:0": 300_000,
+    "amazon.nova-lite-v1:0": 300_000,
+    "amazon.nova-micro-v1:0": 128_000,
+    "anthropic.claude-3-5-sonnet-20241022-v2:0": 200_000,
+    "anthropic.claude-3-5-haiku-20241022-v1:0": 200_000,
+    "anthropic.claude-3-opus-20240229-v1:0": 200_000,
+    "anthropic.claude-3-sonnet-20240229-v1:0": 200_000,
+    "anthropic.claude-3-haiku-20240307-v1:0": 200_000,
+    "meta.llama3-70b-instruct-v1:0": 128_000,
+    "meta.llama3-8b-instruct-v1:0": 8_000,
+    "mistral.mistral-large-2402-v1:0": 32_000,
+}
+DEFAULT_MODEL_ID = "anthropic.claude-3-5-haiku-20241022-v1:0"
+DEFAULT_CONTEXT_LENGTH = 200_000
+
+
+class BedrockAIAssistant(BaseAIAssistant):
+    """AWS Bedrock AI Assistant.
+
+    Uses standard boto3 credentials (IAM role, environment variables, or
+    ~/.aws/credentials). No additional configuration is required when running
+    on AWS infrastructure with an appropriate IAM role attached.
+
+    Optional configuration (via model_args in assistant config):
+        model_id: Bedrock model ID (default: anthropic.claude-3-5-haiku-20241022-v1:0)
+        region_name: AWS region (default: boto3 default region)
+    """
+
+    @property
+    def name(self) -> str:
+        return "bedrock"
+
+    def _get_context_length_by_model(self, model_name: str) -> int:
+        return BEDROCK_MODEL_CONTEXT_WINDOW_SIZE.get(model_name, DEFAULT_CONTEXT_LENGTH)
+
+    def _get_default_llm_config(self):
+        default_config = super()._get_default_llm_config()
+        if not default_config.get("model_id"):
+            default_config["model_id"] = DEFAULT_MODEL_ID
+        return default_config
+
+    def _get_token_count(self, ai_command: str, prompt: str) -> int:
+        # Approximation: ~4 characters per token
+        return len(prompt) // 4
+
+    def _get_llm(self, ai_command: str, prompt_length: int):
+        config = self._get_llm_config(ai_command)
+        model_id = config.get("model_id", DEFAULT_MODEL_ID)
+        region_name = config.get("region_name")
+
+        kwargs = {"model_id": model_id}
+        if region_name:
+            kwargs["region_name"] = region_name
+
+        return ChatBedrock(**kwargs)

querybook/server/lib/ai_assistant/assistants/ollama_assistant.py

@@ -0,0 +1,50 @@
+import os
+
+import tiktoken
+from langchain_ollama import ChatOllama
+
+from lib.ai_assistant.base_ai_assistant import BaseAIAssistant
+from lib.logger import get_logger
+
+LOG = get_logger(__file__)
+
+DEFAULT_MODEL_NAME = "llama3.2"
+DEFAULT_CONTEXT_LENGTH = 4096
+
+
+class OllamaAIAssistant(BaseAIAssistant):
+    """Ollama AI Assistant for locally hosted models.
+
+    Required environment variable:
+        OLLAMA_BASE_URL: Base URL of the Ollama server (e.g. http://localhost:11434)
+
+    Optional configuration (via model_args in assistant config):
+        model_name: Ollama model to use (default: llama3.2)
+        base_url: Alternative to OLLAMA_BASE_URL env var
+    """
+
+    @property
+    def name(self) -> str:
+        return "ollama"
+
+    def _get_context_length_by_model(self, model_name: str) -> int:
+        return DEFAULT_CONTEXT_LENGTH
+
+    def _get_default_llm_config(self):
+        default_config = super()._get_default_llm_config()
+        if not default_config.get("model_name"):
+            default_config["model_name"] = DEFAULT_MODEL_NAME
+        return default_config
+
+    def _get_token_count(self, ai_command: str, prompt: str) -> int:
+        encoding = tiktoken.get_encoding("cl100k_base")
+        return len(encoding.encode(prompt))
+
+    def _get_llm(self, ai_command: str, prompt_length: int):
+        config = self._get_llm_config(ai_command)
+        model = config.get("model_name", DEFAULT_MODEL_NAME)
+        base_url = os.environ.get("OLLAMA_BASE_URL") or config.get(
+            "base_url", "http://localhost:11434"
+        )
+
+        return ChatOllama(model=model, base_url=base_url)

requirements/ai/langchain.txt

@@ -1,3 +1,5 @@
 langchain==0.3.24
 langchain-openai==0.3.14
+langchain-ollama==0.3.3
+langchain-aws==0.2.22
 opensearch-py==2.8.0