feat: Add Gemini support for language models (#3632)

Introduce support for the Gemini model in the language model
configuration, allowing users to utilize Gemini alongside existing
models. Update the configuration to include Gemini-specific settings and
ensure compatibility with the overall architecture.

Author: chloedia

core/pyproject.toml

@@ -23,6 +23,8 @@ dependencies = [
     "markupsafe>=2.1.5",
     "megaparse-sdk>=0.1.11",
     "langchain-mistralai>=0.2.3",
+    "langchain-google-genai>=2.1.3",
+    "langchain-xai>=0.2.3",
     "fasttext-langdetect>=1.0.5",
     "langfuse>=2.57.0",
 ]

core/quivr_core/llm/llm_endpoint.py

@@ -7,8 +7,10 @@
 import tiktoken
 from langchain_anthropic import ChatAnthropic
 from langchain_core.language_models.chat_models import BaseChatModel
+from langchain_google_genai import ChatGoogleGenerativeAI
 from langchain_mistralai import ChatMistralAI
 from langchain_openai import AzureChatOpenAI, ChatOpenAI
+from langchain_xai import ChatXAI
 from pydantic import SecretStr
 
 from quivr_core.brain.info import LLMInfo
@@ -206,7 +208,14 @@ def get_config(self):
 
     @classmethod
     def from_config(cls, config: LLMEndpointConfig = LLMEndpointConfig()):
-        _llm: Union[AzureChatOpenAI, ChatOpenAI, ChatAnthropic, ChatMistralAI]
+        _llm: Union[
+            AzureChatOpenAI,
+            ChatOpenAI,
+            ChatAnthropic,
+            ChatMistralAI,
+            ChatGoogleGenerativeAI,
+            ChatXAI,
+        ]
         try:
             if config.supplier == DefaultModelSuppliers.AZURE:
                 # Parse the URL
@@ -255,6 +264,27 @@ def from_config(cls, config: LLMEndpointConfig = LLMEndpointConfig()):
                     base_url=config.llm_base_url,
                     temperature=config.temperature,
                 )
+            elif config.supplier == DefaultModelSuppliers.GEMINI:
+                _llm = ChatGoogleGenerativeAI(
+                    model=config.model,
+                    api_key=SecretStr(config.llm_api_key)
+                    if config.llm_api_key
+                    else None,
+                    base_url=config.llm_base_url,
+                    max_tokens=config.max_output_tokens,
+                    temperature=config.temperature,
+                )
+            elif config.supplier == DefaultModelSuppliers.GROQ:
+                _llm = ChatXAI(
+                    model=config.model,
+                    api_key=SecretStr(config.llm_api_key)
+                    if config.llm_api_key
+                    else None,
+                    base_url=config.llm_base_url,
+                    max_tokens=config.max_output_tokens,
+                    temperature=config.temperature,
+                )
+
             else:
                 _llm = ChatOpenAI(
                     model=config.model,

core/quivr_core/rag/entities/config.py

@@ -72,6 +72,7 @@ class DefaultModelSuppliers(str, Enum):
     META = "meta"
     MISTRAL = "mistral"
     GROQ = "groq"
+    GEMINI = "gemini"
 
 
 class LLMConfig(QuivrBaseConfig):
@@ -98,6 +99,11 @@ class LLMModelConfig:
                 max_output_tokens=100000,
                 tokenizer_hub="Quivr/gpt-4o",
             ),
+            "o4-mini": LLMConfig(
+                max_context_tokens=200000,
+                max_output_tokens=100000,
+                tokenizer_hub="Quivr/gpt-4o",
+            ),
             "o1-mini": LLMConfig(
                 max_context_tokens=128000,
                 max_output_tokens=65536,
@@ -139,6 +145,11 @@ class LLMModelConfig:
             ),
         },
         DefaultModelSuppliers.ANTHROPIC: {
+            "claude-3-7-sonnet": LLMConfig(
+                max_context_tokens=200000,
+                max_output_tokens=8192,
+                tokenizer_hub="Quivr/claude-tokenizer",
+            ),
             "claude-3-5-sonnet": LLMConfig(
                 max_context_tokens=200000,
                 max_output_tokens=8192,
@@ -209,6 +220,16 @@ class LLMModelConfig:
             "code-llama": LLMConfig(
                 max_context_tokens=16384, tokenizer_hub="Quivr/llama-code-tokenizer"
             ),
+            "deepseek-r1-distill-llama-70b": LLMConfig(
+                max_context_tokens=128000,
+                max_output_tokens=32768,
+                tokenizer_hub="Quivr/Meta-Llama-3.1-Tokenizer",
+            ),
+            "meta-llama/llama-4-maverick-17b-128e-instruct": LLMConfig(
+                max_context_tokens=128000,
+                max_output_tokens=32768,
+                tokenizer_hub="Quivr/Meta-Llama-3.1-Tokenizer",
+            ),
         },
         DefaultModelSuppliers.MISTRAL: {
             "mistral-large": LLMConfig(
@@ -230,6 +251,13 @@ class LLMModelConfig:
                 max_context_tokens=32000, tokenizer_hub="Quivr/mistral-tokenizer-v3"
             ),
         },
+        DefaultModelSuppliers.GEMINI: {
+            "gemini-2.5": LLMConfig(
+                max_context_tokens=128000,
+                max_output_tokens=4096,
+                tokenizer_hub="Quivr/gemini-tokenizer",
+            ),
+        },
     }
 
     @classmethod

core/quivr_core/rag/quivr_rag_langgraph.py

@@ -945,6 +945,7 @@ def generate_zendesk_rag(self, state: AgentState) -> AgentState:
 
         msg = prompt_template.format_prompt(**inputs)
         llm = self.bind_tools_to_llm(self.generate_zendesk_rag.__name__)
+
         response = llm.invoke(msg)
 
         return {**state, "messages": [response]}