diff --git a/plugins/ai_assistant_plugin/__init__.py b/plugins/ai_assistant_plugin/__init__.py index 0791ff47c..3778e9283 100644 --- a/plugins/ai_assistant_plugin/__init__.py +++ b/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()] diff --git a/querybook/server/lib/ai_assistant/assistants/bedrock_assistant.py b/querybook/server/lib/ai_assistant/assistants/bedrock_assistant.py new file mode 100644 index 000000000..fa46593ce --- /dev/null +++ b/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) diff --git a/querybook/server/lib/ai_assistant/assistants/ollama_assistant.py b/querybook/server/lib/ai_assistant/assistants/ollama_assistant.py new file mode 100644 index 000000000..b888b66e6 --- /dev/null +++ b/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) diff --git a/requirements/ai/langchain.txt b/requirements/ai/langchain.txt index 7a2505ba8..a68d5d78b 100644 --- a/requirements/ai/langchain.txt +++ b/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