khoj 1.24.2.dev16__py3-none-any.whl → 1.25.1.dev34__py3-none-any.whl

khoj/processor/conversation/prompts.py

@@ -45,6 +45,13 @@ Instructions:\n{bio}
 """.strip()
 )
 
+# To make Gemini be more verbose and match language of user's query.
+# Prompt forked from https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models
+gemini_verbose_language_personality = """
+All questions should be answered comprehensively with details, unless the user requests a concise response specifically.
+Respond in the same language as the query. Use markdown to format your responses.
+""".strip()
+
 ## General Conversation
 ## --
 general_conversation = PromptTemplate.from_template(
@@ -404,6 +411,10 @@ Tell the user exactly what the document says in response to their query, while a
 extract_relevant_summary = PromptTemplate.from_template(
     """
 {personality_context}
+
+Conversation History:
+{chat_history}
+
 Target Query: {query}
 
 Document Contents:
@@ -415,10 +426,10 @@ Collate only relevant information from the document to answer the target query.
 
 personality_context = PromptTemplate.from_template(
     """
-    Here's some additional context about you:
-    {personality}
+Here's some additional context about you:
+{personality}
 
-    """
+"""
 )
 
 pick_relevant_output_mode = PromptTemplate.from_template(

khoj/processor/conversation/utils.py

@@ -223,7 +223,7 @@ def truncate_messages(
 ) -> list[ChatMessage]:
     """Truncate messages to fit within max prompt size supported by model"""
 
-    default_tokenizer = "hf-internal-testing/llama-tokenizer"
+    default_tokenizer = "gpt-4o"
 
     try:
         if loaded_model:
@@ -240,13 +240,9 @@ def truncate_messages(
         else:
             encoder = download_model(model_name).tokenizer()
     except:
-        if default_tokenizer in state.pretrained_tokenizers:
-            encoder = state.pretrained_tokenizers[default_tokenizer]
-        else:
-            encoder = AutoTokenizer.from_pretrained(default_tokenizer)
-            state.pretrained_tokenizers[default_tokenizer] = encoder
+        encoder = tiktoken.encoding_for_model(default_tokenizer)
         logger.debug(
-            f"Fallback to default chat model tokenizer: {tokenizer_name}.\nConfigure tokenizer for unsupported model: {model_name} in Khoj settings to improve context stuffing."
+            f"Fallback to default chat model tokenizer: {default_tokenizer}.\nConfigure tokenizer for model: {model_name} in Khoj settings to improve context stuffing."
         )
 
     # Extract system message from messages

khoj/processor/embeddings.py

@@ -13,7 +13,7 @@ from tenacity import (
 )
 from torch import nn
 
-from khoj.utils.helpers import get_device, merge_dicts
+from khoj.utils.helpers import get_device, merge_dicts, timer
 from khoj.utils.rawconfig import SearchResponse
 
 logger = logging.getLogger(__name__)
@@ -37,7 +37,8 @@ class EmbeddingsModel:
         self.model_name = model_name
         self.inference_endpoint = embeddings_inference_endpoint
         self.api_key = embeddings_inference_endpoint_api_key
-        self.embeddings_model = SentenceTransformer(self.model_name, **self.model_kwargs)
+        with timer(f"Loaded embedding model {self.model_name}", logger):
+            self.embeddings_model = SentenceTransformer(self.model_name, **self.model_kwargs)
 
     def inference_server_enabled(self) -> bool:
         return self.api_key is not None and self.inference_endpoint is not None
@@ -101,7 +102,8 @@ class CrossEncoderModel:
         self.inference_endpoint = cross_encoder_inference_endpoint
         self.api_key = cross_encoder_inference_endpoint_api_key
         self.model_kwargs = merge_dicts(model_kwargs, {"device": get_device()})
-        self.cross_encoder_model = CrossEncoder(model_name=self.model_name, **self.model_kwargs)
+        with timer(f"Loaded cross-encoder model {self.model_name}", logger):
+            self.cross_encoder_model = CrossEncoder(model_name=self.model_name, **self.model_kwargs)
 
     def inference_server_enabled(self) -> bool:
         return self.api_key is not None and self.inference_endpoint is not None
@@ -112,6 +114,7 @@ class CrossEncoderModel:
             payload = {"inputs": {"query": query, "passages": [hit.additional[key] for hit in hits]}}
             headers = {"Authorization": f"Bearer {self.api_key}", "Content-Type": "application/json"}
             response = requests.post(target_url, json=payload, headers=headers)
+            response.raise_for_status()
             return response.json()["scores"]
 
         cross_inp = [[query, hit.additional[key]] for hit in hits]

khoj/processor/image/generate.py

@@ -25,7 +25,6 @@ async def text_to_image(
     location_data: LocationData,
     references: List[Dict[str, Any]],
     online_results: Dict[str, Any],
-    subscribed: bool = False,
     send_status_func: Optional[Callable] = None,
     uploaded_image_url: Optional[str] = None,
     agent: Agent = None,
@@ -66,8 +65,8 @@ async def text_to_image(
         note_references=references,
         online_results=online_results,
         model_type=text_to_image_config.model_type,
-        subscribed=subscribed,
         uploaded_image_url=uploaded_image_url,
+        user=user,
         agent=agent,
     )
 

khoj/processor/tools/online_search.py

@@ -10,14 +10,22 @@ import aiohttp
 from bs4 import BeautifulSoup
 from markdownify import markdownify
 
-from khoj.database.models import Agent, KhojUser
+from khoj.database.adapters import ConversationAdapters
+from khoj.database.models import Agent, KhojUser, WebScraper
+from khoj.processor.conversation import prompts
 from khoj.routers.helpers import (
     ChatEvent,
     extract_relevant_info,
     generate_online_subqueries,
     infer_webpage_urls,
 )
-from khoj.utils.helpers import is_internet_connected, is_none_or_empty, timer
+from khoj.utils.helpers import (
+    is_env_var_true,
+    is_internal_url,
+    is_internet_connected,
+    is_none_or_empty,
+    timer,
+)
 from khoj.utils.rawconfig import LocationData
 
 logger = logging.getLogger(__name__)
@@ -25,12 +33,11 @@ logger = logging.getLogger(__name__)
 SERPER_DEV_API_KEY = os.getenv("SERPER_DEV_API_KEY")
 SERPER_DEV_URL = "https://google.serper.dev/search"
 
-JINA_READER_API_URL = "https://r.jina.ai/"
 JINA_SEARCH_API_URL = "https://s.jina.ai/"
 JINA_API_KEY = os.getenv("JINA_API_KEY")
 
-OLOSTEP_API_KEY = os.getenv("OLOSTEP_API_KEY")
-OLOSTEP_API_URL = "https://agent.olostep.com/olostep-p2p-incomingAPI"
+FIRECRAWL_USE_LLM_EXTRACT = is_env_var_true("FIRECRAWL_USE_LLM_EXTRACT")
+
 OLOSTEP_QUERY_PARAMS = {
     "timeout": 35,  # seconds
     "waitBeforeScraping": 1,  # seconds
@@ -53,7 +60,6 @@ async def search_online(
     conversation_history: dict,
     location: LocationData,
     user: KhojUser,
-    subscribed: bool = False,
     send_status_func: Optional[Callable] = None,
     custom_filters: List[str] = [],
     uploaded_image_url: str = None,
@@ -84,33 +90,36 @@ async def search_online(
         search_results = await asyncio.gather(*search_tasks)
         response_dict = {subquery: search_result for subquery, search_result in search_results}
 
-    # Gather distinct web page data from organic results of each subquery without an instant answer.
+    # Gather distinct web pages from organic results for subqueries without an instant answer.
     # Content of web pages is directly available when Jina is used for search.
-    webpages = {
-        (organic.get("link"), subquery, organic.get("content"))
-        for subquery in response_dict
-        for organic in response_dict[subquery].get("organic", [])[:MAX_WEBPAGES_TO_READ]
-        if "answerBox" not in response_dict[subquery]
-    }
+    webpages: Dict[str, Dict] = {}
+    for subquery in response_dict:
+        if "answerBox" in response_dict[subquery]:
+            continue
+        for organic in response_dict[subquery].get("organic", [])[:MAX_WEBPAGES_TO_READ]:
+            link = organic.get("link")
+            if link in webpages:
+                webpages[link]["queries"].add(subquery)
+            else:
+                webpages[link] = {"queries": {subquery}, "content": organic.get("content")}
 
     # Read, extract relevant info from the retrieved web pages
     if webpages:
-        webpage_links = set([link for link, _, _ in webpages])
-        logger.info(f"Reading web pages at: {list(webpage_links)}")
+        logger.info(f"Reading web pages at: {webpages.keys()}")
         if send_status_func:
-            webpage_links_str = "\n- " + "\n- ".join(list(webpage_links))
+            webpage_links_str = "\n- " + "\n- ".join(webpages.keys())
             async for event in send_status_func(f"**Reading web pages**: {webpage_links_str}"):
                 yield {ChatEvent.STATUS: event}
     tasks = [
-        read_webpage_and_extract_content(subquery, link, content, subscribed=subscribed, agent=agent)
-        for link, subquery, content in webpages
+        read_webpage_and_extract_content(data["queries"], link, data["content"], user=user, agent=agent)
+        for link, data in webpages.items()
     ]
     results = await asyncio.gather(*tasks)
 
     # Collect extracted info from the retrieved web pages
-    for subquery, webpage_extract, url in results:
+    for subqueries, url, webpage_extract in results:
         if webpage_extract is not None:
-            response_dict[subquery]["webpages"] = {"link": url, "snippet": webpage_extract}
+            response_dict[subqueries.pop()]["webpages"] = {"link": url, "snippet": webpage_extract}
 
     yield response_dict
 
@@ -141,7 +150,6 @@ async def read_webpages(
     conversation_history: dict,
     location: LocationData,
     user: KhojUser,
-    subscribed: bool = False,
     send_status_func: Optional[Callable] = None,
     uploaded_image_url: str = None,
     agent: Agent = None,
@@ -158,29 +166,66 @@ async def read_webpages(
         webpage_links_str = "\n- " + "\n- ".join(list(urls))
         async for event in send_status_func(f"**Reading web pages**: {webpage_links_str}"):
             yield {ChatEvent.STATUS: event}
-    tasks = [read_webpage_and_extract_content(query, url, subscribed=subscribed, agent=agent) for url in urls]
+    tasks = [read_webpage_and_extract_content({query}, url, user=user, agent=agent) for url in urls]
     results = await asyncio.gather(*tasks)
 
     response: Dict[str, Dict] = defaultdict(dict)
     response[query]["webpages"] = [
-        {"query": q, "link": url, "snippet": web_extract} for q, web_extract, url in results if web_extract is not None
+        {"query": qs.pop(), "link": url, "snippet": extract} for qs, url, extract in results if extract is not None
     ]
     yield response
 
 
+async def read_webpage(
+    url, scraper_type=None, api_key=None, api_url=None, subqueries=None, agent=None
+) -> Tuple[str | None, str | None]:
+    if scraper_type == WebScraper.WebScraperType.FIRECRAWL and FIRECRAWL_USE_LLM_EXTRACT:
+        return None, await query_webpage_with_firecrawl(url, subqueries, api_key, api_url, agent)
+    elif scraper_type == WebScraper.WebScraperType.FIRECRAWL:
+        return await read_webpage_with_firecrawl(url, api_key, api_url), None
+    elif scraper_type == WebScraper.WebScraperType.OLOSTEP:
+        return await read_webpage_with_olostep(url, api_key, api_url), None
+    elif scraper_type == WebScraper.WebScraperType.JINA:
+        return await read_webpage_with_jina(url, api_key, api_url), None
+    else:
+        return await read_webpage_at_url(url), None
+
+
 async def read_webpage_and_extract_content(
-    subquery: str, url: str, content: str = None, subscribed: bool = False, agent: Agent = None
-) -> Tuple[str, Union[None, str], str]:
-    try:
-        if is_none_or_empty(content):
-            with timer(f"Reading web page at '{url}' took", logger):
-                content = await read_webpage_with_olostep(url) if OLOSTEP_API_KEY else await read_webpage_with_jina(url)
-        with timer(f"Extracting relevant information from web page at '{url}' took", logger):
-            extracted_info = await extract_relevant_info(subquery, content, subscribed=subscribed, agent=agent)
-        return subquery, extracted_info, url
-    except Exception as e:
-        logger.error(f"Failed to read web page at '{url}' with {e}")
-        return subquery, None, url
+    subqueries: set[str], url: str, content: str = None, user: KhojUser = None, agent: Agent = None
+) -> Tuple[set[str], str, Union[None, str]]:
+    # Select the web scrapers to use for reading the web page
+    web_scrapers = await ConversationAdapters.aget_enabled_webscrapers()
+    # Only use the direct web scraper for internal URLs
+    if is_internal_url(url):
+        web_scrapers = [scraper for scraper in web_scrapers if scraper.type == WebScraper.WebScraperType.DIRECT]
+
+    # Fallback through enabled web scrapers until we successfully read the web page
+    extracted_info = None
+    for scraper in web_scrapers:
+        try:
+            # Read the web page
+            if is_none_or_empty(content):
+                with timer(f"Reading web page with {scraper.type} at '{url}' took", logger, log_level=logging.INFO):
+                    content, extracted_info = await read_webpage(
+                        url, scraper.type, scraper.api_key, scraper.api_url, subqueries, agent
+                    )
+
+            # Extract relevant information from the web page
+            if is_none_or_empty(extracted_info):
+                with timer(f"Extracting relevant information from web page at '{url}' took", logger):
+                    extracted_info = await extract_relevant_info(subqueries, content, user=user, agent=agent)
+
+            # If we successfully extracted information, break the loop
+            if not is_none_or_empty(extracted_info):
+                break
+        except Exception as e:
+            logger.warning(f"Failed to read web page with {scraper.type} at '{url}' with {e}")
+            # If this is the last web scraper in the list, log an error
+            if scraper.name == web_scrapers[-1].name:
+                logger.error(f"All web scrapers failed for '{url}'")
+
+    return subqueries, url, extracted_info
 
 
 async def read_webpage_at_url(web_url: str) -> str:
@@ -197,23 +242,23 @@ async def read_webpage_at_url(web_url: str) -> str:
             return markdownify(body)
 
 
-async def read_webpage_with_olostep(web_url: str) -> str:
-    headers = {"Authorization": f"Bearer {OLOSTEP_API_KEY}"}
+async def read_webpage_with_olostep(web_url: str, api_key: str, api_url: str) -> str:
+    headers = {"Authorization": f"Bearer {api_key}"}
     web_scraping_params: Dict[str, Union[str, int, bool]] = OLOSTEP_QUERY_PARAMS.copy()  # type: ignore
     web_scraping_params["url"] = web_url
 
     async with aiohttp.ClientSession() as session:
-        async with session.get(OLOSTEP_API_URL, params=web_scraping_params, headers=headers) as response:
+        async with session.get(api_url, params=web_scraping_params, headers=headers) as response:
             response.raise_for_status()
             response_json = await response.json()
             return response_json["markdown_content"]
 
 
-async def read_webpage_with_jina(web_url: str) -> str:
-    jina_reader_api_url = f"{JINA_READER_API_URL}/{web_url}"
+async def read_webpage_with_jina(web_url: str, api_key: str, api_url: str) -> str:
+    jina_reader_api_url = f"{api_url}/{web_url}"
     headers = {"Accept": "application/json", "X-Timeout": "30"}
-    if JINA_API_KEY:
-        headers["Authorization"] = f"Bearer {JINA_API_KEY}"
+    if api_key:
+        headers["Authorization"] = f"Bearer {api_key}"
 
     async with aiohttp.ClientSession() as session:
         async with session.get(jina_reader_api_url, headers=headers) as response:
@@ -222,6 +267,54 @@ async def read_webpage_with_jina(web_url: str) -> str:
             return response_json["data"]["content"]
 
 
+async def read_webpage_with_firecrawl(web_url: str, api_key: str, api_url: str) -> str:
+    firecrawl_api_url = f"{api_url}/v1/scrape"
+    headers = {"Content-Type": "application/json", "Authorization": f"Bearer {api_key}"}
+    params = {"url": web_url, "formats": ["markdown"], "excludeTags": ["script", ".ad"]}
+
+    async with aiohttp.ClientSession() as session:
+        async with session.post(firecrawl_api_url, json=params, headers=headers) as response:
+            response.raise_for_status()
+            response_json = await response.json()
+            return response_json["data"]["markdown"]
+
+
+async def query_webpage_with_firecrawl(
+    web_url: str, queries: set[str], api_key: str, api_url: str, agent: Agent = None
+) -> str:
+    firecrawl_api_url = f"{api_url}/v1/scrape"
+    headers = {"Content-Type": "application/json", "Authorization": f"Bearer {api_key}"}
+    schema = {
+        "type": "object",
+        "properties": {
+            "relevant_extract": {"type": "string"},
+        },
+        "required": [
+            "relevant_extract",
+        ],
+    }
+
+    personality_context = (
+        prompts.personality_context.format(personality=agent.personality) if agent and agent.personality else ""
+    )
+    system_prompt = f"""
+{prompts.system_prompt_extract_relevant_information}
+
+{personality_context}
+User Query: {", ".join(queries)}
+
+Collate only relevant information from the website to answer the target query and in the provided JSON schema.
+""".strip()
+
+    params = {"url": web_url, "formats": ["extract"], "extract": {"systemPrompt": system_prompt, "schema": schema}}
+
+    async with aiohttp.ClientSession() as session:
+        async with session.post(firecrawl_api_url, json=params, headers=headers) as response:
+            response.raise_for_status()
+            response_json = await response.json()
+            return response_json["data"]["extract"]["relevant_extract"]
+
+
 async def search_with_jina(query: str, location: LocationData) -> Tuple[str, Dict[str, List[Dict]]]:
     encoded_query = urllib.parse.quote(query)
     jina_search_api_url = f"{JINA_SEARCH_API_URL}/{encoded_query}"

khoj/routers/api.py

@@ -395,7 +395,7 @@ async def extract_references_and_questions(
     # Infer search queries from user message
     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()
+        conversation_config = await ConversationAdapters.aget_default_conversation_config(user)
         vision_enabled = conversation_config.vision_enabled
 
         if conversation_config.model_type == ChatModelOptions.ModelType.OFFLINE:

khoj/routers/api_agents.py

@@ -35,6 +35,7 @@ class ModifyAgentBody(BaseModel):
     files: Optional[List[str]] = []
     input_tools: Optional[List[str]] = []
     output_modes: Optional[List[str]] = []
+    slug: Optional[str] = None
 
 
 @api_agents.get("", response_class=Response)
@@ -161,7 +162,7 @@ async def delete_agent(
 
 
 @api_agents.post("", response_class=Response)
-@requires(["authenticated"])
+@requires(["authenticated", "premium"])
 async def create_agent(
     request: Request,
     common: CommonQueryParams,
@@ -192,6 +193,7 @@ async def create_agent(
         body.files,
         body.input_tools,
         body.output_modes,
+        body.slug,
     )
 
     agents_packet = {
@@ -213,7 +215,7 @@ async def create_agent(
 
 
 @api_agents.patch("", response_class=Response)
-@requires(["authenticated"])
+@requires(["authenticated", "premium"])
 async def update_agent(
     request: Request,
     common: CommonQueryParams,
@@ -233,7 +235,7 @@ async def update_agent(
             status_code=400,
         )
 
-    selected_agent = await AgentAdapters.aget_agent_by_name(body.name, user)
+    selected_agent = await AgentAdapters.aget_agent_by_slug(body.slug, user)
 
     if not selected_agent:
         return Response(
@@ -253,6 +255,7 @@ async def update_agent(
         body.files,
         body.input_tools,
         body.output_modes,
+        body.slug,
     )
 
     agents_packet = {