khoj 1.33.3.dev32__py3-none-any.whl

khoj/processor/conversation/anthropic/utils.py

@@ -0,0 +1,217 @@
+import logging
+from threading import Thread
+from typing import Dict, List
+
+import anthropic
+from langchain.schema import ChatMessage
+from tenacity import (
+    before_sleep_log,
+    retry,
+    stop_after_attempt,
+    wait_exponential,
+    wait_random_exponential,
+)
+
+from khoj.processor.conversation.utils import (
+    ThreadedGenerator,
+    commit_conversation_trace,
+    get_image_from_url,
+)
+from khoj.utils import state
+from khoj.utils.helpers import (
+    get_chat_usage_metrics,
+    in_debug_mode,
+    is_none_or_empty,
+    is_promptrace_enabled,
+)
+
+logger = logging.getLogger(__name__)
+
+anthropic_clients: Dict[str, anthropic.Anthropic] = {}
+
+
+DEFAULT_MAX_TOKENS_ANTHROPIC = 3000
+
+
+@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 anthropic_completion_with_backoff(
+    messages,
+    system_prompt,
+    model_name,
+    temperature=0,
+    api_key=None,
+    model_kwargs=None,
+    max_tokens=None,
+    response_type="text",
+    tracer={},
+) -> str:
+    if api_key not in anthropic_clients:
+        client: anthropic.Anthropic = anthropic.Anthropic(api_key=api_key)
+        anthropic_clients[api_key] = client
+    else:
+        client = anthropic_clients[api_key]
+
+    formatted_messages = [{"role": message.role, "content": message.content} for message in messages]
+    if response_type == "json_object":
+        # Prefill model response with '{' to make it output a valid JSON object
+        formatted_messages += [{"role": "assistant", "content": "{"}]
+
+    aggregated_response = "{" if response_type == "json_object" else ""
+    max_tokens = max_tokens or DEFAULT_MAX_TOKENS_ANTHROPIC
+
+    final_message = None
+    model_kwargs = model_kwargs or dict()
+    if system_prompt:
+        model_kwargs["system"] = system_prompt
+
+    with client.messages.stream(
+        messages=formatted_messages,
+        model=model_name,  # type: ignore
+        temperature=temperature,
+        timeout=20,
+        max_tokens=max_tokens,
+        **(model_kwargs),
+    ) as stream:
+        for text in stream.text_stream:
+            aggregated_response += text
+        final_message = stream.get_final_message()
+
+    # Calculate cost of chat
+    input_tokens = final_message.usage.input_tokens
+    output_tokens = final_message.usage.output_tokens
+    tracer["usage"] = get_chat_usage_metrics(model_name, input_tokens, output_tokens, tracer.get("usage"))
+
+    # Save conversation trace
+    tracer["chat_model"] = model_name
+    tracer["temperature"] = temperature
+    if is_promptrace_enabled():
+        commit_conversation_trace(messages, aggregated_response, tracer)
+
+    return aggregated_response
+
+
+@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 anthropic_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,
+    tracer={},
+):
+    g = ThreadedGenerator(compiled_references, online_results, completion_func=completion_func)
+    t = Thread(
+        target=anthropic_llm_thread,
+        args=(g, messages, system_prompt, model_name, temperature, api_key, max_prompt_size, model_kwargs, tracer),
+    )
+    t.start()
+    return g
+
+
+def anthropic_llm_thread(
+    g, messages, system_prompt, model_name, temperature, api_key, max_prompt_size=None, model_kwargs=None, tracer={}
+):
+    try:
+        if api_key not in anthropic_clients:
+            client: anthropic.Anthropic = anthropic.Anthropic(api_key=api_key)
+            anthropic_clients[api_key] = client
+        else:
+            client: anthropic.Anthropic = anthropic_clients[api_key]
+
+        formatted_messages: List[anthropic.types.MessageParam] = [
+            anthropic.types.MessageParam(role=message.role, content=message.content) for message in messages
+        ]
+
+        aggregated_response = ""
+        final_message = None
+        with client.messages.stream(
+            messages=formatted_messages,
+            model=model_name,  # type: ignore
+            temperature=temperature,
+            system=system_prompt,
+            timeout=20,
+            max_tokens=DEFAULT_MAX_TOKENS_ANTHROPIC,
+            **(model_kwargs or dict()),
+        ) as stream:
+            for text in stream.text_stream:
+                aggregated_response += text
+                g.send(text)
+            final_message = stream.get_final_message()
+
+        # Calculate cost of chat
+        input_tokens = final_message.usage.input_tokens
+        output_tokens = final_message.usage.output_tokens
+        tracer["usage"] = get_chat_usage_metrics(model_name, input_tokens, output_tokens, tracer.get("usage"))
+
+        # Save conversation trace
+        tracer["chat_model"] = model_name
+        tracer["temperature"] = temperature
+        if is_promptrace_enabled():
+            commit_conversation_trace(messages, aggregated_response, tracer)
+    except Exception as e:
+        logger.error(f"Error in anthropic_llm_thread: {e}", exc_info=True)
+    finally:
+        g.close()
+
+
+def format_messages_for_anthropic(messages: list[ChatMessage], system_prompt=None):
+    """
+    Format messages for Anthropic
+    """
+    # Extract system prompt
+    system_prompt = system_prompt or ""
+    for message in messages.copy():
+        if message.role == "system":
+            system_prompt += message.content
+            messages.remove(message)
+    system_prompt = None if is_none_or_empty(system_prompt) else system_prompt
+
+    # Anthropic requires the first message to be a 'user' message
+    if len(messages) == 1:
+        messages[0].role = "user"
+    elif len(messages) > 1 and messages[0].role == "assistant":
+        messages = messages[1:]
+
+    # Convert image urls to base64 encoded images in Anthropic message format
+    for message in messages:
+        if isinstance(message.content, list):
+            content = []
+            # Sort the content. Anthropic models prefer that text comes after images.
+            message.content.sort(key=lambda x: 0 if x["type"] == "image_url" else 1)
+            for idx, part in enumerate(message.content):
+                if part["type"] == "text":
+                    content.append({"type": "text", "text": part["text"]})
+                elif part["type"] == "image_url":
+                    image = get_image_from_url(part["image_url"]["url"], type="b64")
+                    # Prefix each image with text block enumerating the image number
+                    # This helps the model reference the image in its response. Recommended by Anthropic
+                    content.extend(
+                        [
+                            {
+                                "type": "text",
+                                "text": f"Image {idx + 1}:",
+                            },
+                            {
+                                "type": "image",
+                                "source": {"type": "base64", "media_type": image.type, "data": image.content},
+                            },
+                        ]
+                    )
+            message.content = content
+
+    return messages, system_prompt

khoj/processor/conversation/google/gemini_chat.py

@@ -0,0 +1,253 @@
+import logging
+from datetime import datetime, timedelta
+from typing import Dict, List, Optional
+
+import pyjson5
+from langchain.schema import ChatMessage
+
+from khoj.database.models import Agent, ChatModel, KhojUser
+from khoj.processor.conversation import prompts
+from khoj.processor.conversation.google.utils import (
+    format_messages_for_gemini,
+    gemini_chat_completion_with_backoff,
+    gemini_completion_with_backoff,
+)
+from khoj.processor.conversation.utils import (
+    clean_json,
+    construct_structured_message,
+    generate_chatml_messages_with_context,
+    messages_to_print,
+)
+from khoj.utils.helpers import (
+    ConversationCommand,
+    is_none_or_empty,
+    truncate_code_context,
+)
+from khoj.utils.rawconfig import FileAttachment, LocationData
+from khoj.utils.yaml import yaml_dump
+
+logger = logging.getLogger(__name__)
+
+
+def extract_questions_gemini(
+    text,
+    model: Optional[str] = "gemini-1.5-flash",
+    conversation_log={},
+    api_key=None,
+    temperature=0,
+    max_tokens=None,
+    location_data: LocationData = None,
+    user: KhojUser = None,
+    query_images: Optional[list[str]] = None,
+    vision_enabled: bool = False,
+    personality_context: Optional[str] = None,
+    query_files: str = None,
+    tracer: dict = {},
+):
+    """
+    Infer search queries to retrieve relevant notes to answer user query
+    """
+    # Extract Past User Message and Inferred Questions from Conversation Log
+    location = f"{location_data}" if location_data else "Unknown"
+    username = prompts.user_name.format(name=user.get_full_name()) if user and user.get_full_name() else ""
+
+    # Extract Past User Message and Inferred Questions from Conversation Log
+    chat_history = "".join(
+        [
+            f'User: {chat["intent"]["query"]}\nAssistant: {{"queries": {chat["intent"].get("inferred-queries") or list([chat["intent"]["query"]])}}}\nA: {chat["message"]}\n\n'
+            for chat in conversation_log.get("chat", [])[-4:]
+            if chat["by"] == "khoj"
+        ]
+    )
+
+    # Get dates relative to today for prompt creation
+    today = datetime.today()
+    current_new_year = today.replace(month=1, day=1)
+    last_new_year = current_new_year.replace(year=today.year - 1)
+
+    system_prompt = prompts.extract_questions_anthropic_system_prompt.format(
+        current_date=today.strftime("%Y-%m-%d"),
+        day_of_week=today.strftime("%A"),
+        current_month=today.strftime("%Y-%m"),
+        last_new_year=last_new_year.strftime("%Y"),
+        last_new_year_date=last_new_year.strftime("%Y-%m-%d"),
+        current_new_year_date=current_new_year.strftime("%Y-%m-%d"),
+        yesterday_date=(today - timedelta(days=1)).strftime("%Y-%m-%d"),
+        location=location,
+        username=username,
+        personality_context=personality_context,
+    )
+
+    prompt = prompts.extract_questions_anthropic_user_message.format(
+        chat_history=chat_history,
+        text=text,
+    )
+
+    prompt = construct_structured_message(
+        message=prompt,
+        images=query_images,
+        model_type=ChatModel.ModelType.GOOGLE,
+        vision_enabled=vision_enabled,
+        attached_file_context=query_files,
+    )
+
+    messages = []
+
+    messages.append(ChatMessage(content=prompt, role="user"))
+    messages.append(ChatMessage(content=system_prompt, role="system"))
+
+    response = gemini_send_message_to_model(
+        messages, api_key, model, response_type="json_object", temperature=temperature, tracer=tracer
+    )
+
+    # Extract, Clean Message from Gemini's Response
+    try:
+        response = clean_json(response)
+        response = pyjson5.loads(response)
+        response = [q.strip() for q in response["queries"] if q.strip()]
+        if not isinstance(response, list) or not response:
+            logger.error(f"Invalid response for constructing subqueries: {response}")
+            return [text]
+        return response
+    except:
+        logger.warning(f"Gemini returned invalid JSON. Falling back to using user message as search query.\n{response}")
+        questions = [text]
+    logger.debug(f"Extracted Questions by Gemini: {questions}")
+    return questions
+
+
+def gemini_send_message_to_model(
+    messages,
+    api_key,
+    model,
+    response_type="text",
+    temperature=0,
+    model_kwargs=None,
+    tracer={},
+):
+    """
+    Send message to model
+    """
+    messages, system_prompt = format_messages_for_gemini(messages)
+
+    model_kwargs = {}
+
+    # Sometimes, this causes unwanted behavior and terminates response early. Disable for now while it's flaky.
+    # if response_type == "json_object":
+    #     model_kwargs["response_mime_type"] = "application/json"
+
+    # Get Response from Gemini
+    return gemini_completion_with_backoff(
+        messages=messages,
+        system_prompt=system_prompt,
+        model_name=model,
+        api_key=api_key,
+        temperature=temperature,
+        model_kwargs=model_kwargs,
+        tracer=tracer,
+    )
+
+
+def converse_gemini(
+    references,
+    user_query,
+    online_results: Optional[Dict[str, Dict]] = None,
+    code_results: Optional[Dict[str, Dict]] = None,
+    conversation_log={},
+    model: Optional[str] = "gemini-1.5-flash",
+    api_key: Optional[str] = None,
+    temperature: float = 0.2,
+    completion_func=None,
+    conversation_commands=[ConversationCommand.Default],
+    max_prompt_size=None,
+    tokenizer_name=None,
+    location_data: LocationData = None,
+    user_name: str = None,
+    agent: Agent = None,
+    query_images: Optional[list[str]] = None,
+    vision_available: bool = False,
+    query_files: str = None,
+    generated_files: List[FileAttachment] = None,
+    generated_asset_results: Dict[str, Dict] = {},
+    program_execution_context: List[str] = None,
+    tracer={},
+):
+    """
+    Converse with user using Google's Gemini
+    """
+    # Initialize Variables
+    current_date = datetime.now()
+
+    if agent and agent.personality:
+        system_prompt = prompts.custom_personality.format(
+            name=agent.name,
+            bio=agent.personality,
+            current_date=current_date.strftime("%Y-%m-%d"),
+            day_of_week=current_date.strftime("%A"),
+        )
+    else:
+        system_prompt = prompts.personality.format(
+            current_date=current_date.strftime("%Y-%m-%d"),
+            day_of_week=current_date.strftime("%A"),
+        )
+
+    system_prompt += f"{system_prompt}\n\n{prompts.gemini_verbose_language_personality}"
+    if location_data:
+        location_prompt = prompts.user_location.format(location=f"{location_data}")
+        system_prompt = f"{system_prompt}\n{location_prompt}"
+
+    if user_name:
+        user_name_prompt = prompts.user_name.format(name=user_name)
+        system_prompt = f"{system_prompt}\n{user_name_prompt}"
+
+    # Get Conversation Primer appropriate to Conversation Type
+    if conversation_commands == [ConversationCommand.Notes] and is_none_or_empty(references):
+        completion_func(chat_response=prompts.no_notes_found.format())
+        return iter([prompts.no_notes_found.format()])
+    elif conversation_commands == [ConversationCommand.Online] and is_none_or_empty(online_results):
+        completion_func(chat_response=prompts.no_online_results_found.format())
+        return iter([prompts.no_online_results_found.format()])
+
+    context_message = ""
+    if not is_none_or_empty(references):
+        context_message = f"{prompts.notes_conversation.format(query=user_query, references=yaml_dump(references))}\n\n"
+    if ConversationCommand.Online in conversation_commands or ConversationCommand.Webpage in conversation_commands:
+        context_message += f"{prompts.online_search_conversation.format(online_results=yaml_dump(online_results))}\n\n"
+    if ConversationCommand.Code in conversation_commands and not is_none_or_empty(code_results):
+        context_message += (
+            f"{prompts.code_executed_context.format(code_results=truncate_code_context(code_results))}\n\n"
+        )
+    context_message = context_message.strip()
+
+    # Setup Prompt with Primer or Conversation History
+    messages = generate_chatml_messages_with_context(
+        user_query,
+        context_message=context_message,
+        conversation_log=conversation_log,
+        model_name=model,
+        max_prompt_size=max_prompt_size,
+        tokenizer_name=tokenizer_name,
+        query_images=query_images,
+        vision_enabled=vision_available,
+        model_type=ChatModel.ModelType.GOOGLE,
+        query_files=query_files,
+        generated_files=generated_files,
+        generated_asset_results=generated_asset_results,
+        program_execution_context=program_execution_context,
+    )
+
+    messages, system_prompt = format_messages_for_gemini(messages, system_prompt)
+    logger.debug(f"Conversation Context for Gemini: {messages_to_print(messages)}")
+
+    # Get Response from Google AI
+    return gemini_chat_completion_with_backoff(
+        messages=messages,
+        compiled_references=references,
+        online_results=online_results,
+        model_name=model,
+        temperature=temperature,
+        api_key=api_key,
+        system_prompt=system_prompt,
+        completion_func=completion_func,
+        tracer=tracer,
+    )