khoj 1.21.5__py3-none-any.whl → 1.21.7.dev7__py3-none-any.whl

khoj/processor/conversation/offline/chat_model.py

@@ -7,7 +7,7 @@ from typing import Any, Iterator, List, Union
 from langchain.schema import ChatMessage
 from llama_cpp import Llama
 
-from khoj.database.models import Agent, KhojUser
+from khoj.database.models import Agent, ChatModelOptions, KhojUser
 from khoj.processor.conversation import prompts
 from khoj.processor.conversation.offline.utils import download_model
 from khoj.processor.conversation.utils import (
@@ -76,7 +76,11 @@ def extract_questions_offline(
     )
 
     messages = generate_chatml_messages_with_context(
-        example_questions, model_name=model, loaded_model=offline_chat_model, max_prompt_size=max_prompt_size
+        example_questions,
+        model_name=model,
+        loaded_model=offline_chat_model,
+        max_prompt_size=max_prompt_size,
+        model_type=ChatModelOptions.ModelType.OFFLINE,
     )
 
     state.chat_lock.acquire()
@@ -201,6 +205,7 @@ def converse_offline(
         loaded_model=offline_chat_model,
         max_prompt_size=max_prompt_size,
         tokenizer_name=tokenizer_name,
+        model_type=ChatModelOptions.ModelType.OFFLINE,
     )
 
     truncated_messages = "\n".join({f"{message.content[:70]}..." for message in messages})

khoj/processor/conversation/openai/gpt.py

@@ -5,13 +5,16 @@ from typing import Dict, Optional
 
 from langchain.schema import ChatMessage
 
-from khoj.database.models import Agent, KhojUser
+from khoj.database.models import Agent, ChatModelOptions, KhojUser
 from khoj.processor.conversation import prompts
 from khoj.processor.conversation.openai.utils import (
     chat_completion_with_backoff,
     completion_with_backoff,
 )
-from khoj.processor.conversation.utils import generate_chatml_messages_with_context
+from khoj.processor.conversation.utils import (
+    construct_structured_message,
+    generate_chatml_messages_with_context,
+)
 from khoj.utils.helpers import ConversationCommand, is_none_or_empty
 from khoj.utils.rawconfig import LocationData
 
@@ -24,9 +27,10 @@ def extract_questions(
     conversation_log={},
     api_key=None,
     api_base_url=None,
-    temperature=0.7,
     location_data: LocationData = None,
     user: KhojUser = None,
+    uploaded_image_url: Optional[str] = None,
+    vision_enabled: bool = False,
 ):
     """
     Infer search queries to retrieve relevant notes to answer user query
@@ -63,18 +67,18 @@ def extract_questions(
         location=location,
         username=username,
     )
-    messages = [ChatMessage(content=prompt, role="user")]
 
-    # Get Response from GPT
-    response = completion_with_backoff(
-        messages=messages,
-        model=model,
-        temperature=temperature,
-        api_base_url=api_base_url,
-        model_kwargs={"response_format": {"type": "json_object"}},
-        openai_api_key=api_key,
+    prompt = construct_structured_message(
+        message=prompt,
+        image_url=uploaded_image_url,
+        model_type=ChatModelOptions.ModelType.OPENAI,
+        vision_enabled=vision_enabled,
     )
 
+    messages = [ChatMessage(content=prompt, role="user")]
+
+    response = send_message_to_model(messages, api_key, model, response_type="json_object", api_base_url=api_base_url)
+
     # Extract, Clean Message from GPT's Response
     try:
         response = response.strip()
@@ -123,6 +127,8 @@ def converse(
     location_data: LocationData = None,
     user_name: str = None,
     agent: Agent = None,
+    image_url: Optional[str] = None,
+    vision_available: bool = False,
 ):
     """
     Converse with user using OpenAI's ChatGPT
@@ -178,6 +184,9 @@ def converse(
         model_name=model,
         max_prompt_size=max_prompt_size,
         tokenizer_name=tokenizer_name,
+        uploaded_image_url=image_url,
+        vision_enabled=vision_available,
+        model_type=ChatModelOptions.ModelType.OPENAI,
     )
     truncated_messages = "\n".join({f"{message.content[:70]}..." for message in messages})
     logger.debug(f"Conversation Context for GPT: {truncated_messages}")

khoj/processor/conversation/utils.py

@@ -12,7 +12,7 @@ from llama_cpp.llama import Llama
 from transformers import AutoTokenizer
 
 from khoj.database.adapters import ConversationAdapters
-from khoj.database.models import ClientApplication, KhojUser
+from khoj.database.models import ChatModelOptions, ClientApplication, KhojUser
 from khoj.processor.conversation.offline.utils import download_model, infer_max_tokens
 from khoj.utils import state
 from khoj.utils.helpers import is_none_or_empty, merge_dicts
@@ -101,12 +101,16 @@ def save_to_conversation_log(
     client_application: ClientApplication = None,
     conversation_id: int = None,
     automation_id: str = None,
+    uploaded_image_url: str = None,
 ):
     user_message_time = user_message_time or datetime.now().strftime("%Y-%m-%d %H:%M:%S")
     updated_conversation = message_to_log(
         user_message=q,
         chat_response=chat_response,
-        user_message_metadata={"created": user_message_time},
+        user_message_metadata={
+            "created": user_message_time,
+            "uploadedImageData": uploaded_image_url,
+        },
         khoj_message_metadata={
             "context": compiled_references,
             "intent": {"inferred-queries": inferred_queries, "type": intent_type},
@@ -133,6 +137,13 @@ Khoj: "{inferred_queries if ("text-to-image" in intent_type) else chat_response}
     )
 
 
+# Format user and system messages to chatml format
+def construct_structured_message(message, image_url, model_type, vision_enabled):
+    if image_url and vision_enabled and model_type == ChatModelOptions.ModelType.OPENAI:
+        return [{"type": "text", "text": message}, {"type": "image_url", "image_url": {"url": image_url}}]
+    return message
+
+
 def generate_chatml_messages_with_context(
     user_message,
     system_message=None,
@@ -141,6 +152,9 @@ def generate_chatml_messages_with_context(
     loaded_model: Optional[Llama] = None,
     max_prompt_size=None,
     tokenizer_name=None,
+    uploaded_image_url=None,
+    vision_enabled=False,
+    model_type="",
 ):
     """Generate messages for ChatGPT with context from previous conversation"""
     # Set max prompt size from user config or based on pre-configured for model and machine specs
@@ -154,24 +168,35 @@ def generate_chatml_messages_with_context(
     lookback_turns = max_prompt_size // 750
 
     # Extract Chat History for Context
-    chat_logs = []
+    chatml_messages: List[ChatMessage] = []
     for chat in conversation_log.get("chat", []):
-        chat_notes = f'\n\n Notes:\n{chat.get("context")}' if chat.get("context") else "\n"
-        chat_logs += [chat["message"] + chat_notes]
+        message_notes = f'\n\n Notes:\n{chat.get("context")}' if chat.get("context") else "\n"
+        role = "user" if chat["by"] == "you" else "assistant"
+
+        message_content = chat["message"] + message_notes
 
-    rest_backnforths: List[ChatMessage] = []
-    # Extract in reverse chronological order
-    for user_msg, assistant_msg in zip(chat_logs[-2::-2], chat_logs[::-2]):
-        if len(rest_backnforths) >= 2 * lookback_turns:
+        if chat.get("uploadedImageData") and vision_enabled:
+            message_content = construct_structured_message(
+                message_content, chat.get("uploadedImageData"), model_type, vision_enabled
+            )
+
+        reconstructed_message = ChatMessage(content=message_content, role=role)
+
+        chatml_messages.insert(0, reconstructed_message)
+
+        if len(chatml_messages) >= 2 * lookback_turns:
             break
-        rest_backnforths += reciprocal_conversation_to_chatml([user_msg, assistant_msg])[::-1]
 
-    # Format user and system messages to chatml format
     messages = []
     if not is_none_or_empty(user_message):
-        messages.append(ChatMessage(content=user_message, role="user"))
-    if len(rest_backnforths) > 0:
-        messages += rest_backnforths
+        messages.append(
+            ChatMessage(
+                content=construct_structured_message(user_message, uploaded_image_url, model_type, vision_enabled),
+                role="user",
+            )
+        )
+    if len(chatml_messages) > 0:
+        messages += chatml_messages
     if not is_none_or_empty(system_message):
         messages.append(ChatMessage(content=system_message, role="system"))
 

khoj/processor/tools/online_search.py

@@ -56,6 +56,7 @@ async def search_online(
     subscribed: bool = False,
     send_status_func: Optional[Callable] = None,
     custom_filters: List[str] = [],
+    uploaded_image_url: str = None,
 ):
     query += " ".join(custom_filters)
     if not is_internet_connected():
@@ -64,7 +65,9 @@ async def search_online(
         return
 
     # Breakdown the query into subqueries to get the correct answer
-    subqueries = await generate_online_subqueries(query, conversation_history, location, user)
+    subqueries = await generate_online_subqueries(
+        query, conversation_history, location, user, uploaded_image_url=uploaded_image_url
+    )
     response_dict = {}
 
     if subqueries:
@@ -138,13 +141,14 @@ async def read_webpages(
     user: KhojUser,
     subscribed: bool = False,
     send_status_func: Optional[Callable] = None,
+    uploaded_image_url: str = None,
 ):
     "Infer web pages to read from the query and extract relevant information from them"
     logger.info(f"Inferring web pages to read")
     if send_status_func:
         async for event in send_status_func(f"**Inferring web pages to read**"):
             yield {ChatEvent.STATUS: event}
-    urls = await infer_webpage_urls(query, conversation_history, location, user)
+    urls = await infer_webpage_urls(query, conversation_history, location, user, uploaded_image_url)
 
     logger.info(f"Reading web pages at: {urls}")
     if send_status_func:

khoj/routers/api.py

@@ -331,6 +331,7 @@ async def extract_references_and_questions(
     conversation_commands: List[ConversationCommand] = [ConversationCommand.Default],
     location_data: LocationData = None,
     send_status_func: Optional[Callable] = None,
+    uploaded_image_url: Optional[str] = None,
 ):
     user = request.user.object if request.user.is_authenticated else None
 
@@ -370,6 +371,7 @@ async def extract_references_and_questions(
     with timer("Extracting search queries took", logger):
         # If we've reached here, either the user has enabled offline chat or the openai model is enabled.
         conversation_config = await ConversationAdapters.aget_default_conversation_config()
+        vision_enabled = conversation_config.vision_enabled
 
         if conversation_config.model_type == ChatModelOptions.ModelType.OFFLINE:
             using_offline_chat = True
@@ -403,6 +405,8 @@ async def extract_references_and_questions(
                 conversation_log=meta_log,
                 location_data=location_data,
                 user=user,
+                uploaded_image_url=uploaded_image_url,
+                vision_enabled=vision_enabled,
             )
         elif conversation_config.model_type == ChatModelOptions.ModelType.ANTHROPIC:
             api_key = conversation_config.openai_config.api_key

khoj/routers/api_chat.py

@@ -1,4 +1,5 @@
 import asyncio
+import base64
 import json
 import logging
 import time
@@ -46,11 +47,13 @@ from khoj.routers.helpers import (
     update_telemetry_state,
     validate_conversation_config,
 )
+from khoj.routers.storage import upload_image_to_bucket
 from khoj.utils import state
 from khoj.utils.helpers import (
     AsyncIteratorWrapper,
     ConversationCommand,
     command_descriptions,
+    convert_image_to_webp,
     get_device,
     is_none_or_empty,
 )
@@ -517,21 +520,26 @@ async def set_conversation_title(
     )
 
 
-@api_chat.get("")
+class ChatRequestBody(BaseModel):
+    q: str
+    n: Optional[int] = 7
+    d: Optional[float] = None
+    stream: Optional[bool] = False
+    title: Optional[str] = None
+    conversation_id: Optional[int] = None
+    city: Optional[str] = None
+    region: Optional[str] = None
+    country: Optional[str] = None
+    timezone: Optional[str] = None
+    image: Optional[str] = None
+
+
+@api_chat.post("")
 @requires(["authenticated"])
 async def chat(
     request: Request,
     common: CommonQueryParams,
-    q: str,
-    n: int = 7,
-    d: float = None,
-    stream: Optional[bool] = False,
-    title: Optional[str] = None,
-    conversation_id: Optional[int] = None,
-    city: Optional[str] = None,
-    region: Optional[str] = None,
-    country: Optional[str] = None,
-    timezone: Optional[str] = None,
+    body: ChatRequestBody,
     rate_limiter_per_minute=Depends(
         ApiUserRateLimiter(requests=60, subscribed_requests=60, window=60, slug="chat_minute")
     ),
@@ -539,7 +547,20 @@ async def chat(
         ApiUserRateLimiter(requests=600, subscribed_requests=600, window=60 * 60 * 24, slug="chat_day")
     ),
 ):
-    async def event_generator(q: str):
+    # Access the parameters from the body
+    q = body.q
+    n = body.n
+    d = body.d
+    stream = body.stream
+    title = body.title
+    conversation_id = body.conversation_id
+    city = body.city
+    region = body.region
+    country = body.country
+    timezone = body.timezone
+    image = body.image
+
+    async def event_generator(q: str, image: str):
         start_time = time.perf_counter()
         ttft = None
         chat_metadata: dict = {}
@@ -550,6 +571,17 @@ async def chat(
         q = unquote(q)
         nonlocal conversation_id
 
+        uploaded_image_url = None
+        if image:
+            decoded_string = unquote(image)
+            base64_data = decoded_string.split(",", 1)[1]
+            image_bytes = base64.b64decode(base64_data)
+            webp_image_bytes = convert_image_to_webp(image_bytes)
+            try:
+                uploaded_image_url = upload_image_to_bucket(webp_image_bytes, request.user.object.id)
+            except:
+                uploaded_image_url = None
+
         async def send_event(event_type: ChatEvent, data: str | dict):
             nonlocal connection_alive, ttft
             if not connection_alive or await request.is_disconnected():
@@ -637,7 +669,7 @@ async def chat(
 
         if conversation_commands == [ConversationCommand.Default] or is_automated_task:
             conversation_commands = await aget_relevant_information_sources(
-                q, meta_log, is_automated_task, subscribed=subscribed
+                q, meta_log, is_automated_task, subscribed=subscribed, uploaded_image_url=uploaded_image_url
             )
             conversation_commands_str = ", ".join([cmd.value for cmd in conversation_commands])
             async for result in send_event(
@@ -645,7 +677,7 @@ async def chat(
             ):
                 yield result
 
-            mode = await aget_relevant_output_modes(q, meta_log, is_automated_task)
+            mode = await aget_relevant_output_modes(q, meta_log, is_automated_task, uploaded_image_url)
             async for result in send_event(ChatEvent.STATUS, f"**Decided Response Mode:** {mode.value}"):
                 yield result
             if mode not in conversation_commands:
@@ -693,7 +725,9 @@ async def chat(
                     ):
                         yield result
 
-                    response = await extract_relevant_summary(q, contextual_data, subscribed=subscribed)
+                    response = await extract_relevant_summary(
+                        q, contextual_data, subscribed=subscribed, uploaded_image_url=uploaded_image_url
+                    )
                     response_log = str(response)
                     async for result in send_llm_response(response_log):
                         yield result
@@ -711,6 +745,7 @@ async def chat(
                 intent_type="summarize",
                 client_application=request.user.client_app,
                 conversation_id=conversation_id,
+                uploaded_image_url=uploaded_image_url,
             )
             return
 
@@ -753,6 +788,7 @@ async def chat(
                 conversation_id=conversation_id,
                 inferred_queries=[query_to_run],
                 automation_id=automation.id,
+                uploaded_image_url=uploaded_image_url,
             )
             async for result in send_llm_response(llm_response):
                 yield result
@@ -771,6 +807,7 @@ async def chat(
             conversation_commands,
             location,
             partial(send_event, ChatEvent.STATUS),
+            uploaded_image_url=uploaded_image_url,
         ):
             if isinstance(result, dict) and ChatEvent.STATUS in result:
                 yield result[ChatEvent.STATUS]
@@ -807,6 +844,7 @@ async def chat(
                     subscribed,
                     partial(send_event, ChatEvent.STATUS),
                     custom_filters,
+                    uploaded_image_url=uploaded_image_url,
                 ):
                     if isinstance(result, dict) and ChatEvent.STATUS in result:
                         yield result[ChatEvent.STATUS]
@@ -823,7 +861,13 @@ async def chat(
         if ConversationCommand.Webpage in conversation_commands:
             try:
                 async for result in read_webpages(
-                    defiltered_query, meta_log, location, user, subscribed, partial(send_event, ChatEvent.STATUS)
+                    defiltered_query,
+                    meta_log,
+                    location,
+                    user,
+                    subscribed,
+                    partial(send_event, ChatEvent.STATUS),
+                    uploaded_image_url=uploaded_image_url,
                 ):
                     if isinstance(result, dict) and ChatEvent.STATUS in result:
                         yield result[ChatEvent.STATUS]
@@ -869,6 +913,7 @@ async def chat(
                 online_results=online_results,
                 subscribed=subscribed,
                 send_status_func=partial(send_event, ChatEvent.STATUS),
+                uploaded_image_url=uploaded_image_url,
             ):
                 if isinstance(result, dict) and ChatEvent.STATUS in result:
                     yield result[ChatEvent.STATUS]
@@ -898,6 +943,7 @@ async def chat(
                 conversation_id=conversation_id,
                 compiled_references=compiled_references,
                 online_results=online_results,
+                uploaded_image_url=uploaded_image_url,
             )
             content_obj = {
                 "intentType": intent_type,
@@ -924,6 +970,7 @@ async def chat(
             conversation_id,
             location,
             user_name,
+            uploaded_image_url,
         )
 
         # Send Response
@@ -949,9 +996,9 @@ async def chat(
 
     ## Stream Text Response
     if stream:
-        return StreamingResponse(event_generator(q), media_type="text/plain")
+        return StreamingResponse(event_generator(q, image=image), media_type="text/plain")
     ## Non-Streaming Text Response
     else:
-        response_iterator = event_generator(q)
+        response_iterator = event_generator(q, image=image)
         response_data = await read_chat_stream(response_iterator)
         return Response(content=json.dumps(response_data), media_type="application/json", status_code=200)