khoj 1.22.3__py3-none-any.whl → 1.22.4.dev6__py3-none-any.whl

khoj/processor/conversation/google/utils.py

@@ -0,0 +1,93 @@
+import logging
+from threading import Thread
+
+import google.generativeai as genai
+from tenacity import (
+    before_sleep_log,
+    retry,
+    stop_after_attempt,
+    wait_exponential,
+    wait_random_exponential,
+)
+
+from khoj.processor.conversation.utils import ThreadedGenerator
+
+logger = logging.getLogger(__name__)
+
+
+DEFAULT_MAX_TOKENS_GEMINI = 8192
+
+
+@retry(
+    wait=wait_random_exponential(min=1, max=10),
+    stop=stop_after_attempt(2),
+    before_sleep=before_sleep_log(logger, logging.DEBUG),
+    reraise=True,
+)
+def gemini_completion_with_backoff(
+    messages, system_prompt, model_name, temperature=0, api_key=None, model_kwargs=None, max_tokens=None
+) -> str:
+    genai.configure(api_key=api_key)
+    max_tokens = max_tokens or DEFAULT_MAX_TOKENS_GEMINI
+    model_kwargs = model_kwargs or dict()
+    model_kwargs["temperature"] = temperature
+    model_kwargs["max_output_tokens"] = max_tokens
+    model = genai.GenerativeModel(model_name, generation_config=model_kwargs, system_instruction=system_prompt)
+
+    formatted_messages = [{"role": message.role, "parts": [message.content]} for message in messages]
+    # all messages up to the last are considered to be part of the chat history
+    chat_session = model.start_chat(history=formatted_messages[0:-1])
+    # the last message is considered to be the current prompt
+    aggregated_response = chat_session.send_message(formatted_messages[-1]["parts"][0])
+    return aggregated_response.text
+
+
+@retry(
+    wait=wait_exponential(multiplier=1, min=4, max=10),
+    stop=stop_after_attempt(2),
+    before_sleep=before_sleep_log(logger, logging.DEBUG),
+    reraise=True,
+)
+def gemini_chat_completion_with_backoff(
+    messages,
+    compiled_references,
+    online_results,
+    model_name,
+    temperature,
+    api_key,
+    system_prompt,
+    max_prompt_size=None,
+    completion_func=None,
+    model_kwargs=None,
+):
+    g = ThreadedGenerator(compiled_references, online_results, completion_func=completion_func)
+    t = Thread(
+        target=gemini_llm_thread,
+        args=(g, messages, system_prompt, model_name, temperature, api_key, max_prompt_size, model_kwargs),
+    )
+    t.start()
+    return g
+
+
+def gemini_llm_thread(
+    g, messages, system_prompt, model_name, temperature, api_key, max_prompt_size=None, model_kwargs=None
+):
+    try:
+        genai.configure(api_key=api_key)
+        max_tokens = max_prompt_size or DEFAULT_MAX_TOKENS_GEMINI
+        model_kwargs = model_kwargs or dict()
+        model_kwargs["temperature"] = temperature
+        model_kwargs["max_output_tokens"] = max_tokens
+        model_kwargs["stop_sequences"] = ["Notes:\n["]
+        model = genai.GenerativeModel(model_name, generation_config=model_kwargs, system_instruction=system_prompt)
+
+        formatted_messages = [{"role": message.role, "parts": [message.content]} for message in messages]
+        # all messages up to the last are considered to be part of the chat history
+        chat_session = model.start_chat(history=formatted_messages[0:-1])
+        # the last message is considered to be the current prompt
+        for chunk in chat_session.send_message(formatted_messages[-1]["parts"][0], stream=True):
+            g.send(chunk.text)
+    except Exception as e:
+        logger.error(f"Error in gemini_llm_thread: {e}", exc_info=True)
+    finally:
+        g.close()

khoj/processor/conversation/openai/gpt.py

@@ -14,6 +14,7 @@ from khoj.processor.conversation.openai.utils import (
 from khoj.processor.conversation.utils import (
     construct_structured_message,
     generate_chatml_messages_with_context,
+    remove_json_codeblock,
 )
 from khoj.utils.helpers import ConversationCommand, is_none_or_empty
 from khoj.utils.rawconfig import LocationData
@@ -85,6 +86,7 @@ def extract_questions(
     # Extract, Clean Message from GPT's Response
     try:
         response = response.strip()
+        response = remove_json_codeblock(response)
         response = json.loads(response)
         response = [q.strip() for q in response["queries"] if q.strip()]
         if not isinstance(response, list) or not response:

khoj/processor/conversation/openai/utils.py

@@ -45,15 +45,28 @@ def completion_with_backoff(
         openai_clients[client_key] = client
 
     formatted_messages = [{"role": message.role, "content": message.content} for message in messages]
+    stream = True
+
+    # Update request parameters for compatability with o1 model series
+    # Refer: https://platform.openai.com/docs/guides/reasoning/beta-limitations
+    if model.startswith("o1"):
+        stream = False
+        temperature = 1
+        model_kwargs.pop("stop", None)
+        model_kwargs.pop("response_format", None)
 
     chat = client.chat.completions.create(
-        stream=True,
+        stream=stream,
         messages=formatted_messages,  # type: ignore
         model=model,  # type: ignore
         temperature=temperature,
         timeout=20,
         **(model_kwargs or dict()),
     )
+
+    if not stream:
+        return chat.choices[0].message.content
+
     aggregated_response = ""
     for chunk in chat:
         if len(chunk.choices) == 0:
@@ -112,9 +125,18 @@ def llm_thread(g, messages, model_name, temperature, openai_api_key=None, api_ba
             client: openai.OpenAI = openai_clients[client_key]
 
         formatted_messages = [{"role": message.role, "content": message.content} for message in messages]
+        stream = True
+
+        # Update request parameters for compatability with o1 model series
+        # Refer: https://platform.openai.com/docs/guides/reasoning/beta-limitations
+        if model_name.startswith("o1"):
+            stream = False
+            temperature = 1
+            model_kwargs.pop("stop", None)
+            model_kwargs.pop("response_format", None)
 
         chat = client.chat.completions.create(
-            stream=True,
+            stream=stream,
             messages=formatted_messages,
             model=model_name,  # type: ignore
             temperature=temperature,
@@ -122,14 +144,17 @@ def llm_thread(g, messages, model_name, temperature, openai_api_key=None, api_ba
             **(model_kwargs or dict()),
         )
 
-        for chunk in chat:
-            if len(chunk.choices) == 0:
-                continue
-            delta_chunk = chunk.choices[0].delta
-            if isinstance(delta_chunk, str):
-                g.send(delta_chunk)
-            elif delta_chunk.content:
-                g.send(delta_chunk.content)
+        if not stream:
+            g.send(chat.choices[0].message.content)
+        else:
+            for chunk in chat:
+                if len(chunk.choices) == 0:
+                    continue
+                delta_chunk = chunk.choices[0].delta
+                if isinstance(delta_chunk, str):
+                    g.send(delta_chunk)
+                elif delta_chunk.content:
+                    g.send(delta_chunk.content)
     except Exception as e:
         logger.error(f"Error in llm_thread: {e}", exc_info=True)
     finally:

khoj/processor/conversation/prompts.py

@@ -13,8 +13,8 @@ You were created by Khoj Inc. with the following capabilities:
 - You *CAN* generate images, look-up real-time information from the internet, set reminders and answer questions based on the user's notes.
 - Say "I don't know" or "I don't understand" if you don't know what to say or if you don't know the answer to a question.
 - Make sure to use the specific LaTeX math mode delimiters for your response. LaTex math mode specific delimiters as following
-    - inline math mode : `\\(` and `\\)`
-    - display math mode: insert linebreak after opening `$$`, `\\[` and before closing `$$`, `\\]`
+    - inline math mode : \\( and \\)
+    - display math mode: insert linebreak after opening $$, \\[ and before closing $$, \\]
 - Ask crisp follow-up questions to get additional context, when the answer cannot be inferred from the provided notes or past conversations.
 - Sometimes the user will share personal information that needs to be remembered, like an account ID or a residential address. These can be acknowledged with a simple "Got it" or "Okay".
 - Provide inline references to quotes from the user's notes or any web pages you refer to in your responses in markdown format. For example, "The farmer had ten sheep. [1](https://example.com)". *ALWAYS CITE YOUR SOURCES AND PROVIDE REFERENCES*. Add them inline to directly support your claim.

khoj/processor/conversation/utils.py

@@ -1,4 +1,3 @@
-import json
 import logging
 import math
 import queue
@@ -24,6 +23,8 @@ model_to_prompt_size = {
     "gpt-4-0125-preview": 20000,
     "gpt-4-turbo-preview": 20000,
     "gpt-4o-mini": 20000,
+    "o1-preview": 20000,
+    "o1-mini": 20000,
     "TheBloke/Mistral-7B-Instruct-v0.2-GGUF": 3500,
     "NousResearch/Hermes-2-Pro-Mistral-7B-GGUF": 3500,
     "bartowski/Meta-Llama-3.1-8B-Instruct-GGUF": 20000,
@@ -220,8 +221,9 @@ def truncate_messages(
     try:
         if loaded_model:
             encoder = loaded_model.tokenizer()
-        elif model_name.startswith("gpt-"):
-            encoder = tiktoken.encoding_for_model(model_name)
+        elif model_name.startswith("gpt-") or model_name.startswith("o1"):
+            # as tiktoken doesn't recognize o1 model series yet
+            encoder = tiktoken.encoding_for_model("gpt-4o" if model_name.startswith("o1") else model_name)
         elif tokenizer_name:
             if tokenizer_name in state.pretrained_tokenizers:
                 encoder = state.pretrained_tokenizers[tokenizer_name]
@@ -278,10 +280,19 @@ def truncate_messages(
         )
 
     if system_message:
-        system_message.role = "user" if "gemma-2" in model_name else "system"
+        # Default system message role is system.
+        # Fallback to system message role of user for models that do not support this role like gemma-2 and openai's o1 model series.
+        system_message.role = "user" if "gemma-2" in model_name or model_name.startswith("o1") else "system"
     return messages + [system_message] if system_message else messages
 
 
 def reciprocal_conversation_to_chatml(message_pair):
     """Convert a single back and forth between user and assistant to chatml format"""
     return [ChatMessage(content=message, role=role) for message, role in zip(message_pair, ["user", "assistant"])]
+
+
+def remove_json_codeblock(response):
+    """Remove any markdown json codeblock formatting if present. Useful for non schema enforceable models"""
+    if response.startswith("```json") and response.endswith("```"):
+        response = response[7:-3]
+    return response

khoj/processor/tools/online_search.py

@@ -7,6 +7,7 @@ from collections import defaultdict
 from typing import Callable, Dict, List, Optional, Tuple, Union
 
 import aiohttp
+import requests
 from bs4 import BeautifulSoup
 from markdownify import markdownify
 
@@ -94,7 +95,7 @@ async def search_online(
 
     # Read, extract relevant info from the retrieved web pages
     if webpages:
-        webpage_links = [link for link, _, _ in webpages]
+        webpage_links = set([link for link, _, _ in webpages])
         logger.info(f"Reading web pages at: {list(webpage_links)}")
         if send_status_func:
             webpage_links_str = "\n- " + "\n- ".join(list(webpage_links))

khoj/routers/api.py

@@ -31,6 +31,7 @@ from khoj.database.models import ChatModelOptions, KhojUser, SpeechToTextModelOp
 from khoj.processor.conversation.anthropic.anthropic_chat import (
     extract_questions_anthropic,
 )
+from khoj.processor.conversation.google.gemini_chat import extract_questions_gemini
 from khoj.processor.conversation.offline.chat_model import extract_questions_offline
 from khoj.processor.conversation.offline.whisper import transcribe_audio_offline
 from khoj.processor.conversation.openai.gpt import extract_questions
@@ -419,6 +420,18 @@ async def extract_references_and_questions(
                 location_data=location_data,
                 user=user,
             )
+        elif conversation_config.model_type == ChatModelOptions.ModelType.GOOGLE:
+            api_key = conversation_config.openai_config.api_key
+            chat_model = conversation_config.chat_model
+            inferred_queries = extract_questions_gemini(
+                defiltered_query,
+                model=chat_model,
+                api_key=api_key,
+                conversation_log=meta_log,
+                location_data=location_data,
+                max_tokens=conversation_config.max_prompt_size,
+                user=user,
+            )
 
     # Collate search results as context for GPT
     with timer("Searching knowledge base took", logger):

khoj/routers/helpers.py

@@ -76,6 +76,10 @@ from khoj.processor.conversation.anthropic.anthropic_chat import (
     anthropic_send_message_to_model,
     converse_anthropic,
 )
+from khoj.processor.conversation.google.gemini_chat import (
+    converse_gemini,
+    gemini_send_message_to_model,
+)
 from khoj.processor.conversation.offline.chat_model import (
     converse_offline,
     send_message_to_model_offline,
@@ -84,6 +88,7 @@ from khoj.processor.conversation.openai.gpt import converse, send_message_to_mod
 from khoj.processor.conversation.utils import (
     ThreadedGenerator,
     generate_chatml_messages_with_context,
+    remove_json_codeblock,
     save_to_conversation_log,
 )
 from khoj.processor.speech.text_to_speech import is_eleven_labs_enabled
@@ -136,7 +141,7 @@ async def is_ready_to_chat(user: KhojUser):
         await ConversationAdapters.aget_default_conversation_config()
     )
 
-    if user_conversation_config and user_conversation_config.model_type == "offline":
+    if user_conversation_config and user_conversation_config.model_type == ChatModelOptions.ModelType.OFFLINE:
         chat_model = user_conversation_config.chat_model
         max_tokens = user_conversation_config.max_prompt_size
         if state.offline_chat_processor_config is None:
@@ -146,7 +151,14 @@ async def is_ready_to_chat(user: KhojUser):
 
     if (
         user_conversation_config
-        and (user_conversation_config.model_type == "openai" or user_conversation_config.model_type == "anthropic")
+        and (
+            user_conversation_config.model_type
+            in [
+                ChatModelOptions.ModelType.OPENAI,
+                ChatModelOptions.ModelType.ANTHROPIC,
+                ChatModelOptions.ModelType.GOOGLE,
+            ]
+        )
         and user_conversation_config.openai_config
     ):
         return True
@@ -287,9 +299,7 @@ async def aget_relevant_information_sources(
 
     try:
         response = response.strip()
-        # Remove any markdown json codeblock formatting if present (useful for gemma-2)
-        if response.startswith("```json"):
-            response = response[7:-3]
+        response = remove_json_codeblock(response)
         response = json.loads(response)
         response = [q.strip() for q in response["source"] if q.strip()]
         if not isinstance(response, list) or not response or len(response) == 0:
@@ -342,7 +352,9 @@ async def aget_relevant_output_modes(
         response = await send_message_to_model_wrapper(relevant_mode_prompt, response_type="json_object")
 
     try:
-        response = json.loads(response.strip())
+        response = response.strip()
+        response = remove_json_codeblock(response)
+        response = json.loads(response)
 
         if is_none_or_empty(response):
             return ConversationCommand.Text
@@ -422,9 +434,7 @@ async def generate_online_subqueries(
     # Validate that the response is a non-empty, JSON-serializable list
     try:
         response = response.strip()
-        # Remove any markdown json codeblock formatting if present (useful for gemma-2)
-        if response.startswith("```json") and response.endswith("```"):
-            response = response[7:-3]
+        response = remove_json_codeblock(response)
         response = json.loads(response)
         response = [q.strip() for q in response["queries"] if q.strip()]
         if not isinstance(response, list) or not response or len(response) == 0:
@@ -607,9 +617,10 @@ async def send_message_to_model_wrapper(
         else conversation_config.max_prompt_size
     )
     tokenizer = conversation_config.tokenizer
+    model_type = conversation_config.model_type
     vision_available = conversation_config.vision_enabled
 
-    if conversation_config.model_type == "offline":
+    if model_type == ChatModelOptions.ModelType.OFFLINE:
         if state.offline_chat_processor_config is None or state.offline_chat_processor_config.loaded_model is None:
             state.offline_chat_processor_config = OfflineChatProcessorModel(chat_model, max_tokens)
 
@@ -633,7 +644,7 @@ async def send_message_to_model_wrapper(
             response_type=response_type,
         )
 
-    elif conversation_config.model_type == "openai":
+    elif model_type == ChatModelOptions.ModelType.OPENAI:
         openai_chat_config = conversation_config.openai_config
         api_key = openai_chat_config.api_key
         api_base_url = openai_chat_config.api_base_url
@@ -657,7 +668,7 @@ async def send_message_to_model_wrapper(
         )
 
         return openai_response
-    elif conversation_config.model_type == "anthropic":
+    elif model_type == ChatModelOptions.ModelType.ANTHROPIC:
         api_key = conversation_config.openai_config.api_key
         truncated_messages = generate_chatml_messages_with_context(
             user_message=message,
@@ -666,6 +677,7 @@ async def send_message_to_model_wrapper(
             max_prompt_size=max_tokens,
             tokenizer_name=tokenizer,
             vision_enabled=vision_available,
+            uploaded_image_url=uploaded_image_url,
             model_type=conversation_config.model_type,
         )
 
@@ -674,6 +686,21 @@ async def send_message_to_model_wrapper(
             api_key=api_key,
             model=chat_model,
         )
+    elif model_type == ChatModelOptions.ModelType.GOOGLE:
+        api_key = conversation_config.openai_config.api_key
+        truncated_messages = generate_chatml_messages_with_context(
+            user_message=message,
+            system_message=system_message,
+            model_name=chat_model,
+            max_prompt_size=max_tokens,
+            tokenizer_name=tokenizer,
+            vision_enabled=vision_available,
+            uploaded_image_url=uploaded_image_url,
+        )
+
+        return gemini_send_message_to_model(
+            messages=truncated_messages, api_key=api_key, model=chat_model, response_type=response_type
+        )
     else:
         raise HTTPException(status_code=500, detail="Invalid conversation config")
 
@@ -692,7 +719,7 @@ def send_message_to_model_wrapper_sync(
     max_tokens = conversation_config.max_prompt_size
     vision_available = conversation_config.vision_enabled
 
-    if conversation_config.model_type == "offline":
+    if conversation_config.model_type == ChatModelOptions.ModelType.OFFLINE:
         if state.offline_chat_processor_config is None or state.offline_chat_processor_config.loaded_model is None:
             state.offline_chat_processor_config = OfflineChatProcessorModel(chat_model, max_tokens)
 
@@ -714,7 +741,7 @@ def send_message_to_model_wrapper_sync(
             response_type=response_type,
         )
 
-    elif conversation_config.model_type == "openai":
+    elif conversation_config.model_type == ChatModelOptions.ModelType.OPENAI:
         api_key = conversation_config.openai_config.api_key
         truncated_messages = generate_chatml_messages_with_context(
             user_message=message,
@@ -730,7 +757,7 @@ def send_message_to_model_wrapper_sync(
 
         return openai_response
 
-    elif conversation_config.model_type == "anthropic":
+    elif conversation_config.model_type == ChatModelOptions.ModelType.ANTHROPIC:
         api_key = conversation_config.openai_config.api_key
         truncated_messages = generate_chatml_messages_with_context(
             user_message=message,
@@ -746,6 +773,22 @@ def send_message_to_model_wrapper_sync(
             api_key=api_key,
             model=chat_model,
         )
+
+    elif conversation_config.model_type == ChatModelOptions.ModelType.GOOGLE:
+        api_key = conversation_config.openai_config.api_key
+        truncated_messages = generate_chatml_messages_with_context(
+            user_message=message,
+            system_message=system_message,
+            model_name=chat_model,
+            max_prompt_size=max_tokens,
+            vision_enabled=vision_available,
+        )
+
+        return gemini_send_message_to_model(
+            messages=truncated_messages,
+            api_key=api_key,
+            model=chat_model,
+        )
     else:
         raise HTTPException(status_code=500, detail="Invalid conversation config")
 
@@ -811,7 +854,7 @@ def generate_chat_response(
                 agent=agent,
             )
 
-        elif conversation_config.model_type == "openai":
+        elif conversation_config.model_type == ChatModelOptions.ModelType.OPENAI:
             openai_chat_config = conversation_config.openai_config
             api_key = openai_chat_config.api_key
             chat_model = conversation_config.chat_model
@@ -834,7 +877,7 @@ def generate_chat_response(
                 vision_available=vision_available,
             )
 
-        elif conversation_config.model_type == "anthropic":
+        elif conversation_config.model_type == ChatModelOptions.ModelType.ANTHROPIC:
             api_key = conversation_config.openai_config.api_key
             chat_response = converse_anthropic(
                 compiled_references,
@@ -851,6 +894,23 @@ def generate_chat_response(
                 user_name=user_name,
                 agent=agent,
             )
+        elif conversation_config.model_type == ChatModelOptions.ModelType.GOOGLE:
+            api_key = conversation_config.openai_config.api_key
+            chat_response = converse_gemini(
+                compiled_references,
+                q,
+                online_results,
+                meta_log,
+                model=conversation_config.chat_model,
+                api_key=api_key,
+                completion_func=partial_completion,
+                conversation_commands=conversation_commands,
+                max_prompt_size=conversation_config.max_prompt_size,
+                tokenizer_name=conversation_config.tokenizer,
+                location_data=location_data,
+                user_name=user_name,
+                agent=agent,
+            )
 
         metadata.update({"chat_model": conversation_config.chat_model})
 
@@ -1217,11 +1277,6 @@ def scheduled_chat(
             token = token[0].token
         headers["Authorization"] = f"Bearer {token}"
 
-    # Log request details
-    logger.info(f"POST URL: {url}")
-    logger.info(f"Headers: {headers}")
-    logger.info(f"Payload: {json_payload}")
-
     # Call the chat API endpoint with authenticated user token and query
     raw_response = requests.post(url, headers=headers, json=json_payload, allow_redirects=False)
 
@@ -1231,14 +1286,6 @@ def scheduled_chat(
         logger.info(f"Redirecting to {redirect_url}")
         raw_response = requests.post(redirect_url, headers=headers, json=json_payload)
 
-    # Log response details
-    logger.info(f"Response status code: {raw_response.status_code}")
-    logger.info(f"Response headers: {raw_response.headers}")
-    logger.info(f"Response text: {raw_response.text}")
-    if raw_response.history:
-        for resp in raw_response.history:
-            logger.info(f"Redirected from {resp.url} with status code {resp.status_code}")
-
     # Stop if the chat API call was not successful
     if raw_response.status_code != 200:
         logger.error(f"Failed to run schedule chat: {raw_response.text}, user: {user}, query: {query_to_run}")

{khoj-1.22.3.dist-info → khoj-1.22.4.dev6.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: khoj
-Version: 1.22.3
+Version: 1.22.4.dev6
 Summary: Your Second Brain
 Project-URL: Homepage, https://khoj.dev
 Project-URL: Documentation, https://docs.khoj.dev
@@ -36,6 +36,7 @@ Requires-Dist: django==5.0.8
 Requires-Dist: docx2txt==0.8
 Requires-Dist: einops==0.8.0
 Requires-Dist: fastapi>=0.110.0
+Requires-Dist: google-generativeai==0.7.2
 Requires-Dist: httpx==0.25.0
 Requires-Dist: huggingface-hub>=0.22.2
 Requires-Dist: itsdangerous==2.1.2